[train] enforce fp32 master weights in every training entry point

Optimizer steps applied in-place to bf16 parameters round away updates
below ~half an ulp of each weight's magnitude: after 1000 DMD2 steps,
205/211 of the H3 student's norm-scale params were bit-identical to the
base checkpoint, and bf16 exp_avg_sq starves the same way. Guard every
path:

- fastvideo/train: Trainer.run refuses to start unless all trainable
  role params are fp32; opt out via
  training.model.allow_low_precision_master_weights
- fastvideo/training (legacy): same check at the top of the three
  train() funnels (base / distillation / self-forcing), opt out via
  FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS=1. Shipped recipes
  already pass --dit_precision fp32 and are unaffected.
- AdamWBeta1Zero: keep exp_avg_sq and update math in fp32 when the
  param storage is not fp32 (fp32 path stays bitwise-identical)

Also carries the checkpointing_start_step / reset_lr_on_resume config
plumbing for the previous commit (same dataclass/parse files).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

fastvideo/train/trainer.py

@@ -11,6 +11,7 @@ import torch
 from tqdm.auto import tqdm
 
 from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.logger import init_logger
 from fastvideo.train.callbacks.callback import CallbackDict
 from fastvideo.train.methods.base import LogScalar, TrainingMethod
 from fastvideo.train.utils.tracking import build_tracker
@@ -19,6 +20,41 @@ if TYPE_CHECKING:
     from fastvideo.train.utils.training_config import (
         TrainingConfig, )
 
+logger = init_logger(__name__)
+
+
+def _verify_master_weight_precision(method: TrainingMethod, tc: TrainingConfig) -> None:
+    """Refuse to train on low-precision master weights unless opted in.
+
+    Optimizer steps applied in-place to bf16/fp16 parameters round away
+    updates below ~half an ulp of each weight's magnitude; O(1)-magnitude
+    parameters (norm gains) freeze entirely at typical distillation learning
+    rates, and ``zeros_like``-allocated optimizer state inherits the same
+    starved dtype. fp32 sharded masters (``training.dit_precision: fp32``)
+    keep compute in bf16 via FSDP's param_dtype while fixing both.
+    """
+    if bool(getattr(tc.model, "allow_low_precision_master_weights", False)):
+        return
+    offenders: list[str] = []
+    for role, model in getattr(method, "_role_models", {}).items():
+        if not getattr(model, "_trainable", False):
+            continue
+        transformer = getattr(model, "transformer", None)
+        if transformer is None:
+            continue
+        for name, param in transformer.named_parameters():
+            if param.requires_grad and param.dtype != torch.float32:
+                offenders.append(f"{role}:{name} ({param.dtype})")
+                break
+    if offenders:
+        raise RuntimeError("Trainable master weights are not fp32: "
+                           f"{offenders}. bf16/fp16 parameter storage silently rounds away "
+                           "optimizer updates below ~half an ulp per weight (norm-scale "
+                           "parameters freeze completely). Set training.dit_precision: fp32 "
+                           "(fp32 sharded masters; compute stays bf16 via FSDP param_dtype), "
+                           "or acknowledge the effect explicitly with "
+                           "training.model.allow_low_precision_master_weights: true.")
+
 
 def _coerce_log_scalar(
     value: Any,
@@ -114,6 +150,7 @@ class Trainer:
         )
 
         method.set_tracker(self.tracker)
+        _verify_master_weight_precision(method, tc)
         method.on_train_start()
         self.callbacks.on_train_start(
             method,
@@ -127,6 +164,11 @@ class Trainer:
             resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint)))
             if resumed_step is not None:
                 start_step = int(resumed_step)
+                if bool(getattr(tc.checkpoint, "reset_lr_on_resume", False)):
+                    # The DCP load above restored the checkpoint's optimizer
+                    # LRs and scheduler base_lrs; re-apply the YAML's values.
+                    method.apply_configured_lrs()
+                    logger.info("reset_lr_on_resume: re-applied configured learning rates at step %s", start_step)
         self.callbacks.on_validation_begin(
             method,
             iteration=start_step,

fastvideo/train/utils/config.py

@@ -422,6 +422,8 @@ def _build_training_config(
             resume_from_checkpoint=str(ck.get("resume_from_checkpoint", "") or ""),
             training_state_checkpointing_steps=int(ck.get("training_state_checkpointing_steps", 0) or 0),
             checkpoints_total_limit=int(ck.get("checkpoints_total_limit", 0) or 0),
+            checkpointing_start_step=int(ck.get("checkpointing_start_step", 0) or 0),
+            reset_lr_on_resume=bool(ck.get("reset_lr_on_resume", False)),
         ),
         tracker=TrackerConfig(
             trackers=list(tr.get("trackers", []) or []),
@@ -439,6 +441,7 @@ def _build_training_config(
             precondition_outputs=bool(m.get("precondition_outputs", False)),
             moba_config=dict(m.get("moba_config", {}) or {}),
             enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
+            allow_low_precision_master_weights=bool(m.get("allow_low_precision_master_weights", False)),
         ),
         pipeline_config=pipeline_config,
         model_path=model_path,

fastvideo/train/utils/optimizer.py

@@ -60,14 +60,30 @@ class AdamWBeta1Zero(torch.optim.Optimizer):
                 state = self.state[p]
                 if not state:
                     state["step"] = 0
-                    state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Low-precision params must not starve the second moment:
+                    # bf16 v loses g^2 increments below ~0.4% of its running
+                    # magnitude. Keep v in fp32 whenever p is not fp32 (the
+                    # fp32-master path is unchanged and stays bitwise-equal
+                    # to torch.optim.AdamW).
+                    if p.dtype == torch.float32:
+                        state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(p, dtype=torch.float32)
                 state["step"] += 1
                 exp_avg_sq = state["exp_avg_sq"]
                 p.mul_(1 - lr * weight_decay)
-                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
-                bias_correction2_sqrt = (1 - beta2**state["step"])**0.5
-                denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
-                p.addcdiv_(grad, denom, value=-lr)
+                if p.dtype == torch.float32:
+                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+                    bias_correction2_sqrt = (1 - beta2**state["step"])**0.5
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+                    p.addcdiv_(grad, denom, value=-lr)
+                else:
+                    grad_f = grad.float()
+                    exp_avg_sq.mul_(beta2).addcmul_(grad_f, grad_f, value=1 - beta2)
+                    bias_correction2_sqrt = (1 - beta2**state["step"])**0.5
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+                    update = grad_f.div_(denom).mul_(-lr)
+                    p.add_(update.to(p.dtype))
         return loss
 
 

fastvideo/train/utils/training_config.py

@@ -58,6 +58,11 @@ class CheckpointConfig:
     resume_from_checkpoint: str = ""
     training_state_checkpointing_steps: int = 0
     checkpoints_total_limit: int = 0
+    checkpointing_start_step: int = 0
+    # DCP checkpoints restore optimizer param-group LRs and scheduler
+    # base_lrs, silently overriding the YAML on resume. Set true to re-apply
+    # the configured learning rates after loading (LR-change experiments).
+    reset_lr_on_resume: bool = False
 
 
 @dataclass(slots=True)
@@ -77,6 +82,13 @@ class ModelTrainingConfig:
     precondition_outputs: bool = False
     moba_config: dict = field(default_factory=dict)
     enable_gradient_checkpointing_type: str | None = None
+    # Optimizer steps applied to bf16/fp16 parameter storage round away
+    # updates smaller than ~half an ulp of each weight's magnitude —
+    # O(1)-magnitude parameters (norm gains) freeze entirely at typical
+    # distillation LRs. The trainer refuses to start unless master weights
+    # are fp32 (training.dit_precision: fp32) or this explicit opt-in
+    # acknowledges the effect (memory-constrained topologies).
+    allow_low_precision_master_weights: bool = False
 
 
 @dataclass(slots=True)

fastvideo/training/distillation_pipeline.py

@@ -33,7 +33,8 @@ from fastvideo.training.activation_checkpoint import (apply_activation_checkpoin
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
                                                get_scheduler, load_distillation_checkpoint,
-                                               save_distillation_checkpoint, shift_timestep)
+                                               save_distillation_checkpoint, shift_timestep,
+                                               verify_master_weight_precision)
 from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory)
 
 try:
@@ -1282,6 +1283,12 @@ class DistillationPipeline(TrainingPipeline):
         """Main training loop with distillation-specific logging."""
         assert self.training_args.seed is not None, "seed must be set"
         seed = self.training_args.seed
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+            "fake_score_transformer": getattr(self, "fake_score_transformer", None),
+            "fake_score_transformer_2": getattr(self, "fake_score_transformer_2", None),
+        })
 
         # Set the same seed within each SP group to ensure reproducibility
         if self.sp_world_size > 1:

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -21,7 +21,7 @@ from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfF
 from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines import TrainingBatch
 from fastvideo.training.distillation_pipeline import DistillationPipeline
-from fastvideo.training.training_utils import (EMA_FSDP, save_distillation_checkpoint)
+from fastvideo.training.training_utils import (EMA_FSDP, save_distillation_checkpoint, verify_master_weight_precision)
 from fastvideo.utils import is_vsa_available, set_random_seed
 from fastvideo.profiler import profile_region
 
@@ -797,6 +797,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         """Main training loop with self-forcing specific logging."""
         assert self.training_args.seed is not None, "seed must be set"
         seed = self.training_args.seed
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+            "fake_score_transformer": getattr(self, "fake_score_transformer", None),
+            "fake_score_transformer_2": getattr(self, "fake_score_transformer_2", None),
+        })
 
         # Set the same seed within each SP group to ensure reproducibility
         if self.sp_world_size > 1:

fastvideo/training/training_pipeline.py

@@ -42,7 +42,8 @@ from fastvideo.training.activation_checkpoint import (apply_activation_checkpoin
 from fastvideo.training.trackers import (DummyTracker, TrackerType, initialize_trackers, Trackers)
 from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases,
                                                compute_density_for_timestep_sampling, count_trainable, get_scheduler,
-                                               get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint)
+                                               get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint,
+                                               verify_master_weight_precision)
 from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict)
 
 try:
@@ -519,6 +520,10 @@ class TrainingPipeline(LoRAPipeline, ABC):
         logger.info('rank: %s: start training', self.global_rank, local_main_process_only=False)
         if not self.post_init_called:
             self.post_init()
+        verify_master_weight_precision({
+            "transformer": self.transformer,
+            "transformer_2": getattr(self, "transformer_2", None),
+        })
         num_trainable_params = count_trainable(self.transformer)
         logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3))
 

fastvideo/training/training_utils.py

@@ -23,6 +23,36 @@ logger = init_logger(__name__)
 _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
 
 
+def verify_master_weight_precision(modules: dict[str, torch.nn.Module | None]) -> None:
+    """Refuse to train on low-precision master weights unless opted in.
+
+    Optimizer steps applied in-place to bf16/fp16 parameters round away
+    updates below ~half an ulp of each weight's magnitude; O(1)-magnitude
+    parameters (norm gains) freeze entirely at typical fine-tune/distillation
+    learning rates, and ``zeros_like``-allocated optimizer state inherits the
+    same starved dtype. ``--dit_precision fp32`` gives fp32 sharded masters
+    while compute stays bf16 via FSDP's ``param_dtype``.
+    """
+    if os.environ.get("FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS", "0") == "1":
+        return
+    offenders: list[str] = []
+    for module_name, module in modules.items():
+        if module is None:
+            continue
+        for name, param in module.named_parameters():
+            if param.requires_grad and param.dtype != torch.float32:
+                offenders.append(f"{module_name}:{name} ({param.dtype})")
+                break
+    if offenders:
+        raise RuntimeError("Trainable master weights are not fp32: "
+                           f"{offenders}. bf16/fp16 parameter storage silently rounds away "
+                           "optimizer updates below ~half an ulp per weight (norm-scale "
+                           "parameters freeze completely). Pass --dit_precision fp32 "
+                           "(fp32 sharded masters; compute stays bf16 via FSDP param_dtype), "
+                           "or acknowledge the effect explicitly with "
+                           "FASTVIDEO_ALLOW_LOW_PRECISION_MASTER_WEIGHTS=1.")
+
+
 def gather_state_dict_on_cpu_rank0(
     model,
     device: torch.device | None = None,