nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix] EMA shadow on resume and EMA under MoE path (#1441)

Browse Source 
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Kaiqin Kong
2026-06-09 00:55:43 +00:00
committed by GitHub
co-authored by mergify[bot]
parent 2e35b0c6bd
commit 9ea77d37f3
3 changed files with 60 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/training/distillation_pipeline.py
+37 -16
View File
@@ -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)
fastvideo/training/self_forcing_distillation_pipeline.py
+2 -9
View File
@@ -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)
fastvideo/training/training_utils.py
+21 -22
View File
@@ -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)