Merge remote-tracking branch 'origin/main' into feat/kaiqin/mg_solaris

2026-03-10 15:10:34 -07:00
parent 369970e600 d92858659d
commit 72ff437d8a
88 changed files with 11611 additions and 107 deletions
@@ -0,0 +1,70 @@
+# DFSFT (Diffusion-Forcing SFT): Wan 2.1 T2V 1.3B Causal
+#
+# - Student: trainable causal Wan model
+# - Training: inhomogeneous timesteps per chunk (diffusion forcing)
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
+  chunk_size: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_causal_dfsft
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_wan_r
+    run_name: wan2.1_causal_dfsft
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json
+    every_steps: 50
+    sampling_steps: [40]
+    guidance_scale: 6.0
+    num_frames: 69
+
+pipeline:
+  flow_shift: 8
@@ -0,0 +1,94 @@
+# DMD2 distillation: Wan 2.1 T2V 1.3B (teacher 50-step -> student 4-step).
+#
+# - Teacher: frozen pretrained Wan 2.1 T2V 1.3B
+# - Student: trainable, initialized from the same pretrained weights
+# - Critic:  trainable, initialized from the same pretrained weights
+# - Validation: 4-step SDE sampling
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  # Critic optimizer (required — no fallback to training.optimizer)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_dmd2_4steps
+    training_state_checkpointing_steps: 20
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_wan
+    run_name: wan2.1_dmd2_4steps
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_dmd_pipeline.WanDMDPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    guidance_scale: 6.0
+  ema:                                                                                                                     
+    _target_: fastvideo.train.callbacks.ema.EMACallback                                                                    
+    decay: 0.98                                                                                                          
+    start_iter: 0        # delay EMA updates until this iteration 
+
+pipeline:
+  flow_shift: 8
@@ -0,0 +1,194 @@
+# ==============================================================================
+# Full configuration reference for fastvideo.train
+#
+# Legend:
+#   [TYPED]    — parsed into a typed dataclass; fields are validated with
+#                defaults. Unknown keys are silently ignored.
+#   [FREE]     — free-form dict passed as-is to the target class / method.
+#                Keys depend on the _target_ class constructor / method_config.
+#   [RESOLVED] — parsed by PipelineConfig.from_kwargs(); auto-populated from
+#                the model's config files. Only scalar overrides are useful.
+# ==============================================================================
+
+# ------------------------------------------------------------------------------
+# models: [FREE]
+#
+# Each role is instantiated via _target_(*, training_config=..., **kwargs).
+# Keys here are constructor kwargs of the _target_ class (e.g. WanModel).
+# You can define any role name (student, teacher, critic, etc.).
+# ------------------------------------------------------------------------------
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel    # required
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers     # required: HF repo or local path
+    trainable: true                                   # default: true
+    disable_custom_init_weights: false                # default: false
+    flow_shift: 3.0                                   # default: 3.0
+    enable_gradient_checkpointing_type: null           # default: null (falls back to training.model)
+
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+# ------------------------------------------------------------------------------
+# method: [FREE]
+#
+# Instantiated via _target_(*, cfg=RunConfig, role_models=...).
+# All keys besides _target_ are available in self.method_config (a plain dict).
+# Keys depend entirely on the method class.
+# ------------------------------------------------------------------------------
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method  # required
+
+  # --- DMD2-specific keys (read from self.method_config) ---
+  rollout_mode: simulate              # required: "simulate" or "data_latent"
+  generator_update_interval: 5        # default: 1
+  dmd_denoising_steps: [1000, 750, 500, 250]  # SDE timestep schedule
+
+  # Critic optimizer (all required — no fallback)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+  # CFG conditioning policy (optional)
+  # cfg_uncond:
+  #   on_missing: error               # "error" or "ignore"
+  #   text: keep                      # "keep", "zero", "drop", "negative_prompt"
+  #   image: keep                     # "keep", "zero", "drop"
+  #   action: keep                    # "keep", "zero", "drop"
+
+  # --- FineTuneMethod keys (if using finetune instead) ---
+  # _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+# ------------------------------------------------------------------------------
+# training: [TYPED] -> TrainingConfig
+#
+# Every field below has a typed default. Unknown keys are ignored.
+# ------------------------------------------------------------------------------
+training:
+
+  # --- training.distributed [TYPED] -> DistributedConfig ---
+  distributed:
+    num_gpus: 8                       # default: 1
+    tp_size: 1                        # default: 1
+    sp_size: 1                        # default: 1 (defaults to num_gpus in loader)
+    hsdp_replicate_dim: 1             # default: 1
+    hsdp_shard_dim: 8                 # default: -1 (defaults to num_gpus in loader)
+    pin_cpu_memory: false             # default: false
+
+  # --- training.data [TYPED] -> DataConfig ---
+  data:
+    data_path: data/my_dataset        # default: ""
+    train_batch_size: 1               # default: 1
+    dataloader_num_workers: 4         # default: 0
+    training_cfg_rate: 0.1            # default: 0.0
+    seed: 1000                        # default: 0
+    num_height: 448                   # default: 0
+    num_width: 832                    # default: 0
+    num_latent_t: 20                  # default: 0
+    num_frames: 77                    # default: 0
+
+  # --- training.optimizer [TYPED] -> OptimizerConfig ---
+  # Note: only for the student optimizer. Critic optimizer is in method config.
+  optimizer:
+    learning_rate: 2.0e-6             # default: 0.0
+    betas: [0.9, 0.999]              # default: [0.9, 0.999]
+    weight_decay: 0.01                # default: 0.0
+    lr_scheduler: constant            # default: "constant"
+    lr_warmup_steps: 0                # default: 0
+    lr_num_cycles: 0                  # default: 0
+    lr_power: 0.0                     # default: 0.0
+    min_lr_ratio: 0.5                 # default: 0.5
+
+  # --- training.loop [TYPED] -> TrainingLoopConfig ---
+  loop:
+    max_train_steps: 10000            # default: 0
+    gradient_accumulation_steps: 1    # default: 1
+
+  # --- training.checkpoint [TYPED] -> CheckpointConfig ---
+  checkpoint:
+    output_dir: outputs/my_run        # default: ""
+    resume_from_checkpoint: ""        # default: "" (or use --resume-from-checkpoint CLI)
+    training_state_checkpointing_steps: 1000  # default: 0 (disabled)
+    checkpoints_total_limit: 3        # default: 0 (keep all)
+
+  # --- training.tracker [TYPED] -> TrackerConfig ---
+  tracker:
+    trackers: []                      # default: [] (auto-adds "wandb" if project_name is set)
+    project_name: my_project          # default: "fastvideo"
+    run_name: my_run                  # default: ""
+
+  # --- training.vsa [TYPED] -> VSAConfig ---
+  vsa:
+    sparsity: 0.0                     # default: 0.0 (0.0 = disabled)
+
+  # --- training.model [TYPED] -> ModelTrainingConfig ---
+  model:
+    weighting_scheme: uniform         # default: "uniform"
+    logit_mean: 0.0                   # default: 0.0
+    logit_std: 1.0                    # default: 1.0
+    mode_scale: 1.0                   # default: 1.0
+    precondition_outputs: false       # default: false
+    moba_config: {}                   # default: {}
+    enable_gradient_checkpointing_type: full  # default: null ("full" or null)
+
+  # --- training top-level [TYPED] ---
+  dit_precision: fp32                 # default: "fp32" (master weight precision)
+  # model_path: ...                   # default: "" (auto-derived from models.student.init_from)
+
+# ------------------------------------------------------------------------------
+# callbacks: [FREE]
+#
+# Each callback is instantiated via _target_(*, **kwargs).
+# The callback name (e.g. "grad_clip") is arbitrary — only _target_ matters.
+# training_config is injected automatically (not from YAML).
+# ------------------------------------------------------------------------------
+callbacks:
+
+  # --- GradNormClipCallback ---
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback  # optional if using default registry
+    max_grad_norm: 1.0                # default: 0.0 (0.0 = disabled)
+    log_grad_norms: false             # default: false
+
+  # --- EMACallback ---
+  # ema:
+  #   _target_: fastvideo.train.callbacks.ema.EMACallback
+  #   decay: 0.9999                   # default: 0.9999
+  #   start_iter: 0                   # default: 0
+
+  # --- ValidationCallback ---
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback  # optional if using default registry
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline  # required
+    dataset_file: path/to/validation.json     # required
+    every_steps: 100                  # default: 100
+    sampling_steps: [4]               # default: [40]
+    guidance_scale: 5.0               # default: null (uses model default)
+    num_frames: null                  # default: null (derived from training.data)
+    output_dir: null                  # default: null (falls back to training.checkpoint.output_dir)
+    sampling_timesteps: null          # default: null (explicit timestep list for streaming)
+
+# ------------------------------------------------------------------------------
+# pipeline: [RESOLVED] -> PipelineConfig
+#
+# Parsed by PipelineConfig.from_kwargs(). Most fields are auto-populated from
+# the model's config files (vae_config, dit_config, text_encoder_configs, etc.).
+# Only scalar overrides are typically needed here.
+# ------------------------------------------------------------------------------
+pipeline:
+  flow_shift: 3                       # default: null (model-specific)
+  # flow_shift_sr: null               # default: null (super-resolution shift)
+  # embedded_cfg_scale: 6.0           # default: 6.0
+  # is_causal: false                  # default: false
+  # vae_tiling: true                  # default: true
+  # vae_sp: true                      # default: true
+  # disable_autocast: false           # default: false
@@ -0,0 +1,74 @@
+# V3 config: Wan 2.1 T2V 1.3B finetune with VSA (phase 3.4, 0.9 sparsity).
+#
+# Uses _target_-based instantiation — each model role is an independent
+# class instance; the method class is resolved directly from the YAML.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/phase3.4_wan2.1_finetune_vsa_0.9_v3
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: distillation_wan_r
+    run_name: phase3.4_wan_finetune_vsa_0.9_v3
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  vsa:
+    sparsity: 0.9
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [50]
+    guidance_scale: 5.0
+
+pipeline:
+  flow_shift: 3
@@ -0,0 +1,98 @@
+# Self-Forcing distillation: Wan 2.1 T2V 1.3B Causal
+#
+# - Teacher: frozen pretrained Wan 2.1 T2V 1.3B
+# - Student: trainable causal Wan model
+# - Critic:  trainable, initialized from pretrained weights
+# - Training: streaming rollout with SDE sampling
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    transformer_override_safetensor: /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  chunk_size: 3
+  student_sample_type: sde
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1e-5
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_causal_self_forcing
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: self-forcing-wan
+    run_name: wan2.1_causal_self_forcing
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline.WanCausalDMDPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 69
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,105 @@
+# Self-Forcing distillation: Wan 2.1 T2V 1.3B Causal with VSA (0.9 sparsity)
+#
+# - Teacher: frozen pretrained Wan 2.1 T2V 14B
+# - Student: trainable causal Wan model with VSA sparse attention
+# - Critic:  trainable, initialized from pretrained weights
+# - Training: streaming rollout with SDE sampling + VSA 0.9 sparsity
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    transformer_override_safetensor: /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  chunk_size: 3
+  student_sample_type: sde
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+  attn_kind: vsa
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1e-5
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_causal_self_forcing_vsa_0.9
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: self-forcing-wan
+    run_name: wan2.1_causal_self_forcing_vsa_0.9
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  vsa:
+    sparsity: 0.9
+    decay_rate: 0.03
+    decay_interval_steps: 1
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline.WanCausalDMDPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 69
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Launch training from a YAML config.
+#
+# Usage:
+#   bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
+#
+# Examples:
+#   bash examples/train/run.sh examples/train/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml --dry-run
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#       --training.distributed.num_gpus 4 \
+#       --training.optimizer.learning_rate 1e-5
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
+#
+# Logs are written to logs/<config_name>_<timestamp>.log (and also printed to stdout).
+
+set -euo pipefail
+
+CONFIG="${1:?Usage: $0 <config.yaml> [extra flags...]}"
+shift
+
+# ── GPU / node settings ──────────────────────────────────────────
+NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}"
+NUM_GPUS="${NUM_GPUS:-1}"
+NNODES="${NNODES:-1}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29501}"
+export TOKENIZERS_PARALLELISM=false
+# ── W&B ──────────────────────────────────────────────────────────
+export WANDB_API_KEY="${WANDB_API_KEY:-}"
+export WANDB_MODE="${WANDB_MODE:-online}"
+
+# ── Log file ─────────────────────────────────────────────────────
+CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+LOG_DIR="${LOG_DIR:-examples/train}"
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/${CONFIG_NAME}_${TIMESTAMP}.log"
+
+echo "=== Train Training ==="
+echo "Config:      ${CONFIG}"
+echo "Num GPUs:    ${NUM_GPUS}"
+echo "Num Nodes:   ${NNODES}"
+echo "Node Rank:   ${NODE_RANK}"
+echo "Master:      ${MASTER_ADDR}:${MASTER_PORT}"
+echo "Extra args:  $*"
+echo "Log file:    ${LOG_FILE}"
+echo "=============================="
+
+python -m torch.distributed.run \
+    --nnodes "${NNODES}" \
+    --node_rank "${NODE_RANK}" \
+    --nproc_per_node "${NUM_GPUS}" \
+    --master_addr "${MASTER_ADDR}" \
+    --master_port "${MASTER_PORT}" \
+    fastvideo/train/entrypoint/train.py \
+    --config "${CONFIG}" \
+    "$@" \
+    2>&1 | tee "${LOG_FILE}"
@@ -0,0 +1,114 @@
+#!/usr/bin/env bash
+# Submit a multi-node Slurm training job.
+#
+# Usage:
+#   bash examples/train/run_slurm.sh <config.yaml> <num_nodes> [--dotted.key value ...]
+#
+# Examples:
+#   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 2
+#   bash examples/train/run_slurm.sh examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml 4 \
+#       --training.optimizer.learning_rate 1e-5
+#   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 8 \
+#       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
+#
+# Environment variables (override defaults):
+#   PARTITION       Slurm partition           (default: main)
+#   NUM_GPUS        GPUs per node             (default: 8)
+#   CPUS_PER_TASK   CPUs per task             (default: 128)
+#   MEM             Memory per node           (default: 1440G)
+#   JOB_NAME        Slurm job name            (default: derived from config)
+#   OUTPUT_DIR      Directory for slurm logs  (default: slurm_logs)
+#   MASTER_PORT     Rendezvous port           (default: 29500)
+#   EXCLUDE         Nodes to exclude          (default: "")
+#   WANDB_API_KEY   W&B API key               (default: "")
+#   WANDB_MODE      W&B mode                  (default: online)
+
+set -euo pipefail
+
+CONFIG="${1:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
+NUM_NODES="${2:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
+shift 2
+
+# ── Defaults ──────────────────────────────────────────────────────
+PARTITION="${PARTITION:-main}"
+NUM_GPUS="${NUM_GPUS:-8}"
+CPUS_PER_TASK="${CPUS_PER_TASK:-128}"
+MEM="${MEM:-1440G}"
+MASTER_PORT="${MASTER_PORT:-29500}"
+EXCLUDE="${EXCLUDE:-}"
+WANDB_API_KEY="${WANDB_API_KEY:-}"
+WANDB_MODE="${WANDB_MODE:-online}"
+
+TOTAL_GPUS=$(( NUM_NODES * NUM_GPUS ))
+
+CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
+JOB_NAME="${JOB_NAME:-${CONFIG_NAME}}"
+OUTPUT_DIR="${OUTPUT_DIR:-logs/slurm}"
+mkdir -p "${OUTPUT_DIR}"
+
+# ── Build sbatch args ─────────────────────────────────────────────
+SBATCH_ARGS=(
+    --job-name="${JOB_NAME}"
+    --partition="${PARTITION}"
+    --nodes="${NUM_NODES}"
+    --ntasks="${NUM_NODES}"
+    --ntasks-per-node=1
+    --gres="gpu:${NUM_GPUS}"
+    --cpus-per-task="${CPUS_PER_TASK}"
+    --mem="${MEM}"
+    --output="${OUTPUT_DIR}/${JOB_NAME}_%j.out"
+    --error="${OUTPUT_DIR}/${JOB_NAME}_%j.err"
+    --exclusive
+)
+
+if [[ -n "${EXCLUDE}" ]]; then
+    SBATCH_ARGS+=(--exclude="${EXCLUDE}")
+fi
+
+# ── Collect extra overrides for the training script ───────────────
+EXTRA_ARGS=("$@")
+
+echo "=== Slurm Training Submission ==="
+echo "Config:      ${CONFIG}"
+echo "Nodes:       ${NUM_NODES}"
+echo "GPUs/node:   ${NUM_GPUS}"
+echo "Total GPUs:  ${TOTAL_GPUS}"
+echo "Partition:   ${PARTITION}"
+echo "Job name:    ${JOB_NAME}"
+echo "Extra args:  ${EXTRA_ARGS[*]:-}"
+echo "================================="
+
+# ── Submit ────────────────────────────────────────────────────────
+sbatch "${SBATCH_ARGS[@]}" <<EOF
+#!/bin/bash
+set -e -x
+
+# ── Environment ───────────────────────────────────────────────────
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export TRITON_CACHE_DIR=/tmp/triton_cache_\${SLURM_PROCID}
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="${WANDB_API_KEY}"
+export WANDB_MODE="${WANDB_MODE}"
+
+# ── Rendezvous ────────────────────────────────────────────────────
+export MASTER_PORT=${MASTER_PORT}
+nodes=( \$(scontrol show hostnames \$SLURM_JOB_NODELIST) )
+export MASTER_ADDR=\${nodes[0]}
+export NODE_RANK=\$SLURM_PROCID
+
+echo "MASTER_ADDR: \$MASTER_ADDR"
+echo "NODE_RANK:   \$NODE_RANK"
+
+# ── Launch ────────────────────────────────────────────────────────
+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/train/entrypoint/train.py \\
+    --config ${CONFIG} \\
+    --training.distributed.num_gpus ${TOTAL_GPUS} \\
+    ${EXTRA_ARGS[*]:-}
+EOF
@@ -2,35 +2,15 @@
  "data": [
    {
      "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.",
-      "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
    },
    {
      "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.",
-      "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
    },
    {
      "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.",
-      "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
    },
    {
      "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.",
-      "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
    }
  ]
 }