126 lines
5.7 KiB
Markdown
126 lines
5.7 KiB
Markdown
# `fastvideo/train/` — Modular Training Framework
|
||
|
||
YAML-driven trainer composed from interchangeable **methods × models × callbacks**. Preferred location for new training code. (See sibling `fastvideo/training/AGENTS.md` for the legacy stack.)
|
||
|
||
## Layout
|
||
|
||
```
|
||
train/
|
||
├── trainer.py # Core training loop coordinator
|
||
├── README.md # User-facing overview (legacy → new diff)
|
||
├── entrypoint/ # train.py + dcp_to_diffusers.py CLI entrypoints
|
||
├── methods/
|
||
│ ├── base.py # TrainingMethod ABC
|
||
│ ├── fine_tuning/ # FineTuneMethod, DiffusionForcingSFTMethod
|
||
│ ├── distribution_matching/ # DMD2Method, SelfForcingMethod
|
||
│ ├── knowledge_distillation/ # KDMethod, KDCausalMethod
|
||
│ ├── consistency_model/ # Consistency-model training methods
|
||
│ └── rl/ # RL methods, sampling helpers, rewards
|
||
├── models/
|
||
│ ├── base.py # ModelBase / CausalModelBase wrappers
|
||
│ ├── wan/, hunyuan/, cosmos/ # Per-family training wrappers
|
||
├── callbacks/ # callback.py base + ema, grad_clip, validation
|
||
└── utils/
|
||
├── training_config.py # Hierarchical YAML config dataclasses
|
||
├── builder.py # Build trainer from config
|
||
├── checkpoint.py # DCP save/load
|
||
├── optimizer.py # AdamW / fused optimizer factory
|
||
├── tracking.py # build_tracker (W&B / TensorBoard)
|
||
└── dataloader.py # StatefulDataLoader wiring
|
||
```
|
||
|
||
## Composition Model
|
||
|
||
```
|
||
Trainer = Method × Model × [Callback...] × Config
|
||
```
|
||
|
||
- A **Method** owns the loss + optimizer step (`compute_loss`, `step_post_grad`).
|
||
- A **Model** owns the forward + parameter-grouping (`forward`, `trainable_parameters`).
|
||
- **Callbacks** subscribe to lifecycle hooks (`on_train_start`,
|
||
`on_training_step_end`, `on_before_optimizer_step`, `on_validation_begin`,
|
||
`on_validation_end`, `on_train_end`) and compose freely.
|
||
- **Config** is a Pydantic-style hierarchical YAML resolved by
|
||
`utils/training_config.py`. Dotted-key CLI overrides go through
|
||
`parse_overrides`.
|
||
|
||
## Adding a New Model Plugin
|
||
|
||
1. Subclass `ModelBase` (or `CausalModelBase`) in `models/<family>/`.
|
||
2. Wrap the existing inference DiT from `fastvideo/models/dits/`. Do not
|
||
reimplement.
|
||
3. Expose `trainable_parameters()` so the optimizer factory can group them.
|
||
4. Register in `utils/builder.py` if the trainer dispatches by name.
|
||
|
||
## Adding a New Method
|
||
|
||
1. Subclass `TrainingMethod` in `methods/<family>/`.
|
||
2. Implement `compute_loss(batch, model_outputs) -> dict`.
|
||
3. If the method needs a teacher / second model, expose a `build_extras(cfg)`
|
||
classmethod — never instantiate inside `__init__`.
|
||
|
||
## RL/RLHF Boundaries
|
||
|
||
- RL methods live under `methods/rl/`, subclass `TrainingMethod`, and own
|
||
reward collection, advantage computation, policy loss, KL/reference terms,
|
||
and optimizer cadence.
|
||
- Keep model-family forward logic and latent decoding in `ModelBase`
|
||
wrappers. RL methods must use `ModelBase.decode_latents` rather than reach
|
||
into VAE normalization internals or inference pipelines.
|
||
- Shared generation, grouped prompt sampling, validation helpers, and scheduler
|
||
math belong under `methods/rl/common/`.
|
||
- Use `common/sampling.py` for generation and `common/prompt_sampling.py`
|
||
for reusable patterns such as DiffusionNFT K-repeat.
|
||
- Return `manages_optimization() == True` only for a genuinely nonstandard
|
||
outer/inner loop, implemented through
|
||
`managed_train_step(data_stream, iteration)`. Existing callbacks,
|
||
checkpointing, tracking, and validation must continue to work.
|
||
|
||
### Sampling and configuration
|
||
|
||
- Put method knobs under `method` and sampler knobs under
|
||
`method.sampling`.
|
||
- A diffusers-style scheduler owns both its schedule and `step()` solver.
|
||
Use `trajectory` only for higher-level ODE versus explicit re-noise
|
||
behavior; missing timesteps means to ask the scheduler for its defaults.
|
||
- Avoid fixed timestep lists in examples unless reproducing a known baseline.
|
||
- Do not put scheduler or trajectory policy in model configs.
|
||
|
||
### Reward models
|
||
|
||
- Put reusable rewards under `methods/rl/rewards/` and export their builders
|
||
from that package's `__init__.py`. Keep method-specific aggregation and
|
||
advantage logic out of reward classes.
|
||
- Rewards receive decoded media and return one scalar per prompt/sample.
|
||
Support `[B, C, H, W]` images and `[B, C, T, H, W]` videos when practical.
|
||
Each reward must explicitly choose whether it evaluates the first frame,
|
||
sampled frames, or the full video.
|
||
- Preserve upstream attribution and SPDX headers for ported reward code.
|
||
|
||
### RL tests
|
||
|
||
- Add fake-model tests for sampler/method behavior and fake-scorer injection
|
||
for reward aggregation.
|
||
- Test config parsing, tensor layouts, weighted reward metrics, prompt grouping,
|
||
and managed-optimization opt-in without loading heavyweight checkpoints.
|
||
- Confirm existing non-RL methods remain on the default Trainer path with
|
||
`manages_optimization() == False`.
|
||
|
||
## Configs
|
||
|
||
YAML lives under `examples/train/`. Schemas in `utils/training_config.py`. Add
|
||
new fields with explicit defaults; agents and humans both rely on the dataclass
|
||
to discover knobs.
|
||
|
||
## Anti-Patterns
|
||
|
||
- Importing from `fastvideo.training.*` (legacy stack). Forbidden cross-import.
|
||
- Adding a new training method as a fork of an existing pipeline file. Compose
|
||
via `Method` instead.
|
||
- Logging via stdlib `logging` — use `init_logger(__name__)`.
|
||
- Mutating the global state of a model in a callback. Callbacks operate on the
|
||
trainer state object passed in.
|
||
- Binding RL methods to inference pipeline classes.
|
||
- Hardcoding scheduler timesteps or embedding reward-model code inside one RL
|
||
method.
|