--- title: anyflow_pretrain ---

AnyFlow pretrain (flow-map central-difference) training method.

Stage 1 of the AnyFlow two-stage recipe. Trains a single student network u_θ(x_t, t, r) to predict the average velocity from time t back to time r via the central-difference target

target = (eps - x_0) - ((t - r) / N) * dF/dt

where N = num_train_timesteps and dF/dt is estimated from the student's own forward at (t ± δ, r) (with one-sided fallback near the schedule endpoints).

Per-batch (t, r) sampling follows the AnyFlow paper:

Reference: trainer_wan_anyflow_pretrain.py in NVlabs/AnyFlow at commit 549236a.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod

TrainingMethod(*, cfg: Any, role_models: dict[str, ModelBase])

Base training method (algorithm layer).

Subclasses own their role models (student, teacher, critic, …) as plain attributes and manage optimizers directly — no RoleManager or RoleHandle.

The constructor receives role_models (a dict[str, ModelBase]) and a cfg object. It calls init_preprocessors on the student and builds self.role_modules for FSDP wrapping.

A single shared CUDA RNG generator (cuda_generator) is created in :meth:on_train_start. All torch.randn / torch.randint calls in methods and models must use this generator instead of relying on global RNG state.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.set_tracker

set_tracker(tracker: Any) -> None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.single_train_step

single_train_step(batch: dict[str, Any], iteration: int) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.get_optimizers

get_optimizers(iteration: int) -> Sequence[torch.optim.Optimizer]

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.get_lr_schedulers

get_lr_schedulers(iteration: int) -> Sequence[Any]

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.checkpoint_state

checkpoint_state() -> dict[str, Any]

Return DCP-ready checkpoint state for all trainable roles.

Keys follow the convention: roles.<role>.<module>, optimizers.<role>, schedulers.<role>, random_state.*.

EMA state is managed by the EMACallback and is checkpointed through the callback state mechanism.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.backward

backward(loss_map: dict[str, torch.Tensor], outputs: dict[str, Any], *, grad_accum_rounds: int = 1) -> None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.optimizers_schedulers_step

optimizers_schedulers_step(iteration: int) -> None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.optimizers_zero_grad

optimizers_zero_grad(iteration: int) -> None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.seed_optimizer_state_for_resume

seed_optimizer_state_for_resume() -> None

Seed optimizer state so DCP can load saved state.

A fresh optimizer has empty state (exp_avg, exp_avg_sq, step are only created on the first optimizer.step()). DCP needs matching entries to load into; without them the saved optimizer state is silently dropped.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.manages_optimization

manages_optimization() -> bool

Whether the method owns backward/optimizer stepping internally.

Most methods return loss tensors and let :class:Trainer handle gradient accumulation, callbacks, optimizer stepping, and scheduler stepping. RL-style methods such as DiffusionNFT need to preserve their own sample-then-inner-train loop, so they can provide a specific managed_train_step.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.managed_train_step

managed_train_step(data_stream: Any, iteration: int) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]

Run one method-managed step.

Subclasses that return True from :meth:manages_optimization should override this. The fallback consumes one dataloader batch and delegates to single_train_step so tests can exercise the hook with tiny fake methods.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.on_validation_begin

on_validation_begin(iteration: int = 0) -> dict[str, LogScalar]

Run method-owned validation, if any.

Pipeline-style validation should remain in callbacks. Methods that intentionally avoid inference pipelines, such as RL methods with their own sampler/reward loop, can override this hook and return metrics for the trainer to log at iteration.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.get_grad_clip_targets

get_grad_clip_targets(iteration: int) -> dict[str, torch.nn.Module]

Return modules whose gradients should be clipped.

Override in subclasses to add/conditionally include modules (e.g. critic, conditionally student). Default: student transformer.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.apply_configured_lrs

apply_configured_lrs() -> None

Re-apply the config's learning rates to live optimizers/schedulers.

Called after a checkpoint resume when training.checkpoint.reset_lr_on_resume is set: DCP restores param-group lr/initial_lr and scheduler base_lrs from the checkpoint, which would otherwise silently override an LR change in the YAML. Methods with per-role LRs override this.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.inference_checkpoint_modules

inference_checkpoint_modules(role: str = 'student') -> dict[str, torch.nn.Module]

Return the complete modules that form an inference checkpoint.

This intentionally does not reuse :meth:checkpoint_state: resumable state includes optimizers and every trainable role, while inference exports select one explicit role and must retain frozen parameters too. Methods with a different deployable role contract may override this hook; EMA is never selected implicitly.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.inference_checkpoint_base_model_path

inference_checkpoint_base_model_path(role: str = 'student') -> str

Return the immutable model directory supplying non-trained parts.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.get_grad_sanitize_targets

get_grad_sanitize_targets(iteration: int) -> dict[str, torch.nn.Module]

Return active modules whose gradients must be made finite.

Most methods sanitize and clip the same modules. Methods that leave an optimizer role unclipped can override this independently.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.set_is_last_backward

set_is_last_backward(iteration: int, is_last_backward: bool) -> None

Mark an FSDP2 accumulation boundary without changing reductions.

Intermediate backwards still reduce-scatter their gradients. Only FSDP's pending-reduction waits and backward-prefetch cleanup are deferred. Use the active gradient targets, including roles whose gradients are sanitized but deliberately not clipped. Non-FSDP modules are unchanged.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.set_requires_gradient_sync

set_requires_gradient_sync(iteration: int, requires_gradient_sync: bool) -> None

Control FSDP gradient communication for accumulation rounds.

FSDP2 normally reduces gradients after every backward. During gradient accumulation only the final microbatch needs that communication; earlier rounds can accumulate local gradients. The active gradient targets include roles that are deliberately not clipped.

Non-FSDP modules do not expose set_requires_gradient_sync and are intentionally left unchanged.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.supports_gradient_accumulation_no_sync

supports_gradient_accumulation_no_sync(iteration: int) -> bool

Whether every active role safely retains unsharded gradients.

FSDP2 no-sync holds full, unreduced gradients between accumulation rounds. Model plugins must opt in explicitly; this prevents the optimization from adding roughly one full model's FP32 gradients to memory for ordinary full-weight training.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.TrainingMethod.on_train_start

on_train_start() -> None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase

ModelBase(*, trainable: bool = True, lora: LoraConfig | dict[str, Any] | None = None, attention_backend: AttentionBackendEnum | str | None = None)

Per-role model instance.

Every role (student, teacher, critic, …) gets its own ModelBase instance. Each instance owns its own transformer and noise_scheduler. Heavyweight resources (VAE, dataloader, RNG seeds) are loaded lazily via :meth:init_preprocessors, which the method calls only on the student.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.attention_backend_name

attention_backend_name: str | None

Explicit per-role backend name, or None for global/default.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.transformer_module_type

transformer_module_type: str

Physical checkpoint component that owns this role's transformer.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.device

device: torch.device

The local CUDA device for this rank.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.init_preprocessors

init_preprocessors(training_config: TrainingConfig) -> None

Load VAE, build dataloader, seed RNGs.

Called only on the student by the method's __init__. Default is a no-op so teacher/critic instances skip this.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.on_train_start

on_train_start() -> None

Called once before the training loop begins.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.eager_optimizer_state_parameters

eager_optimizer_state_parameters() -> tuple[torch.nn.Parameter, ...]

Parameters whose optimizer state must exist before their first grad.

Most models return an empty tuple. Models with intermittently used trainable parameters can override this so strict distributed checkpoints retain a stable optimizer-state schema across steps.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.decode_latents

decode_latents(latents_b_t_c_h_w: torch.Tensor) -> torch.Tensor

Decode [B, T, C, H, W] latents to [B, C, T, H, W] media.

RL reward methods call this hook instead of reaching into model-specific VAE normalization details.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.num_train_timesteps

num_train_timesteps: int

Return the scheduler's training timestep horizon.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.shift_and_clamp_timestep

shift_and_clamp_timestep(timestep: torch.Tensor) -> torch.Tensor

Apply model/pipeline timestep shifting and clamp.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.prepare_batch

prepare_batch(raw_batch: dict[str, Any], *, generator: torch.Generator, latents_source: Literal['data', 'zeros'] = 'data') -> TrainingBatch

Convert a dataloader batch into forward primitives.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.add_noise

add_noise(clean_latents: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor

Apply forward-process noise at timestep.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.add_noise_for_batch

add_noise_for_batch(clean_latents: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch) -> torch.Tensor

Apply noise with optional immutable geometry carried by batch.

Video-only and fixed-shape models keep their existing behavior. Joint packed models override this hook when splitting the tensor requires batch-local shape context.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.predict_noise

predict_noise(noisy_latents: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch, *, conditional: bool, cfg_uncond: dict[str, Any] | None = None, attn_kind: Literal['dense', 'vsa'] = 'dense') -> NoisePrediction

Predict video flow or an ordered (video, audio) flow pair.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.predict_x0

predict_x0(noisy_latents: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch, *, conditional: bool, cfg_uncond: dict[str, Any] | None = None, attn_kind: Literal['dense', 'vsa'] = 'dense') -> torch.Tensor

Convert a video-only flow prediction to clean video latents.

This helper owns one noisy video tensor and one video scheduler. Joint video/audio callers apply their modality-specific conversions where both noisy tensors and both schedulers are available.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.ModelBase.backward

backward(loss: torch.Tensor, ctx: Any, *, grad_accum_rounds: int) -> None

Backward that may restore forward-context.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.get_optional_float

get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> float | None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.get_optional_int

get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | None

fastvideo.train.methods.distribution_matching.anyflow_pretrain.build_optimizer_and_scheduler

build_optimizer_and_scheduler(*, params: list[torch.nn.Parameter], optimizer_config: OptimizerConfig, loop_config: TrainingLoopConfig, learning_rate: float, betas: tuple[float, float], scheduler_name: str, weight_decay: float | None = None) -> tuple[torch.optim.Optimizer, object]

Build an AdamW optimizer and LR scheduler.

Returns (optimizer, lr_scheduler) so the caller can store them as method-level attributes.

fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod

AnyFlowPretrainMethod(*, cfg: Any, role_models: dict[str, ModelBase])

AnyFlow flow-map pretrain method.

Single-student training; no teacher or critic. The student must implement predict_velocity_with_r(noisy, t, r, batch, ...) — typically a WanModel with r_embedder=True in its arch config.

Source
class AnyFlowPretrainMethod(TrainingMethod):
    """AnyFlow flow-map pretrain method.

    Single-student training; no teacher or critic. The student must
    implement ``predict_velocity_with_r(noisy, t, r, batch, ...)`` —
    typically a ``WanModel`` with ``r_embedder=True`` in its arch config.
    """

    def __init__(
        self,
        *,
        cfg: Any,
        role_models: dict[str, ModelBase],
    ) -> None:
        super().__init__(cfg=cfg, role_models=role_models)
        if "student" not in role_models:
            raise ValueError("AnyFlowPretrainMethod requires role 'student'")
        if not self.student._trainable:
            raise ValueError("AnyFlowPretrainMethod requires student to be trainable")

        mcfg = self.method_config
        self._diffusion_ratio = float(
            get_optional_float(mcfg, "diffusion_ratio", where="method.diffusion_ratio") or 0.5)
        self._consistency_ratio = float(
            get_optional_float(mcfg, "consistency_ratio", where="method.consistency_ratio") or 0.25)
        if self._diffusion_ratio + self._consistency_ratio > 1.0:
            raise ValueError("method.diffusion_ratio + method.consistency_ratio must "
                             f"be <= 1, got {self._diffusion_ratio} + "
                             f"{self._consistency_ratio}")

        # δ: finite-difference step in absolute train-timestep units.
        epsilon = get_optional_int(mcfg, "epsilon", where="method.epsilon")
        self._fd_epsilon = float(epsilon) if epsilon is not None else 5.0

        # Loss weighting scheme (uniform / gaussian / beta08).
        raw_weight_type = mcfg.get("weight_type", "beta08")
        if not isinstance(raw_weight_type, str):
            raise ValueError("method.weight_type must be a string, got "
                             f"{type(raw_weight_type).__name__}")
        weight_type = raw_weight_type.strip().lower()
        if weight_type not in {"uniform", "gaussian", "beta08"}:
            raise ValueError("method.weight_type must be one of "
                             "{uniform, gaussian, beta08}, "
                             f"got {raw_weight_type!r}")
        self._weight_type = weight_type

        # Guidance fused into the training target (default 1.0 = unused).
        fg = get_optional_float(mcfg, "fuse_guidance_scale", where="method.fuse_guidance_scale")
        self._fuse_guidance_scale = float(fg) if fg is not None else 1.0
        if self._fuse_guidance_scale <= 0.0:
            raise ValueError("method.fuse_guidance_scale must be positive, "
                             f"got {self._fuse_guidance_scale}")

        # Flow-map scheduler — uses pipeline_config.flow_shift if present
        # and falls back to method.shift (and finally 1.0).
        shift = float(getattr(self.training_config.pipeline_config, "flow_shift", 0.0) or 0.0)
        if shift <= 0.0:
            shift_override = get_optional_float(mcfg, "shift", where="method.shift")
            shift = float(shift_override) if shift_override is not None else 1.0
        self._shift = shift

        # Lazy-imported to avoid circular imports on package load.
        from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
            FlowMapEulerDiscreteScheduler, )
        self._flow_map_scheduler = FlowMapEulerDiscreteScheduler(
            num_train_timesteps=int(self.student.num_train_timesteps),
            shift=self._shift,
        )

        self.student.init_preprocessors(self.training_config)
        self._init_optimizer_and_scheduler()

    @property
    def _optimizer_dict(self) -> dict[str, torch.optim.Optimizer]:
        return {"student": self._student_optimizer}

    @property
    def _lr_scheduler_dict(self) -> dict[str, Any]:
        return {"student": self._student_lr_scheduler}

    def get_optimizers(
        self,
        iteration: int,
    ) -> Sequence[torch.optim.Optimizer]:
        del iteration
        return [self._student_optimizer]

    def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
        del iteration
        return [self._student_lr_scheduler]

    def single_train_step(
        self,
        batch: dict[str, Any],
        iteration: int,
    ) -> tuple[
            dict[str, torch.Tensor],
            dict[str, Any],
            dict[str, LogScalar],
    ]:
        del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.

        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
            diff_mean = per_sample[diff_mask].mean() if diff_mask.any() else 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
        params = [p for p in self.student.transformer.parameters() if p.requires_grad]
        (
            self._student_optimizer,
            self._student_lr_scheduler,
        ) = build_optimizer_and_scheduler(
            params=params,
            optimizer_config=tc.optimizer,
            loop_config=tc.loop,
            learning_rate=float(tc.optimizer.learning_rate),
            betas=tc.optimizer.betas,
            scheduler_name=str(tc.optimizer.lr_scheduler),
        )

fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.get_optimizers

get_optimizers(iteration: int) -> Sequence[torch.optim.Optimizer]
Source
def get_optimizers(
    self,
    iteration: int,
) -> Sequence[torch.optim.Optimizer]:
    del iteration
    return [self._student_optimizer]

fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.get_lr_schedulers

get_lr_schedulers(iteration: int) -> Sequence[Any]
Source
def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
    del iteration
    return [self._student_lr_scheduler]

fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.single_train_step

single_train_step(batch: dict[str, Any], iteration: int) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]
Source
def single_train_step(
    self,
    batch: dict[str, Any],
    iteration: int,
) -> tuple[
        dict[str, torch.Tensor],
        dict[str, Any],
        dict[str, LogScalar],
]:
    del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.

    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
        diff_mean = per_sample[diff_mask].mean() if diff_mask.any() else 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

fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.backward

backward(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.

Source
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,
    )