Merge pull request #64 from VEGA1215/wan2.2-packing-github

[Task] Wan2.2 data packing training support

Author: nullnonenilNULL

docs/source/wan_tutorial/index.rst

@@ -0,0 +1,8 @@
+Diffusion Training Tutorial
+============================
+
+.. toctree::
+   :maxdepth: 2
+
+   quick_start_wan_training
+   wan_packing

docs/source/wan_tutorial/quick_start_wan_training.md

@@ -99,7 +99,7 @@ Each `.pth` file contains the following three keys:
 Edit `examples/wan/convert_wan2.2.sh` (section `hg2mcore`):
 - `--checkpoint_path` → source HF folder (`high_noise_model` / `low_noise_model`)
 - `--save_path` → target Megatron checkpoint folder
-- `--tp`, `--pp`, `--num_layers`, `--num_checkpoints` → match your conversion setup
+- `--num_layers`, `--num_checkpoints` → match your conversion setup
 
 Run from `examples/wan` because the script invokes conversion utilities with relative paths:
 ```bash
@@ -137,6 +137,7 @@ CP_RING_DEGREE = CP_SIZE / CP_ULYSSES_DEGREE
 - `DATASET_PATH` → output path from Section 1 (e.g. `./data/preprocessed`)
 - `--context-parallel-size 4`
 - `--context-parallel-ulysses-degree 2`
+- Optional packing: add `--packing-sft-data` to enable WAN sample packing, and tune `--packing-buffer-size` for the packing buffer size.
 
 **Step-2** Start  
 - Single-node:  

docs/source/wan_tutorial/wan_packing.md

@@ -0,0 +1,52 @@
+# Wan2.2 Packing Training
+
+Wan2.2 packing concatenates multiple variable-length video samples into one packed training sequence. It keeps per-sample attention and loss boundaries through THD `PackedSeqParams`, so packed samples do not attend to each other and padding tokens do not contribute to loss.
+
+## When to Use
+
+Enable packing when your Wan2.2 training set contains videos or text prompts with noticeably different lengths. Packing reduces wasted computation from padding and is compatible with context parallel training.
+
+Supported context-parallel modes:
+- No CP: `CP_SIZE=1`
+- Ring CP: `CP_SIZE>1`, `CP_ULYSSES_DEGREE=1`
+- Ulysses CP: `CP_SIZE=CP_ULYSSES_DEGREE`
+- Hybrid Ring + Ulysses: `CP_SIZE>CP_ULYSSES_DEGREE>1`
+
+## Data Requirements
+
+Use the same preprocessed Wan dataset format as normal training. Each sample should provide:
+- `input_latents`: video latent tensor
+- `y`: optional image-conditioning latent
+- `context`: text embedding
+- `seed`: sample seed used for deterministic noise and timestep generation
+- `grid_sizes`: optional latent patch grid; when missing, LoongForge derives it from `input_latents`
+
+Packing supports variable-length samples. For CP training, each sample in a packed bin is padded to the per-sample CP split boundary before the bin is concatenated.
+
+## How to Enable
+
+Add the packing flags to the Wan pretrain script:
+
+```bash
+--packing-sft-data
+--packing-buffer-size 512
+```
+
+Example launch:
+
+```bash
+cd examples/wan
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+CP_SIZE=4 \
+CP_ULYSSES_DEGREE=2 \
+bash pretrain_wan2.2_i2v_a14b.sh
+```
+
+`--packing-buffer-size` controls how many samples are buffered before forming packed bins. Larger buffers can improve packing density but use more host memory.
+
+## Notes and Limitations
+
+- Packing currently uses `micro_batch_size=1`; the validator enforces this when packing is enabled.
+- The packed attention path uses THD metadata and may not be bitwise identical to non-packed dense attention, but loss should stay numerically close.
+- Keep `seq_length` large enough for one packed bin. In CP mode, LoongForge aligns the effective sequence length to the required CP split boundary.
+- For accuracy checks, compare the first several training iterations against a packing-off run with the same data order and do not change `train-iters` between the two runs.

examples/wan/pretrain_wan2.2_i2v_a14b.sh

@@ -102,6 +102,8 @@ TRAINING_ARGS=(
     --no-load-rng
     --no-strict-fsdp-dtensor-load
     --finetune
+    # --packing-sft-data
+    # --packing-buffer-size ${PACKING_BUFFER_SIZE:-512}
     ${RECOMPUTE_ARGS[@]}
 )