From 467150f68fd1d54aebf1e230a7d9e83ecd213e14 Mon Sep 17 00:00:00 2001 From: bubbliiiing <3323290568@qq.com> Date: Mon, 19 Oct 2026 20:23:02 +0800 Subject: [PATCH] Update sparkwan --- .../README_TRAIN_DISTILL_wan2.1.md | 18 +++++----- .../README_TRAIN_DISTILL_wan2.1_zh-CN.md | 18 +++++----- .../README_TRAIN_DISTILL_wan2.2.md | 35 ++++++++++--------- .../README_TRAIN_DISTILL_wan2.2_zh-CN.md | 35 ++++++++++--------- .../train_distill_wan2.2_high.sh | 1 + .../train_distill_wan2.2_low.sh | 1 + 6 files changed, 58 insertions(+), 50 deletions(-) diff --git a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md index a315c80..c930757 100644 --- a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md +++ b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1.md @@ -130,7 +130,7 @@ The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` ### 3.2 Stage 1 — Sparse Warm-up -Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launcher caps it at `--max_train_steps=400`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2. +Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launcher caps it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2. Run [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh): @@ -156,9 +156,9 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --train_batch_size=1 \ --gradient_accumulation_steps=1 \ --num_train_epochs=100 \ - --max_train_steps=400 \ - --checkpointing_steps=100 \ - --learning_rate=2e-05 \ + --max_train_steps=2000 \ + --checkpointing_steps=200 \ + --learning_rate=1e-4 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -188,7 +188,7 @@ Stage-1 keys: | `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). | | `--max_train_steps` | Warm-up budget; keep it short. | -On completion the warm-up product is `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors`. +On completion the warm-up product is `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`. ### 3.3 Stage 2 — DMD Distill @@ -201,7 +201,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 warm-up product. Env-overridable; --student_topk below MUST match Stage 1. -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ @@ -278,7 +278,7 @@ Setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill | `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` | | `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.10` | | `--sparse_student` | (Stage 2) build the generator as the sparse student | - | -| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-400/diffusion_pytorch_model.safetensors` | +| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` | | `--denoising_step_indices_list` | Discrete few-step grid the generator rolls out; validation reuses it (`num_inference_steps` = its length) | `1000 750 500 250` | | `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic Euler steps (LightX2V-style), no x0 re-noising | - | | `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - | @@ -293,7 +293,7 @@ Setting `--loss_scale_div > 0` layers the official SparkDiffusion **CrossDistill | `--div_rf_boundary` | Phase-0 right boundary `rf_b` (official `backward_timesteps[0]`) | `0.933781` | | `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` | | `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` | -| `--checkpointing_steps` | Save checkpoint every N steps | `200` (distill) / `100` (warm-up) | +| `--checkpointing_steps` | Save checkpoint every N steps | `200` | | `--max_grad_norm` | Gradient clipping threshold | `0.05` | | `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` | @@ -323,7 +323,7 @@ DeepSpeed Zero-2 shards optimizer states and gradients but **not** parameters. F export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \ diff --git a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md index e4027ba..5a41429 100644 --- a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md +++ b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.1_zh-CN.md @@ -130,7 +130,7 @@ modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_ ### 3.2 Stage 1 — 稀疏暖启 -Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=400`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。 +Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。 用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.1.sh](./train_finetune_sparse_wan2.1.sh) 运行： @@ -156,9 +156,9 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --train_batch_size=1 \ --gradient_accumulation_steps=1 \ --num_train_epochs=100 \ - --max_train_steps=400 \ - --checkpointing_steps=100 \ - --learning_rate=2e-05 \ + --max_train_steps=2000 \ + --checkpointing_steps=200 \ + --learning_rate=1e-4 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -188,7 +188,7 @@ Stage-1 关键参数： | `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 | | `--max_train_steps` | 暖启预算；保持简短。 | -完成后暖启产物为 `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors`。 +完成后暖启产物为 `output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。 ### 3.3 Stage 2 — DMD 蒸馏 @@ -201,7 +201,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 暖启产物。可用环境变量覆盖；下方 --student_topk 必须与 Stage 1 一致。 -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ @@ -278,7 +278,7 @@ bash scripts/sparkdiffusion/train_distill_wan2.1.sh | `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` | | `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.10` | | `--sparse_student` | （Stage 2）把 generator 建成稀疏学生 | - | -| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-400/diffusion_pytorch_model.safetensors` | +| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` | | `--denoising_step_indices_list` | generator rollout 的离散少步网格；validation 复用同一网格（`num_inference_steps` = 其长度） | `1000 750 500 250` | | `--flow_euler_rollout` | 让 generator 在 flow/速度空间用确定性 Euler 步 rollout（LightX2V 式），不做 x0 重新加噪 | - | | `--uniform_sampling` | 在训练调度上均匀采样 DMD 噪声档位 | - | @@ -293,7 +293,7 @@ bash scripts/sparkdiffusion/train_distill_wan2.1.sh | `--div_rf_boundary` | phase-0 右边界 `rf_b`（官方 `backward_timesteps[0]`） | `0.933781` | | `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` | | `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` | -| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200`（蒸馏）/ `100`（暖启） | +| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` | | `--max_grad_norm` | 梯度裁剪阈值 | `0.05` | | `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` | @@ -323,7 +323,7 @@ DeepSpeed Zero-2 只分片优化器态与梯度、**不**分片参数。对三 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.1_1.3b_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/sparkdiffusion/train_distill_wan2.1.py \ diff --git a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md index fdbafd4..b288491 100644 --- a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md +++ b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2.md @@ -130,7 +130,7 @@ The base checkpoint seeds all three Stage-2 transformers (the sparse `generator` ### 3.2 Stage 1 — Sparse Warm-up -Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launchers cap it at `--max_train_steps=400` (≈2,000 videos at 8 GPUs × batch 1). Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2. +Stage 1 inserts the RoLa sparse self-attention and briefly warms it up with the step-local flow-matching objective. It is a **short** stage — a brief warm-up is enough to establish a usable prior for distillation, so the launchers cap it at `--max_train_steps=2000`. Do **not** judge this stage by its own sample quality; the sharpness is established in Stage 2. Run the high-noise expert with [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh): @@ -152,9 +152,9 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --train_batch_size=1 \ --gradient_accumulation_steps=1 \ --num_train_epochs=100 \ - --max_train_steps=400 \ - --checkpointing_steps=100 \ - --learning_rate=2e-05 \ + --max_train_steps=2000 \ + --checkpointing_steps=200 \ + --learning_rate=1e-4 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -191,7 +191,7 @@ Stage-1 keys: | `--trainable_modules` | `"attn_op"` = Stage A (train only the sparse branch); `"."` = Stage B (full sparse finetune, load Stage A via `--transformer_path`). | | `--max_train_steps` | Warm-up budget; keep it short. | -On completion the warm-up product is `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors`. +On completion the warm-up product is `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`. ### 3.3 Stage 2 — CrossDistill @@ -204,7 +204,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 warm-up product. -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ @@ -223,7 +223,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -236,6 +236,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ @@ -293,11 +294,11 @@ On the **high**-noise expert, setting `--loss_scale_div > 0` layers the official | `--train_mode` | SparkDiffusion supports `normal` (T2V) only | `normal` | | `--student_topk` | Sparse attention keep-ratio (sparsity = 1 − topk); Stage 1 and Stage 2 must match | `0.05` | | `--sparse_student` | (Stage 2) build generator + init_rollout as the sparse student | - | -| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-400/diffusion_pytorch_model.safetensors` | +| `--generator_transformer_path` | (Stage 2) Stage-1 warm-up checkpoint loaded into the generator | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` | | `--uniform_sampling` | Sample the DMD noise level uniformly over the training schedule | - | | `--flow_euler_rollout` | Roll the generator in flow/velocity space with deterministic rf-native Euler steps (no x0 re-noising) | - | | `--gen_update_interval` | Update the generator once every N steps (the critic updates every step) | `5` | -| `--joint_backward_timesteps` | Joint few-step rf grid (descending), split at the config `boundary` into per-expert segments by `--boundary_type`; auto-defaulted, and validation reuses it | `0.933781 0.875 0.608979` | +| `--joint_backward_timesteps` | Joint few-step rf grid (descending), split at the config `boundary` into per-expert segments by `--boundary_type`; passed explicitly in the launchers (equals the built-in default), and validation reuses it | `0.933781 0.875 0.608979` | | `--div_mode` | CrossDistill diversity mode: `pcm_random` (random phase-0 offset) or `pcm_fixed` (high expert only) | `pcm_random` | | `--loss_scale_div` | Weight of the CrossDistill PCM diversity term (per-element-mean space); `0` disables it | `10.0` (high) / `0` (low) | | `--dcm_total_steps` | Subdivision of the diversity term's continuous rf grid (`dw = 1 / N`) | `48` | @@ -307,12 +308,12 @@ On the **high**-noise expert, setting `--loss_scale_div > 0` layers the official | `--div_teacher_guidance` | CFG scale of the diversity teacher rollout (independent of `--real_guidance_scale`) | `5.0` | | `--real_guidance_scale` | CFG scale of the frozen dense teacher | `5.0` | | `--init_rollout_path` | (low expert) trained high student used to roll the rollout's boundary start latent | high-distill ckpt | -| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `4e-07` | -| `--checkpointing_steps` | Save checkpoint every N steps | `200` (distill) / `100` (warm-up) | +| `--learning_rate` / `--learning_rate_critic` | Generator / critic learning rate | `2e-06` / `2e-06` | +| `--checkpointing_steps` | Save checkpoint every N steps | `200` | | `--max_grad_norm` | Gradient clipping threshold | `0.05` | | `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` | -> 💡 The per-expert few-step grid is **auto-derived**: `--joint_backward_timesteps` (the joint rf grid) is split at the config `boundary` by `--boundary_type`, and `--init_rollout_timesteps` (the low expert's frozen high-segment relay) is filled from the knots above the boundary. Leave both at their defaults; validation reuses the same grid, so `num_inference_steps` = `len(--joint_backward_timesteps) + 1` (`4` here). +> 💡 The per-expert few-step grid is **auto-derived**: `--joint_backward_timesteps` (the joint rf grid, passed explicitly in the launchers) is split at the config `boundary` by `--boundary_type`, and `--init_rollout_timesteps` (the low expert's frozen high-segment relay) is filled from the knots above the boundary. Leave `--init_rollout_timesteps` at its default; validation reuses the same grid, so `num_inference_steps` = `len(--joint_backward_timesteps) + 1` (`4` here). ### 3.5 FSDP Backend and Memory Notes @@ -339,7 +340,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 warm-up product. -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -357,7 +358,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -370,6 +371,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ @@ -394,7 +396,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 warm-up product. -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -412,7 +414,7 @@ accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_ --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -425,6 +427,7 @@ accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_ --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ diff --git a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md index 251cee7..988cb4a 100644 --- a/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md +++ b/scripts/sparkdiffusion/README_TRAIN_DISTILL_wan2.2_zh-CN.md @@ -130,7 +130,7 @@ modelscope download --model Wan-AI/Wan2.2-T2V-A14B --local_dir models/Diffusion_ ### 3.2 Stage 1 — 稀疏暖启 -Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=400`（8 卡 × batch 1 时约 2,000 个视频）。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。 +Stage 1 插入 RoLa 稀疏自注意力，并用 step-local flow-matching 目标短暂暖启。这是一个**短**阶段——短暂暖启即足以为蒸馏建立可用先验，因此启动脚本把它限制在 `--max_train_steps=2000`。**不要**用该阶段自身的采样质量来评判它；清晰度是在 Stage 2 建立的。 用 [scripts/sparkdiffusion/train_finetune_sparse_wan2.2_high.sh](./train_finetune_sparse_wan2.2_high.sh) 运行高噪专家： @@ -152,9 +152,9 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --train_batch_size=1 \ --gradient_accumulation_steps=1 \ --num_train_epochs=100 \ - --max_train_steps=400 \ - --checkpointing_steps=100 \ - --learning_rate=2e-05 \ + --max_train_steps=2000 \ + --checkpointing_steps=200 \ + --learning_rate=1e-4 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -191,7 +191,7 @@ Stage-1 关键参数： | `--trainable_modules` | `"attn_op"` = Stage A（只训练稀疏分支）；`"."` = Stage B（稀疏网全量微调，用 `--transformer_path` 载入 Stage A）。 | | `--max_train_steps` | 暖启预算；保持简短。 | -完成后暖启产物为 `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors`。 +完成后暖启产物为 `output_dir_sparkdiffusion_wan2.2_14b_{high,low}_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors`。 ### 3.3 Stage 2 — CrossDistill @@ -204,7 +204,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 暖启产物。 -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" NCCL_DEBUG=INFO accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ @@ -223,7 +223,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -236,6 +236,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ @@ -293,11 +294,11 @@ bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh | `--train_mode` | SparkDiffusion 仅支持 `normal`（T2V） | `normal` | | `--student_topk` | 稀疏注意力保留比例（稀疏度 = 1 − topk）；Stage 1 与 Stage 2 必须一致 | `0.05` | | `--sparse_student` | （Stage 2）把 generator + init_rollout 建成稀疏学生 | - | -| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-400/diffusion_pytorch_model.safetensors` | +| `--generator_transformer_path` | （Stage 2）载入 generator 的 Stage-1 暖启 checkpoint | `.../checkpoint-2000/diffusion_pytorch_model.safetensors` | | `--uniform_sampling` | 在训练调度上均匀采样 DMD 的噪声水平 | - | | `--flow_euler_rollout` | 在 flow/velocity 空间用确定性的 rf-native Euler 步进滚动 generator（不做 x0 再加噪） | - | | `--gen_update_interval` | 每 N 步更新一次 generator（critic 每步都更新） | `5` | -| `--joint_backward_timesteps` | 联合少步 rf 网格（降序），按 config 的 `boundary` 由 `--boundary_type` 切分成各专家的段；自动取默认值，validation 复用它 | `0.933781 0.875 0.608979` | +| `--joint_backward_timesteps` | 联合少步 rf 网格（降序），按 config 的 `boundary` 由 `--boundary_type` 切分成各专家的段；在启动脚本中显式传入（与内置默认值一致），validation 复用它 | `0.933781 0.875 0.608979` | | `--div_mode` | CrossDistill 多样性模式：`pcm_random`（随机 phase-0 偏移）或 `pcm_fixed`（仅高噪专家） | `pcm_random` | | `--loss_scale_div` | CrossDistill PCM 多样性项的权重（逐元素均值空间）；`0` 表示关闭 | `10.0`（高噪）/ `0`（低噪） | | `--dcm_total_steps` | 多样性项连续 rf 网格的等分数（`dw = 1 / N`） | `48` | @@ -307,12 +308,12 @@ bash scripts/sparkdiffusion/train_distill_wan2.2_low.sh | `--div_teacher_guidance` | 多样性 teacher rollout 的 CFG 尺度（独立于 `--real_guidance_scale`） | `5.0` | | `--real_guidance_scale` | 冻结 dense teacher 的 CFG 尺度 | `5.0` | | `--init_rollout_path` | （低噪专家）用于滚动 rollout 分界起点 latent 的、已训练好的高噪学生 | 高噪蒸馏 ckpt | -| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `4e-07` | -| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200`（蒸馏）/ `100`（暖启） | +| `--learning_rate` / `--learning_rate_critic` | 生成器 / critic 学习率 | `2e-06` / `2e-06` | +| `--checkpointing_steps` | 每 N 步保存 checkpoint | `200` | | `--max_grad_norm` | 梯度裁剪阈值 | `0.05` | | `--resume_from_checkpoint` | 续训路径或 `"latest"` | `latest` | -> 💡 各专家的少步网格是**自动推导**的：`--joint_backward_timesteps`（联合 rf 网格）按 config 的 `boundary` 由 `--boundary_type` 切分，`--init_rollout_timesteps`（低噪专家冻结的高噪段中继）由分界以上的节点填充。两者都保持默认值即可；validation 复用同一网格，故 `num_inference_steps` = `len(--joint_backward_timesteps) + 1`（此处为 `4`）。 +> 💡 各专家的少步网格是**自动推导**的：`--joint_backward_timesteps`（联合 rf 网格，在启动脚本中显式传入）按 config 的 `boundary` 由 `--boundary_type` 切分，`--init_rollout_timesteps`（低噪专家冻结的高噪段中继）由分界以上的节点填充。`--init_rollout_timesteps` 保持默认值即可；validation 复用同一网格，故 `num_inference_steps` = `len(--joint_backward_timesteps) + 1`（此处为 `4`）。 ### 3.5 FSDP Backend 与显存说明 @@ -339,7 +340,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 暖启产物。 -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -357,7 +358,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -370,6 +371,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ @@ -394,7 +396,7 @@ export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-T2V-A14B" export DATASET_NAME="datasets/X-Fun-Videos-Demo/" export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" # Stage-1 暖启产物。 -export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-400/diffusion_pytorch_model.safetensors" +export SPARSE_STUDENT_CKPT="output_dir_sparkdiffusion_wan2.2_14b_high_sparse_rola/checkpoint-2000/diffusion_pytorch_model.safetensors" # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. # export NCCL_IB_DISABLE=1 # export NCCL_P2P_DISABLE=1 @@ -412,7 +414,7 @@ accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_ --num_train_epochs=100 \ --checkpointing_steps=200 \ --learning_rate=2e-06 \ - --learning_rate_critic=4e-07 \ + --learning_rate_critic=2e-06 \ --lr_scheduler="constant_with_warmup" \ --lr_warmup_steps=100 \ --seed=42 \ @@ -425,6 +427,7 @@ accelerate launch --mixed_precision="bf16" scripts/sparkdiffusion/train_distill_ --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ diff --git a/scripts/sparkdiffusion/train_distill_wan2.2_high.sh b/scripts/sparkdiffusion/train_distill_wan2.2_high.sh index cf447ce..1a876f2 100644 --- a/scripts/sparkdiffusion/train_distill_wan2.2_high.sh +++ b/scripts/sparkdiffusion/train_distill_wan2.2_high.sh @@ -62,6 +62,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --div_mode pcm_random \ --loss_scale_div 10.0 \ --dcm_total_steps 48 \ diff --git a/scripts/sparkdiffusion/train_distill_wan2.2_low.sh b/scripts/sparkdiffusion/train_distill_wan2.2_low.sh index 38b2fc3..c2232cd 100644 --- a/scripts/sparkdiffusion/train_distill_wan2.2_low.sh +++ b/scripts/sparkdiffusion/train_distill_wan2.2_low.sh @@ -65,6 +65,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR --uniform_sampling \ --flow_euler_rollout \ --gen_update_interval 5 \ + --joint_backward_timesteps 0.933781 0.875 0.608979 \ --boundary_type="low" \ --real_guidance_scale=5.0 \ --init_rollout_path=$INIT_ROLLOUT_CKPT \