+
+ | DeepSeek-V4-Flash |
+ GB200 MXFP8 |
+ 128 |
+ 1/1/64/1/1 |
+ 1/2048/4096 |
+ BSHD; paged stash; full CG; HybridEP |
+ 646.4 |
+
| DeepSeek-V3 |
+ H100 BF16 |
+ 1024 |
+ 1/16/64/1/1 |
+ 1/8192/4096 |
+ BF16 baseline |
+ |
+
+
+ | H100 FP8 |
+ 1024 |
+ 2/8/64/1/1 |
+ 1/8192/4096 |
+ DeepEP; EP overlap |
+ |
+
+
| B200 MXFP8 |
256 |
1/8/32/1/1 |
1/2048/4096 |
DeepEP; EP overlap |
+ 776.7 |
| GB200 MXFP8 |
@@ -34,6 +61,7 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
1/4/64/1/1 |
1/8192/4096 |
HybridEP; partial CG; EP overlap; offload |
+ 1127.6 |
| GB300 MXFP8 |
@@ -41,28 +69,40 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
1/4/64/1/1 |
1/8192/4096 |
HybridEP; partial CG; EP overlap |
+ 1357.7 |
- | H100 BF16 |
- 1024 |
- 1/16/64/1/1 |
- 1/8192/4096 |
- BF16 baseline |
+ Qwen3-235B-A22B |
+ H100 BF16 |
+ 256 |
+ 2/8/32/1/1 |
+ 1/2048/4096 |
+ router/preprocess CG; HybridEP; EP overlap |
+ 214.3 |
- | H100 FP8 |
- 1024 |
- 2/8/64/1/1 |
- 1/8192/4096 |
- DeepEP; EP overlap |
+ B200 MXFP8 |
+ 128 |
+ 2/4/16/1/1 |
+ 3/3072/4096 |
+ DeepEP; EP overlap; VPP12 |
+ 633.9 |
+
+
+ | B300 MXFP8 partial CG |
+ 128 |
+ 2/1/16/1/1 |
+ 3/9216/4096 |
+ partial CG; HybridEP; EP overlap |
+ 922.6 |
- | Qwen3-235B-A22B |
GB200 MXFP8 full CG |
128 |
1/1/64/1/1 |
1/8192/4096 |
paged stash; full CG; HybridEP; EP overlap |
+ 1159.5 |
| GB200 MXFP8 partial CG |
@@ -70,6 +110,7 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
1/1/64/1/1 |
1/8192/4096 |
partial CG; HybridEP; EP overlap |
+ 1018 |
| GB300 MXFP8 full CG |
@@ -77,13 +118,48 @@ This directory contains self-contained MoE training recipes. Each YAML file incl
1/1/64/1/1 |
1/8192/4096 |
paged stash; full CG; HybridEP; EP overlap |
+ 1323.1 |
- | H100 BF16 |
- 256 |
- 2/8/32/1/1 |
- 1/2048/4096 |
- router/preprocess CG; HybridEP; EP overlap |
+ Qwen3-30B-A3B |
+ H100 FP8 |
+ 32 |
+ 1/1/8/1/1 |
+ 1/256/4096 |
+ FP8 blockwise; router/preprocess CG |
+ 232.4 |
+
+
+ | H100 BF16 |
+ 32 |
+ 1/1/8/1/1 |
+ 1/256/4096 |
+ BF16 baseline |
+ 234.1 |
+
+
+ | GB200 BF16 |
+ 16 |
+ 1/1/16/1/1 |
+ 4/512/4096 |
+ BF16 baseline |
+ 528 |
+
+
+ | GB200 MXFP8 partial CG |
+ 16 |
+ 1/1/16/1/1 |
+ 4/512/4096 |
+ MXFP8; partial CG |
+ 514 |
+
+
+ | GB200 MXFP8 paged stash |
+ 16 |
+ 1/1/16/1/1 |
+ 4/512/4096 |
+ MXFP8; paged stash; full CG |
+ 646.8 |
diff --git a/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_SL4K_128GPU_TP1PP1EP64.yaml b/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_SL4K_128GPU_TP1PP1EP64.yaml
new file mode 100644
index 00000000000..9fe34ff9ca3
--- /dev/null
+++ b/examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_SL4K_128GPU_TP1PP1EP64.yaml
@@ -0,0 +1,242 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.04-py3
+ dockerfile: |
+ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: dsv4-gb200-torch2604
+ #
+ # DeepSeek-V4 training container for GB200 (arm64).
+ #
+
+ FROM nvcr.io/nvidia/pytorch:26.04-py3 AS base
+
+ ENV SHELL=/bin/bash
+
+ # System packages + yq
+ RUN bash -ex <<"EOF"
+ rm -rf /opt/megatron-lm
+ apt-get update
+ apt-get install -y --no-install-recommends \
+ sudo gdb bash-builtins git zsh autojump tmux curl gettext libfabric-dev
+ wget https://github.com/mikefarah/yq/releases/download/v4.27.5/yq_linux_arm64 -O /usr/bin/yq
+ chmod +x /usr/bin/yq
+ apt-get clean
+ rm -rf /var/lib/apt/lists/*
+ EOF
+
+ # Python deps (mcore + dev + test + one-logger + cutlass-dsl pin)
+ RUN unset PIP_CONSTRAINT && pip install --no-cache-dir \
+ debugpy dm-tree torch_tb_profiler einops wandb \
+ sentencepiece tokenizers transformers==4.57.1 torchvision ftfy modelcards datasets tqdm pydantic omegaconf \
+ nvidia-pytriton py-spy yapf darker \
+ tiktoken flask-restful \
+ nltk wrapt pytest pytest_asyncio pytest-cov pytest_mock pytest-random-order \
+ black==24.4.2 isort==5.13.2 flake8==7.1.0 pylint==3.2.6 coverage mypy \
+ one-logger --index-url https://sc-hw-artf.nvidia.com/artifactory/api/pypi/hwinf-mlwfo-pypi/simple \
+ setuptools==69.5.1 nvidia-cutlass-dsl==4.5.2
+
+ # TransformerEngine pinned to release_v2.9-based commit with CPU/quantization fixes
+ ARG TE_COMMIT="3bca93857a9103ee7869d57c464936547573860a"
+ RUN pip install --no-cache-dir flash-attn-4==4.0.0b4 nvidia-mathdx==25.1.1 && \
+ unset PIP_CONSTRAINT && \
+ NVTE_CUDA_ARCHS="100a;103a" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch \
+ pip install --no-build-isolation --no-cache-dir \
+ "git+https://github.com/NVIDIA/TransformerEngine.git@${TE_COMMIT}"
+
+ # HybridEP
+ WORKDIR /home/
+ RUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \
+ cd DeepEP && git checkout 1b8f467965bb818bf2f6511e06993f5607e1721f && \
+ TORCH_CUDA_ARCH_LIST="10.0" pip install --no-build-isolation .
+
+ # Fast Hadamard Transform (used by DSA indexer)
+ WORKDIR /home/
+ RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git && \
+ cd fast-hadamard-transform && \
+ pip install --no-build-isolation .
+
+ # Emerging-Optimizers (Muon)
+ WORKDIR /home/
+ RUN git clone https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git && \
+ cd Emerging-Optimizers && \
+ pip install --no-build-isolation .
+
+ # FlashMLA (DSA kernels)
+ WORKDIR /opt/
+ RUN git clone --branch nv_dev https://github.com/deepseek-ai/FlashMLA.git && \
+ cd FlashMLA && \
+ FLASH_MLA_DISABLE_SM90=1 \
+ NVCC_THREADS=16 \
+ CFLAGS="-I/usr/local/cuda/include/cccl" \
+ CXXFLAGS="-I/usr/local/cuda/include/cccl" \
+ pip install --no-build-isolation .
+
+ # cudnn_frontend
+ RUN pip install apache-tvm-ffi && \
+ pip install --force-reinstall --no-deps --no-build-isolation git+https://github.com/NVIDIA/cudnn-frontend.git && \
+ pip install --force-reinstall 'nvidia-cutlass-dsl[cu13]==4.5.2'
+
+ RUN unset PIP_CONSTRAINT && \
+ pip install --no-cache-dir nvidia-resiliency-ext==0.6.0
+
+ # Cleanup
+ RUN rm -rf /root/.cache /tmp/*
+ WORKDIR /home/
+
+ENV_VARS:
+ TORCH_NCCL_AVOID_RECORD_STREAMS: '0'
+ NVTE_ALLOW_NONDETERMINISTIC_ALGO: '1'
+ PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True,graph_capture_record_stream_reuse:True
+ NCCL_NVLS_ENABLE: '0'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ PYTHONWARNINGS: ignore
+ NCCL_DEBUG: VERSION
+ NCCL_GRAPH_REGISTER: '0'
+ NVTE_CUTEDSL_FUSED_GROUPED_MLP: '1'
+ NVTE_CPU_OFFLOAD_V1: '1'
+ NUM_OF_TOKENS_PER_CHUNK_COMBINE_API: '128'
+ NUM_OF_STAGES_DISPATCH_API: '10'
+ NUM_OF_IN_FLIGHT_S2G_DISPATCH_API: '8'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: unsloth/DeepSeek-V3
+ num_layers: 43
+ hidden_size: 4096
+ num_attention_heads: 64
+ kv_channels: 512
+ max_position_embeddings: 4096
+ normalization: RMSNorm
+ norm_epsilon: 1e-6
+ swiglu: true
+ disable_bias_linear: true
+ untie_embeddings_and_output_weights: true
+ position_embedding_type: rope
+ rotary_base: 10000
+ make_vocab_size_divisible_by: 3232
+ multi_latent_attention: true
+ q_lora_rank: 1024
+ qk_pos_emb_head_dim: 64
+ v_head_dim: 512
+ rotary_scaling_factor: 4
+ mscale: 1.0
+ mscale_all_dim: 1.0
+ qk_layernorm: true
+ o_groups: 8
+ o_lora_rank: 1024
+ original_max_position_embeddings: 65536
+ experimental_attention_variant: dsv4_hybrid
+ csa_window_size: 128
+ csa_compress_ratios: ([0,0,4]+[128,4]*20+[0])
+ csa_compress_rotary_base: 40000
+ dsa_indexer_n_heads: 64
+ dsa_indexer_head_dim: 128
+ dsa_indexer_topk: 512
+ dsa_indexer_loss_coeff: 1e-2
+ dsa_indexer_use_sparse_loss: true
+ num_experts: 256
+ moe_n_hash_layers: 3
+ moe_ffn_hidden_size: 2048
+ moe_shared_expert_intermediate_size: 2048
+ moe_router_load_balancing_type: seq_aux_loss
+ moe_router_topk: 6
+ moe_aux_loss_coeff: 1e-4
+ moe_router_topk_scaling_factor: 1.5
+ moe_router_score_function: sqrtsoftplus
+ moe_router_enable_expert_bias: true
+ moe_router_bias_update_rate: 1e-3
+ activation_func_clamp_value: 10.0
+ enable_hyper_connections: true
+ num_residual_streams: 4
+ mhc_sinkhorn_iterations: 20
+ use_fused_mhc: true
+ mtp_num_layers: 1
+ mtp_loss_scaling_factor: 0.1
+ attention_dropout: 0.0
+ hidden_dropout: 0.0
+ mock_data: true
+ seq_length: 4096
+ moe_router_force_load_balancing: true
+ tensor_model_parallel_size: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 64
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ overlap_grad_reduce: true
+ overlap_param_gather: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_hybridep_num_sms: 32
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ recompute_granularity: selective
+ recompute_modules:
+ - mla_up_proj
+ use_transformer_engine_op_fuser: true
+ moe_mlp_glu_interleave_size: 32
+ moe_paged_stash: true
+ moe_paged_stash_page_size: 64
+ moe_paged_stash_buffer_size_factor_cuda: 1.2
+ moe_paged_stash_buffer_size_factor_cpu: 0.0
+ moe_pad_experts_for_cuda_graph_inference: true
+ moe_expert_rank_capacity_factor: 1.5
+ cuda_graph_impl: local
+ cuda_graph_scope: full_iteration
+ use_mcore_models: true
+ use_flash_attn: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 1
+ global_batch_size: 2048
+ train_samples: 585937500
+ exit_duration_in_mins: 220
+ no_save_optim: true
+ no_check_for_nan_in_loss_and_grad: true
+ cross_entropy_loss_fusion: true
+ cross_entropy_fusion_impl: native
+ no_create_attention_mask_in_dataloader: true
+ manual_gc: true
+ manual_gc_interval: 10
+ 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_recipe: mxfp8
+ fp8_format: e4m3
+ fp8_param_gather: true
+ reuse_grad_buf_for_mxfp8_param_ag: true
+ use_precision_aware_optimizer: true
+ main_grads_dtype: fp32
+ main_params_dtype: fp32
+ exp_avg_dtype: bf16
+ exp_avg_sq_dtype: bf16
+ moe_router_padding_for_quantization: true
+ init_method_std: 0.02
+ eval_iters: 32
+ eval_interval: 200
+ finetune: false
+ no_load_optim: true
+ no_load_rng: true
+ auto_detect_ckpt_format: true
+ load: ${LOAD_PATH}
+ save_interval: 500
+ dist_ckpt_strictness: log_all
+ log_throughput: true
+ log_interval: 1
+ logging_level: 20
+ log_timers_to_tensorboard: true
+ log_memory_to_tensorboard: true
+ log_validation_ppl_to_tensorboard: true
+ tensorboard_dir: ${OUTPUT_PATH}/tensorboard
+ wandb_exp_name: DeepSeek-V4-Flash-GB200-MXFP8-TP1PP1EP64-GBS2048SEQLEN4096
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_235b/b200/mxfp8_128GPU_TP2PP4EP16.yaml b/examples/moe_recipes/qwen3_235b/b200/mxfp8_128GPU_TP2PP4EP16.yaml
new file mode 100644
index 00000000000..f3aed1dbe86
--- /dev/null
+++ b/examples/moe_recipes/qwen3_235b/b200/mxfp8_128GPU_TP2PP4EP16.yaml
@@ -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
diff --git a/examples/moe_recipes/qwen3_235b/b300/mxfp8_128GPU_TP2PP1EP16_partial_cg.yaml b/examples/moe_recipes/qwen3_235b/b300/mxfp8_128GPU_TP2PP1EP16_partial_cg.yaml
new file mode 100644
index 00000000000..f7567f21def
--- /dev/null
+++ b/examples/moe_recipes/qwen3_235b/b300/mxfp8_128GPU_TP2PP1EP16_partial_cg.yaml
@@ -0,0 +1,201 @@
+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: 1
+ expert_model_parallel_size: 16
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ 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
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ moe_router_padding_for_quantization: true
+ moe_hybridep_num_sms: 32
+ cuda_graph_impl: transformer_engine
+ cuda_graph_scope:
+ - attn
+ - moe_router
+ cuda_graph_warmup_steps: 2
+ te_rng_tracker: true
+ use_mcore_models: true
+ use_flash_attn: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 3
+ global_batch_size: 9216
+ 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-B300-MXFP8-TP2PP1EP16-MBS3GBS9216
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_30b/gb200/bf16_16GPU_TP1PP1EP16.yaml b/examples/moe_recipes/qwen3_30b/gb200/bf16_16GPU_TP1PP1EP16.yaml
new file mode 100644
index 00000000000..b4a678b9223
--- /dev/null
+++ b/examples/moe_recipes/qwen3_30b/gb200/bf16_16GPU_TP1PP1EP16.yaml
@@ -0,0 +1,181 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
+ dockerfile: |
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: gb200-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"
+
+ # 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, and cuDNN frontend with CuTe DSL 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"
+
+ # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes.
+ RUN pip install --no-cache-dir --no-deps \
+ 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'
+ HF_HUB_OFFLINE: '1'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ CUDA_DEVICE_MAX_CONNECTIONS: '1'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: Qwen/Qwen3-30B-A3B
+ num_layers: 48
+ hidden_size: 2048
+ ffn_hidden_size: 6144
+ num_attention_heads: 32
+ kv_channels: 128
+ max_position_embeddings: 40960
+ 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: 768
+ 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: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 16
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ use_mcore_models: true
+ use_flash_attn: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 4
+ global_batch_size: 512
+ 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: 0.00012
+ min_lr: 1.2e-05
+ lr_decay_style: cosine
+ lr_decay_samples: 255126953
+ lr_warmup_samples: 162761
+ weight_decay: 0.1
+ clip_grad: 1.0
+ adam_beta1: 0.9
+ adam_beta2: 0.95
+ bf16: true
+ 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: 500
+ 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-30B-GB200-BF16-TP1PP1EP16-GBS512
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_paged_stash.yaml b/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_paged_stash.yaml
new file mode 100644
index 00000000000..694373bbaad
--- /dev/null
+++ b/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_paged_stash.yaml
@@ -0,0 +1,195 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
+ dockerfile: |
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: gb200-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"
+
+ # 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, and cuDNN frontend with CuTe DSL 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"
+
+ # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes.
+ RUN pip install --no-cache-dir --no-deps \
+ 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,graph_capture_record_stream_reuse:True
+ NCCL_NVLS_ENABLE: '0'
+ NCCL_GRAPH_REGISTER: 0
+ HF_HUB_OFFLINE: '1'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ CUDA_DEVICE_MAX_CONNECTIONS: '1'
+ NVTE_CUTEDSL_FUSED_GROUPED_MLP: '1'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: Qwen/Qwen3-30B-A3B
+ num_layers: 48
+ hidden_size: 2048
+ ffn_hidden_size: 6144
+ num_attention_heads: 32
+ kv_channels: 128
+ max_position_embeddings: 40960
+ 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: 768
+ 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: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 16
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ use_transformer_engine_op_fuser: true
+ moe_paged_stash: true
+ moe_expert_rank_capacity_factor: 1.5
+ moe_paged_stash_page_size: 64
+ moe_paged_stash_buffer_size_factor_cuda: 1.1
+ cuda_graph_impl: local
+ cuda_graph_scope: full_iteration
+ moe_pad_experts_for_cuda_graph_inference: true
+ moe_mlp_glu_interleave_size: 32
+ use_mcore_models: true
+ use_flash_attn: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 4
+ global_batch_size: 512
+ train_samples: 268554688
+ exit_duration_in_mins: 230
+ no_create_attention_mask_in_dataloader: true
+ no_check_for_nan_in_loss_and_grad: true
+ cross_entropy_loss_fusion: true
+ cross_entropy_fusion_impl: te
+ manual_gc: true
+ manual_gc_interval: 5
+ lr: 0.00012
+ min_lr: 1.2e-05
+ lr_decay_style: cosine
+ lr_decay_samples: 255126953
+ lr_warmup_samples: 162761
+ weight_decay: 0.1
+ clip_grad: 1.0
+ adam_beta1: 0.9
+ adam_beta2: 0.95
+ bf16: true
+ fp8_format: e4m3
+ fp8_recipe: mxfp8
+ 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: 500
+ 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: 20
+ tensorboard_dir: ${OUTPUT_PATH}/tensorboard
+ wandb_exp_name: Qwen3-30B-GB200-MXFP8-PagedStash-TP1PP1EP16-GBS512
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_partial_cg.yaml b/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_partial_cg.yaml
new file mode 100644
index 00000000000..9a03735e782
--- /dev/null
+++ b/examples/moe_recipes/qwen3_30b/gb200/mxfp8_16GPU_TP1PP1EP16_partial_cg.yaml
@@ -0,0 +1,190 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
+ dockerfile: |
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: gb200-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"
+
+ # 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, and cuDNN frontend with CuTe DSL 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"
+
+ # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes.
+ RUN pip install --no-cache-dir --no-deps \
+ 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'
+ HF_HUB_OFFLINE: '1'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ CUDA_DEVICE_MAX_CONNECTIONS: '1'
+ NCCL_GRAPH_REGISTER: '0'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: Qwen/Qwen3-30B-A3B
+ num_layers: 48
+ hidden_size: 2048
+ ffn_hidden_size: 6144
+ num_attention_heads: 32
+ kv_channels: 128
+ max_position_embeddings: 40960
+ 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: 768
+ 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: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 16
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ external_cuda_graph: true
+ cuda_graph_scope:
+ - attn
+ - moe_router
+ - moe_preprocess
+ te_rng_tracker: true
+ use_mcore_models: true
+ use_flash_attn: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 4
+ global_batch_size: 512
+ 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: 0.00012
+ min_lr: 1.2e-05
+ lr_decay_style: cosine
+ lr_decay_samples: 255126953
+ lr_warmup_samples: 162761
+ weight_decay: 0.1
+ clip_grad: 1.0
+ adam_beta1: 0.9
+ adam_beta2: 0.95
+ bf16: true
+ fp8_format: e4m3
+ fp8_recipe: mxfp8
+ 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: 500
+ 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-30B-GB200-MXFP8-PartialCG-TP1PP1EP16-GBS512
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_30b/h100/bf16_32GPU_TP1PP1EP8.yaml b/examples/moe_recipes/qwen3_30b/h100/bf16_32GPU_TP1PP1EP8.yaml
new file mode 100644
index 00000000000..a55ac6434c5
--- /dev/null
+++ b/examples/moe_recipes/qwen3_30b/h100/bf16_32GPU_TP1PP1EP8.yaml
@@ -0,0 +1,176 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
+ dockerfile: |
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: h100-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="90a"
+
+ # 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 (standard, not FA4 which requires SM100+)
+ RUN pip install --no-cache-dir flash-attn
+
+ # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes.
+ RUN pip install --no-cache-dir --no-deps \
+ 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 cf78085241ebfdd809da8f169b41fa08e589b316
+ RDMA_CORE_HOME=/workspace/rdma-core/build HYBRID_EP_MULTINODE=1 \
+ TORCH_CUDA_ARCH_LIST="9.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'
+ HF_HUB_OFFLINE: '1'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ CUDA_DEVICE_MAX_CONNECTIONS: '1'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: Qwen/Qwen3-30B-A3B
+ num_layers: 48
+ hidden_size: 2048
+ ffn_hidden_size: 6144
+ num_attention_heads: 32
+ kv_channels: 128
+ max_position_embeddings: 40960
+ 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: 768
+ 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: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 8
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ overlap_grad_reduce: true
+ overlap_param_gather: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ use_mcore_models: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 1
+ global_batch_size: 256
+ 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: 0.00012
+ min_lr: 1.2e-05
+ lr_decay_style: cosine
+ lr_decay_samples: 255126953
+ lr_warmup_samples: 162761
+ weight_decay: 0.1
+ clip_grad: 1.0
+ adam_beta1: 0.9
+ adam_beta2: 0.95
+ bf16: true
+ init_method_std: 0.02
+ eval_iters: 32
+ eval_interval: 500
+ finetune: true
+ auto_detect_ckpt_format: true
+ load: ${LOAD_PATH}
+ save_interval: 500
+ 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
+ tensorboard_dir: ${OUTPUT_PATH}/tensorboard
+ wandb_exp_name: Qwen3-30B-TP1PP1EP8-GBS256
+ enable_experimental: true
diff --git a/examples/moe_recipes/qwen3_30b/h100/fp8_32GPU_TP1PP1EP8.yaml b/examples/moe_recipes/qwen3_30b/h100/fp8_32GPU_TP1PP1EP8.yaml
new file mode 100644
index 00000000000..2954da5387f
--- /dev/null
+++ b/examples/moe_recipes/qwen3_30b/h100/fp8_32GPU_TP1PP1EP8.yaml
@@ -0,0 +1,190 @@
+DEPENDENCIES:
+ pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3
+ dockerfile: |
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ # IMAGE_NAME: h100-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="90a"
+
+ # 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 (standard, not FA4 which requires SM100+)
+ RUN pip install --no-cache-dir flash-attn
+
+ # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes.
+ RUN pip install --no-cache-dir --no-deps \
+ 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 cf78085241ebfdd809da8f169b41fa08e589b316
+ RDMA_CORE_HOME=/workspace/rdma-core/build HYBRID_EP_MULTINODE=1 \
+ TORCH_CUDA_ARCH_LIST="9.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'
+ HF_HUB_OFFLINE: '1'
+ NVTE_FUSED_ATTN: '1'
+ NVTE_NORM_FWD_USE_CUDNN: '1'
+ NVTE_NORM_BWD_USE_CUDNN: '1'
+ CUDA_DEVICE_MAX_CONNECTIONS: '1'
+ NCCL_GRAPH_REGISTER: '0'
+ARGS:
+ tokenizer_type: HuggingFaceTokenizer
+ tokenizer_model: Qwen/Qwen3-30B-A3B
+ num_layers: 48
+ hidden_size: 2048
+ ffn_hidden_size: 6144
+ num_attention_heads: 32
+ kv_channels: 128
+ max_position_embeddings: 40960
+ 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: 768
+ 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: 1
+ pipeline_model_parallel_size: 1
+ expert_model_parallel_size: 8
+ context_parallel_size: 1
+ expert_tensor_parallel_size: 1
+ use_distributed_optimizer: true
+ sequence_parallel: true
+ overlap_grad_reduce: true
+ overlap_param_gather: true
+ moe_token_dispatcher_type: flex
+ moe_flex_dispatcher_backend: hybridep
+ moe_grouped_gemm: true
+ moe_permute_fusion: true
+ moe_router_fusion: true
+ moe_router_dtype: fp32
+ cuda_graph_impl: transformer_engine
+ cuda_graph_scope:
+ - attn
+ - moe_router
+ - moe_preprocess
+ use_mcore_models: true
+ transformer_impl: transformer_engine
+ micro_batch_size: 1
+ global_batch_size: 256
+ 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: 0.00012
+ min_lr: 1.2e-05
+ lr_decay_style: cosine
+ lr_decay_samples: 255126953
+ lr_warmup_samples: 162761
+ weight_decay: 0.1
+ clip_grad: 1.0
+ adam_beta1: 0.9
+ adam_beta2: 0.95
+ bf16: true
+ fp8_recipe: blockwise
+ fp8_format: e4m3
+ fp8_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
+ load: ${LOAD_PATH}
+ save_interval: 500
+ 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
+ tensorboard_dir: ${OUTPUT_PATH}/tensorboard
+ wandb_exp_name: Qwen3-30B-FP8-PartialCG-TP1PP1EP8-GBS256
+ enable_experimental: true
diff --git a/examples/multimodal/Dockerfile b/examples/multimodal/Dockerfile
index d7c4fd41af5..46db1319566 100644
--- a/examples/multimodal/Dockerfile
+++ b/examples/multimodal/Dockerfile
@@ -37,7 +37,7 @@ RUN uv pip install --system --no-cache --break-system-packages \
flask-restful \
wandb \
bitstring \
- filetype
+ filetype \
setuptools
# Install CLIP from GitHub
diff --git a/examples/multimodal/layer_specs.py b/examples/multimodal/layer_specs.py
index c51fb69f496..caff5ac7e0b 100644
--- a/examples/multimodal/layer_specs.py
+++ b/examples/multimodal/layer_specs.py
@@ -1,4 +1,6 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
+from functools import partial
+
import torch
from megatron.core.extensions.transformer_engine import HAVE_TE
@@ -112,7 +114,7 @@ def get_layer_spec_te(is_vit=False, padding=False) -> ModuleSpec:
submodules=SelfAttentionSubmodules(
linear_qkv=not_none(TELayerNormColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
q_layernorm=IdentityOp,
k_layernorm=IdentityOp,
),
@@ -158,7 +160,7 @@ def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec:
submodules=SelfAttentionSubmodules(
linear_qkv=not_none(TELayerNormColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
),
),
self_attn_bda=get_bias_dropout_add,
@@ -170,8 +172,8 @@ def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec:
mlp_layer=ModuleSpec(
module=MLPLayer,
submodules=TransformerLayerSubmodules(
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
@@ -184,10 +186,10 @@ def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec:
)
-def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
+def get_mlp_module_spec(use_te: bool = True):
# Dense MLP w/ or w/o TE modules.
- return ModuleSpec(
- module=MLP,
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TEColumnParallelLinear) if use_te else ColumnParallelLinear,
linear_fc2=not_none(TERowParallelLinear) if use_te else RowParallelLinear,
@@ -195,9 +197,9 @@ def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
)
-def get_norm_mlp_module_spec_te() -> ModuleSpec:
- return ModuleSpec(
- module=MLP,
+def get_norm_mlp_module_spec_te():
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
diff --git a/examples/multimodal/nvlm/internvit.py b/examples/multimodal/nvlm/internvit.py
index 0018bb5ccb9..d38ac64c16b 100644
--- a/examples/multimodal/nvlm/internvit.py
+++ b/examples/multimodal/nvlm/internvit.py
@@ -160,10 +160,10 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata={}):
return super().sharded_state_dict(prefix, sharded_offsets, metadata)
-def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
+def get_mlp_module_spec(use_te: bool = True):
# Dense MLP w/ or w/o TE modules.
- return ModuleSpec(
- module=MLP,
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=TEColumnParallelLinear if use_te else ColumnParallelLinear,
linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
diff --git a/examples/multimodal/radio/radio_g.py b/examples/multimodal/radio/radio_g.py
index 9883d58db61..a3d0317b03b 100644
--- a/examples/multimodal/radio/radio_g.py
+++ b/examples/multimodal/radio/radio_g.py
@@ -1,12 +1,11 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
from functools import partial
-import torch
-
from examples.multimodal.layer_scaling import (
LayerScalingTransformerLayer,
get_bias_dropout_add_layer_scaling,
)
+from megatron.core.extensions.transformer_engine import HAVE_TE
from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
from megatron.core.transformer.dot_product_attention import DotProductAttention
@@ -14,9 +13,8 @@
from megatron.core.transformer.identity_op import IdentityOp
from megatron.core.transformer.mlp import MLP, MLPSubmodules
from megatron.core.transformer.spec_utils import ModuleSpec
-from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules
from megatron.core.typed_torch import not_none
-from megatron.core.extensions.transformer_engine import HAVE_TE
if HAVE_TE:
from megatron.core.extensions.transformer_engine import (
@@ -51,10 +49,10 @@
LNImpl = WrappedTorchNorm
-def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
+def get_mlp_module_spec(use_te: bool = True):
# Dense MLP w/ or w/o TE modules.
- return ModuleSpec(
- module=MLP,
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TEColumnParallelLinear) if use_te else ColumnParallelLinear,
linear_fc2=not_none(TERowParallelLinear) if use_te else RowParallelLinear,
@@ -62,9 +60,9 @@ def get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
)
-def get_norm_mlp_module_spec_te() -> ModuleSpec:
- return ModuleSpec(
- module=MLP,
+def get_norm_mlp_module_spec_te():
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
@@ -125,7 +123,7 @@ def get_radio_g_layer_spec_te() -> ModuleSpec:
submodules=SelfAttentionSubmodules(
linear_qkv=not_none(TELayerNormColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
q_layernorm=IdentityOp,
k_layernorm=IdentityOp,
),
diff --git a/examples/multimodal/train.py b/examples/multimodal/train.py
index 98536f72d1e..82927c61793 100644
--- a/examples/multimodal/train.py
+++ b/examples/multimodal/train.py
@@ -26,9 +26,9 @@
get_tensor_model_parallel_rank,
is_pipeline_last_stage,
)
-from megatron.core.utils import nvtx_range_pop, nvtx_range_push
+from megatron.core.utils import get_batch_on_this_cp_rank, nvtx_range_pop, nvtx_range_push
from megatron.training import get_args, get_timers, get_tokenizer, pretrain
-from megatron.training.utils import get_batch_on_this_cp_rank, is_last_rank
+from megatron.training.utils import is_last_rank
def get_batch(data_iterator, image_token_index, img_seq_len):
diff --git a/examples/multimodal_dev/README.md b/examples/multimodal_dev/README.md
new file mode 100644
index 00000000000..e4e6c53ceb4
--- /dev/null
+++ b/examples/multimodal_dev/README.md
@@ -0,0 +1,222 @@
+# multimodal_dev — Standalone Multimodal Training
+
+Standalone, model-agnostic training entry point for multimodal
+vision-language models built on Megatron-Core (FSDP + EP).
+
+## Directory Structure
+
+```
+multimodal_dev/
+├── pretrain_multimodal.py # Training entry point (model-agnostic)
+├── forward_step.py # Forward step, TP broadcast, loss computation
+├── arguments.py # Multimodal CLI arguments
+├── data/
+│ └── mock.py # Mock dataset for end-to-end testing
+├── models/
+│ ├── __init__.py # MODEL_REGISTRY — central model registry
+│ ├── base.py # MultimodalModel base class (vision encoder + GPTModel)
+│ └── qwen35_vl/ # Qwen3.5-VL architecture
+│ ├── factory.py # Factory functions for pretrain entry point
+│ ├── model.py # Qwen35VLModel (MRoPE, vision encoder wiring)
+│ ├── configuration.py # TransformerConfig builders and constants
+│ ├── specs.py # Layer spec builders (hybrid attention, ViT)
+│ ├── mrope.py # 3D MRoPE position ID computation
+│ └── vision_encoder.py# ViT encoder (patch embed, merger, RoPE)
+└── scripts/ # Launch scripts (torchrun, Slurm)
+```
+
+## Quick Start
+
+```bash
+torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \
+ --model-arch qwen35_vl \
+ --dataset-provider mock \
+ ... # other Megatron args (--num-layers, --hidden-size, etc.)
+```
+
+## Checkpoint Conversion (HF → Megatron-FSDP DTensor)
+
+Convert a HuggingFace release to a Megatron-FSDP DTensor checkpoint via
+[Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) before
+pretraining from pretrained weights.
+
+### Setup
+
+Clone Bridge and pin its `3rdparty/Megatron-LM` submodule to this branch:
+
+```bash
+git clone --recurse-submodules https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
+cd Megatron-Bridge/3rdparty/Megatron-LM
+git remote add wplf https://github.com/wplf/Megatron-LM.git
+git fetch wplf feat/qwen35-vl-example
+git checkout feat/qwen35-vl-example
+cd ../..
+```
+
+### Convert
+
+Single 8×GPU node, EP=8 / TP=CP=1; substitute any Qwen3.5 variant for
+`--hf-model`:
+
+```bash
+PYTHONPATH=./src:./3rdparty/Megatron-LM/ \
+ torchrun --nproc_per_node=8 \
+ examples/conversion/mfsdp/convert_checkpoints_fsdp.py import \
+ --hf-model Qwen/Qwen3.5-35B-A3B \
+ --megatron-path ${WORKSPACE}/models/Qwen/Qwen3.5-35B-A3B-fsdp \
+ --ckpt-format fsdp_dtensor \
+ --ep 8
+```
+
+HF weights are auto-fetched on first run via `huggingface_hub`. Adjust
+`--tp` / `--cp` / `--ep` to match the training topology (must satisfy
+`WORLD_SIZE % (TP*CP*EP) == 0`).
+
+### Output
+
+```
+${WORKSPACE}/models/Qwen/Qwen3.5-35B-A3B-fsdp/
+├── iter_0000000/
+│ ├── __0_0.distcp .. __7_0.distcp # FSDP DTensor shards, one per rank (~18 GB each for 35B-A3B)
+│ ├── .metadata
+│ ├── run_config.yaml
+│ └── train_state.pt
+├── latest_checkpointed_iteration.txt
+└── latest_train_state.pt
+```
+
+### Bridge dependency
+
+Requires
+[NVIDIA-NeMo/Megatron-Bridge#3987](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/3987)
+(skip tokenizer save). Without that fix the checkpoint is still written
+correctly but the script exits non-zero after save with
+`AttributeError: 'TokenizerConfig' object has no attribute 'make_vocab_size_divisible_by'`
+against this branch's `megatron.core.tokenizers.utils.build_tokenizer`.
+
+## Architecture
+
+`pretrain_multimodal.py` is **model-agnostic**. All model-specific logic
+is delegated to factory functions registered in `MODEL_REGISTRY`
+(`models/__init__.py`). The entry point handles only generic concerns:
+
+- Building `language_config` from Megatron CLI args
+- Constructing `vision_config` via the registry
+- Applying vision recompute and dtype propagation
+- Routing to model and dataset factories
+
+The `forward_step` is also model-agnostic — it uses the model's
+`compute_position_ids()` method polymorphically and passes a standard
+batch dict.
+
+## Adding a New Model Architecture
+
+Adding a new model (e.g. `llava_next`) requires **no changes** to
+`pretrain_multimodal.py` or `forward_step.py`. Follow these steps:
+
+### Step 1 — Create the model package
+
+```
+multimodal_dev/models/llava_next/
+├── __init__.py
+├── factory.py # Required: factory functions
+├── configuration.py # Vision/language TransformerConfig builders
+├── model.py # Model class (subclass MultimodalModel)
+├── specs.py # Layer spec builders
+└── vision_encoder.py # Vision encoder (if custom)
+```
+
+### Step 2 — Implement factory functions
+
+Create `factory.py` with up to three functions:
+
+```python
+# models/llava_next/factory.py
+
+def post_language_config(language_config, args):
+ """(Optional) Mutate language_config with model-specific fields."""
+ # e.g. language_config.some_field = value
+ pass
+
+def set_vision_flops_metadata(args, language_config, vision_config):
+ """(Optional) Set vision FLOPs metadata on args."""
+ args.count_vision_model_flops = True
+ args.vision_flops_variant = "llava_next"
+ # ... set dimension fields for FLOPs calculation
+
+def build_model(args, language_config, vision_config, **kwargs):
+ """(Required) Build and return the complete model instance."""
+ from .model import LlavaNextModel
+ from .specs import get_llava_next_language_spec
+
+ language_spec = get_llava_next_language_spec(
+ config=language_config,
+ vp_stage=kwargs.get("vp_stage", None),
+ pp_rank=None,
+ )
+ return LlavaNextModel(
+ language_config=language_config,
+ language_spec=language_spec,
+ vision_config=vision_config,
+ # ... model-specific args
+ )
+```
+
+### Step 3 — Register in `MODEL_REGISTRY`
+
+Add an entry in `models/__init__.py`:
+
+```python
+from multimodal_dev.models.llava_next.configuration import (
+ get_llava_next_vision_config,
+)
+from multimodal_dev.models.llava_next.factory import (
+ build_model as _build_llava_next_model,
+ post_language_config as _llava_next_post_language_config,
+ set_vision_flops_metadata as _llava_next_vision_flops,
+)
+
+MODEL_REGISTRY["llava_next"] = {
+ "model_factory_fn": _build_llava_next_model, # required
+ "vision_config_fn": get_llava_next_vision_config, # required
+ "post_language_config_fn": _llava_next_post_language_config, # optional
+ "vision_flops_fn": _llava_next_vision_flops, # optional
+ "dataset_providers": { # optional
+ "mock": "multimodal_dev.data.llava_mock.train_valid_test_datasets_provider",
+ },
+}
+```
+
+### Step 4 — (Optional) Add a dataset provider
+
+Create a dataset module under `data/` if the model needs custom data
+preprocessing. The provider function signature is:
+
+```python
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+ """Return (train_dataset, val_dataset, test_dataset)."""
+ ...
+```
+
+Register it in the `dataset_providers` dict of the registry entry.
+Providers can be either direct callables or dotted import path strings
+(resolved lazily at runtime).
+
+### Step 5 — Launch
+
+```bash
+torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \
+ --model-arch llava_next \
+ --dataset-provider mock \
+ ...
+```
+
+## Registry Entry Reference
+
+| Field | Required | Signature |
+|-------|----------|-----------|
+| `model_factory_fn` | Yes | `(args, language_config, vision_config, **kwargs) -> MegatronModule` |
+| `vision_config_fn` | Yes | `(num_layers_override=None) -> TransformerConfig` |
+| `post_language_config_fn` | No | `(language_config, args) -> None` |
+| `vision_flops_fn` | No | `(args, language_config, vision_config) -> None` |
+| `dataset_providers` | No | `Dict[str, str \| callable]` |
diff --git a/examples/multimodal_dev/__init__.py b/examples/multimodal_dev/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/examples/multimodal_dev/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/examples/multimodal_dev/arguments.py b/examples/multimodal_dev/arguments.py
new file mode 100644
index 00000000000..35655831bfb
--- /dev/null
+++ b/examples/multimodal_dev/arguments.py
@@ -0,0 +1,100 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Extra CLI arguments for multimodal_dev standalone training."""
+
+
+def add_multimodal_args(parser):
+ """Add multimodal-specific arguments to the Megatron argument parser."""
+ group = parser.add_argument_group(
+ "Multimodal", "Multimodal model arguments",
+ )
+
+ group.add_argument(
+ "--model-arch",
+ type=str,
+ default="qwen35_vl",
+ help="Model architecture. Available: qwen35_vl",
+ )
+ group.add_argument(
+ "--model-variant",
+ type=str,
+ default="proxy",
+ help="Model variant (size). E.g. proxy, 9b, 397b_a17b",
+ )
+ group.add_argument(
+ "--dataset-provider",
+ type=str,
+ default="mock",
+ help="Dataset provider: mock",
+ )
+ group.add_argument(
+ "--image-token-id",
+ type=int,
+ default=248056,
+ help="Token ID for image placeholder tokens",
+ )
+ group.add_argument(
+ "--image-size",
+ type=int,
+ default=224,
+ help="Image size (height and width) for mock data",
+ )
+ group.add_argument(
+ "--total-seq-length",
+ type=int,
+ default=1024,
+ help="Total sequence length for mock data",
+ )
+ group.add_argument(
+ "--image-seq-length",
+ type=int,
+ default=256,
+ help="Number of image tokens in mock data",
+ )
+ group.add_argument(
+ "--vision-num-layers",
+ type=int,
+ default=None,
+ help=(
+ "Override for vision backbone depth. "
+ "Useful for proxy perf runs."
+ ),
+ )
+ group.add_argument(
+ "--hf-processor-path",
+ type=str,
+ default=None,
+ help=(
+ "HuggingFace processor path for real VLM datasets "
+ "(e.g. Qwen/Qwen2.5-VL-7B-Instruct)"
+ ),
+ )
+ group.add_argument(
+ "--recompute-vision",
+ action="store_true",
+ default=False,
+ help=(
+ "Enable full activation recomputation for vision encoder layers. "
+ "Uses uniform method and recomputes every layer. "
+ "Independent of the decoder --recompute-* flags."
+ ),
+ )
+ group.add_argument(
+ "--use-packed-sequence",
+ action="store_true",
+ default=False,
+ help=(
+ "Pack variable-length sequences into THD format to eliminate "
+ "padding waste."
+ ),
+ )
+ group.add_argument(
+ "--use-vanilla-collate-fn",
+ action="store_true",
+ default=False,
+ help=(
+ "Use vanilla collate function to collate the data."
+ ),
+ )
+
+ return parser
diff --git a/examples/multimodal_dev/data/__init__.py b/examples/multimodal_dev/data/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/examples/multimodal_dev/data/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/examples/multimodal_dev/data/cord_v2.py b/examples/multimodal_dev/data/cord_v2.py
new file mode 100644
index 00000000000..69fd4c13ec4
--- /dev/null
+++ b/examples/multimodal_dev/data/cord_v2.py
@@ -0,0 +1,397 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""CORD-V2 VLM dataset for multimodal_dev training.
+
+Single-turn image-text dataset using a HuggingFace ``AutoProcessor`` for
+tokenization and image preprocessing. This module is the reference
+implementation for the CORD-V2 receipt-OCR dataset. No multi-turn support —
+each sample is one image + question → answer pair.
+
+Each image is preprocessed via ``qwen_vl_utils.process_vision_info`` and
+fed to the processor with Qwen-VL's recommended ``min_pixels`` /
+``max_pixels`` budget, so the per-sample patch grid varies with aspect
+ratio. The run script must therefore pass ``--use-vanilla-collate-fn``
+(which the example launcher does) so the dataloader does not try to stack
+variable-shape tensors.
+
+Usage::
+
+ torchrun ... pretrain_multimodal.py \\
+ --model-arch qwen35_vl --dataset-provider cord_v2 \\
+ --hf-processor-path Qwen/Qwen3.5-397B-A17B \\
+ --total-seq-length 4096 --use-vanilla-collate-fn
+
+Adding another VLM dataset
+--------------------------
+
+The dataset layer mirrors the model layer's registry pattern: each dataset
+ships its own module and a ``train_valid_test_datasets_provider`` factory,
+and the model's registry entry maps a ``--dataset-provider`` name to that
+factory's dotted path. To add a new dataset (e.g. NLVR2):
+
+1. Create ``examples/multimodal_dev/data/