[feat] AnyFlow pretrain: central-difference target + full single_train_step
Implements AnyFlow's central-difference loss path end-to-end:
1. WanModel gains predict_velocity_with_r(noisy, t, r, batch, ...) — mirrors
predict_noise's forward_context + autocast plumbing but injects
r_timestep into the transformer kwargs.
2. anyflow_pretrain._central_difference_dF_dt — symmetric finite difference
over (t ± delta) with the sample also moved along the flow trajectory by
v_pred * (delta / num_train_timesteps), mirroring AnyFlow's reference
trainer_wan_anyflow_pretrain.py::compute_central_difference. Wrapped in
torch.no_grad so the two extra forwards stay out of the backward graph.
3. AnyFlowPretrainMethod.single_train_step:
- sample (t, r) ∈ [0, 1] via _sample_pair_timesteps
- apply scheduler shift + scale to absolute units
- rebuild noisy with the flow-map scheduler's add_noise
- run conditional + (optional) unconditional student forwards
- apply guidance distillation when fuse_guidance_scale != 1
- compute target = (eps - x0) - (t - r) * dF/dt
- per-sample MSE * per-timestep weight
- stop-grad scale balance so non-diffusion losses match the diffusion
branch's magnitude
- emit metrics for diffusion/consistency fractions + scale weight mean
4. AnyFlowPretrainMethod.backward overrides TrainingMethod.backward to
route the loss through self.student.backward inside the correct
forward_context, matching DMD2Method's pattern.
This commit is contained in:
Enderfga
committed by
SolitaryThinker
SolitaryThinker
parent
a8b99d2113
commit
f6d78fa389
@@ -433,3 +433,107 @@ def test_sample_pair_timesteps_rejects_negative_ratios() -> None:
|
||||
device=torch.device("cpu"),
|
||||
generator=None,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 7: central-difference target.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _StubStudent:
|
||||
"""Stand-in student for unit-testing the central-difference helper.
|
||||
|
||||
The "velocity prediction" is a closed-form function of x and t
|
||||
(no actual neural network) so we can compare against the analytical
|
||||
derivative.
|
||||
"""
|
||||
|
||||
def __init__(self, alpha: float = 0.3) -> None:
|
||||
self.alpha = float(alpha)
|
||||
|
||||
def predict_velocity_with_r(
|
||||
self,
|
||||
noisy: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
r: torch.Tensor,
|
||||
batch,
|
||||
*,
|
||||
conditional: bool = True,
|
||||
attn_kind: str = "dense",
|
||||
cfg_uncond=None,
|
||||
) -> torch.Tensor:
|
||||
del batch, conditional, attn_kind, cfg_uncond, r
|
||||
# f(x, t) = x + alpha * (t / 1000) → dF/dt = alpha / 1000.
|
||||
view = [-1] + [1] * (noisy.ndim - 1)
|
||||
return noisy + self.alpha * (t.view(*view).float() / 1000.0)
|
||||
|
||||
|
||||
def test_central_difference_dF_dt_linear_function() -> None:
|
||||
"""For f(x, t) = x + alpha * (t / N), the central-difference estimate
|
||||
of dF/dt must be alpha / N (in absolute t-units that's alpha / N
|
||||
velocity per t-unit, and our helper divides the difference by
|
||||
2*delta -> exactly alpha / N)."""
|
||||
from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
|
||||
_central_difference_dF_dt, )
|
||||
|
||||
student = _StubStudent(alpha=0.3)
|
||||
x = torch.randn(2, 1, 4, 4, 4)
|
||||
latents = torch.zeros_like(x)
|
||||
noise = torch.zeros_like(x) # v_pred = noise - latents = 0
|
||||
t = torch.tensor([500.0, 250.0])
|
||||
r = torch.tensor([100.0, 0.0])
|
||||
|
||||
dF = _central_difference_dF_dt(
|
||||
student=student,
|
||||
batch=None,
|
||||
noisy=x,
|
||||
latents=latents,
|
||||
noise=noise,
|
||||
t=t,
|
||||
r=r,
|
||||
delta=5.0,
|
||||
num_train_timesteps=1000.0,
|
||||
)
|
||||
|
||||
expected = torch.full_like(x, 0.3 / 1000.0)
|
||||
torch.testing.assert_close(dF, expected, rtol=1e-5, atol=1e-5)
|
||||
|
||||
|
||||
def test_central_difference_dF_dt_guidance_scaling() -> None:
|
||||
"""With guidance_scale != 1, the result must be divided by guidance."""
|
||||
from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
|
||||
_central_difference_dF_dt, )
|
||||
|
||||
student = _StubStudent(alpha=0.6)
|
||||
x = torch.zeros(1, 1, 4, 4, 4)
|
||||
latents = torch.zeros_like(x)
|
||||
noise = torch.zeros_like(x)
|
||||
t = torch.tensor([500.0])
|
||||
r = torch.tensor([100.0])
|
||||
|
||||
dF_g1 = _central_difference_dF_dt(
|
||||
student=student, batch=None, noisy=x, latents=latents, noise=noise,
|
||||
t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=1.0)
|
||||
dF_g3 = _central_difference_dF_dt(
|
||||
student=student, batch=None, noisy=x, latents=latents, noise=noise,
|
||||
t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=3.0)
|
||||
|
||||
torch.testing.assert_close(dF_g1 / 3.0, dF_g3, rtol=1e-5, atol=1e-5)
|
||||
|
||||
|
||||
def test_central_difference_dF_dt_rejects_zero_delta() -> None:
|
||||
from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
|
||||
_central_difference_dF_dt, )
|
||||
|
||||
with pytest.raises(ValueError, match="delta must be positive"):
|
||||
_central_difference_dF_dt(
|
||||
student=_StubStudent(),
|
||||
batch=None,
|
||||
noisy=torch.zeros(1, 1, 4, 4, 4),
|
||||
latents=torch.zeros(1, 1, 4, 4, 4),
|
||||
noise=torch.zeros(1, 1, 4, 4, 4),
|
||||
t=torch.tensor([500.0]),
|
||||
r=torch.tensor([100.0]),
|
||||
delta=0.0,
|
||||
num_train_timesteps=1000.0,
|
||||
)
|
||||
|
||||
@@ -92,6 +92,58 @@ def _sample_pair_timesteps(
|
||||
return t, r, is_diffusion, is_consistency
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _central_difference_dF_dt(
|
||||
*,
|
||||
student: Any,
|
||||
batch: Any,
|
||||
noisy: torch.Tensor,
|
||||
latents: torch.Tensor,
|
||||
noise: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
r: torch.Tensor,
|
||||
delta: float,
|
||||
num_train_timesteps: float,
|
||||
attn_kind: str = "dense",
|
||||
guidance_scale: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Estimate ``dF/dt`` for the AnyFlow central-difference target.
|
||||
|
||||
Computes a symmetric finite difference of the velocity prediction in
|
||||
*absolute train-timestep units*:
|
||||
|
||||
dF/dt ≈ [u_θ(x_{t+δ}, t+δ, r) - u_θ(x_{t-δ}, t-δ, r)] / (2 * δ * guidance)
|
||||
|
||||
The sample is also moved along the flow trajectory by the same
|
||||
finite step (``v_pred * (δ / N)``) to match AnyFlow's reference
|
||||
formulation in ``trainer_wan_anyflow_pretrain.py::compute_central_difference``.
|
||||
Wrapped in ``torch.no_grad`` so the two extra forwards never enter
|
||||
the backward graph.
|
||||
"""
|
||||
if delta <= 0.0:
|
||||
raise ValueError(f"delta must be positive, got {delta}")
|
||||
if guidance_scale <= 0.0:
|
||||
raise ValueError(
|
||||
f"guidance_scale must be positive, got {guidance_scale}")
|
||||
|
||||
v_pred = noise - latents # ground-truth flow velocity
|
||||
delta_x = delta / float(num_train_timesteps)
|
||||
|
||||
t_plus = t + delta
|
||||
noisy_plus = noisy + v_pred * delta_x
|
||||
f_plus = student.predict_velocity_with_r(
|
||||
noisy_plus, t_plus, r, batch,
|
||||
conditional=True, attn_kind=attn_kind)
|
||||
|
||||
t_minus = t - delta
|
||||
noisy_minus = noisy - v_pred * delta_x
|
||||
f_minus = student.predict_velocity_with_r(
|
||||
noisy_minus, t_minus, r, batch,
|
||||
conditional=True, attn_kind=attn_kind)
|
||||
|
||||
return (f_plus - f_minus) / (2.0 * delta * guidance_scale)
|
||||
|
||||
|
||||
class AnyFlowPretrainMethod(TrainingMethod):
|
||||
"""AnyFlow flow-map pretrain method.
|
||||
|
||||
@@ -206,7 +258,155 @@ class AnyFlowPretrainMethod(TrainingMethod):
|
||||
dict[str, LogScalar],
|
||||
]:
|
||||
del iteration # AnyFlow pretrain has no iteration-dependent dispatch.
|
||||
raise NotImplementedError("filled in Task 7")
|
||||
|
||||
training_batch = self.student.prepare_batch(
|
||||
batch,
|
||||
generator=self.cuda_generator,
|
||||
latents_source="data",
|
||||
)
|
||||
latents = training_batch.latents # [B, T, C, H, W] (post-permute in prepare_batch).
|
||||
if latents is None or latents.ndim != 5:
|
||||
raise RuntimeError(
|
||||
"AnyFlow pretrain expects TrainingBatch.latents of shape "
|
||||
"[B, T, C, H, W] after prepare_batch; got "
|
||||
f"{None if latents is None else tuple(latents.shape)}")
|
||||
device = latents.device
|
||||
dtype = latents.dtype
|
||||
batch_size = int(latents.shape[0])
|
||||
|
||||
# AnyFlow (t, r) sampling — overrides the timestep drawn by
|
||||
# WanModel._sample_timesteps inside prepare_batch.
|
||||
t_norm, r_norm, is_diffusion, is_consistency = _sample_pair_timesteps(
|
||||
batch_size=batch_size,
|
||||
diffusion_ratio=self._diffusion_ratio,
|
||||
consistency_ratio=self._consistency_ratio,
|
||||
device=device,
|
||||
generator=self.cuda_generator,
|
||||
)
|
||||
|
||||
sched = self._flow_map_scheduler
|
||||
n_train = float(self.student.num_train_timesteps)
|
||||
t = (sched.apply_shift(t_norm) * n_train).to(device=device, dtype=dtype)
|
||||
r = (sched.apply_shift(r_norm) * n_train).to(device=device, dtype=dtype)
|
||||
|
||||
# Fresh noise drawn from the method's RNG; ignore the noise that
|
||||
# prepare_batch attached (it pairs with the discarded timestep).
|
||||
noise = torch.randn(
|
||||
latents.shape,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
generator=self.cuda_generator,
|
||||
)
|
||||
noisy = sched.add_noise(latents, noise, t)
|
||||
|
||||
# Keep training_batch coherent with the new (t, noisy): downstream
|
||||
# forward_context uses these fields.
|
||||
training_batch.timesteps = t
|
||||
training_batch.noise = noise
|
||||
training_batch.noisy_model_input = noisy.permute(0, 2, 1, 3, 4)
|
||||
|
||||
# Student velocity prediction at (t, r).
|
||||
noise_pred = self.student.predict_velocity_with_r(
|
||||
noisy, t, r, training_batch,
|
||||
conditional=True,
|
||||
attn_kind="dense",
|
||||
)
|
||||
|
||||
# Optional guidance distillation — fuse CFG into the training target so
|
||||
# the resulting checkpoint can be sampled at guidance_scale=1.0.
|
||||
if self._fuse_guidance_scale != 1.0:
|
||||
with torch.no_grad():
|
||||
noise_pred_uncond = self.student.predict_velocity_with_r(
|
||||
noisy, t, r, training_batch,
|
||||
conditional=False,
|
||||
attn_kind="dense",
|
||||
)
|
||||
g = float(self._fuse_guidance_scale)
|
||||
noise_pred = (noise_pred - (1.0 - g) * noise_pred_uncond) / g
|
||||
|
||||
dF_dt = _central_difference_dF_dt(
|
||||
student=self.student,
|
||||
batch=training_batch,
|
||||
noisy=noisy,
|
||||
latents=latents,
|
||||
noise=noise,
|
||||
t=t,
|
||||
r=r,
|
||||
delta=self._fd_epsilon,
|
||||
num_train_timesteps=n_train,
|
||||
attn_kind="dense",
|
||||
guidance_scale=self._fuse_guidance_scale,
|
||||
)
|
||||
|
||||
# AnyFlow target: target = (eps - x_0) - (t - r) * dF/dt
|
||||
# dF/dt is in (velocity per absolute t-unit); (t - r) is in absolute units;
|
||||
# the product cancels back to velocity units, matching noise_pred.
|
||||
view = [batch_size] + [1] * (latents.ndim - 1)
|
||||
target = (noise - latents) - (t - r).view(*view) * dF_dt
|
||||
|
||||
# Per-sample squared error, then per-timestep weight, then scale-balance
|
||||
# so the non-diffusion branches stay on the same magnitude as the
|
||||
# diffusion branch (matches AnyFlow's stop-grad rescaling).
|
||||
per_sample = torch.mean(
|
||||
((noise_pred.float() - target.float()) ** 2).reshape(
|
||||
batch_size, -1),
|
||||
dim=-1,
|
||||
)
|
||||
weight = sched.get_train_weight(t, weight_type=self._weight_type)
|
||||
per_sample = per_sample * weight
|
||||
|
||||
with torch.no_grad():
|
||||
diff_mask = is_diffusion
|
||||
if diff_mask.any():
|
||||
diff_mean = per_sample[diff_mask].mean()
|
||||
else:
|
||||
diff_mean = per_sample.mean()
|
||||
non_diff_mask = ~diff_mask
|
||||
if non_diff_mask.any():
|
||||
scale = diff_mean / (per_sample[non_diff_mask] + 1e-5)
|
||||
else:
|
||||
scale = torch.tensor(1.0, device=device, dtype=per_sample.dtype)
|
||||
non_diff_idx = torch.nonzero(non_diff_mask, as_tuple=False).flatten()
|
||||
if non_diff_idx.numel() > 0:
|
||||
per_sample = per_sample.clone()
|
||||
per_sample[non_diff_idx] = per_sample[non_diff_idx] * scale
|
||||
|
||||
total_loss = per_sample.mean()
|
||||
|
||||
loss_map = {"total_loss": total_loss}
|
||||
metrics: dict[str, LogScalar] = {
|
||||
"diffusion_fraction": float(is_diffusion.float().mean()),
|
||||
"consistency_fraction": float(is_consistency.float().mean()),
|
||||
"scale_weight_mean": float(scale.mean()) if isinstance(
|
||||
scale, torch.Tensor) else float(scale),
|
||||
}
|
||||
outputs = {
|
||||
"student_ctx": (
|
||||
training_batch.timesteps,
|
||||
training_batch.attn_metadata,
|
||||
),
|
||||
}
|
||||
return loss_map, outputs, metrics
|
||||
|
||||
def backward(
|
||||
self,
|
||||
loss_map: dict[str, torch.Tensor],
|
||||
outputs: dict[str, Any],
|
||||
*,
|
||||
grad_accum_rounds: int = 1,
|
||||
) -> None:
|
||||
"""Route the loss backward through the student's forward_context
|
||||
so attn metadata stays attached during gradient computation."""
|
||||
student_ctx = outputs.get("student_ctx")
|
||||
if student_ctx is None:
|
||||
super().backward(
|
||||
loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
|
||||
return
|
||||
self.student.backward(
|
||||
loss_map["total_loss"],
|
||||
student_ctx,
|
||||
grad_accum_rounds=grad_accum_rounds,
|
||||
)
|
||||
|
||||
def _init_optimizer_and_scheduler(self) -> None:
|
||||
tc = self.training_config
|
||||
|
||||
@@ -358,6 +358,54 @@ class WanModel(ModelBase):
|
||||
pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
|
||||
return pred_noise
|
||||
|
||||
def predict_velocity_with_r(
|
||||
self,
|
||||
noisy_latents: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
r_timestep: torch.Tensor,
|
||||
batch: TrainingBatch,
|
||||
*,
|
||||
conditional: bool,
|
||||
cfg_uncond: dict[str, Any] | None = None,
|
||||
attn_kind: Literal["dense", "vsa"] = "dense",
|
||||
) -> torch.Tensor:
|
||||
"""AnyFlow forward: predict average velocity from ``t`` back to ``r``.
|
||||
|
||||
Same plumbing as :meth:`predict_noise` but injects ``r_timestep``
|
||||
into the transformer kwargs. The transformer must have been
|
||||
constructed with an arch config that sets ``r_embedder=True`` for
|
||||
the dual-timestep branch to be active — otherwise ``r_timestep``
|
||||
is silently ignored by the embedder and the forward reduces to
|
||||
the single-timestep path.
|
||||
"""
|
||||
device_type = self.device.type
|
||||
dtype = noisy_latents.dtype
|
||||
if conditional:
|
||||
text_dict = batch.conditional_dict
|
||||
if text_dict is None:
|
||||
raise RuntimeError("Missing conditional_dict in "
|
||||
"TrainingBatch")
|
||||
else:
|
||||
text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
|
||||
|
||||
if attn_kind == "dense":
|
||||
attn_metadata = batch.attn_metadata
|
||||
elif attn_kind == "vsa":
|
||||
attn_metadata = batch.attn_metadata_vsa
|
||||
else:
|
||||
raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
|
||||
|
||||
with torch.autocast(device_type, dtype=dtype), set_forward_context(
|
||||
current_timestep=batch.timesteps,
|
||||
attn_metadata=attn_metadata,
|
||||
):
|
||||
input_kwargs = (
|
||||
self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
|
||||
input_kwargs["r_timestep"] = r_timestep
|
||||
transformer = self._get_transformer(timestep)
|
||||
pred_velocity = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
|
||||
return pred_velocity
|
||||
|
||||
def backward(
|
||||
self,
|
||||
loss: torch.Tensor,
|
||||
|
||||
Reference in New Issue
Block a user