--- title: anyflow ---

AnyFlow on-policy distillation method.

Stage 2 of the AnyFlow two-stage recipe. Continues from a pretrained flow-map student and refines it via distribution-matching distillation (DMD2) where the student is rolled out for student_sample_steps Euler-flow steps from pure noise. One randomly-chosen step in the rollout is gradient-enabled (and broadcast across ranks so every worker agrees on which step to gradient-enable); the rest run under torch.no_grad.

Inherits DMD2Method for the alternating student / critic update machinery and the existing DMD VSD-with-fake-score loss. Overrides _student_rollout to drive the multi-step Euler-flow rollout with r = t_next (mean-velocity sampling — matches the AnyFlow paper's WanAnyFlowPipeline.training_rollout with use_mean_velocity=True).

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

fastvideo.train.methods.distribution_matching.anyflow.DMD2Method

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

DMD2 distillation algorithm (method layer).

Owns role model instances directly: - self.student — trainable student :class:ModelBase - self.teacher — frozen teacher :class:ModelBase - self.critic — trainable critic :class:ModelBase

fastvideo.train.methods.distribution_matching.anyflow.DMD2Method.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.DMD2Method.backward

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

fastvideo.train.methods.distribution_matching.anyflow.DMD2Method.get_optimizers

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

fastvideo.train.methods.distribution_matching.anyflow.DMD2Method.get_lr_schedulers

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

fastvideo.train.methods.distribution_matching.anyflow.DMD2Method.get_grad_clip_targets

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

fastvideo.train.methods.distribution_matching.anyflow.get_optional_float

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

fastvideo.train.methods.distribution_matching.anyflow.get_optional_int

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

fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod

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

AnyFlow on-policy distillation (multi-step rollout).

Source
class AnyFlowMethod(DMD2Method):
    """AnyFlow on-policy distillation (multi-step rollout)."""

    def __init__(
        self,
        *,
        cfg: Any,
        role_models: dict[str, Any],
    ) -> None:
        super().__init__(cfg=cfg, role_models=role_models)
        mcfg = self.method_config

        student_sample_steps = get_optional_int(mcfg, "student_sample_steps", where="method.student_sample_steps")
        if student_sample_steps is None:
            student_sample_steps = 4
        if int(student_sample_steps) <= 0:
            raise ValueError("method.student_sample_steps must be positive, "
                             f"got {student_sample_steps}")
        self._student_sample_steps = int(student_sample_steps)

        use_mean_velocity_raw = mcfg.get("use_mean_velocity", True)
        if not isinstance(use_mean_velocity_raw, bool):
            raise ValueError("method.use_mean_velocity must be a bool, "
                             f"got {type(use_mean_velocity_raw).__name__}")
        self._use_mean_velocity = bool(use_mean_velocity_raw)

        # Optional pinned rollout schedule (descending, absolute t-units).
        # Falls back to dmd_denoising_steps when absent.
        raw_t_list = mcfg.get("t_list_override", None)
        if raw_t_list is None:
            self._t_list_override: list[float] | None = None
        else:
            if not isinstance(raw_t_list, list) or not raw_t_list:
                raise ValueError("method.t_list_override must be a non-empty list of "
                                 f"floats when set, got {raw_t_list!r}")
            t_list = [float(x) for x in raw_t_list]
            for i in range(len(t_list) - 1):
                if t_list[i] < t_list[i + 1]:
                    raise ValueError("method.t_list_override must be descending, "
                                     f"got {t_list!r}")
            self._t_list_override = t_list

        # Scoring conditioning: AnyFlow scores against r=0 for the DMD branch.
        score_r_raw = mcfg.get("dmd_score_r_value", 0.0)
        try:
            self._dmd_score_r = float(score_r_raw)
        except (TypeError, ValueError) as exc:
            raise ValueError("method.dmd_score_r_value must be numeric, "
                             f"got {score_r_raw!r}") from exc

        # Optional teacher guidance scale for the DMD loss (carry over from
        # DMD2Method's behavior; default 1.0).
        guidance = get_optional_float(mcfg, "real_score_guidance_scale", where="method.real_score_guidance_scale")
        self._real_score_guidance = float(guidance) if guidance is not None else 1.0

    # ------------------------------------------------------------------
    # Rollout schedule

    def _get_rollout_schedule(self, *, device: torch.device) -> torch.Tensor:
        """Build the descending timestep schedule used by the on-policy
        rollout. Length is ``num_steps + 1`` so ``num_steps`` Euler steps
        consume the full range.

        Order of precedence:
          1. ``method.t_list_override`` — used verbatim (absolute units).
          2. ``method.dmd_denoising_steps`` (inherited from DMD2) appended
             with a final 0 boundary if the last entry isn't already 0.
        """
        if self._t_list_override is not None:
            return torch.tensor(self._t_list_override, device=device, dtype=torch.float32)

        steps = self._get_denoising_step_list(device).to(dtype=torch.float32)
        if float(steps[-1].item()) != 0.0:
            zero = torch.zeros(1, device=device, dtype=torch.float32)
            steps = torch.cat([steps, zero], dim=0)
        return steps

    def _broadcast_grad_step_index(
        self,
        num_steps: int,
        *,
        device: torch.device,
    ) -> int:
        """Pick the rollout step that gets gradient enabled. In distributed
        runs the choice is broadcast from rank 0 so every worker agrees."""
        if num_steps <= 0:
            raise ValueError("num_steps must be positive")
        if dist.is_initialized() and dist.get_rank() != 0:
            idx_tensor = torch.empty((1, ), dtype=torch.long, device=device)
        else:
            idx_tensor = torch.randint(0,
                                       num_steps, (1, ),
                                       device=device,
                                       dtype=torch.long,
                                       generator=self.cuda_generator)
        if dist.is_initialized():
            dist.broadcast(idx_tensor, src=0)
        return int(idx_tensor.item())

    # ------------------------------------------------------------------
    # Rollout

    def _student_rollout(
        self,
        batch: Any,
        *,
        with_grad: bool,
    ) -> torch.Tensor:
        """Multi-step Euler-flow rollout from pure noise.

        Returns the predicted clean latent ``x_0`` after the chosen
        gradient step (or the final ``x`` after the last step if
        ``with_grad`` is False — used by the critic path).
        """
        latents = batch.latents
        if latents is None or latents.ndim != 5:
            raise RuntimeError("AnyFlow on-policy rollout requires TrainingBatch.latents "
                               "of shape [B, T, C, H, W] for shape templating")
        device = latents.device
        dtype = latents.dtype

        schedule = self._get_rollout_schedule(device=device)
        num_entries = int(schedule.numel())
        num_steps = num_entries - 1
        if num_steps <= 0:
            raise RuntimeError("rollout schedule must have at least two entries "
                               f"(got {num_entries})")
        if num_steps > self._student_sample_steps:
            # Trim to the configured cap, keeping the last (=0) boundary.
            schedule = torch.cat([schedule[:self._student_sample_steps], schedule[-1:]], dim=0)
            num_steps = self._student_sample_steps

        grad_step = self._broadcast_grad_step_index(num_steps, device=device) if with_grad else -1

        attn_kind: Literal["dense", "vsa"] = "vsa"
        n_train = float(self.student.num_train_timesteps)

        x = torch.randn(latents.shape, device=device, dtype=dtype, generator=self.cuda_generator)
        last_pred_x0: torch.Tensor | None = None
        batch_size = int(latents.shape[0])

        for i in range(num_steps):
            t_cur = schedule[i].expand(batch_size)
            t_next = schedule[i + 1].expand(batch_size)
            r = t_next if self._use_mean_velocity else t_cur

            enable_grad = bool(with_grad) and (i == grad_step)
            with torch.set_grad_enabled(enable_grad):
                v = self.student.predict_velocity_with_r(
                    x,
                    t_cur,
                    r,
                    batch,
                    conditional=True,
                    cfg_uncond=self._cfg_uncond,
                    attn_kind=attn_kind,
                )
            # Euler step in absolute units: x ← x - ((t_cur - t_next) / N) * v.
            view = [-1] + [1] * (x.ndim - 1)
            dt = ((t_cur - t_next) / n_train).view(*view)
            x = x - dt * v

            if enable_grad:
                # We treat the rollout output (post-step) as a predicted
                # clean latent — AnyFlow's last Euler step lands at t=0.
                last_pred_x0 = x

        if last_pred_x0 is None:
            # No gradient step taken (with_grad=False path).
            last_pred_x0 = x

        if hasattr(batch, "dmd_latent_vis_dict"):
            batch.dmd_latent_vis_dict["generator_timestep"] = (schedule[-1].detach().clone())
        return last_pred_x0