Skip to content

[Dev] Add MoE recipe performance summary #5289

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
yaox12 merged 2 commits into from
Jun 15, 2026
Merged

[Dev] Add MoE recipe performance summary #5289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1feba7f
docs: add MoE recipe performance summary
Jun 12, 2026
2d37044
Merge branch 'dev' into denliu/moe-recipes-release-performance
Jun 15, 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
68 changes: 50 additions & 18 deletions examples/moe_recipes/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
<th>TP/PP/EP/CP/ETP</th>
<th>MBS/GBS/SL</th>
<th>Features</th>
<th>Median TFLOP/s/GPU</th>
</tr>
</thead>
<tbody>
Expand All @@ -27,71 +28,97 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
<td>1/1/64/1/1</td>
<td>1/2048/4096</td>
<td>BSHD; paged stash; full CG; HybridEP</td>
<td>646.4</td>
</tr>
<tr>
<td rowspan="5">DeepSeek-V3</td>
<td><a href="deepseek_v3/h100/bf16_1024GPU_TP1PP16EP64.yaml">H100 BF16</a></td>
<td>1024</td>
<td>1/16/64/1/1</td>
<td>1/8192/4096</td>
<td>BF16 baseline</td>
<td></td>
</tr>
<tr>
<td><a href="deepseek_v3/h100/fp8_1024GPU_TP2PP8EP64.yaml">H100 FP8</a></td>
<td>1024</td>
<td>2/8/64/1/1</td>
<td>1/8192/4096</td>
<td>DeepEP; EP overlap</td>
<td></td>
</tr>
<tr>
<td><a href="deepseek_v3/b200/mxfp8_256GPU_TP1PP8EP32.yaml">B200 MXFP8</a></td>
<td>256</td>
<td>1/8/32/1/1</td>
<td>1/2048/4096</td>
<td>DeepEP; EP overlap</td>
<td>776.7</td>
</tr>
<tr>
<td><a href="deepseek_v3/gb200/mxfp8_256GPU_TP1PP4EP64.yaml">GB200 MXFP8</a></td>
<td>256</td>
<td>1/4/64/1/1</td>
<td>1/8192/4096</td>
<td>HybridEP; partial CG; EP overlap; offload</td>
<td>1127.6</td>
</tr>
<tr>
<td><a href="deepseek_v3/gb300/mxfp8_256GPU_TP1PP4EP64.yaml">GB300 MXFP8</a></td>
<td>256</td>
<td>1/4/64/1/1</td>
<td>1/8192/4096</td>
<td>HybridEP; partial CG; EP overlap</td>
<td>1357.7</td>
</tr>
<tr>
<td><a href="deepseek_v3/h100/bf16_1024GPU_TP1PP16EP64.yaml">H100 BF16</a></td>
<td>1024</td>
<td>1/16/64/1/1</td>
<td>1/8192/4096</td>
<td>BF16 baseline</td>
<td rowspan="6">Qwen3-235B-A22B</td>
<td><a href="qwen3_235b/h100/bf16_256GPU_TP2PP8EP32.yaml">H100 BF16</a></td>
<td>256</td>
<td>2/8/32/1/1</td>
<td>1/2048/4096</td>
<td>router/preprocess CG; HybridEP; EP overlap</td>
<td>214.3</td>
</tr>
<tr>
<td><a href="deepseek_v3/h100/fp8_1024GPU_TP2PP8EP64.yaml">H100 FP8</a></td>
<td>1024</td>
<td>2/8/64/1/1</td>
<td>1/8192/4096</td>
<td>DeepEP; EP overlap</td>
<td><a href="qwen3_235b/b200/mxfp8_128GPU_TP2PP4EP16.yaml">B200 MXFP8</a></td>
<td>128</td>
<td>2/4/16/1/1</td>
<td>3/3072/4096</td>
<td>DeepEP; EP overlap; VPP12</td>
<td>633.9</td>
</tr>
<tr>
<td><a href="qwen3_235b/b300/mxfp8_128GPU_TP2PP1EP16_partial_cg.yaml">B300 MXFP8 partial CG</a></td>
<td>128</td>
<td>2/1/16/1/1</td>
<td>3/9216/4096</td>
<td>partial CG; HybridEP; EP overlap</td>
<td>922.6</td>
</tr>
<tr>
<td rowspan="4">Qwen3-235B-A22B</td>
<td><a href="qwen3_235b/gb200/mxfp8_128GPU_TP1PP1EP64_paged_stash_fullcg_overlap.yaml">GB200 MXFP8 full CG</a></td>
<td>128</td>
<td>1/1/64/1/1</td>
<td>1/8192/4096</td>
<td>paged stash; full CG; HybridEP; EP overlap</td>
<td>1159.5</td>
</tr>
<tr>
<td><a href="qwen3_235b/gb200/mxfp8_128GPU_TP1PP1EP64_partial_cg_overlap.yaml">GB200 MXFP8 partial CG</a></td>
<td>128</td>
<td>1/1/64/1/1</td>
<td>1/8192/4096</td>
<td>partial CG; HybridEP; EP overlap</td>
<td>1018</td>
</tr>
<tr>
<td><a href="qwen3_235b/gb300/mxfp8_128GPU_TP1PP1EP64_paged_stash_full_cg.yaml">GB300 MXFP8 full CG</a></td>
<td>128</td>
<td>1/1/64/1/1</td>
<td>1/8192/4096</td>
<td>paged stash; full CG; HybridEP; EP overlap</td>
</tr>
<tr>
<td><a href="qwen3_235b/h100/bf16_256GPU_TP2PP8EP32.yaml">H100 BF16</a></td>
<td>256</td>
<td>2/8/32/1/1</td>
<td>1/2048/4096</td>
<td>router/preprocess CG; HybridEP; EP overlap</td>
<td>1323.1</td>
</tr>
<tr>
<td rowspan="5">Qwen3-30B-A3B</td>
Expand All @@ -100,34 +127,39 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
<td>1/1/8/1/1</td>
<td>1/256/4096</td>
<td>FP8 blockwise; router/preprocess CG</td>
<td>232.4</td>
</tr>
<tr>
<td><a href="qwen3_30b/h100/bf16_32GPU_TP1PP1EP8.yaml">H100 BF16</a></td>
<td>32</td>
<td>1/1/8/1/1</td>
<td>1/256/4096</td>
<td>BF16 baseline</td>
<td>234.1</td>
</tr>
<tr>
<td><a href="qwen3_30b/gb200/bf16_16GPU_TP1PP1EP16.yaml">GB200 BF16</a></td>
<td>16</td>
<td>1/1/16/1/1</td>
<td>4/512/4096</td>
<td>BF16 baseline</td>
<td>528</td>
</tr>
<tr>
<td><a href="qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_partial_cg.yaml">GB200 MXFP8 partial CG</a></td>
<td>16</td>
<td>1/1/16/1/1</td>
<td>4/512/4096</td>
<td>MXFP8; partial CG</td>
<td>514</td>
</tr>
<tr>
<td><a href="qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_paged_stash.yaml">GB200 MXFP8 paged stash</a></td>
<td>16</td>
<td>1/1/16/1/1</td>
<td>4/512/4096</td>
<td>MXFP8; paged stash; full CG</td>
<td>646.8</td>
</tr>
</tbody>
</table>
Expand Down
197 changes: 197 additions & 0 deletions examples/moe_recipes/qwen3_235b/b200/mxfp8_128GPU_TP2PP4EP16.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,197 @@
DEPENDENCIES:
pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
dockerfile: |
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
# IMAGE_NAME: b300-torch2603

ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:26.03-py3
FROM ${FROM_IMAGE_NAME} AS base

ENV SHELL=/bin/bash
ENV DEBIAN_FRONTEND=noninteractive

ARG YQ_VERSION=4.27.5
ARG TE_TAG=release_v2.14
ARG NVTE_CUDA_ARCHS="100a;103a"

# Install system dependencies
RUN bash -ex <<"EOF"
rm -rf /opt/megatron-lm
apt-get update
apt-get install -y --no-install-recommends git curl gettext sudo
ARCH=$(uname -m)
case "${ARCH}" in
"x86_64") YQ_ARCH=amd64 ;;
"aarch64") YQ_ARCH=arm64 ;;
*) echo "Unsupported architecture: ${ARCH}" && exit 1 ;;
esac
wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/bin/yq
chmod +x /usr/bin/yq
apt-get clean
rm -rf /var/lib/apt/lists/*
EOF

# Install Megatron-Core dependencies (mlm + dev + test groups from pyproject.toml)
RUN bash -ex <<"EOF"
unset PIP_CONSTRAINT
pip install \
sentencepiece tiktoken transformers accelerate \
wandb einops tqdm datasets nvtx \
flask-restful flask[async] fastapi hypercorn \
nltk wrapt pydantic pyyaml omegaconf \
pytest==8.3.5 pytest-mock pytest-cov pytest-random-order pytest-asyncio \
coverage tensorboard
EOF

# Install TransformerEngine
RUN unset PIP_CONSTRAINT && \
NVTE_CUDA_ARCHS="${NVTE_CUDA_ARCHS}" NVTE_FRAMEWORK=pytorch \
pip install --no-build-isolation --no-cache-dir \
"git+https://github.com/nvidia/TransformerEngine.git@${TE_TAG}"

# Install Flash Attention 4, CUTLASS DSL, cuDNN frontend, and FLA for
# Qwen3.5 Gated Delta Net support.
# Keep cudnn-frontend at 1.22.1: 1.23.0 removes the c_dtype kwarg from
# grouped_gemm_quant_wrapper_sm100, which TE 2.14 still passes.
RUN pip install --no-cache-dir flash-attn-4==4.0.0b4 "nvidia-cutlass-dsl[cu13]==4.4.2" && \
pip install --no-cache-dir --no-deps "nvidia-cudnn-frontend==1.22.1" && \
pip install --no-cache-dir fla-core==0.4.2 flash-linear-attention==0.4.2

# =========================
# Option 1: HybridEP
# =========================
FROM base AS hybridep

RUN bash -ex <<"EOF"
cd /workspace
git clone https://github.com/linux-rdma/rdma-core.git
cd rdma-core && git checkout tags/v60.0 && sh build.sh
apt-get update
apt-get install -y --no-install-recommends libnvidia-ml-dev
git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git
cd DeepEP && git checkout 7febc6e25660af0f54d95dd781ecdcd62265ecca
RDMA_CORE_HOME=/workspace/rdma-core/build HYBRID_EP_MULTINODE=1 \
TORCH_CUDA_ARCH_LIST="10.0" MAX_JOBS=8 \
pip install --no-build-isolation .
apt-get purge -y libnvidia-ml-dev
apt-get autoremove -y
rm -rf /root/.cache /tmp/* /var/lib/apt/lists/*
EOF

WORKDIR /workspace/
ENV_VARS:
NVTE_ALLOW_NONDETERMINISTIC_ALGO: '1'
PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
NCCL_NVLS_ENABLE: '0'
NVTE_FUSED_ATTN: '1'
NVTE_NORM_FWD_USE_CUDNN: '1'
NVTE_NORM_BWD_USE_CUDNN: '1'
CUDA_DEVICE_MAX_CONNECTIONS: '32'
ARGS:
tokenizer_type: HuggingFaceTokenizer
tokenizer_model: Qwen/Qwen3-235B-A22B
num_layers: 94
hidden_size: 4096
ffn_hidden_size: 12288
num_attention_heads: 64
kv_channels: 128
max_position_embeddings: 4096
group_query_attention: true
num_query_groups: 4
qk_layernorm: true
normalization: RMSNorm
norm_epsilon: 1e-6
swiglu: true
disable_bias_linear: true
untie_embeddings_and_output_weights: true
position_embedding_type: rope
rotary_percent: 1.0
rotary_base: 1000000
make_vocab_size_divisible_by: 1187
num_experts: 128
moe_ffn_hidden_size: 1536
moe_router_load_balancing_type: aux_loss
moe_router_topk: 8
moe_router_pre_softmax: false
moe_aux_loss_coeff: 1e-3
attention_dropout: 0.0
hidden_dropout: 0.0
mock_data: true
seq_length: 4096
moe_router_force_load_balancing: true
tensor_model_parallel_size: 2
pipeline_model_parallel_size: 4
expert_model_parallel_size: 16
context_parallel_size: 1
expert_tensor_parallel_size: 1
num_layers_per_virtual_pipeline_stage: 2
use_distributed_optimizer: true
sequence_parallel: true
account_for_embedding_in_pipeline_split: true
account_for_loss_in_pipeline_split: true
delay_wgrad_compute: true
overlap_moe_expert_parallel_comm: true
overlap_p2p_comm: true
moe_token_dispatcher_type: flex
moe_flex_dispatcher_backend: deepep
moe_grouped_gemm: true
moe_permute_fusion: true
moe_router_fusion: true
moe_router_dtype: fp32
moe_router_padding_for_quantization: true
moe_deepep_num_sms: 20
use_mcore_models: true
use_flash_attn: true
transformer_impl: transformer_engine
micro_batch_size: 3
global_batch_size: 3072
train_samples: 268554688
exit_duration_in_mins: 230
no_create_attention_mask_in_dataloader: true
cross_entropy_loss_fusion: true
cross_entropy_fusion_impl: te
manual_gc: true
manual_gc_interval: 5
lr: 3.9e-06
min_lr: 3.9e-07
lr_warmup_init: 3.9e-07
lr_decay_style: cosine
lr_decay_samples: 584765624
lr_warmup_samples: 1536000
weight_decay: 0.1
clip_grad: 1.0
adam_beta1: 0.9
adam_beta2: 0.95
bf16: true
fp8_format: e4m3
fp8_recipe: mxfp8
fp8_param_gather: true
reuse_grad_buf_for_mxfp8_param_ag: true
overlap_grad_reduce: true
overlap_param_gather: true
use_precision_aware_optimizer: true
main_grads_dtype: fp32
main_params_dtype: fp32
exp_avg_dtype: bf16
exp_avg_sq_dtype: bf16
init_method_std: 0.02
eval_iters: 32
eval_interval: 500
finetune: true
auto_detect_ckpt_format: true
no_load_rng: true
no_load_optim: true
load: ${LOAD_PATH}
save_interval: 100
dist_ckpt_strictness: log_all
log_throughput: true
log_interval: 1
log_timers_to_tensorboard: true
log_memory_to_tensorboard: true
log_num_zeros_in_grad: true
log_params_norm: true
log_validation_ppl_to_tensorboard: true
logging_level: 40
tensorboard_dir: ${OUTPUT_PATH}/tensorboard
wandb_exp_name: Qwen3-235B-A22B-B200-MXFP8-TP2PP4EP16-MBS3GBS3072
enable_experimental: true
Loading
Loading