--- title: causal_cd ---

Causal consistency distillation method (algorithm layer).

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler

SelfForcingFlowMatchScheduler(num_inference_steps = 100, num_train_timesteps = 1000, shift = 3.0, sigma_max = 1.0, sigma_min = 0.003 / 1.002, inverse_timesteps = False, extra_one_step = False, reverse_sigmas = False, training = False)

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.set_timesteps

set_timesteps(num_inference_steps = 100, denoising_strength = 1.0, training = False, return_dict = False, **kwargs)

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.step

step(model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final = False, return_dict = False, **kwargs)

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high

calculate_alpha_beta_high(sigma, sigma_bound)

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.add_noise

add_noise(original_samples, noise, timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.add_noise_high

add_noise_high(original_samples, noise, timestep, boundary_timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.training_target

training_target(sample, noise, timestep)

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.training_weight

training_weight(timestep)

Input

Output: the corresponding weighting [B*T]

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor

fastvideo.train.methods.consistency_model.causal_cd.SelfForcingFlowMatchScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.TrainingMethod.set_tracker

set_tracker(tracker: Any) -> None

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.TrainingMethod.get_optimizers

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

fastvideo.train.methods.consistency_model.causal_cd.TrainingMethod.get_lr_schedulers

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

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.TrainingMethod.backward

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

fastvideo.train.methods.consistency_model.causal_cd.TrainingMethod.optimizers_schedulers_step

optimizers_schedulers_step(iteration: int) -> None

fastvideo.train.methods.consistency_model.causal_cd.TrainingMethod.optimizers_zero_grad

optimizers_zero_grad(iteration: int) -> None

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.TrainingMethod.on_train_start

on_train_start() -> None

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.ModelBase.attention_backend_name

attention_backend_name: str | None

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

fastvideo.train.methods.consistency_model.causal_cd.ModelBase.transformer_module_type

transformer_module_type: str

Physical checkpoint component that owns this role's transformer.

fastvideo.train.methods.consistency_model.causal_cd.ModelBase.device

device: torch.device

The local CUDA device for this rank.

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.ModelBase.on_train_start

on_train_start() -> None

Called once before the training loop begins.

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.ModelBase.num_train_timesteps

num_train_timesteps: int

Return the scheduler's training timestep horizon.

fastvideo.train.methods.consistency_model.causal_cd.ModelBase.shift_and_clamp_timestep

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

Apply model/pipeline timestep shifting and clamp.

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.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.consistency_model.causal_cd.ModelBase.backward

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

Backward that may restore forward-context.

fastvideo.train.methods.consistency_model.causal_cd.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.consistency_model.causal_cd.CausalConsistencyDistillationMethod

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

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

        for role in ("student", "teacher", "ema"):
            if role not in role_models:
                raise ValueError(f"Causal-CD requires role {role!r} "
                                 "(student trainable; teacher + ema frozen, "
                                 "both initialized from the student's "
                                 "checkpoint)")
        if not self.student._trainable:
            raise ValueError("Causal-CD requires student to be trainable")
        self.teacher = role_models["teacher"]
        self.ema_model = role_models["ema"]

        self._attn_kind = self._infer_attn_kind()
        self._guidance_scale = float(self.method_config.get("guidance_scale", 3.0))
        self._discrete_cd_n = int(self.method_config.get("discrete_cd_N", 48))
        if self._discrete_cd_n < 2:
            raise ValueError("method.discrete_cd_N must be >= 2")
        self._ema_decay = float(self.method_config.get("ema_decay", 0.99))
        self._ema_start_step = int(self.method_config.get("ema_start_step", 200))
        shift = getattr(self.training_config.pipeline_config, "flow_shift", None)
        self._flow_shift = float(shift) if shift else 5.0

        self.student.init_preprocessors(self.training_config)
        self._sf_scheduler = SelfForcingFlowMatchScheduler(
            num_inference_steps=self._discrete_cd_n,
            num_train_timesteps=int(self.student.num_train_timesteps),
            shift=self._flow_shift,
            sigma_min=0.0,
            sigma_max=1.0,
            extra_one_step=True,
            training=False,
        )
        self._init_optimizers_and_schedulers()

    # ------------------------------------------------------------------

    @property
    def _optimizer_dict(self) -> dict[str, Any]:
        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) -> list[torch.optim.Optimizer]:
        del iteration
        return [self._student_optimizer]

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

    def checkpoint_state(self) -> dict[str, Any]:
        # The EMA role is frozen (so the base class skips it) but mutated by
        # _update_ema every step; without persisting it a resume reloads the
        # EMA from init_from and the consistency target snaps back to the
        # base checkpoint. Mirrors DiffusionNFT's frozen "old" role.
        states = super().checkpoint_state()
        states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer)
        return states

    # ------------------------------------------------------------------

    def single_train_step(
        self,
        batch: dict[str, Any],
        iteration: int,
    ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]:
        del iteration
        training_batch = self.student.prepare_batch(
            batch,
            generator=self.cuda_generator,
            latents_source="data",
        )
        clean_latents = training_batch.latents
        if not torch.is_tensor(clean_latents) or clean_latents.ndim != 5:
            raise ValueError("Causal-CD expects [B, T, C, H, W] latents")

        batch_size, num_latents = int(clean_latents.shape[0]), int(clean_latents.shape[1])
        device = clean_latents.device

        sigmas = self._sf_scheduler.sigmas.to(device)
        timesteps = self._sf_scheduler.timesteps.to(device)
        idx = torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).squeeze(0)
        t, t_next = timesteps[idx], timesteps[idx + 1]
        sigma_t, sigma_t_next = sigmas[idx], sigmas[idx + 1]
        t_pf = t * torch.ones(batch_size, num_latents, device=device)
        t_next_pf = t_next * torch.ones(batch_size, num_latents, device=device)

        noise = torch.randn(
            clean_latents.shape,
            generator=self.cuda_generator,
            device=device,
            dtype=clean_latents.dtype,
        )
        latent_t = (1.0 - sigma_t) * clean_latents + sigma_t * noise

        # Set before any forward: predict_noise feeds batch.timesteps into
        # set_forward_context (VSA sparsity gating), so the teacher CFG
        # passes below must not see the stale timesteps from prepare_batch.
        training_batch.timesteps = t_pf

        with torch.no_grad():
            v_cond = self._predict_flow(self.teacher,
                                        latent_t,
                                        t_pf,
                                        training_batch,
                                        conditional=True,
                                        clean_x=clean_latents)
            v_uncond = self._predict_flow(self.teacher,
                                          latent_t,
                                          t_pf,
                                          training_batch,
                                          conditional=False,
                                          clean_x=clean_latents)
            v_pred = v_uncond + self._guidance_scale * (v_cond - v_uncond)
            dt = ((t - t_next) / float(self.student.num_train_timesteps))
            latent_t_next = latent_t - dt * v_pred

        flow_student = self._predict_flow(self.student,
                                          latent_t,
                                          t_pf,
                                          training_batch,
                                          conditional=True,
                                          clean_x=clean_latents)
        x0_t = latent_t - sigma_t * flow_student

        with torch.no_grad():
            flow_ema = self._predict_flow(self.ema_model,
                                          latent_t_next,
                                          t_next_pf,
                                          training_batch,
                                          conditional=True,
                                          clean_x=clean_latents)
            x0_t_next = latent_t_next - sigma_t_next * flow_ema

        loss = F.mse_loss(x0_t.float(), x0_t_next.float())

        loss_map = {"total_loss": loss, "causal_cd_loss": loss}
        attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata)
        outputs: dict[str, Any] = {"_fv_backward": (t_pf, attn_metadata)}
        metrics: dict[str, LogScalar] = {}
        return loss_map, outputs, metrics

    # ------------------------------------------------------------------

    def backward(
        self,
        loss_map: dict[str, torch.Tensor],
        outputs: dict[str, Any],
        *,
        grad_accum_rounds: int = 1,
    ) -> None:
        grad_accum_rounds = max(1, int(grad_accum_rounds))
        ctx = outputs.get("_fv_backward")
        if ctx is None:
            super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
            return
        self.student.backward(loss_map["total_loss"], ctx, grad_accum_rounds=grad_accum_rounds)

    def optimizers_schedulers_step(self, iteration: int) -> None:
        super().optimizers_schedulers_step(iteration)
        if iteration >= self._ema_start_step:
            self._update_ema()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _predict_flow(
        self,
        model: ModelBase,
        latents: torch.Tensor,
        timestep: torch.Tensor,
        batch: Any,
        *,
        conditional: bool,
        clean_x: torch.Tensor,
    ) -> torch.Tensor:
        return model.predict_noise(latents,
                                   timestep,
                                   batch,
                                   conditional=conditional,
                                   cfg_uncond=None,
                                   attn_kind=self._attn_kind,
                                   clean_x=clean_x)

    @torch.no_grad()
    def _update_ema(self) -> None:
        decay = self._ema_decay
        for ema_p, p in zip(self.ema_model.transformer.parameters(), self.student.transformer.parameters(),
                            strict=True):
            ema_p.mul_(decay).add_(p.detach().to(ema_p.dtype), alpha=1.0 - decay)

    def _init_optimizers_and_schedulers(self) -> None:
        tc = self.training_config
        student_lr = float(tc.optimizer.learning_rate)
        if student_lr <= 0.0:
            raise ValueError("training.learning_rate must be > 0 for causal-cd")
        student_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=student_params,
            optimizer_config=tc.optimizer,
            loop_config=tc.loop,
            learning_rate=student_lr,
            betas=tc.optimizer.betas,
            scheduler_name=str(tc.optimizer.lr_scheduler),
        )

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.get_optimizers

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

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.get_lr_schedulers

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

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.checkpoint_state

checkpoint_state() -> dict[str, Any]
Source
def checkpoint_state(self) -> dict[str, Any]:
    # The EMA role is frozen (so the base class skips it) but mutated by
    # _update_ema every step; without persisting it a resume reloads the
    # EMA from init_from and the consistency target snaps back to the
    # base checkpoint. Mirrors DiffusionNFT's frozen "old" role.
    states = super().checkpoint_state()
    states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer)
    return states

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.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
    training_batch = self.student.prepare_batch(
        batch,
        generator=self.cuda_generator,
        latents_source="data",
    )
    clean_latents = training_batch.latents
    if not torch.is_tensor(clean_latents) or clean_latents.ndim != 5:
        raise ValueError("Causal-CD expects [B, T, C, H, W] latents")

    batch_size, num_latents = int(clean_latents.shape[0]), int(clean_latents.shape[1])
    device = clean_latents.device

    sigmas = self._sf_scheduler.sigmas.to(device)
    timesteps = self._sf_scheduler.timesteps.to(device)
    idx = torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).squeeze(0)
    t, t_next = timesteps[idx], timesteps[idx + 1]
    sigma_t, sigma_t_next = sigmas[idx], sigmas[idx + 1]
    t_pf = t * torch.ones(batch_size, num_latents, device=device)
    t_next_pf = t_next * torch.ones(batch_size, num_latents, device=device)

    noise = torch.randn(
        clean_latents.shape,
        generator=self.cuda_generator,
        device=device,
        dtype=clean_latents.dtype,
    )
    latent_t = (1.0 - sigma_t) * clean_latents + sigma_t * noise

    # Set before any forward: predict_noise feeds batch.timesteps into
    # set_forward_context (VSA sparsity gating), so the teacher CFG
    # passes below must not see the stale timesteps from prepare_batch.
    training_batch.timesteps = t_pf

    with torch.no_grad():
        v_cond = self._predict_flow(self.teacher,
                                    latent_t,
                                    t_pf,
                                    training_batch,
                                    conditional=True,
                                    clean_x=clean_latents)
        v_uncond = self._predict_flow(self.teacher,
                                      latent_t,
                                      t_pf,
                                      training_batch,
                                      conditional=False,
                                      clean_x=clean_latents)
        v_pred = v_uncond + self._guidance_scale * (v_cond - v_uncond)
        dt = ((t - t_next) / float(self.student.num_train_timesteps))
        latent_t_next = latent_t - dt * v_pred

    flow_student = self._predict_flow(self.student,
                                      latent_t,
                                      t_pf,
                                      training_batch,
                                      conditional=True,
                                      clean_x=clean_latents)
    x0_t = latent_t - sigma_t * flow_student

    with torch.no_grad():
        flow_ema = self._predict_flow(self.ema_model,
                                      latent_t_next,
                                      t_next_pf,
                                      training_batch,
                                      conditional=True,
                                      clean_x=clean_latents)
        x0_t_next = latent_t_next - sigma_t_next * flow_ema

    loss = F.mse_loss(x0_t.float(), x0_t_next.float())

    loss_map = {"total_loss": loss, "causal_cd_loss": loss}
    attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata)
    outputs: dict[str, Any] = {"_fv_backward": (t_pf, attn_metadata)}
    metrics: dict[str, LogScalar] = {}
    return loss_map, outputs, metrics

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.backward

backward(loss_map: dict[str, torch.Tensor], outputs: dict[str, Any], *, grad_accum_rounds: int = 1) -> None
Source
def backward(
    self,
    loss_map: dict[str, torch.Tensor],
    outputs: dict[str, Any],
    *,
    grad_accum_rounds: int = 1,
) -> None:
    grad_accum_rounds = max(1, int(grad_accum_rounds))
    ctx = outputs.get("_fv_backward")
    if ctx is None:
        super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
        return
    self.student.backward(loss_map["total_loss"], ctx, grad_accum_rounds=grad_accum_rounds)

fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod.optimizers_schedulers_step

optimizers_schedulers_step(iteration: int) -> None
Source
def optimizers_schedulers_step(self, iteration: int) -> None:
    super().optimizers_schedulers_step(iteration)
    if iteration >= self._ema_start_step:
        self._update_ema()