Skip to content

Commit 9eaace1

Browse files
authored
Checkpoint conversion for Wan docs (#81)
Added instructions for converting HuggingFace checkpoints to Megatron format and vice versa, including necessary commands and notes on exported checkpoints.
1 parent be9f8ca commit 9eaace1

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
@@ -153,6 +153,38 @@ uv run --group megatron-bridge python -m torch.distributed.run --nproc-per-node
153153

154154
**Note**: Current inference path is single-GPU. Parallel inference is not yet supported.
155155

156+
157+
---
158+
159+
### 🔄 Checkpoint Converting (optional)
160+
161+
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.
162+
163+
Follow these steps to convert your checkpoints:
164+
```
165+
# Download the HF checkpoint locally
166+
huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
167+
--local-dir /root/.cache/huggingface/wan2.1 \
168+
--local-dir-use-symlinks False
169+
170+
# Import a HuggingFace model to Megatron format
171+
python examples/megatron/recipes/wan/conversion/convert_checkpoints.py import \
172+
--hf-model /root/.cache/huggingface/wan2.1 \
173+
--megatron-path /workspace/checkpoints/megatron_checkpoints/wan_1_3b
174+
175+
# Export a Megatron checkpoint to HuggingFace format
176+
python examples/megatron/recipes/wan/conversion/convert_checkpoints.py export \
177+
--hf-model /root/.cache/huggingface/wan2.1 \
178+
--megatron-path /workspace/checkpoints/megatron_checkpoints/wan_1_3b/iter_0000000 \
179+
--hf-path /workspace/checkpoints/hf_checkpoints/wan_1_3b_hf
180+
181+
```
182+
183+
**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.).
184+
To assemble a functional inference directory:
185+
- Duplicate the original HF checkpoint directory.
186+
- Replace the `./transformer` folder in that directory with your newly exported `/transformer` folder.
187+
156188
---
157189

158190
### ⚡ Parallelism Support

0 commit comments

Comments
 (0)