diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 0b8e82fb9..4af537796 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -248,13 +248,7 @@ class DistillationPipeline(TrainingPipeline): ema_enabled = (self.training_args.ema_decay is not None) and (self.training_args.ema_decay > 0.0) if ema_enabled and (self.training_args.ema_start_step <= 0): # Only eager-construct from the cold init weights when averaging starts at step 0. - self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay) - logger.info("Initialized generator EMA with decay=%s", self.training_args.ema_decay) - - # Initialize EMA for transformer_2 if it exists - if self.transformer_2 is not None: - self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay) - logger.info("Initialized generator EMA_2 with decay=%s", self.training_args.ema_decay) + self._build_generator_emas(context="eager init, ema_start_step<=0") elif ema_enabled: # Defer construction to the lazy block in the train loop, which builds the EMA AT # ema_start_step from the already-trained weights. Eager-constructing here would anchor @@ -904,11 +898,45 @@ class DistillationPipeline(TrainingPipeline): training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss return training_batch + def _build_generator_emas(self, context: str = "") -> None: + # Idempotently construct whichever generator EMA shadows are missing, per-expert and + # decoupled. Safe to call repeatedly; no-op once both exist or when EMA is disabled. + if (self.training_args.ema_decay is None) or (self.training_args.ema_decay <= 0.0): + return + suffix = f" [{context}]" if context else "" + if self.generator_ema is None: + self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay) + logger.info("Built generator EMA (decay=%s)%s", self.training_args.ema_decay, suffix) + if self.transformer_2 is not None and self.generator_ema_2 is None: + self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay) + logger.info("Built generator EMA_2 (decay=%s)%s", self.training_args.ema_decay, suffix) + + def _build_deferred_ema_for_resume(self) -> None: + # Build a deferred EMA before checkpoint load only when a saved shard exists, so the + # shadow reloads instead of being skipped and rebuilt fresh. Gating on shard existence + # matters: building when none exists would reintroduce cold-init contamination. + if (self.training_args.ema_decay is None) or (self.training_args.ema_decay <= 0.0): + return + + ema_shard_dir = os.path.join(self.training_args.resume_from_checkpoint, "ema_local_shard") + + if (self.generator_ema is None + and os.path.exists(os.path.join(ema_shard_dir, f"generator_ema_rank{self.global_rank}.pt"))): + self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay) + logger.info("Pre-built generator EMA for resume from existing shard") + + if (self.transformer_2 is not None and self.generator_ema_2 is None + and os.path.exists(os.path.join(ema_shard_dir, f"generator_ema_2_rank{self.global_rank}.pt"))): + self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay) + logger.info("Pre-built generator EMA_2 for resume from existing shard") + def _resume_from_checkpoint(self) -> None: """Resume training from checkpoint with distillation models.""" logger.info("Loading distillation checkpoint from %s", self.training_args.resume_from_checkpoint) + self._build_deferred_ema_for_resume() + resumed_step = load_distillation_checkpoint( self.transformer, self.fake_score_transformer, @@ -1317,15 +1345,8 @@ class DistillationPipeline(TrainingPipeline): self.current_trainstep = step training_batch.current_vsa_sparsity = current_vsa_sparsity - if (step >= self.training_args.ema_start_step) and \ - (self.generator_ema is None) and (self.training_args.ema_decay > 0): - self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay) - logger.info("Created generator EMA at step %s with decay=%s", step, self.training_args.ema_decay) - - # Create EMA for transformer_2 if it exists - if self.transformer_2 is not None and self.generator_ema_2 is None: - self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay) - logger.info("Created generator EMA_2 at step %s with decay=%s", step, self.training_args.ema_decay) + if step >= self.training_args.ema_start_step: + self._build_generator_emas(context=f"lazy @ step {step}") with torch.autocast("cuda", dtype=torch.bfloat16): training_batch = self.train_one_step(training_batch) diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 9975f15bd..9e6b0b26f 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -852,15 +852,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): self.current_trainstep = step training_batch.current_vsa_sparsity = current_vsa_sparsity - if (step >= self.training_args.ema_start_step) and \ - (self.generator_ema is None) and (self.training_args.ema_decay > 0): - self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay) - logger.info("Created generator EMA at step %s with decay=%s", step, self.training_args.ema_decay) - - # Create EMA for transformer_2 if it exists - if self.transformer_2 is not None and self.generator_ema_2 is None: - self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay) - logger.info("Created generator EMA_2 at step %s with decay=%s", step, self.training_args.ema_decay) + if step >= self.training_args.ema_start_step: + self._build_generator_emas(context=f"lazy @ step {step}") with torch.autocast("cuda", dtype=torch.bfloat16): training_batch = self.train_one_step(training_batch) diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 84ece8194..2cc1a0a87 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -613,18 +613,10 @@ def load_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Load EMA separately if saved in rank0_full mode + # Load EMA separately if generator_ema is not None: try: - if getattr(generator_ema, "mode", None) == "rank0_full": - ema_path = os.path.join(checkpoint_path, "ema", "generator_ema.pt") - if rank == 0 and os.path.exists(ema_path): - ema_state = torch.load(ema_path, map_location="cpu") - generator_ema.load_state_dict(ema_state) - logger.info("rank: %s, generator EMA (rank0_full) loaded from %s", rank, ema_path) - elif rank == 0: - logger.info("rank: %s, generator EMA file not found at %s; skipping", rank, ema_path) - elif getattr(generator_ema, "mode", None) == "local_shard": + if getattr(generator_ema, "mode", None) == "local_shard": ema_path = os.path.join(checkpoint_path, "ema_local_shard", f"generator_ema_rank{rank}.pt") if os.path.exists(ema_path): ema_state = torch.load(ema_path, map_location="cpu") @@ -632,9 +624,28 @@ def load_distillation_checkpoint( logger.info("rank: %s, generator EMA shard (local_shard) loaded from %s", rank, ema_path) else: logger.info("rank: %s, generator EMA shard file not found at %s; skipping", rank, ema_path) + else: + logger.info("rank: %s, generator EMA mode %s not supported for resume; skipping", rank, + getattr(generator_ema, "mode", None)) except Exception as e: logger.warning("rank: %s, failed to load generator EMA: %s", rank, str(e)) + # Load EMA_2 from its shard (symmetric with generator_ema above) + if generator_ema_2 is not None: + try: + if getattr(generator_ema_2, "mode", None) == "local_shard": + ema_2_path = os.path.join(checkpoint_path, "ema_local_shard", f"generator_ema_2_rank{rank}.pt") + if os.path.exists(ema_2_path): + generator_ema_2.load_state_dict(torch.load(ema_2_path, map_location="cpu")) + logger.info("rank: %s, generator EMA_2 shard (local_shard) loaded from %s", rank, ema_2_path) + else: + logger.info("rank: %s, generator EMA_2 shard file not found at %s; skipping", rank, ema_2_path) + else: + logger.info("rank: %s, generator EMA_2 mode %s not supported for resume; skipping", rank, + getattr(generator_ema_2, "mode", None)) + except Exception as e: + logger.warning("rank: %s, failed to load generator EMA_2: %s", rank, str(e)) + # Load generator_2 distributed checkpoint (MoE support) if generator_transformer_2 is not None: generator_2_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "generator_2") @@ -665,18 +676,6 @@ def load_distillation_checkpoint( rank, end_time - begin_time, local_main_process_only=False) - - # Load EMA_2 state if available and generator_ema_2 is provided - if generator_ema_2 is not None: - try: - ema_2_state = generator_2_states.get("ema") - if ema_2_state is not None: - generator_ema_2.load_state_dict(ema_2_state) - logger.info("rank: %s, generator_2 EMA state loaded successfully", rank) - else: - logger.info("rank: %s, no EMA_2 state found in checkpoint", rank) - except Exception as e: - logger.warning("rank: %s, failed to load EMA_2 state: %s", rank, str(e)) else: logger.info("rank: %s, generator_2 checkpoint not found, skipping", rank)