diff --git a/.gitignore b/.gitignore index dd1ea30e5..5010b2cd3 100644 --- a/.gitignore +++ b/.gitignore @@ -33,6 +33,7 @@ env **.txt *.log weights/ +logs/ # SSIM test outputs fastvideo/tests/ssim/generated_videos/ @@ -87,8 +88,3 @@ preprocess_output_text/ .claude/ .codex/ openspec/ - -preprocess_output/ -finetune_i2v_output/ -finetune_df_output/ -sf_train_output/ diff --git a/examples/train/configs/dfsft_matrixgame2.0_i2v_1.3B.yaml b/examples/train/configs/dfsft_matrixgame2.0_i2v_1.3B.yaml deleted file mode 100644 index 42e4598b6..000000000 --- a/examples/train/configs/dfsft_matrixgame2.0_i2v_1.3B.yaml +++ /dev/null @@ -1,61 +0,0 @@ -# MatrixGame causal diffusion-forcing SFT with the new trainer framework. -# -# This mirrors the legacy matrixgame_ar_diffusion_pipeline, but uses the -# modular train/ architecture. By default it keeps the legacy latent-wise -# timestep behavior. - -models: - student: - _target_: fastvideo.train.models.matrixgame.matrixgame_causal.MatrixGameCausalModel - init_from: ../mg_models/Solaris-6K - trainable: true - -method: - _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod - chunk_size: 3 - share_timestep_within_chunk: false - -training: - distributed: - num_gpus: 1 - sp_size: 1 - tp_size: 1 - hsdp_replicate_dim: 1 - hsdp_shard_dim: 1 - - data: - data_path: /mnt/weka/home/hao.zhang/kaiqin/solaris/datasets/vpt/vpt/train_81f/preprocessed_0306 - dataloader_num_workers: 1 - train_batch_size: 1 - training_cfg_rate: 0.0 - seed: 1000 - num_latent_t: 21 - num_height: 352 - num_width: 640 - num_frames: 81 - - optimizer: - learning_rate: 2.0e-6 - betas: [0.9, 0.999] - weight_decay: 1.0e-4 - lr_scheduler: constant - lr_warmup_steps: 0 - - loop: - max_train_steps: 1500000 - gradient_accumulation_steps: 1 - - checkpoint: - output_dir: checkpoints/matrixgame_dfsft_v3 - training_state_checkpointing_steps: 100 - checkpoints_total_limit: 3 - - tracker: - project_name: mg_dfsft_solaris - run_name: matrixgame_dfsft_v3 - - model: - enable_gradient_checkpointing_type: full - -pipeline: - flow_shift: 3 diff --git a/examples/train/configs/finetune_matrixgame2.0_i2v_1.3B.yaml b/examples/train/configs/finetune_matrixgame2.0_i2v_1.3B.yaml deleted file mode 100644 index 4a46db76e..000000000 --- a/examples/train/configs/finetune_matrixgame2.0_i2v_1.3B.yaml +++ /dev/null @@ -1,58 +0,0 @@ -# MatrixGame I2V finetune with the new trainer framework. -# -# This config mirrors the legacy matrixgame_training_pipeline and -# uses a MatrixGame-specific model plugin with the generic FineTuneMethod. - -models: - student: - _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel - init_from: ../mg_models/Solaris-6K - trainable: true - -method: - _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod - -training: - distributed: - num_gpus: 1 - sp_size: 1 - tp_size: 1 - hsdp_replicate_dim: 1 - hsdp_shard_dim: 1 - - data: - data_path: /mnt/weka/home/hao.zhang/kaiqin/solaris/datasets/vpt/vpt/train_81f/preprocessed_0306 - dataloader_num_workers: 1 - train_batch_size: 1 - training_cfg_rate: 0.1 - seed: 1000 - num_latent_t: 21 - num_height: 352 - num_width: 640 - num_frames: 81 - - optimizer: - learning_rate: 2.0e-5 - betas: [0.9, 0.999] - weight_decay: 1.0e-4 - lr_scheduler: constant - lr_warmup_steps: 0 - - loop: - max_train_steps: 1500000 - gradient_accumulation_steps: 1 - - checkpoint: - output_dir: checkpoints/matrixgame_finetune_v3 - training_state_checkpointing_steps: 100 - checkpoints_total_limit: 3 - - tracker: - project_name: mg_finetune_solaris - run_name: matrixgame_finetune_v3 - - model: - enable_gradient_checkpointing_type: full - -pipeline: - flow_shift: 3 diff --git a/examples/train/configs/kd_wan2.1_causal_t2v_1.3B.yaml b/examples/train/configs/kd_wan2.1_causal_t2v_1.3B.yaml new file mode 100644 index 000000000..0dc9940eb --- /dev/null +++ b/examples/train/configs/kd_wan2.1_causal_t2v_1.3B.yaml @@ -0,0 +1,89 @@ +# Knowledge Distillation: Wan 2.1 Causal T2V 1.3B +# +# Trains a causal student (1.3B) with per-frame block-quantized MSE loss +# on teacher (14B) ODE trajectories. The resulting checkpoint serves as +# the ode_init weight initialization for downstream Self-Forcing training. +# +# The same teacher_path_cache can be shared with kd_wan2.1_t2v_1.3B.yaml +# since the cache stores teacher trajectories independent of the student. +# +# After training, export the ode_init checkpoint: +# python -m fastvideo.train.entrypoint.dcp_to_diffusers \ +# --role student \ +# --output_dir outputs/wan2.1_causal_kd_ode_init/diffusers \ +# --checkpoint_dir outputs/wan2.1_causal_kd_ode_init/latest +# +# The exported model.safetensors can then be passed as: +# transformer_override_safetensor in a Self-Forcing config. +# +# Remove the 'teacher' block once the cache is complete to avoid +# loading the 14B model on subsequent runs. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: # remove once cache is complete + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers + trainable: false + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.kd.KDCausalMethod + teacher_path_cache: data/kd_cache/wan14b_4step + t_list: [999, 937, 833, 624, 0] # student inference schedule (num_steps+1) + student_sample_steps: 4 + teacher_inference_steps: 48 # teacher ODE steps; nearest indices selected for t_list + teacher_guidance_scale: 3.5 # CFG scale during teacher path generation + num_frames_per_block: 3 # frames sharing the same noise level (per-frame KD) + +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: 2 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 20 + num_height: 448 + num_width: 832 + num_frames: 77 + + optimizer: + learning_rate: 7e-5 + betas: [0.9, 0.999] + weight_decay: 0.0 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 10000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_causal_kd_ode_init + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + + tracker: + project_name: kd-wan + run_name: wan2.1_causal_kd_ode_init + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/kd_wan2.1_t2v_1.3B.yaml b/examples/train/configs/kd_wan2.1_t2v_1.3B.yaml new file mode 100644 index 000000000..079434f39 --- /dev/null +++ b/examples/train/configs/kd_wan2.1_t2v_1.3B.yaml @@ -0,0 +1,85 @@ +# Knowledge Distillation: Wan 2.1 T2V 1.3B +# +# Trains student (1.3B) with MSE loss on teacher (14B) ODE trajectories. +# The resulting checkpoint serves as the ode_init weight initialization +# for downstream Self-Forcing training. +# +# After training, export the ode_init checkpoint: +# python -m fastvideo.train.entrypoint.dcp_to_diffusers \ +# --role student \ +# --output_dir outputs/wan2.1_kd_ode_init/diffusers \ +# --checkpoint_dir outputs/wan2.1_kd_ode_init/latest +# +# The exported model.safetensors can then be passed as: +# transformer_override_safetensor in a Self-Forcing config. +# +# Remove the 'teacher' block once the cache is complete to avoid +# loading the 14B model on subsequent runs. + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: # remove once cache is complete + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers + trainable: false + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.kd.KDMethod + teacher_path_cache: data/kd_cache/wan14b_4step + t_list: [999, 937, 833, 624, 0] # student inference schedule (num_steps+1) + student_sample_steps: 4 + teacher_inference_steps: 48 # teacher ODE steps; nearest indices selected for t_list + teacher_guidance_scale: 3.5 # CFG scale during teacher path generation + +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: 2 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 20 + num_height: 448 + num_width: 832 + num_frames: 77 + + optimizer: + learning_rate: 7e-5 + betas: [0.9, 0.999] + weight_decay: 0.0 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 10000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_kd_ode_init + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + + tracker: + project_name: kd-wan + run_name: wan2.1_kd_ode_init + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/kdtest/decode_kd_cache.py b/examples/train/configs/kdtest/decode_kd_cache.py new file mode 100644 index 000000000..e44519689 --- /dev/null +++ b/examples/train/configs/kdtest/decode_kd_cache.py @@ -0,0 +1,92 @@ +"""Decode 'real' latents from a KD cache to videos. + +Uses fastvideo's AutoencoderKLWan and its built-in denormalization +(latent * std + mean, matching normalize_dit_input's inverse). + +Usage: + python examples/train/configs/kdtest/decode_kd_cache.py \ + --cache_dir data/kd_test_cache_small \ + --out_dir data/kd_test_cache_small/video \ + --fps 16 +""" + +import argparse +from pathlib import Path + +import imageio +import numpy as np +import torch +from diffusers import AutoencoderKLWan + +MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + + +def denormalize(latents: torch.Tensor, vae: AutoencoderKLWan) -> torch.Tensor: + """Inverse of normalize_dit_input('wan', ...). + + normalize_dit_input: normalized = (latent - mean) * (1/std) + inverse: latent = normalized * std + mean + """ + cfg = vae.config + mean = torch.tensor(cfg.latents_mean, dtype=latents.dtype, + device=latents.device).view(1, -1, 1, 1, 1) + std = torch.tensor(cfg.latents_std, dtype=latents.dtype, + device=latents.device).view(1, -1, 1, 1, 1) + return latents * std + mean + + +def decode_real(pt_path: Path, vae: AutoencoderKLWan, + device: torch.device) -> np.ndarray: + """Load one .pt cache file and decode its 'real' latent to uint8 [T,H,W,3].""" + d = torch.load(pt_path, weights_only=True, map_location="cpu") + real = d["real"].float() # [T, C, H, W] (normalized) + + # [T, C, H, W] → [1, C, T, H, W] + latent = real.permute(1, 0, 2, 3).unsqueeze(0).to(device) + latent = denormalize(latent, vae) + + with torch.no_grad(): + video = vae.decode(latent).sample # [1, 3, T_px, H_px, W_px] + + video = video.squeeze(0).permute(1, 2, 3, 0) # [T, H, W, 3] + video = ((video.clamp(-1, 1) + 1) / 2 * 255).cpu().float().numpy() + return video.astype(np.uint8) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cache_dir", default="data/kd_test_cache_small") + parser.add_argument("--out_dir", default="data/kd_test_cache_small/video") + parser.add_argument("--model_id", default=MODEL_ID) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + samples_dir = Path(args.cache_dir) / "samples" + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print(f"Loading VAE from {args.model_id} ...") + vae = AutoencoderKLWan.from_pretrained( + args.model_id, subfolder="vae", torch_dtype=torch.float32) + vae.eval().to(device) + + pts = sorted(samples_dir.glob("*.pt")) + print(f"Found {len(pts)} samples in {samples_dir}") + + for pt in pts: + out_path = out_dir / (pt.stem + ".mp4") + if out_path.exists() and not args.overwrite: + print(f" skip {pt.name} (exists)") + continue + video_np = decode_real(pt, vae, device) + imageio.mimsave(str(out_path), video_np, fps=args.fps) + print(f" {pt.name} → {out_path.name} {video_np.shape}") + + print(f"Done. Videos in {out_dir}/") + + +if __name__ == "__main__": + main() diff --git a/examples/train/configs/kdtest/kd_causal.yaml b/examples/train/configs/kdtest/kd_causal.yaml new file mode 100644 index 000000000..67cc1f843 --- /dev/null +++ b/examples/train/configs/kdtest/kd_causal.yaml @@ -0,0 +1,91 @@ +# KD Test: Causal Wan 1.3B ODE-init +# +# Smoke-test for KDCausalMethod on a tiny 16-sample dataset. +# Uses 1.3B as teacher (instead of 14B) so everything fits on 1 GPU. +# +# Step 1 — Run KD training: +# torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \ +# --config examples/train/configs/kdtest/kd_causal.yaml +# +# Step 2 — Export ode_init checkpoint: +# python -m fastvideo.train.entrypoint.dcp_to_diffusers \ +# --role student \ +# --checkpoint_dir outputs/kdtest/kd_causal/latest \ +# --output_dir outputs/kdtest/kd_ode_init +# +# Step 3 — Run Self-Forcing with the exported ode_init: +# torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \ +# --config examples/train/configs/kdtest/self_forcing_causal.yaml +# +# Success criteria: +# - Cache generates 16 .pt files under data/kd_test_cache/samples/ +# - kd_loss decreases steadily and approaches ~0 by step 200 + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: # using 1.3B for single-GPU test; swap to 14B for real runs + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.kd.KDCausalMethod + teacher_path_cache: data/kd_test_cache_small + t_list: [995, 937, 833, 625, 0] + student_sample_steps: 4 + teacher_inference_steps: 48 + teacher_guidance_scale: 3.5 + num_frames_per_block: 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_small + dataloader_num_workers: 0 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 42 + num_latent_t: 20 + num_height: 448 + num_width: 832 + num_frames: 77 + + optimizer: + learning_rate: 7e-5 + betas: [0.9, 0.999] + weight_decay: 0.0 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 300 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/kdtest/kd_causal + training_state_checkpointing_steps: 300 + checkpoints_total_limit: 1 + + tracker: + project_name: kd-test + run_name: kd_causal_smoke + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/kdtest/self_forcing_causal.yaml b/examples/train/configs/kdtest/self_forcing_causal.yaml new file mode 100644 index 000000000..eea00f200 --- /dev/null +++ b/examples/train/configs/kdtest/self_forcing_causal.yaml @@ -0,0 +1,111 @@ +# Self-Forcing Test: Causal Wan 1.3B +# +# Smoke-test for SelfForcingMethod on a tiny 16-sample dataset. +# Run this AFTER kd_causal.yaml + dcp_to_diffusers export. +# +# Prerequisite — export ode_init from the KD run: +# python -m fastvideo.train.entrypoint.dcp_to_diffusers \ +# --role student \ +# --checkpoint_dir outputs/kdtest/kd_causal/latest \ +# --output_dir outputs/kdtest/kd_ode_init +# +# Then run this config: +# torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \ +# --config examples/train/configs/kdtest/self_forcing_causal.yaml +# +# Success criteria: +# - generator_loss and critic_loss both finite at step 0 +# - Validation videos logged to tracker every 20 steps +# - No crash through max_train_steps=100 + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers + trainable: true + transformer_override_safetensor: outputs/kdtest/kd_ode_init/transformer/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_small + dataloader_num_workers: 0 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 42 + 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: 100 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/kdtest/self_forcing_causal + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 1 + + tracker: + project_name: kd-test + run_name: self_forcing_causal_smoke + + 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/train/configs/kdtest/validation.json + every_steps: 20 + sampling_steps: [4] + sampling_timesteps: [1000, 750, 500, 250] + num_frames: 69 + guidance_scale: 6.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/kdtest/usage.md b/examples/train/configs/kdtest/usage.md new file mode 100644 index 000000000..e8086ceba --- /dev/null +++ b/examples/train/configs/kdtest/usage.md @@ -0,0 +1,8 @@ +# Step 1 — KD training + bash examples/train/run.sh examples/train/configs/kdtest/kd_causal.yaml + + # Step 2 — Export ode_init +python -m fastvideo.train.entrypoint.dcp_to_diffusers --role student --checkpoint outputs/kdtest/kd_causal/checkpoint-300 --output-dir outputs/kdtest/kd_ode_init + + # Step 3 — Self-Forcing + bash examples/train/run.sh examples/train/configs/kdtest/self_forcing_causal.yaml \ No newline at end of file diff --git a/examples/train/configs/kdtest/validation.json b/examples/train/configs/kdtest/validation.json new file mode 100644 index 000000000..8784d8a6d --- /dev/null +++ b/examples/train/configs/kdtest/validation.json @@ -0,0 +1,20 @@ +{ + "data": [ + { + "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights.", + "video_path": null, + "num_inference_steps": 4, + "height": 448, + "width": 832, + "num_frames": 69 + }, + { + "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The scene is captured from a ground-level angle, giving a low and intimate perspective.", + "video_path": null, + "num_inference_steps": 4, + "height": 448, + "width": 832, + "num_frames": 69 + } + ] +} diff --git a/examples/train/configs/self_forcing_matrixgame2.0_i2v_1.3B.yaml b/examples/train/configs/self_forcing_matrixgame2.0_i2v_1.3B.yaml deleted file mode 100644 index 31087a207..000000000 --- a/examples/train/configs/self_forcing_matrixgame2.0_i2v_1.3B.yaml +++ /dev/null @@ -1,87 +0,0 @@ -# MatrixGame self-forcing distillation with the new trainer framework. -# -# This is the minimal train/ config wiring the existing MatrixGame model -# plugins into SelfForcingMethod. It does not try to replicate every legacy -# pipeline detail; it only matches the current train/ method surface. - -models: - student: - _target_: fastvideo.train.models.matrixgame.matrixgame_causal.MatrixGameCausalModel - init_from: ../mg_models/Solaris-6K - trainable: true - - teacher: - _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel - init_from: ../mg_models/Solaris-6K - trainable: false - disable_custom_init_weights: true - - critic: - _target_: fastvideo.train.models.matrixgame.matrixgame.MatrixGameModel - init_from: ../mg_models/Solaris-6K - 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: 1.0 - dmd_denoising_steps: [1000, 750, 500, 250] - - chunk_size: 3 - student_sample_type: sde - same_step_across_blocks: false - last_step_only: false - context_noise: 0.0 - enable_gradient_in_rollout: true - start_gradient_frame: 0 - - fake_score_learning_rate: 8.0e-6 - fake_score_betas: [0.0, 0.999] - fake_score_lr_scheduler: constant - -training: - distributed: - num_gpus: 1 - sp_size: 1 - tp_size: 1 - hsdp_replicate_dim: 1 - hsdp_shard_dim: 1 - - data: - data_path: /mnt/weka/home/hao.zhang/kaiqin/solaris/datasets/vpt/vpt/train_81f/preprocessed_0306 - dataloader_num_workers: 1 - train_batch_size: 1 - training_cfg_rate: 0.0 - seed: 1000 - num_latent_t: 21 - num_height: 352 - num_width: 640 - num_frames: 81 - - optimizer: - learning_rate: 1.0e-5 - betas: [0.0, 0.999] - weight_decay: 1.0e-4 - lr_scheduler: constant - lr_warmup_steps: 0 - - loop: - max_train_steps: 1500000 - gradient_accumulation_steps: 1 - - checkpoint: - output_dir: checkpoints/matrixgame_self_forcing_v3 - training_state_checkpointing_steps: 100 - checkpoints_total_limit: 3 - - tracker: - project_name: mg_self_forcing - run_name: matrixgame_self_forcing_v3 - - model: - enable_gradient_checkpointing_type: full - -pipeline: - flow_shift: 3