Skip to content

Commit 66abea2

Browse files
authored
Checkpoint conversion for Wan docs
Added instructions for converting HuggingFace checkpoints to Megatron format and vice versa, including necessary commands and notes on exported checkpoints.
1 parent 859c4a7 commit 66abea2

File tree

1 file changed

+32
-0
lines changed

1 file changed

+32
-0
lines changed

docs/megatron/recipes/wan/wan2.1.md

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,38 @@ uv run --group megatron-bridge python -m torch.distributed.run --nproc-per-node
147147

148148
**Note**: Current inference path is single-GPU. Parallel inference is not yet supported.
149149

150+
151+
---
152+
153+
### 🔄 Checkpoint Converting (optional)
154+
155+
If you plan to fine-tune Wan using a pre-trained model, you must first convert the HuggingFace checkpoint (e.g., `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`) into the Megatron format. The provided script supports bidirectional conversion, allowing you to move between HuggingFace and Megatron formats as needed.
156+
157+
Follow these steps to convert your checkpoints:
158+
```
159+
# Download the HF checkpoint locally
160+
huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
161+
--local-dir /root/.cache/huggingface/wan2.1 \
162+
--local-dir-use-symlinks False
163+
164+
# Import a HuggingFace model to Megatron format
165+
python examples/megatron/recipes/wan/conversion/convert_checkpoints.py import \
166+
--hf-model /root/.cache/huggingface/wan2.1 \
167+
--megatron-path /workspace/checkpoints/megatron_checkpoints/wan_1_3b
168+
169+
# Export a Megatron checkpoint to HuggingFace format
170+
python examples/megatron/recipes/wan/conversion/convert_checkpoints.py export \
171+
--hf-model /root/.cache/huggingface/wan2.1 \
172+
--megatron-path /workspace/checkpoints/megatron_checkpoints/wan_1_3b/iter_0000000 \
173+
--hf-path /workspace/checkpoints/hf_checkpoints/wan_1_3b_hf
174+
175+
```
176+
177+
**Note**: The exported checkpoint from Megatron to HuggingFace (`/workspace/checkpoints/hf_checkpoints/wan_1_3b_hf`) contains only the DiT transformer weights. To run inference, you still require the other pipeline components (VAE, text encoders, etc.).
178+
To assemble a functional inference directory:
179+
- Duplicate the original HF checkpoint directory.
180+
- Replace the `./transformer` folder in that directory with your newly exported `/transformer` folder.
181+
150182
---
151183

152184
### ⚡ Parallelism Support

0 commit comments

Comments
 (0)