Skip to content
This repository was archived by the owner on May 20, 2026. It is now read-only.

Commit f8c1c09

Browse files
parthmannanparthmannan
authored
Update perf reproduce instructions (#85)
* Update perf reproduce instructions Signed-off-by: Parth Mannan <pmannan@nvidia.com> * Docs update Signed-off-by: Parth Mannan <pmannan@nvidia.com> * Adding B200 Perf recipe and numbers for 25.11 Signed-off-by: Parth Mannan <pmannan@nvidia.com> * Update seq len formula in README Signed-off-by: Parth Mannan <pmannan@nvidia.com> * Update GB200 recipe with FSDP Signed-off-by: Parth Mannan <pmannan@nvidia.com> * Update Signed-off-by: Parth Mannan <pmannan@nvidia.com> --------- Signed-off-by: Parth Mannan <pmannan@nvidia.com>
1 parent 269973f commit f8c1c09

4 files changed

Lines changed: 74 additions & 9 deletions

File tree

‎docs/performance-summary.md‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,14 +46,20 @@ The performance data includes:
4646

4747
| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | SP | PP | CP | VP | EP | Model TFLOP / sec / GPU |
4848
|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-------------------------|
49-
|Wan 2.1 14B|32|64|1|37440|0|1|0|1|4|0|0|787.59|
49+
|Wan 2.1 14B|32|64|1|37440|1|1|0|1|2|0|0|899.62|
5050

5151

5252
#### System: DGX-GB300
5353

5454
| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | SP | PP | CP | VP | EP | Model TFLOP / sec / GPU |
5555
|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-------------------------|
56-
|Wan 2.1 14B|32|64|1|37440|0|1|0|1|2|0|0|1,022.26|
56+
|Wan 2.1 14B|32|64|1|37440|0|1|0|1|2|0|0|1,030.67|
57+
58+
#### System: DGX-B200
59+
60+
| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | SP | PP | CP | VP | EP | Model TFLOP / sec / GPU |
61+
|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-------------------------|
62+
|Wan 2.1 14B|32|64|1|37440|1|1|0|1|2|0|0|804.02|
5763

5864
#### System: DGX-H100
5965

‎examples/megatron/recipes/wan/README_perf_test.md‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ This guide provides concise steps to set up the environment and run WAN pretrain
55
## Container Launch
66

77
```bash
8-
CONT="nvcr.io/nvidia/nemo:25.09.00"
8+
CONT="nvcr.io/nvidia/nemo:25.11"
99
MOUNT="/lustre/fsw/:/lustre/fsw/"
1010

1111
srun -t 02:00:00 \
@@ -28,18 +28,18 @@ cd /opt/
2828

2929
# DFM (pinned)
3030
git clone --no-checkout https://github.com/NVIDIA-NeMo/DFM.git
31-
git -C DFM checkout 174bb7b34de002ebbbcae1ba8e2b12363c7dee01
31+
git -C DFM checkout 9eaace14995a724c982fe53726a909be2edc93cb
3232
export DFM_PATH=/opt/DFM
3333

3434
# Megatron-Bridge (pinned)
3535
rm -rf /opt/Megatron-Bridge
36-
git clone --no-checkout https://github.com/huvunvidia/Megatron-Bridge.git
37-
git -C Megatron-Bridge checkout 713ab548e4bfee307eb94a7bb3f57c17dbb31b50
36+
git clone --no-checkout https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
37+
git -C Megatron-Bridge checkout 953aabf75c0500180dc14a6a76cf9e7e7c4baec7
3838

3939
# Megatron-LM (pinned)
4040
rm -rf /opt/Megatron-LM
4141
git clone --no-checkout https://github.com/NVIDIA/Megatron-LM.git
42-
git -C Megatron-LM checkout ce8185cbbe04f38beb74360e878450f2e8525885
42+
git -C Megatron-LM checkout 2d398b42fd4237fffb553109563d73ac099751c3
4343

4444
# Python path
4545
export PYTHONPATH="${DFM_PATH}/.:/opt/Megatron-Bridge/.:/opt/Megatron-LM"
@@ -141,7 +141,13 @@ NVTE_FUSED_ATTN=1 torchrun --nproc_per_node=8 examples/megatron/recipes/wan/pret
141141
### Using mock data (optional, for debugging)
142142

143143
- Using `--mock` argument.
144-
- Adjust `video_size` (F_latents, H_latents, W_latents) and `number_packed_samples` of `WanMockDataModuleConfig` in `wan.py`. Total `seq_len = F * H * W * number_packed_samples`.
144+
- Adjust `video_size` (F_latents, H_latents, W_latents) and `number_packed_samples` of `WanMockDataModuleConfig` in `wan.py`. Total `seq_len = (F_latents // patch_temporal) * (H_latents // patch_temporal) * (W_latents // patch_temporal) * number_packed_samples`.
145+
146+
### Reproducing performance recipes
147+
148+
- Please use the appropriate system config recipe in `examples/megatron/recipes/wan/conf/<h100/gb200/gb300>_perf_pretrain_mock.yaml`
149+
- Usage example `examples/megatron/recipes/wan/pretrain_wan.py --mock --training-mode pretrain --config-file examples/megatron/recipes/wan/conf/gb300_perf_pretrain_mock.yaml`
150+
- Note that the FLOPs calculation for Wan 2.1 is not currently supported in Megatron-Bridge. Please use a manual calculator until a fix is made.
145151

146152
## Inference
147153

‎examples/megatron/recipes/wan/conf/b200_perf_pretrain_mock.yaml‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
model:
2+
tensor_model_parallel_size: 1
3+
sequence_parallel: false
4+
pipeline_model_parallel_size: 1
5+
context_parallel_size: 2
6+
crossattn_emb_size: 5120
7+
hidden_size: 5120
8+
ffn_hidden_size: 13824
9+
num_attention_heads: 40
10+
num_layers: 40
11+
qkv_format: thd
12+
seq_length: 2048 # This is not used
13+
14+
train:
15+
global_batch_size: 64
16+
micro_batch_size: 1
17+
eval_iters: 0
18+
19+
scheduler:
20+
lr_decay_style: constant
21+
lr_warmup_iters: 0
22+
23+
optimizer:
24+
lr: 5e-6
25+
min_lr: 5e-6
26+
27+
dataset:
28+
seq_length: 2048 # This is not used
29+
global_batch_size: 64
30+
micro_batch_size: 1
31+
32+
logger:
33+
log_interval: 1
34+
35+
ddp:
36+
use_megatron_fsdp: true
37+
data_parallel_sharding_strategy: "optim_grads_params"
38+
39+
dist:
40+
use_megatron_fsdp: true
41+
42+
checkpoint:
43+
ckpt_format: "fsdp_dtensor"

‎examples/megatron/recipes/wan/conf/gb200_perf_pretrain_mock.yaml‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ model:
22
tensor_model_parallel_size: 1
33
sequence_parallel: false
44
pipeline_model_parallel_size: 1
5-
context_parallel_size: 4
5+
context_parallel_size: 2
66
crossattn_emb_size: 5120
77
hidden_size: 5120
88
ffn_hidden_size: 13824
@@ -31,3 +31,13 @@ dataset:
3131

3232
logger:
3333
log_interval: 1
34+
35+
ddp:
36+
use_megatron_fsdp: true
37+
data_parallel_sharding_strategy: "optim_grads_params"
38+
39+
dist:
40+
use_megatron_fsdp: true
41+
42+
checkpoint:
43+
ckpt_format: "fsdp_dtensor"

0 commit comments

Comments
 (0)