training script testing

README.md

@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/wZPZTLKg" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">

examples/distill/SFWan2.1-I2V/distill_dmd_i2v_1.3B.slurm

@@ -0,0 +1,150 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your_env
+
+# Basic Info
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+
+# Model paths for DMD distillation:
+GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+DATA_DIR=your_data_dir
+VALIDATION_DATASET_FILE=your_validation_dataset_file
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_dmd
+  --output_dir "checkpoints/SFwan_t2v_finetune"
+  --max_train_steps 1000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 64
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 200
+  --validation_sampling_steps "3"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+)
+
+# DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 4e-7
+  --fake_score_betas '0.0,0.999'
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"

fastvideo/configs/models/dits/wanvideo.py

@@ -93,8 +93,8 @@ class WanVideoArchConfig(DiTArchConfig):
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
     # Causal Wan
-    local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention)
-    sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+    local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
+    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
 
     def __post_init__(self):
         super().__post_init__()

fastvideo/configs/pipelines/wan.py

@@ -146,3 +146,5 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
 @dataclass
 class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     is_causal: bool = True
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 750, 500, 250])

fastvideo/configs/sample/registry.py

@@ -39,6 +39,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # Wan2.2
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
+    Wan2_2_TI2V_5B_SamplingParam,
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
 

fastvideo/fastvideo_args.py

@@ -590,6 +590,11 @@ class TrainingArgs(FastVideoArgs):
     # text encoder & vae & diffusion model
     pretrained_model_name_or_path: str = ""
     dit_model_name_or_path: str = ""
+    
+    # DMD model paths - separate paths for each network
+    generator_model_path: str = ""  # path for generator (student) model
+    real_score_model_path: str = ""  # path for real score (teacher) model
+    fake_score_model_path: str = ""  # path for fake score (critic) model
 
     # diffusion setting
     ema_decay: float = 0.0
@@ -644,6 +649,7 @@ class TrainingArgs(FastVideoArgs):
     linear_quadratic_threshold: float = 0.0
     linear_range: float = 0.0
     weight_decay: float = 0.0
+    betas: str = "0.9,0.999"  # betas for optimizer, format: "beta1,beta2"
     use_ema: bool = False
     multi_phased_distill_schedule: str = ""
     pred_decay_weight: float = 0.0
@@ -669,6 +675,7 @@ class TrainingArgs(FastVideoArgs):
     real_score_guidance_scale: float = 3.5
     fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
     fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
+    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
     training_state_checkpointing_steps: int = 0  # for resuming training
     weight_only_checkpointing_steps: int = 0  # for inference
     log_visualization: bool = False
@@ -774,6 +781,17 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--cache-dir",
                             type=str,
                             help="Directory to cache models")
+        
+        # DMD model paths - separate paths for each network
+        parser.add_argument("--generator-model-path",
+                            type=str,
+                            help="Path to generator (student) model for DMD distillation")
+        parser.add_argument("--real-score-model-path",
+                            type=str,
+                            help="Path to real score (teacher) model for DMD distillation")
+        parser.add_argument("--fake-score-model-path",
+                            type=str,
+                            help="Path to fake score (critic) model for DMD distillation")
 
         # Diffusion settings
         parser.add_argument("--ema-decay",
@@ -949,6 +967,10 @@ class TrainingArgs(FastVideoArgs):
                             help="Linear quadratic threshold")
         parser.add_argument("--linear-range", type=float, help="Linear range")
         parser.add_argument("--weight-decay", type=float, help="Weight decay")
+        parser.add_argument("--betas",
+                            type=str,
+                            default=TrainingArgs.betas,
+                            help="Betas for optimizer (format: 'beta1,beta2')")
         parser.add_argument("--use-ema",
                             action=StoreBoolean,
                             help="Whether to use EMA")
@@ -1006,6 +1028,10 @@ class TrainingArgs(FastVideoArgs):
                             type=float,
                             default=TrainingArgs.fake_score_learning_rate,
                             help="Learning rate for fake score transformer")
+        parser.add_argument("--fake-score-betas",
+                            type=str,
+                            default=TrainingArgs.fake_score_betas,
+                            help="Betas for fake score optimizer (format: 'beta1,beta2')")
         parser.add_argument(
             "--fake-score-lr-scheduler",
             type=str,

fastvideo/layers/rotary_embedding.py

@@ -33,6 +33,8 @@ from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
+logger = init_logger(__name__)
+
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
     x1 = x[..., :x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2:]
@@ -297,7 +299,6 @@ def get_nd_rotary_pos_embed(
 
     if start_frame > 0:
         full_grid[0] += start_frame
-        logger.info(full_grid[0])
 
     # Shard the grid if using sequence parallelism (sp_world_size > 1)
     assert shard_dim < len(

fastvideo/models/dits/causal_wanvideo.py

@@ -58,14 +58,6 @@ class CausalWanSelfAttention(nn.Module):
         self.parallel_attention = parallel_attention
         self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
 
-        # layers
-        self.to_q = ReplicatedLinear(dim, dim)
-        self.to_k = ReplicatedLinear(dim, dim)
-        self.to_v = ReplicatedLinear(dim, dim)
-        self.to_out = ReplicatedLinear(dim, dim)
-        self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
         # Scaled dot product attention
         self.attn = LocalAttention(
             num_heads=num_heads,
@@ -96,8 +88,8 @@ class CausalWanSelfAttention(nn.Module):
             cache_start = current_start
 
         cos, sin = freqs_cis
-        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False)
-        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False)
+        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
+        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
 
         if kv_cache is None:
             # Padding for flex attention
@@ -165,9 +157,6 @@ class CausalWanSelfAttention(nn.Module):
             kv_cache["global_end_index"].fill_(current_end)
             kv_cache["local_end_index"].fill_(local_end_index)
 
-        # output
-        x = x.flatten(2)
-        x = self.to_out(x)
         return x
 
 class CausalWanTransformerBlock(nn.Module):
@@ -281,7 +270,7 @@ class CausalWanTransformerBlock(nn.Module):
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
 
-        attn_output, _ = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start)
+        attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start)
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
         attn_output = attn_output.squeeze(1)
@@ -451,6 +440,7 @@ class CausalWanTransformer3DModel(BaseDiT):
                 crossattn_cache: dict = None,
                 current_start: int = 0,
                 cache_start: int = 0,
+                start_frame: int = 0,
                 **kwargs) -> torch.Tensor:
         r"""
         Run the diffusion model with kv caching.
@@ -485,7 +475,7 @@ class CausalWanTransformer3DModel(BaseDiT):
             rope_dim_list,
             dtype=torch.float32 if current_platform.is_mps() else torch.float64,
             rope_theta=10000,
-            start_frame=current_start # Assume that current_start is 0 when kv_cache is None
+            start_frame=start_frame # Assume that start_frame is 0 when kv_cache is None
         )
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
@@ -515,7 +505,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                 causal_kwargs = {
                     "kv_cache": kv_cache[block_index],
                     "current_start": current_start,
-                    "cache_start": cache_start
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
                 }
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
@@ -526,7 +517,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                     "kv_cache": kv_cache[block_index],
                     "crossattn_cache": crossattn_cache[block_index],
                     "current_start": current_start,
-                    "cache_start": cache_start
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
                 }
                 hidden_states = block(hidden_states, encoder_hidden_states,
                                         timestep_proj, freqs_cis,

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -13,8 +13,10 @@ from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
 
 # isort: off
 from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
-                                        CausalDMDDenosingStage, InputValidationStage,
-                                        LatentPreparationStage, TextEncodingStage,
+                                        CausalDMDDenosingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
                                         TimestepPreparationStage)
 # isort: on
 
@@ -65,5 +67,3 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
 
 
 EntryClass = WanCausalDMDPipeline
-
-

fastvideo/pipelines/pipeline_batch_info.py

@@ -173,6 +173,10 @@ class ForwardBatch:
     logging_info: PipelineLoggingInfo = field(
         default_factory=PipelineLoggingInfo)
 
+    # Causal Wan parameters
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 21
+
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 

fastvideo/pipelines/stages/causal_denoising.py

@@ -7,6 +7,7 @@ from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.training.training_utils import pred_noise_to_pred_video
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -24,222 +25,6 @@ except ImportError:
     vsa_available = False
     VideoSparseAttentionBackend = None  # type: ignore
 
-# class CausalDenoisingStage(LoRAPipeline, ComposedPipelineBase):
-#     """
-#     Flow Causal denoising stage.
-#     """
-
-#     def forward(
-#         self,
-#         batch: ForwardBatch,
-#         fastvideo_args: FastVideoArgs,
-#     ) -> ForwardBatch:
-#         # Setup precision and autocast settings
-#         target_dtype = torch.bfloat16
-#         autocast_enabled = (target_dtype != torch.float32
-#                             ) and not fastvideo_args.disable_autocast
-
-#         # Gather args for causal block-wise processing
-#         try:
-#             num_frame_per_block = getattr(fastvideo_args.pipeline_config,
-#                                           "num_frame_per_block", 1)
-#         except Exception:
-#             num_frame_per_block = 1
-#         try:
-#             independent_first_frame = getattr(fastvideo_args.pipeline_config,
-#                                               "independent_first_frame", False)
-#         except Exception:
-#             independent_first_frame = False
-
-#         # Required fields
-#         timesteps = batch.timesteps
-#         if timesteps is None:
-#             raise ValueError("Timesteps must be provided")
-
-#         # Prepare image latents and embeddings for I2V generation
-#         image_embeds = batch.image_embeds
-#         if len(image_embeds) > 0:
-#             assert torch.isnan(image_embeds[0]).sum() == 0
-#             image_embeds = [
-#                 image_embed.to(target_dtype) for image_embed in image_embeds
-#             ]
-
-#         image_kwargs = self.prepare_extra_func_kwargs(
-#             self.transformer.forward,
-#             {
-#                 "encoder_hidden_states_image": image_embeds,
-#                 "mask_strategy": dict_to_3d_list(
-#                     None, t_max=50, l_max=60, h_max=24)
-#             },
-#         )
-
-#         pos_cond_kwargs = self.prepare_extra_func_kwargs(
-#             self.transformer.forward,
-#             {
-#                 "encoder_hidden_states_2": batch.clip_embedding_pos,
-#                 "encoder_attention_mask": batch.prompt_attention_mask,
-#             },
-#         )
-
-#         neg_cond_kwargs = self.prepare_extra_func_kwargs(
-#             self.transformer.forward,
-#             {
-#                 "encoder_hidden_states_2": batch.clip_embedding_neg,
-#                 "encoder_attention_mask": batch.negative_attention_mask,
-#             },
-#         )
-
-#         # Prepare STA parameters
-#         if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-#             self.prepare_sta_param(batch, fastvideo_args)
-
-#         # Get latents and embeddings
-#         latents = batch.latents
-#         assert latents is not None, "latents must be provided"
-#         prompt_embeds = batch.prompt_embeds
-#         assert torch.isnan(prompt_embeds[0]).sum() == 0
-
-#         # Determine blocks
-#         b, c, t, h, w = latents.shape
-#         if not independent_first_frame or (independent_first_frame
-#                                            and batch.image_latent is not None):
-#             if t % num_frame_per_block != 0:
-#                 raise ValueError(
-#                     "num_frames must be divisible by num_frame_per_block for causal denoising"
-#                 )
-#             num_blocks = t // num_frame_per_block
-#             block_sizes = [num_frame_per_block] * num_blocks
-#             start_index = 0
-#         else:
-#             if (t - 1) % num_frame_per_block != 0:
-#                 raise ValueError(
-#                     "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
-#                 )
-#             num_blocks = (t - 1) // num_frame_per_block
-#             block_sizes = [1] + [num_frame_per_block] * num_blocks
-#             start_index = 0
-
-#         # Run causal block-wise denoising loop
-#         with self.progress_bar(total=len(block_sizes) *
-#                                len(timesteps)) as progress_bar:
-#             for current_num_frames in block_sizes:
-#                 # Build the latent input for this block
-#                 current_latents = latents[:, :, start_index:start_index +
-#                                           current_num_frames, :, :]
-
-#                 # Run spatial denoising loop over timesteps
-#                 for i, t_cur in enumerate(timesteps):
-#                     # Expand latents for I2V
-#                     latent_model_input = current_latents.to(target_dtype)
-#                     if batch.image_latent is not None and independent_first_frame and start_index == 0:
-#                         # Prepend the image latent only for the first block when required
-#                         latent_model_input = torch.cat([
-#                             batch.image_latent.to(target_dtype),
-#                             latent_model_input
-#                         ],
-#                                                        dim=2)
-
-#                     latent_model_input = self.scheduler.scale_model_input(
-#                         latent_model_input, t_cur)
-
-#                     # Prepare inputs for transformer
-#                     t_expand = t_cur.repeat(latent_model_input.shape[0])
-#                     guidance_expand = (torch.tensor(
-#                         [fastvideo_args.pipeline_config.embedded_cfg_scale] *
-#                         latent_model_input.shape[0],
-#                         dtype=torch.float32,
-#                         device=get_local_torch_device(),
-#                     ).to(target_dtype) *
-#                                        1000.0 if fastvideo_args.pipeline_config.
-#                                        embedded_cfg_scale is not None else None)
-
-#                     # Predict noise residual
-#                     with torch.autocast(device_type="cuda",
-#                                         dtype=target_dtype,
-#                                         enabled=autocast_enabled):
-#                         if (st_attn_available and self.attn_backend == SlidingTileAttentionBackend) or \
-#                            (vsa_available and self.attn_backend == VideoSparseAttentionBackend):
-#                             self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-#                             )
-#                             if self.attn_metadata_builder_cls is not None:
-#                                 self.attn_metadata_builder = self.attn_metadata_builder_cls(
-#                                 )
-#                                 attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-#                                     current_timestep=i,  # type: ignore
-#                                     raw_latent_shape=(current_num_frames, h,
-#                                                       w),  # type: ignore
-#                                     patch_size=fastvideo_args.pipeline_config.
-#                                     dit_config.patch_size,  # type: ignore
-#                                     STA_param=batch.STA_param,  # type: ignore
-#                                     VSA_sparsity=fastvideo_args.
-#                                     VSA_sparsity,  # type: ignore
-#                                     device=get_local_torch_device(),
-#                                 )
-#                                 assert attn_metadata is not None, "attn_metadata cannot be None"
-#                             else:
-#                                 attn_metadata = None
-#                         else:
-#                             attn_metadata = None
-
-#                         batch.is_cfg_negative = False
-#                         with set_forward_context(current_timestep=i,
-#                                                  attn_metadata=attn_metadata,
-#                                                  forward_batch=batch):
-#                             noise_pred = self.transformer(
-#                                 latent_model_input,
-#                                 prompt_embeds,
-#                                 t_expand,
-#                                 guidance=guidance_expand,
-#                                 **image_kwargs,
-#                                 **pos_cond_kwargs,
-#                             )
-
-#                         if batch.do_classifier_free_guidance:
-#                             batch.is_cfg_negative = True
-#                             with set_forward_context(
-#                                     current_timestep=i,
-#                                     attn_metadata=attn_metadata,
-#                                     forward_batch=batch):
-#                                 noise_pred_uncond = self.transformer(
-#                                     latent_model_input,
-#                                     batch.negative_prompt_embeds,
-#                                     t_expand,
-#                                     guidance=guidance_expand,
-#                                     **image_kwargs,
-#                                     **neg_cond_kwargs,
-#                                 )
-#                             noise_pred_text = noise_pred
-#                             noise_pred = noise_pred_uncond + batch.guidance_scale * (
-#                                 noise_pred_text - noise_pred_uncond)
-
-#                             if batch.guidance_rescale > 0.0:
-#                                 noise_pred = self.rescale_noise_cfg(
-#                                     noise_pred,
-#                                     noise_pred_text,
-#                                     guidance_rescale=batch.guidance_rescale,
-#                                 )
-
-#                         # Compute the previous noisy sample for this block only
-#                         current_latents = self.scheduler.step(
-#                             noise_pred,
-#                             t_cur,
-#                             current_latents,
-#                             return_dict=False)[0]
-
-#                     if progress_bar is not None:
-#                         progress_bar.update()
-
-#                 # After finishing timesteps, write denoised block back
-#                 latents[:, :, start_index:start_index +
-#                         current_num_frames, :, :] = current_latents
-
-#                 # Advance to next block
-#                 start_index += current_num_frames
-
-#         # Update batch with final latents
-#         batch.latents = latents
-#         return batch
-
 
 class CausalDMDDenosingStage(DenoisingStage):
     """
@@ -250,8 +35,8 @@ class CausalDMDDenosingStage(DenoisingStage):
         super().__init__(transformer, scheduler)
         self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
         # KV and cross-attention cache state (initialized on first forward)
-        self.kv_cache1 = None
-        self.crossattn_cache = None
+        self.kv_cache1: list | None = None
+        self.crossattn_cache: list | None = None
         # Model-dependent constants (aligned with causal_inference.py assumptions)
         self.num_transformer_blocks = 30
         self.frame_seq_length = 1560
@@ -273,10 +58,9 @@ class CausalDMDDenosingStage(DenoisingStage):
 
         # Args
         try:
-            num_frame_per_block = getattr(fastvideo_args.pipeline_config,
-                                          "num_frame_per_block", 1)
+            num_frames_per_block = getattr(batch, "num_frames_per_block", 1)
         except Exception:
-            num_frame_per_block = 1
+            num_frames_per_block = 1
         try:
             independent_first_frame = getattr(fastvideo_args.pipeline_config,
                                               "independent_first_frame", False)
@@ -290,7 +74,7 @@ class CausalDMDDenosingStage(DenoisingStage):
             device=get_local_torch_device())
 
         # Image kwargs (kept empty unless caller provides compatible args)
-        image_kwargs = {}
+        image_kwargs: dict = {}
 
         pos_cond_kwargs = self.prepare_extra_func_kwargs(
             self.transformer.forward,
@@ -314,12 +98,13 @@ class CausalDMDDenosingStage(DenoisingStage):
         # Initialize or reset caches
         if self.kv_cache1 is None:
             self._initialize_kv_cache(batch_size=latents.shape[0],
-                                      dtype=latents.dtype,
+                                      dtype=target_dtype,
                                       device=latents.device)
             self._initialize_crossattn_cache(batch_size=latents.shape[0],
-                                             dtype=latents.dtype,
+                                             dtype=target_dtype,
                                              device=latents.device)
         else:
+            assert self.crossattn_cache is not None
             # reset cross-attention cache
             for block_index in range(self.num_transformer_blocks):
                 self.crossattn_cache[block_index][
@@ -372,7 +157,7 @@ class CausalDMDDenosingStage(DenoisingStage):
 
             # process remaining input frames in blocks of num_frame_per_block
             while remaining_frames > 0:
-                block = min(num_frame_per_block, remaining_frames)
+                block = min(num_frames_per_block, remaining_frames)
                 ref_btchw = image_latent[:, :, current_start_frame:
                                          current_start_frame +
                                          block, :, :].to(target_dtype).permute(
@@ -400,20 +185,20 @@ class CausalDMDDenosingStage(DenoisingStage):
         # Determine block sizes
         if not independent_first_frame or (independent_first_frame
                                            and batch.image_latent is not None):
-            if t % num_frame_per_block != 0:
+            if t % num_frames_per_block != 0:
                 raise ValueError(
-                    "num_frames must be divisible by num_frame_per_block for causal DMD denoising"
+                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
                 )
-            num_blocks = t // num_frame_per_block
-            block_sizes = [num_frame_per_block] * num_blocks
+            num_blocks = t // num_frames_per_block
+            block_sizes = [num_frames_per_block] * num_blocks
             start_index = 0
         else:
-            if (t - 1) % num_frame_per_block != 0:
+            if (t - 1) % num_frames_per_block != 0:
                 raise ValueError(
                     "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
                 )
-            num_blocks = (t - 1) // num_frame_per_block
-            block_sizes = [1] + [num_frame_per_block] * num_blocks
+            num_blocks = (t - 1) // num_frames_per_block
+            block_sizes = [1] + [num_frames_per_block] * num_blocks
             start_index = 0
 
         # DMD loop in causal blocks
@@ -468,26 +253,25 @@ class CausalDMDDenosingStage(DenoisingStage):
 
                     with torch.autocast(device_type="cuda",
                                         dtype=target_dtype,
-                                        enabled=autocast_enabled):
-                        with set_forward_context(current_timestep=i,
-                                                 attn_metadata=attn_metadata,
-                                                 forward_batch=batch):
-                            # Run transformer; follow DMD stage pattern
-                            pred_noise_btchw = self.transformer(
-                                latent_model_input.permute(0, 2, 1, 3, 4),
-                                prompt_embeds,
-                                t_expand,
-                                kv_cache=self.kv_cache1,
-                                crossattn_cache=self.crossattn_cache,
-                                current_start=(pos_start_base + start_index) *
-                                self.frame_seq_length,
-                                **image_kwargs,
-                                **pos_cond_kwargs,
-                            ).permute(0, 2, 1, 3, 4)
+                                        enabled=autocast_enabled), \
+                        set_forward_context(current_timestep=i,
+                                            attn_metadata=attn_metadata,
+                                            forward_batch=batch):
+                        # Run transformer; follow DMD stage pattern
+                        pred_noise_btchw = self.transformer(
+                            latent_model_input,
+                            prompt_embeds,
+                            t_expand,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=(pos_start_base + start_index) *
+                            self.frame_seq_length,
+                            start_frame=start_index,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        ).permute(0, 2, 1, 3, 4)
 
                     # Convert pred noise to pred video with FM Euler scheduler utilities
-                    from fastvideo.training.training_utils import (
-                        pred_noise_to_pred_video)
                     pred_video_btchw = pred_noise_to_pred_video(
                         pred_noise=pred_noise_btchw.flatten(0, 1),
                         noise_input_latent=noise_latents.flatten(0, 1),
@@ -507,12 +291,12 @@ class CausalDMDDenosingStage(DenoisingStage):
                                 batch.generator, list) else
                                        batch.generator)).to(self.device)
                         noise_btchw = noise
-                        current_latents_btchw = self.scheduler.add_noise(
+                        noise_latents_btchw = self.scheduler.add_noise(
                             pred_video_btchw.flatten(0, 1),
                             noise_btchw.flatten(0, 1),
                             next_timestep).unflatten(0,
                                                      pred_video_btchw.shape[:2])
-                        current_latents = current_latents_btchw.permute(
+                        current_latents = noise_latents_btchw.permute(
                             0, 2, 1, 3, 4)
                     else:
                         current_latents = pred_video_btchw.permute(
@@ -531,19 +315,22 @@ class CausalDMDDenosingStage(DenoisingStage):
                 t_context = torch.ones([latents.shape[0]],
                                        device=latents.device,
                                        dtype=torch.long) * int(context_noise)
-                context_btchw = current_latents.to(target_dtype).permute(
-                    0, 2, 1, 3, 4)
+                context_bcthw = current_latents.to(target_dtype)
                 with torch.autocast(device_type="cuda",
                                     dtype=target_dtype,
-                                    enabled=autocast_enabled):
+                                    enabled=autocast_enabled), \
+                    set_forward_context(current_timestep=0,
+                                        attn_metadata=attn_metadata,
+                                        forward_batch=batch):
                     _ = self.transformer(
-                        context_btchw,
+                        context_bcthw,
                         prompt_embeds,
                         t_context,
                         kv_cache=self.kv_cache1,
                         crossattn_cache=self.crossattn_cache,
                         current_start=(pos_start_base + start_index) *
                         self.frame_seq_length,
+                        start_frame=start_index,
                         **image_kwargs,
                         **pos_cond_kwargs,
                     )
@@ -552,7 +339,7 @@ class CausalDMDDenosingStage(DenoisingStage):
         batch.latents = latents
         return batch
 
-    def _initialize_kv_cache(self, batch_size, dtype, device):
+    def _initialize_kv_cache(self, batch_size, dtype, device) -> None:
         """
         Initialize a Per-GPU KV cache aligned with the Wan model assumptions.
         """
@@ -580,7 +367,7 @@ class CausalDMDDenosingStage(DenoisingStage):
 
         self.kv_cache1 = kv_cache1
 
-    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+    def _initialize_crossattn_cache(self, batch_size, dtype, device) -> None:
         """
         Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions.
         """

fastvideo/training/distillation_pipeline.py

@@ -37,8 +37,8 @@ from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
     load_distillation_checkpoint, pred_noise_to_pred_video,
-    save_distillation_checkpoint, shift_timestep)
-from fastvideo.utils import is_vsa_available, set_random_seed
+    save_distillation_checkpoint, shift_timestep, EMA_FSDP)
+from fastvideo.utils import is_vsa_available, set_random_seed, load_module_from_path
 
 import wandb  # isort: skip
 
@@ -87,9 +87,25 @@ class DistillationPipeline(TrainingPipeline):
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
-        # self.transformer is the generator model
-        self.real_score_transformer = self.get_module("real_score_transformer")
-        self.fake_score_transformer = self.get_module("fake_score_transformer")
+        if training_args.real_score_model_path:
+            logger.info(f"Loading real score transformer from: {training_args.real_score_model_path}")
+            self.real_score_transformer = load_module_from_path(
+                training_args.real_score_model_path, 
+                "transformer", 
+                training_args
+            )
+        else:
+            self.real_score_transformer = self.get_module("real_score_transformer")
+            
+        if training_args.fake_score_model_path:
+            logger.info(f"Loading fake score transformer from: {training_args.fake_score_model_path}")
+            self.fake_score_transformer = load_module_from_path(
+                training_args.fake_score_model_path, 
+                "transformer", 
+                training_args
+            )
+        else:
+            self.fake_score_transformer = self.get_module("fake_score_transformer")
 
         self.real_score_transformer.requires_grad_(False)
         self.real_score_transformer.eval()
@@ -116,10 +132,13 @@ class DistillationPipeline(TrainingPipeline):
         if fake_score_lr == 0.0:
             fake_score_lr = training_args.learning_rate
 
+        betas_str = training_args.fake_score_betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+        
         self.fake_score_optimizer = torch.optim.AdamW(
             fake_score_params,
             lr=fake_score_lr,
-            betas=(0.9, 0.999),
+            betas=betas,
             weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
@@ -158,6 +177,13 @@ class DistillationPipeline(TrainingPipeline):
 
         self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
 
+        self.generator_ema = None
+        if (self.training_args.ema_decay is not None) and (self.training_args.ema_decay > 0.0):
+            self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
+            logger.info(f"Initialized generator EMA with decay={self.training_args.ema_decay}")
+        else:
+            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
+
     @abstractmethod
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         """Initialize validation pipeline - must be implemented by subclasses."""
@@ -174,6 +200,103 @@ class DistillationPipeline(TrainingPipeline):
 
         return training_batch
 
+    def apply_ema_to_model(self, model):
+        """Apply EMA weights to the model for validation or inference."""
+        if self.generator_ema is not None:
+            with self.generator_ema.apply_to_model(model):
+                return model
+        return model
+
+    def get_ema_model_copy(self):
+        """Get a copy of the model with EMA weights applied."""
+        if self.generator_ema is not None:
+            ema_model = copy.deepcopy(self.transformer)
+            self.generator_ema.copy_to_unwrapped(ema_model)
+            return ema_model
+        return None
+
+    def is_ema_ready(self):
+        """Check if EMA is ready for use (after ema_start_step)."""
+        return (self.generator_ema is not None and 
+                self.current_trainstep >= self.training_args.ema_start_step)
+
+    def save_ema_weights(self, output_dir: str, step: int):
+        """Save EMA weights separately for inference purposes."""
+        if self.generator_ema is None:
+            logger.warning("Cannot save EMA weights: EMA not initialized")
+            return
+        
+        if not self.is_ema_ready():
+            logger.warning("Cannot save EMA weights: EMA not ready yet (step < ema_start_step)")
+            return
+        
+        try:
+            ema_model = self.get_ema_model_copy()
+            if ema_model is None:
+                logger.warning("Failed to create EMA model copy")
+                return
+
+            ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
+            os.makedirs(ema_save_dir, exist_ok=True)
+            
+            # save as diffusers format
+            from fastvideo.training.training_utils import gather_state_dict_on_cpu_rank0, custom_to_hf_state_dict
+            from safetensors.torch import save_file
+            cpu_state = gather_state_dict_on_cpu_rank0(ema_model, device=None)
+            
+            if self.global_rank == 0:
+                weight_path = os.path.join(ema_save_dir, "diffusion_pytorch_model.safetensors")
+                diffusers_state_dict = custom_to_hf_state_dict(
+                    cpu_state, ema_model.reverse_param_names_mapping)
+                save_file(diffusers_state_dict, weight_path)
+                
+                config_dict = ema_model.hf_config
+                if "dtype" in config_dict:
+                    del config_dict["dtype"]
+                config_path = os.path.join(ema_save_dir, "config.json")
+                with open(config_path, "w") as f:
+                    json.dump(config_dict, f, indent=4)
+                
+                logger.info(f"EMA weights saved to {weight_path}")
+            
+            del ema_model
+            
+        except Exception as e:
+            logger.error(f"Failed to save EMA weights: {str(e)}")
+
+    def get_ema_stats(self):
+        """Get EMA statistics for monitoring."""
+        if self.generator_ema is None:
+            return {
+                "ema_enabled": False,
+                "ema_decay": None,
+                "ema_start_step": self.training_args.ema_start_step,
+                "ema_ready": False,
+                "ema_step": self.current_trainstep,
+            }
+        
+        return {
+            "ema_enabled": True,
+            "ema_decay": self.training_args.ema_decay,
+            "ema_start_step": self.training_args.ema_start_step,
+            "ema_ready": self.is_ema_ready(),
+            "ema_step": self.current_trainstep,
+        }
+
+    def reset_ema(self):
+        """Reset EMA to current model weights."""
+        if self.generator_ema is not None:
+            logger.info("Resetting EMA to current model weights")
+            self.generator_ema.update(self.transformer)
+            # Force update to current weights by setting decay to 0 temporarily
+            original_decay = self.generator_ema.decay
+            self.generator_ema.decay = 0.0
+            self.generator_ema.update(self.transformer)
+            self.generator_ema.decay = original_decay
+            logger.info("EMA reset completed")
+        else:
+            logger.warning("Cannot reset EMA: EMA not initialized")
+
     def _build_distill_input_kwargs(
             self, noise_input: torch.Tensor, timestep: torch.Tensor,
             text_dict: dict[str, torch.Tensor] | None,
@@ -587,6 +710,10 @@ class DistillationPipeline(TrainingPipeline):
             self._clip_model_grad_norm_(batch_gen, self.transformer)
             self.optimizer.step()
             self.optimizer.zero_grad(set_to_none=True)
+            
+            if self.generator_ema is not None:
+                self.generator_ema.update(self.transformer)
+            
             avg_dmd_loss = torch.tensor(total_dmd_loss /
                                         gradient_accumulation_steps,
                                         device=self.device)
@@ -637,7 +764,8 @@ class DistillationPipeline(TrainingPipeline):
             self.transformer, self.fake_score_transformer, self.global_rank,
             self.training_args.resume_from_checkpoint, self.optimizer,
             self.fake_score_optimizer, self.train_dataloader, self.lr_scheduler,
-            self.fake_score_lr_scheduler, self.noise_random_generator)
+            self.fake_score_lr_scheduler, self.noise_random_generator,
+            self.generator_ema)
 
         if resumed_step > 0:
             self.init_steps = resumed_step
@@ -668,6 +796,12 @@ class DistillationPipeline(TrainingPipeline):
             sum(p.numel()
                 for p in self.fake_score_transformer.parameters()) / 1e9)
 
+        if self.generator_ema is not None:
+            logger.info("  Generator EMA enabled with decay: %s", self.training_args.ema_decay)
+            logger.info("  Generator EMA start step: %s", self.training_args.ema_start_step)
+        else:
+            logger.info("  Generator EMA disabled")
+
     @torch.no_grad()
     def _log_validation(self, transformer, training_args, global_step) -> None:
         training_args.inference_mode = True
@@ -699,6 +833,17 @@ class DistillationPipeline(TrainingPipeline):
 
         transformer.eval()
 
+        # Optionally use EMA model for validation if available and ready
+        use_ema_for_validation = (self.training_args.use_ema and 
+                                  self.is_ema_ready())
+        if use_ema_for_validation:
+            logger.info("Using EMA model for validation")
+            validation_transformer = self.transformer
+            ema_context = self.generator_ema.apply_to_model(validation_transformer)
+        else:
+            validation_transformer = transformer
+            ema_context = None
+
         validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
         validation_steps = [step for step in validation_steps if step > 0]
@@ -714,50 +859,98 @@ class DistillationPipeline(TrainingPipeline):
             step_videos: list[np.ndarray] = []
             step_captions: list[str] = []
 
-            for validation_batch in validation_dataloader:
-                batch = self._prepare_validation_batch(sampling_param,
-                                                       training_args,
-                                                       validation_batch,
-                                                       num_inference_steps)
+            if ema_context is not None:
+                with ema_context:
+                    for validation_batch in validation_dataloader:
+                        batch = self._prepare_validation_batch(sampling_param,
+                                                               training_args,
+                                                               validation_batch,
+                                                               num_inference_steps)
 
-                negative_prompt = batch.negative_prompt
-                batch_negative = ForwardBatch(
-                    data_type="video",
-                    prompt=negative_prompt,
-                    prompt_embeds=[],
-                    prompt_attention_mask=[],
-                )
-                result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
-                    batch_negative, training_args)
-                self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
-                    0], result_batch.prompt_attention_mask[0]
+                        negative_prompt = batch.negative_prompt
+                        batch_negative = ForwardBatch(
+                            data_type="video",
+                            prompt=negative_prompt,
+                            prompt_embeds=[],
+                            prompt_attention_mask=[],
+                        )
+                        result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                            batch_negative, training_args)
+                        self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                            0], result_batch.prompt_attention_mask[0]
 
-                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
-                            self.global_rank,
-                            self.rank_in_sp_group,
-                            batch.prompt,
-                            local_main_process_only=False)
+                        logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                                    self.global_rank,
+                                    self.rank_in_sp_group,
+                                    batch.prompt,
+                                    local_main_process_only=False)
 
-                assert batch.prompt is not None and isinstance(
-                    batch.prompt, str)
-                step_captions.append(batch.prompt)
+                        assert batch.prompt is not None and isinstance(
+                            batch.prompt, str)
+                        step_captions.append(batch.prompt)
 
-                # Run validation inference
-                with torch.no_grad():
-                    output_batch = self.validation_pipeline.forward(
-                        batch, training_args)
-                samples = output_batch.output
-                if self.rank_in_sp_group != 0:
-                    continue
+                        # Run validation inference
+                        with torch.no_grad():
+                            output_batch = self.validation_pipeline.forward(
+                                batch, training_args)
+                        samples = output_batch.output
+                        if self.rank_in_sp_group != 0:
+                            continue
 
-                # Process outputs
-                video = rearrange(samples, "b c t h w -> t b c h w")
-                frames = []
-                for x in video:
-                    x = torchvision.utils.make_grid(x, nrow=6)
-                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-                    frames.append((x * 255).numpy().astype(np.uint8))
-                step_videos.append(frames)
+                        # Process outputs
+                        video = rearrange(samples, "b c t h w -> t b c h w")
+                        frames = []
+                        for x in video:
+                            x = torchvision.utils.make_grid(x, nrow=6)
+                            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                            frames.append((x * 255).numpy().astype(np.uint8))
+                        step_videos.append(frames)
+            else:
+                # Use original transformer without EMA
+                for validation_batch in validation_dataloader:
+                    batch = self._prepare_validation_batch(sampling_param,
+                                                           training_args,
+                                                           validation_batch,
+                                                           num_inference_steps)
+
+                    negative_prompt = batch.negative_prompt
+                    batch_negative = ForwardBatch(
+                        data_type="video",
+                        prompt=negative_prompt,
+                        prompt_embeds=[],
+                        prompt_attention_mask=[],
+                    )
+                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                        batch_negative, training_args)
+                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+
+                    logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                                self.global_rank,
+                                self.rank_in_sp_group,
+                                batch.prompt,
+                                local_main_process_only=False)
+
+                    assert batch.prompt is not None and isinstance(
+                        batch.prompt, str)
+                    step_captions.append(batch.prompt)
+
+                    # Run validation inference
+                    with torch.no_grad():
+                        output_batch = self.validation_pipeline.forward(
+                            batch, training_args)
+                    samples = output_batch.output
+                    if self.rank_in_sp_group != 0:
+                        continue
+
+                    # Process outputs
+                    video = rearrange(samples, "b c t h w -> t b c h w")
+                    frames = []
+                    for x in video:
+                        x = torchvision.utils.make_grid(x, nrow=6)
+                        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                        frames.append((x * 255).numpy().astype(np.uint8))
+                    step_videos.append(frames)
 
             # Log validation results for this step
             world_group = get_world_group()
@@ -947,6 +1140,11 @@ 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(f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}")
+
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 training_batch = self.train_one_step(training_batch)
 
@@ -965,6 +1163,7 @@ class DistillationPipeline(TrainingPipeline):
                 "fake_score_loss": f"{fake_score_loss:.4f}",
                 "step_time": f"{step_time:.2f}s",
                 "grad_norm": grad_norm,
+                "ema": "✓" if (self.generator_ema is not None and self.is_ema_ready()) else "✗",
             })
             progress_bar.update(1)
 
@@ -992,6 +1191,15 @@ class DistillationPipeline(TrainingPipeline):
                 if use_vsa:
                     log_data["VSA_train_sparsity"] = current_vsa_sparsity
 
+                if self.generator_ema is not None:
+                    log_data["ema_enabled"] = True
+                    log_data["ema_decay"] = self.training_args.ema_decay
+                else:
+                    log_data["ema_enabled"] = False
+
+                ema_stats = self.get_ema_stats()
+                log_data.update(ema_stats)
+
                 if training_batch.dmd_latent_vis_dict:
                     dmd_additional_logs = {
                         "generator_timestep":
@@ -1023,7 +1231,8 @@ class DistillationPipeline(TrainingPipeline):
                     self.global_rank, self.training_args.output_dir, step,
                     self.optimizer, self.fake_score_optimizer,
                     self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator)
+                    self.fake_score_lr_scheduler, self.noise_random_generator,
+                    self.generator_ema)
 
                 if self.transformer:
                     self.transformer.train()
@@ -1040,7 +1249,11 @@ class DistillationPipeline(TrainingPipeline):
                                              self.global_rank,
                                              self.training_args.output_dir,
                                              f"{step}_weight_only",
-                                             only_save_generator_weight=True)
+                                             only_save_generator_weight=True,
+                                             generator_ema=self.generator_ema)
+                
+                if self.training_args.use_ema and self.is_ema_ready():
+                    self.save_ema_weights(self.training_args.output_dir, step)
 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                 if self.training_args.log_visualization:
@@ -1060,7 +1273,10 @@ class DistillationPipeline(TrainingPipeline):
             self.training_args.output_dir, self.training_args.max_train_steps,
             self.optimizer, self.fake_score_optimizer, self.train_dataloader,
             self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator)
+            self.noise_random_generator, self.generator_ema)
+
+        if self.training_args.use_ema and self.is_ema_ready():
+            self.save_ema_weights(self.training_args.output_dir, self.training_args.max_train_steps)
 
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/training/training_pipeline.py

@@ -117,10 +117,14 @@ class TrainingPipeline(LoRAPipeline, ABC):
         params_to_optimize = self.transformer.parameters()
         params_to_optimize = list(
             filter(lambda p: p.requires_grad, params_to_optimize))
+        # Parse betas from string format "beta1,beta2"
+        betas_str = training_args.betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+        
         self.optimizer = torch.optim.AdamW(
             params_to_optimize,
             lr=training_args.learning_rate,
-            betas=(0.9, 0.999),
+            betas=betas,
             weight_decay=training_args.weight_decay,
             eps=1e-8,
         )

fastvideo/training/training_utils.py

@@ -202,6 +202,7 @@ def save_distillation_checkpoint(generator_transformer,
                                  generator_scheduler=None,
                                  fake_score_scheduler=None,
                                  noise_generator=None,
+                                 generator_ema=None,
                                  only_save_generator_weight=False) -> None:
     """
     Save distillation checkpoint with both generator and fake_score models.
@@ -233,6 +234,8 @@ def save_distillation_checkpoint(generator_transformer,
         if generator_scheduler is not None:
             generator_states["scheduler"] = SchedulerWrapper(
                 generator_scheduler)
+        if generator_ema is not None:
+            generator_states["ema"] = generator_ema.state_dict()
 
         generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                          "generator")
@@ -402,7 +405,8 @@ def load_distillation_checkpoint(generator_transformer,
                                  dataloader=None,
                                  generator_scheduler=None,
                                  fake_score_scheduler=None,
-                                 noise_generator=None) -> int:
+                                 noise_generator=None,
+                                 generator_ema=None) -> int:
     """
     Load distillation checkpoint with both generator and fake_score models.
     Returns the step number from which training should resume.
@@ -456,6 +460,18 @@ def load_distillation_checkpoint(generator_transformer,
         end_time - begin_time,
         local_main_process_only=False)
 
+    # Load EMA state if available and generator_ema is provided
+    if generator_ema is not None:
+        try:
+            ema_state = generator_states.get("ema")
+            if ema_state is not None:
+                generator_ema.load_state_dict(ema_state)
+                logger.info("rank: %s, generator EMA state loaded successfully", rank)
+            else:
+                logger.info("rank: %s, no EMA state found in checkpoint", rank)
+        except Exception as e:
+            logger.warning("rank: %s, failed to load EMA state: %s", rank, str(e))
+
     # Load critic distributed checkpoint
     critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                   "critic")
@@ -1299,3 +1315,156 @@ def get_scheduler(
                          num_warmup_steps=num_warmup_steps,
                          num_training_steps=num_training_steps,
                          last_epoch=last_epoch)
+
+
+class EMA_FSDP:
+    """
+    FSDP2-friendly EMA with two modes:
+      - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank.
+        Provides a context manager to temporarily swap EMA weights into the live model for teacher forward.
+      - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only
+        using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward.
+
+    Usage (local_shard for CM teacher):
+      ema = EMA_FSDP(model, decay=0.999, mode="local_shard")
+      for step in ...:
+          ema.update(model)
+      with ema.apply_to_model(model):
+          with torch.no_grad():
+              y_teacher = model(...)
+
+    Usage (rank0_full for export):
+      ema = EMA_FSDP(model, decay=0.999, mode="rank0_full")
+      ema.update(model)
+      ema.state_dict()  # on rank 0
+    """
+    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
+        self.decay = float(decay)
+        self.mode = mode
+        self.shadow: dict[str, torch.Tensor] = {}
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        if self.mode not in {"local_shard", "rank0_full"}:
+            raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}")
+        self._init_shadow(module)
+
+    @staticmethod
+    def _to_local_tensor(t: torch.Tensor) -> torch.Tensor:
+        # DTensor-aware to_local fetch; fall back to raw tensor
+        try:
+            from torch.distributed.tensor import DTensor  # type: ignore
+            if isinstance(t, DTensor):
+                return t.to_local()
+        except Exception:
+            pass
+        return t
+
+    @torch.no_grad()
+    def _init_shadow(self, module):
+        if self.mode == "rank0_full":
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            if self.rank == 0:
+                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
+            else:
+                self.shadow = {}
+            return
+
+        # local_shard: maintain EMA of local shards for requires_grad params
+        self.shadow = {}
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            self.shadow[name] = local.clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, module):
+        d = self.decay
+        if self.mode == "rank0_full":
+            if self.rank != 0:
+                return
+            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
+            for n, v in cpu_state.items():
+                v_cpu = v.detach().float().cpu()
+                if n not in self.shadow:
+                    self.shadow[n] = v_cpu.clone()
+                else:
+                    self.shadow[n].mul_(d).add_(v_cpu, alpha=1.0 - d)
+            return
+
+        # local_shard: update local shard EMA on every rank
+        for name, p in module.named_parameters():
+            if not p.requires_grad:
+                continue
+            local = self._to_local_tensor(p.detach())
+            v_cpu = local.float().cpu()
+            if name not in self.shadow:
+                self.shadow[name] = v_cpu.clone()
+            else:
+                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        if self.mode == "rank0_full":
+            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
+        return {k: v.clone() for k, v in self.shadow.items()}
+
+    def load_state_dict(self, sd: dict[str, torch.Tensor]):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    @torch.no_grad()
+    def copy_to_unwrapped(self, module) -> None:
+        """
+        Copy EMA weights into a non-sharded (unwrapped) module. Intended for export/eval.
+        For mode="rank0_full", only rank 0 has the full EMA state.
+        """
+        if self.mode == "rank0_full" and self.rank != 0:
+            return
+        name_to_param = dict(module.named_parameters())
+        for n, w in self.shadow.items():
+            if n in name_to_param:
+                p = name_to_param[n]
+                p.data.copy_(w.to(dtype=p.dtype, device=p.device))
+
+    class _ApplyEMACtx:
+        def __init__(self, ema: "EMA_FSDP", module):
+            self.ema = ema
+            self.module = module
+            self.saved: dict[str, torch.Tensor] = {}
+
+        def __enter__(self):
+            if self.ema.mode != "local_shard":
+                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if not p.requires_grad:
+                        continue
+                    # Save local shard
+                    p_local = EMA_FSDP._to_local_tensor(p.detach())
+                    if p_local.numel() == 0:
+                        # Nothing to swap on this rank for this param
+                        continue
+                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
+                    if name in self.ema.shadow:
+                        ema_cpu = self.ema.shadow[name]
+                        if ema_cpu.numel() != p_local.numel():
+                            # Shard shape mismatch (e.g., empty shard here), skip
+                            continue
+                        # Copy EMA shard into local param shard
+                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
+            return self.module
+
+        def __exit__(self, exc_type, exc, tb):
+            with torch.no_grad():
+                for name, p in self.module.named_parameters():
+                    if name in self.saved:
+                        p_local = EMA_FSDP._to_local_tensor(p.detach())
+                        if p_local.numel() == 0:
+                            continue
+                        saved_local = self.saved[name]
+                        if saved_local.numel() != p_local.numel():
+                            continue
+                        p_local.copy_(saved_local)
+            self.saved.clear()
+            return False
+
+    def apply_to_model(self, module):
+        return EMA_FSDP._ApplyEMACtx(self, module)

fastvideo/utils.py

@@ -34,6 +34,8 @@ from torch.distributed.fsdp import MixedPrecisionPolicy
 import fastvideo.envs as envs
 from fastvideo.logger import init_logger
 
+from fastvideo.models.loader.component_loader import PipelineComponentLoader
+
 logger = init_logger(__name__)
 
 T = TypeVar("T")
@@ -614,6 +616,47 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
             f"Failed to download or parse model_index.json for {model_name_or_path}: {e}"
         ) from e
 
+def load_module_from_path(model_path: str, module_type: str, training_args: TrainingArgs):
+    """
+    Load a module from a specific path using the same loading logic as the pipeline.
+    
+    Args:
+        model_path: Path to the model
+        module_type: Type of module to load (e.g., "transformer")
+        training_args: Training arguments
+        
+    Returns:
+        The loaded module
+    """
+    logger.info(f"Loading {module_type} from custom path: {model_path}")
+    config = verify_model_config_and_directory(model_path)
+    
+    if module_type not in config:
+        if hasattr(self, '_extra_config_module_map') and module_type in self._extra_config_module_map:
+            extra_module = self._extra_config_module_map[module_type]
+            if extra_module in config:
+                module_type = extra_module
+                logger.info(f"Using {extra_module} for {module_type}")
+            else:
+                raise ValueError(f"Module {module_type} not found in config at {model_path}")
+        else:
+            raise ValueError(f"Module {module_type} not found in config at {model_path}")
+    
+    module_info = config[module_type]
+    if module_info is None:
+        raise ValueError(f"Module {module_type} has null value in config at {model_path}")
+    
+    transformers_or_diffusers, architecture = module_info
+    component_path = os.path.join(model_path, module_type)
+    module = PipelineComponentLoader.load_module(
+        module_name=module_type,
+        component_model_path=component_path,
+        transformers_or_diffusers=transformers_or_diffusers,
+        fastvideo_args=training_args,
+    )
+    
+    logger.info(f"Successfully loaded {module_type} from {component_path}")
+    return module
 
 def update_environment_variables(envs: dict[str, str]):
     for k, v in envs.items():