nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update ema and update taehv

Browse Source
This commit is contained in:
bubbliiiing
2026-10-29 10:59:15 +08:00
parent e8c314fbfa
commit a15b293128
5 changed files with 90 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/taehv/README_TRAIN_SUPER.md
+34 -35
View File
@@ -186,14 +186,16 @@ LPIPS and DINOv2 weights are loaded separately by the loss implementation.
**Wan2.2 TI2V-5B / Fun-2.2VAE (taew2_2_super) Example**:
The launcher uses a 720x1280 canvas, 25 training frames, 81 validation frames,
10,000 updates, and EMA. Set the GPU/process count with `accelerate config`;
four GPUs with batch size 1 and accumulation 1 give an effective batch of 4.
The launcher uses a 720x1280 canvas, 25 training frames, and EMA, with no
`--max_train_steps`: training runs for `--num_train_epochs` (100 by default)
epochs until you kill the processes. Set the GPU/process count with
`accelerate config`; four GPUs with batch size 1 and accumulation 1 give an
effective batch of 4.
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
NCCL_DEBUG=INFO
@@ -206,18 +208,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--fix_sample_size 720 1280 \
--video_sample_strides 1 2 3 \
--video_sample_n_frames=25 \
--validation_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--max_train_steps=10000 \
--checkpointing_steps=1000 \
--validation_steps=1000 \
--validation_max_samples=0 \
--learning_rate=5e-06 \
--lr_scheduler="cosine" \
--lr_warmup_steps=200 \
--lr_min_ratio=0.2 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--l1_loss_weight=0.1 \
--lpips_loss_weight=0.1 \
--dino_loss_weight=0.02 \
@@ -227,16 +234,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--use_ema \
--ema_decay=0.999 \
--selection_lpips_ratio=1.02 \
--selection_grid_ratio=1.05 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0
--selection_grid_ratio=1.05
```
These settings are also provided in `scripts/taehv/train_taehv_super.sh`. Edit its
@@ -253,10 +251,9 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--fix_sample_size 720 1280 \
--video_sample_n_frames=25 \
--video_sample_strides 1 2 3 \
--max_train_steps=10000 \
--learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
--use_ema --ema_decay=0.999 \
--dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
--dataloader_num_workers=8 --checkpointing_steps=1000 \
--output_dir="output_dir_taehv_w2.2_super_scratch"
```
@@ -271,19 +268,20 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--train_data_dir="$DATASET_NAME" \
--train_data_meta="$DATASET_META_NAME" \
--fix_sample_size 720 1280 --video_sample_n_frames=25 \
--video_sample_strides 1 2 3 --max_train_steps=10000 \
--video_sample_strides 1 2 3 \
--learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
--use_ema --ema_decay=0.999 \
--dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
--dataloader_num_workers=8 --checkpointing_steps=1000 \
--output_dir="output_dir_taehv_w2.2_super" \
--resume_from_checkpoint=latest
```
Restore with the same settings used to create the checkpoint, including validation
metadata when supplied. `--max_train_steps` is the total update budget. Keep the
original schedule horizon. The checkpoint restores optimizer, scheduler, EMA, RNG,
step and batch cursor; changing workers or sampling settings can change the resumed
sample sequence. Use a new output directory for a new run.
Restore with the same settings used to create the checkpoint. Without
`--max_train_steps`, the schedule horizon derives from `--num_train_epochs`;
keep the original value when resuming. The checkpoint restores optimizer,
scheduler, EMA, RNG, step and batch cursor; changing workers or sampling
settings can change the resumed sample sequence. Use a new output directory
for a new run.
### 3.3 Training Parameter Reference
@@ -302,7 +300,7 @@ Values below describe the launcher configuration; they are not all Python CLI de
| `--train_batch_size` / `--gradient_accumulation_steps` | Batch per device / gradient accumulation | 1 / 1 |
| `--vae_mini_batch` | Official VAE encoding mini-batch size | 1 |
| `--dataloader_num_workers` | DataLoader workers | 8 |
| `--max_train_steps` | Total optimizer updates; overrides `--num_train_epochs` | 10000 |
| `--num_train_epochs` | Epoch budget used when `--max_train_steps` is not set | 100 |
| `--checkpointing_steps` / `--checkpoints_total_limit` | Save interval / checkpoint count limit | 1000 / `None` |
| `--learning_rate` / `--lr_scheduler` | Peak learning rate / scheduler | `5e-6` / `cosine` |
| `--lr_warmup_steps` / `--lr_min_ratio` | Warmup updates / final-to-peak ratio for cosine decay | 200 / 0.2 |
@@ -317,8 +315,8 @@ Values below describe the launcher configuration; they are not all Python CLI de
**Sample Size Configuration Guide**:
- `--fix_sample_size 720 1280` resizes to cover the canvas and center crops while preserving aspect ratio. Without it, `--video_sample_size` selects a square canvas.
- A 25-frame clip at stride 3 needs at least 73 source frames; 81-frame validation needs at least 81. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log.
- Training uses random temporal starts. Validation starts at frame 0 with stride 1. Data loading retains the original random replacement/retry behavior on decode failures.
- A 25-frame clip at stride 3 needs at least 73 source frames. Short inputs shorten the batch to the shortest 4k+1 clip length, recorded as `frames` in the training log.
- Training uses random temporal starts. Data loading retains the original random replacement/retry behavior on decode failures.
**Loss Configuration Guide**:
@@ -328,9 +326,10 @@ Values below describe the launcher configuration; they are not all Python CLI de
### 3.4 Training Validation
Provide a separate JSON array through `--validation_data_meta` in the training
command, or add that argument to the launcher. Relative video paths resolve under
`--train_data_dir`, so use absolute video paths when validation data has another root.
Validation is optional and disabled in the default launcher. To enable it, add
`--validation_data_meta` (plus the related validation arguments) to the training
command. Relative video paths resolve under `--train_data_dir`, so use absolute
video paths when validation data has another root.
| Parameter | Description | Recommended Value |
|-----------|-------------|-------------------|
@@ -359,7 +358,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
--validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
--fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
--video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \
--video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
--checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
--dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
--output_dir="output_dir_taehv_w2.2_super_smoke"
scripts/taehv/README_TRAIN_SUPER_zh-CN.md
+26 -29
View File
@@ -179,13 +179,14 @@ LPIPS 与 DINOv2 权重由损失实现单独加载。
**Wan2.2 TI2V-5B / Fun-2.2VAE（taew2_2_super）训练示例**：
启动器采用 720x1280 画布、训练 25 帧、验证 81 帧、10,000 步和 EMA。
启动器采用 720x1280 画布、训练 25 帧和 EMA，不设置 `--max_train_steps`：
训练按 `--num_train_epochs`（默认 100）个 epoch 运行，直到手动终止进程。
使用 `accelerate config` 配置 GPU/进程数；四卡、每卡 batch=1、累积=1 时有效 batch 为 4。
```bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
NCCL_DEBUG=INFO
@@ -198,18 +199,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--fix_sample_size 720 1280 \
--video_sample_strides 1 2 3 \
--video_sample_n_frames=25 \
--validation_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--max_train_steps=10000 \
--checkpointing_steps=1000 \
--validation_steps=1000 \
--validation_max_samples=0 \
--learning_rate=5e-06 \
--lr_scheduler="cosine" \
--lr_warmup_steps=200 \
--lr_min_ratio=0.2 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--l1_loss_weight=0.1 \
--lpips_loss_weight=0.1 \
--dino_loss_weight=0.02 \
@@ -219,16 +225,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--use_ema \
--ema_decay=0.999 \
--selection_lpips_ratio=1.02 \
--selection_grid_ratio=1.05 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0
--selection_grid_ratio=1.05
```
相同参数已写入 `scripts/taehv/train_taehv_super.sh`，也可修改其中的训练参数后，
@@ -245,10 +242,9 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--fix_sample_size 720 1280 \
--video_sample_n_frames=25 \
--video_sample_strides 1 2 3 \
--max_train_steps=10000 \
--learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
--use_ema --ema_decay=0.999 \
--dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
--dataloader_num_workers=8 --checkpointing_steps=1000 \
--output_dir="output_dir_taehv_w2.2_super_scratch"
```
@@ -263,16 +259,16 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--train_data_dir="$DATASET_NAME" \
--train_data_meta="$DATASET_META_NAME" \
--fix_sample_size 720 1280 --video_sample_n_frames=25 \
--video_sample_strides 1 2 3 --max_train_steps=10000 \
--video_sample_strides 1 2 3 \
--learning_rate=5e-6 --lr_warmup_steps=200 --lr_min_ratio=0.2 \
--use_ema --ema_decay=0.999 \
--dataloader_num_workers=8 --checkpointing_steps=1000 --validation_steps=1000 \
--dataloader_num_workers=8 --checkpointing_steps=1000 \
--output_dir="output_dir_taehv_w2.2_super" \
--resume_from_checkpoint=latest
```
恢复时保持创建 checkpoint 时的配置，使用过的验证索引也需继续提供。
`--max_train_steps` 表示总更新预算，应保留原 schedule 预算。Checkpoint 恢复
恢复时保持创建 checkpoint 时的配置。未设置 `--max_train_steps` 时，schedule
预算由 `--num_train_epochs` 决定，恢复时应保留原值。Checkpoint 恢复
optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或采样设置可能改变
恢复后的样本顺序。新训练应使用新的输出目录。
@@ -293,7 +289,7 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
| `--train_batch_size` / `--gradient_accumulation_steps` | 每卡 batch / 梯度累积 | 1 / 1 |
| `--vae_mini_batch` | 官方 VAE 编码 mini-batch | 1 |
| `--dataloader_num_workers` | DataLoader workers | 8 |
| `--max_train_steps` | 总优化器更新次数，覆盖 `--num_train_epochs` | 10000 |
| `--num_train_epochs` | 未设置 `--max_train_steps` 时的 epoch 预算 | 100 |
| `--checkpointing_steps` / `--checkpoints_total_limit` | 保存间隔 / checkpoint 数量上限 | 1000 / `None` |
| `--learning_rate` / `--lr_scheduler` | 峰值学习率 / 调度器 | `5e-6` / `cosine` |
| `--lr_warmup_steps` / `--lr_min_ratio` | 预热步数 / cosine 终点与峰值之比 | 200 / 0.2 |
@@ -308,8 +304,8 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
**Sample Size 配置指南**：
- `--fix_sample_size 720 1280` 等比例缩放覆盖画布后中心裁剪；未指定时，`--video_sample_size` 控制正方形画布。
- 25 帧、stride=3 至少需要 73 个源视频帧；81 帧验证至少需要 81 帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。
- 训练随机选择时间起点，验证固定从第 0 帧开始、stride=1。解码失败时沿用原有随机换视频重试逻辑。
- 25 帧、stride=3 至少需要 73 个源视频帧。短视频会使 batch 按最短视频截为 4k+1 帧，实际长度记录在训练日志的 `frames` 字段。
- 训练随机选择时间起点。解码失败时沿用原有随机换视频重试逻辑。
**损失配置指南**：
@@ -319,8 +315,9 @@ optimizer、scheduler、EMA、RNG、step 和 batch 游标；改变 workers 或
### 3.4 训练验证
在训练命令中通过 `--validation_data_meta` 提供独立 JSON 数组，或将该参数加入启动器。
相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他根目录时，可使用绝对视频路径。
验证为可选功能，默认启动器中未开启。需要时在训练命令中通过 `--validation_data_meta`
（及相关验证参数）启用。相对视频路径基于 `--train_data_dir` 解析；验证数据位于其他
根目录时，可使用绝对视频路径。
| 参数 | 说明 | 推荐值 |
|------|------|--------|
@@ -347,7 +344,7 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--train_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
--validation_data_meta="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" \
--fix_sample_size 64 96 --video_sample_n_frames=5 --validation_n_frames=5 \
--video_sample_strides 1 --max_train_steps=2 --lr_warmup_steps=0 \
--video_sample_strides 1 --num_train_epochs=1 --lr_warmup_steps=0 \
--checkpointing_steps=1 --validation_steps=1 --validation_max_samples=1 \
--dataloader_num_workers=0 --use_ema --ema_decay=0.999 \
--output_dir="output_dir_taehv_w2.2_super_smoke"
scripts/taehv/train_taehv_super.py
+2 -6
View File
@@ -614,12 +614,8 @@ def main():
if args.allow_tf32:
torch.backends.cuda.matmul.allow_tf32 = True
if accelerator.is_main_process:
os.makedirs(args.output_dir, exist_ok=True)
if not args.resume_from_checkpoint and (
os.path.exists(os.path.join(args.output_dir, "train.jsonl")) or
any(name.startswith("checkpoint-") for name in os.listdir(args.output_dir))
):
raise FileExistsError("Use a new output_dir or --resume_from_checkpoint")
if args.output_dir is not None:
os.makedirs(args.output_dir, exist_ok=True)
# For mixed precision training, keep trainable parameters and Adam moments FP32.
weight_dtype = torch.float32
scripts/taehv/train_taehv_super.sh
+12 -16
View File
@@ -1,7 +1,7 @@
#!/usr/bin/env bash
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-TI2V-5B"
export DATASET_NAME="${DATASET_NAME:-datasets/X-Fun-Videos-Demo}"
export DATASET_META_NAME="${DATASET_META_NAME:-$DATASET_NAME/metadata_add_width_height.json}"
export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
export TAE_PATH="models/Diffusion_Transformer/taew2_2_super.pth"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
@@ -17,18 +17,23 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--fix_sample_size 720 1280 \
--video_sample_strides 1 2 3 \
--video_sample_n_frames=25 \
--validation_n_frames=81 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--max_train_steps=10000 \
--checkpointing_steps=1000 \
--validation_steps=1000 \
--validation_max_samples=0 \
--learning_rate=5e-06 \
--lr_scheduler="cosine" \
--lr_warmup_steps=200 \
--lr_min_ratio=0.2 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0 \
--l1_loss_weight=0.1 \
--lpips_loss_weight=0.1 \
--dino_loss_weight=0.02 \
@@ -38,13 +43,4 @@ accelerate launch --mixed_precision="bf16" scripts/taehv/train_taehv_super.py \
--use_ema \
--ema_decay=0.999 \
--selection_lpips_ratio=1.02 \
--selection_grid_ratio=1.05 \
--seed=42 \
--output_dir="output_dir_taehv_w2.2_super" \
--mixed_precision="bf16" \
--adam_beta1=0.9 \
--adam_beta2=0.99 \
--adam_weight_decay=1e-2 \
--adam_epsilon=1e-08 \
--vae_mini_batch=1 \
--max_grad_norm=1.0
--selection_grid_ratio=1.05
videox_fun/utils/fsdp_ema.py
+16
View File
@@ -109,6 +109,22 @@ class FSDPEMA:
b_shadow.copy_(b_source.to(device=b_shadow.device, dtype=b_shadow.dtype))
self.optimization_step += 1
@torch.no_grad()
def reinit_from(self, parameters):
r"""Reset the shadow to the current live weights.
Official One-Forcing builds its EMA at ``ema_start_step`` from the live weights of
that step; this reproduces that lazy creation for a shadow instantiated at step 0,
dropping whatever warm-start mass it had accumulated before the start step.
"""
source_params, shadow_params = self._paired_params(parameters)
for p_shadow, p_source in zip(shadow_params, source_params):
if p_shadow.numel() == 0:
continue
p_shadow.copy_(p_source.detach().to(device=p_shadow.device, dtype=p_shadow.dtype))
for b_shadow, b_source in zip(self.module.buffers(), self._source_module.buffers()):
b_shadow.copy_(b_source.to(device=b_shadow.device, dtype=b_shadow.dtype))
@torch.no_grad()
def store(self, parameters):
# Snapshots the live local shards so `copy_to` can be undone with