Skip to content

base

Classes

fastvideo.train.methods.base.TrainingMethod 

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

Bases: Module, ABC

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.

Source code in fastvideo/train/methods/base.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, ModelBase],
) -> None:
    super().__init__()
    self.tracker: Any | None = None
    self._role_models: dict[str, ModelBase] = dict(role_models)

    self.student = role_models["student"]
    self.training_config = cfg.training
    self.method_config: dict[str, Any] = dict(cfg.method)
    self.validation_config: dict[str, Any] = dict(getattr(cfg, "validation", {}) or {})

    # Build nn.ModuleDict for FSDP / checkpoint visibility.
    self.role_modules = torch.nn.ModuleDict()
    for role, model in role_models.items():
        mods: dict[str, torch.nn.Module] = {}
        transformer = getattr(model, "transformer", None)
        if isinstance(transformer, torch.nn.Module):
            mods["transformer"] = transformer
        if mods:
            self.role_modules[role] = torch.nn.ModuleDict(mods)

Methods:

fastvideo.train.methods.base.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.

Source code in fastvideo/train/methods/base.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def checkpoint_state(self) -> 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.
    """
    states: dict[str, Any] = {}

    for role, model in self._role_models.items():
        if not getattr(model, "_trainable", False):
            continue

        modules: dict[str, torch.nn.Module] = {}
        if model.transformer is not None:
            modules["transformer"] = model.transformer

        container = _RoleModuleContainer(modules)

        for module_name, module in modules.items():
            states[f"roles.{role}.{module_name}"] = ModelWrapper(module)

        opt = self._optimizer_dict.get(role)
        if opt is not None:
            states[f"optimizers.{role}"] = OptimizerWrapper(container, opt)

        sched = self._lr_scheduler_dict.get(role)
        if sched is not None:
            states[f"schedulers.{role}"] = SchedulerWrapper(sched)

    return states
fastvideo.train.methods.base.TrainingMethod.get_grad_clip_targets 
get_grad_clip_targets(iteration: int) -> dict[str, Module]

Return modules whose gradients should be clipped.

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

Source code in fastvideo/train/methods/base.py 
240
241
242
243
244
245
246
247
248
249
250
def get_grad_clip_targets(
    self,
    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.
    """
    return {"student": self.student.transformer}
fastvideo.train.methods.base.TrainingMethod.managed_train_step 
managed_train_step(data_stream: Any, iteration: int) -> tuple[dict[str, 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.

Source code in fastvideo/train/methods/base.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def managed_train_step(
    self,
    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.
    """
    return self.single_train_step(next(data_stream), iteration)
fastvideo.train.methods.base.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.

Source code in fastvideo/train/methods/base.py 
200
201
202
203
204
205
206
207
208
209
def manages_optimization(self) -> 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``.
    """
    return False
fastvideo.train.methods.base.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.

Source code in fastvideo/train/methods/base.py 
229
230
231
232
233
234
235
236
237
238
def on_validation_begin(self, 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``.
    """
    del iteration
    return {}
fastvideo.train.methods.base.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.

Source code in fastvideo/train/methods/base.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def seed_optimizer_state_for_resume(self) -> 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.
    """
    for opt in self.get_optimizers(0):
        for group in opt.param_groups:
            for p in group["params"]:
                if not p.requires_grad:
                    continue
                if len(opt.state.get(p, {})) > 0:
                    continue
                opt.state[p] = {
                    "step": torch.tensor(0.0),
                    "exp_avg": torch.zeros_like(p),
                    "exp_avg_sq": torch.zeros_like(p),
                }

Functions: