[Bug] FSDP EMA save failed

[yjhong89]

Describe the bug

Error raised when trying to save EMA

• FastVideo/fastvideo/training/distillation_pipeline.py

  Line 381 in 0082bc6

  ema_model = copy.deepcopy(self.transformer)

• ERROR 12-04 14:11:23 [distillation_pipeline.py:490] Failed to save EMA weights: FSDP does not support deepcopy. Please use state dict for serialization.

Reproduction

Run distill_dmd.sh with my custom preprocessed dataset.

Environment

• CUDA128, PYTORCH271
• H100*8

[yjhong89]

If I changed EMA_FSDP mode to "rank0_full"

• FastVideo/fastvideo/training/distillation_pipeline.py

  Line 266 in 0082bc6

  self.generator_ema = EMA_FSDP(self.transformer,

• Hanging up when generator update (5 step)
  • Not progress after this point

Steps:   0%|                                                                                                                            | 4/4000 [00:31<7:23:16,  6.66s/it, total_loss=0.1850, generator_loss=0.0000, fake_score_loss=0.1850, step_time=4.74s, grad_norm=None, ema=✓, ema2=✗