Skip to content

Update perf reproduce instructions #85

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
parthmannan merged 7 commits into from
Mar 13, 2026
Merged

Update perf reproduce instructions #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2612caa
Update perf reproduce instructions
parthmannan Jan 14, 2026
87eb817
Docs update
parthmannan Jan 14, 2026
f716125
Merge branch 'main' of https://github.com/NVIDIA-NeMo/DFM into pmanna…
parthmannan Mar 9, 2026
7290d77
Adding B200 Perf recipe and numbers for 25.11
parthmannan Mar 9, 2026
da11297
Update seq len formula in README
parthmannan Mar 9, 2026
6375ddb
Update GB200 recipe with FSDP
parthmannan Mar 9, 2026
903eb44
Update
parthmannan Mar 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions docs/performance-summary.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,14 +46,20 @@ The performance data includes:

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


#### System: DGX-GB300

| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | SP | PP | CP | VP | EP | Model TFLOP / sec / GPU |
|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-------------------------|
|Wan 2.1 14B|32|64|1|37440|0|1|0|1|2|0|0|1,022.26|
|Wan 2.1 14B|32|64|1|37440|0|1|0|1|2|0|0|1,030.67|

#### System: DGX-B200

| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | SP | PP | CP | VP | EP | Model TFLOP / sec / GPU |
|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-------------------------|
|Wan 2.1 14B|32|64|1|37440|1|1|0|1|2|0|0|804.02|

#### System: DGX-H100

Expand Down
18 changes: 12 additions & 6 deletions examples/megatron/recipes/wan/README_perf_test.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ This guide provides concise steps to set up the environment and run WAN pretrain
## Container Launch

```bash
CONT="nvcr.io/nvidia/nemo:25.09.00"
CONT="nvcr.io/nvidia/nemo:25.11"
MOUNT="/lustre/fsw/:/lustre/fsw/"

srun -t 02:00:00 \
Expand All @@ -28,18 +28,18 @@ cd /opt/

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

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

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

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

- Using `--mock` argument.
- 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`.
- 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`.

### Reproducing performance recipes

- Please use the appropriate system config recipe in `examples/megatron/recipes/wan/conf/<h100/gb200/gb300>_perf_pretrain_mock.yaml`
- 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`
- 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.
Copy link
Contributor

@huvunvidia huvunvidia Jan 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we share some of our pre-computed FLOPs here, as a reference point for users? It's not trivial calculating the correct FLOPs number.

Copy link
Contributor

@abhinavg4 abhinavg4 Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that might help as well.


## Inference

Expand Down
43 changes: 43 additions & 0 deletions examples/megatron/recipes/wan/conf/b200_perf_pretrain_mock.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
model:
tensor_model_parallel_size: 1
sequence_parallel: false
pipeline_model_parallel_size: 1
context_parallel_size: 2
crossattn_emb_size: 5120
hidden_size: 5120
ffn_hidden_size: 13824
num_attention_heads: 40
num_layers: 40
qkv_format: thd
seq_length: 2048 # This is not used

train:
global_batch_size: 64
micro_batch_size: 1
eval_iters: 0

scheduler:
lr_decay_style: constant
lr_warmup_iters: 0

optimizer:
lr: 5e-6
min_lr: 5e-6

dataset:
seq_length: 2048 # This is not used
global_batch_size: 64
micro_batch_size: 1

logger:
log_interval: 1

ddp:
use_megatron_fsdp: true
data_parallel_sharding_strategy: "optim_grads_params"

dist:
use_megatron_fsdp: true

checkpoint:
ckpt_format: "fsdp_dtensor"
12 changes: 11 additions & 1 deletion examples/megatron/recipes/wan/conf/gb200_perf_pretrain_mock.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ model:
tensor_model_parallel_size: 1
sequence_parallel: false
pipeline_model_parallel_size: 1
context_parallel_size: 4
context_parallel_size: 2
crossattn_emb_size: 5120
hidden_size: 5120
ffn_hidden_size: 13824
Expand Down Expand Up @@ -31,3 +31,13 @@ dataset:

logger:
log_interval: 1

ddp:
use_megatron_fsdp: true
data_parallel_sharding_strategy: "optim_grads_params"

dist:
use_megatron_fsdp: true

checkpoint:
ckpt_format: "fsdp_dtensor"
Loading