fine_tuning ¶

Classes¶

fastvideo.train.methods.fine_tuning.DiffusionForcingSFTMethod ¶

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

Bases: TrainingMethod

Diffusion-forcing SFT (DFSFT): train only student with inhomogeneous timesteps.

Source code in fastvideo/train/methods/fine_tuning/dfsft.py

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("DFSFT requires role 'student'")
    if not self.student._trainable:
        raise ValueError("DFSFT requires student to be trainable")
    self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind())

    self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None))
    self._timestep_index_range = (self._parse_timestep_index_range())
    self._training_weights = self._build_training_weights()

    # Initialize preprocessors on student.
    self.student.init_preprocessors(self.training_config)

    self._init_optimizers_and_schedulers()

fastvideo.train.methods.fine_tuning.FineTuneMethod ¶

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

Bases: TrainingMethod

Supervised finetuning: only student participates.

Source code in fastvideo/train/methods/fine_tuning/finetune.py

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("FineTuneMethod requires role 'student'")
    if not self.student._trainable:
        raise ValueError("FineTuneMethod requires student to be "
                         "trainable")
    self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind())

    # Initialize preprocessors on student.
    self.student.init_preprocessors(self.training_config)

    self._init_optimizers_and_schedulers()

Methods:¶

fastvideo.train.methods.fine_tuning.FineTuneMethod.backward ¶

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

Backpropagate an accumulation-scaled loss through the student model.

Delegating to ModelBase.backward lets each model restore its forward context before the distributed wrapper synchronizes parameter gradients.

Source code in fastvideo/train/methods/fine_tuning/finetune.py

def backward(
    self,
    loss_map: dict[str, torch.Tensor],
    outputs: dict[str, Any],
    *,
    grad_accum_rounds: int = 1,
) -> None:
    """Backpropagate an accumulation-scaled loss through the student model.

    Delegating to ``ModelBase.backward`` lets each model restore its forward
    context before the distributed wrapper synchronizes parameter gradients.
    """
    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.fine_tuning.FineTuneMethod.single_train_step ¶

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

Prepare synchronized targets and compute supervised flow loss.

The returned forward context lets model-specific backward methods restore activation-checkpoint metadata during recomputation.

Source code in fastvideo/train/methods/fine_tuning/finetune.py

def single_train_step(
    self,
    batch: dict[str, Any],
    iteration: int,
) -> tuple[
        dict[str, torch.Tensor],
        dict[str, Any],
        dict[str, LogScalar],
]:
    """Prepare synchronized targets and compute supervised flow loss.

    The returned forward context lets model-specific backward methods
    restore activation-checkpoint metadata during recomputation.
    """
    del iteration
    training_batch = self.student.prepare_batch(
        batch,
        generator=self.cuda_generator,
        latents_source="data",
    )

    if training_batch.latents is None:
        raise RuntimeError("prepare_batch() must set "
                           "TrainingBatch.latents")
    if training_batch.noisy_model_input is None:
        raise RuntimeError("prepare_batch() must set "
                           "TrainingBatch.noisy_model_input")
    if training_batch.noise is None:
        raise RuntimeError("prepare_batch() must set "
                           "TrainingBatch.noise")
    if training_batch.sigmas is None:
        raise RuntimeError("prepare_batch() must set "
                           "TrainingBatch.sigmas")
    if training_batch.timesteps is None:
        raise RuntimeError("prepare_batch() must set "
                           "TrainingBatch.timesteps")

    clean_latents = training_batch.latents
    noisy_latents = (training_batch.noisy_model_input.permute(0, 2, 1, 3, 4))
    noise = training_batch.noise.permute(0, 2, 1, 3, 4)
    sigmas = training_batch.sigmas
    timesteps = training_batch.timesteps

    pred = self.student.predict_noise(
        noisy_latents,
        timesteps,
        training_batch,
        conditional=True,
        attn_kind=self._attn_kind,
    )

    loss_map = _compute_finetune_loss_map(
        pred,
        clean_latents,
        noisy_latents,
        noise,
        sigmas,
        training_batch,
        precondition_outputs=bool(self.training_config.model.precondition_outputs),
    )

    attn_metadata = training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata

    outputs: dict[str, Any] = {
        "_fv_backward": (
            training_batch.timesteps,
            attn_metadata,
        )
    }
    metrics: dict[str, LogScalar] = {}
    return loss_map, outputs, metrics