diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index be6054f77..568e5a156 100644 --- a/fastvideo/train/trainer.py +++ b/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, diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index d396b0258..be787b11d 100644 --- a/fastvideo/train/utils/config.py +++ b/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, diff --git a/fastvideo/train/utils/optimizer.py b/fastvideo/train/utils/optimizer.py index 94fb0d133..b036fd913 100644 --- a/fastvideo/train/utils/optimizer.py +++ b/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 diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 25a546cef..80911659e 100644 --- a/fastvideo/train/utils/training_config.py +++ b/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) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 0933b1861..a6b929ee0 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/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: diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 9e6b0b26f..bedc2d48d 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/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: diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 89779ad4c..ef5a64872 100644 --- a/fastvideo/training/training_pipeline.py +++ b/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)) diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 7eb3847fa..4a8c142d0 100644 --- a/fastvideo/training/training_utils.py +++ b/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,