Merge GDN conv fusion (fused pre-gated-delta-rule kernels)

docker/.ngc_version.dev

@@ -1 +1 @@
-nvcr.io/nvidia/pytorch:26.02-py3
+nvcr.io/nvidia/pytorch:26.04-py3

examples/moe_recipes/README.md

@@ -0,0 +1,127 @@
+# MoE Recipes
+
+This directory contains self-contained MoE training recipes. Each YAML file includes:
+
+- `DEPENDENCIES`: PyTorch base image and Dockerfile content used for the recipe.
+- `ENV_VARS`: Environment variables expected by the runtime.
+- `ARGS`: Megatron-LM training arguments.
+
+## Recipe Index
+
+<table>
+  <thead>
+    <tr>
+      <th>Model</th>
+      <th>Recipe</th>
+      <th>GPUs</th>
+      <th>TP/PP/EP/CP/ETP</th>
+      <th>MBS/GBS/SL</th>
+      <th>Features</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td rowspan="5">DeepSeek-V3</td>
+      <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>
+    </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>
+    </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>
+    </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>
+    </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>
+    </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>
+    </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>
+    </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>
+    </tr>
+    <tr>
+      <td rowspan="5">Qwen3-30B-A3B</td>
+      <td><a href="qwen3_30b/h100/fp8_32GPU_TP1PP1EP8.yaml">H100 FP8</a></td>
+      <td>32</td>
+      <td>1/1/8/1/1</td>
+      <td>1/256/4096</td>
+      <td>FP8 blockwise; router/preprocess CG</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>
+    </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>
+    </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>
+    </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>
+    </tr>
+  </tbody>
+</table>
+
+Legend: TP = tensor parallel, PP = pipeline parallel, EP = expert parallel, CP = context parallel, ETP = expert tensor parallel, MBS = micro-batch size, GBS = global batch size, SL = sequence length, CG = CUDA graph. The tuple orders are `TP/PP/EP/CP/ETP` and `MBS/GBS/SL`.

examples/moe_recipes/deepseek_v3/b200/mxfp8_256GPU_TP1PP8EP32.yaml

@@ -0,0 +1,216 @@
+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'
+  NVTE_FWD_LAYERNORM_SM_MARGIN: '24'
+  NVTE_BWD_LAYERNORM_SM_MARGIN: '24'
+ARGS:
+  tokenizer_type: HuggingFaceTokenizer
+  tokenizer_model: unsloth/DeepSeek-V3
+  num_layers: 61
+  hidden_size: 7168
+  ffn_hidden_size: 18432
+  num_attention_heads: 128
+  kv_channels: 128
+  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: 1536
+  kv_lora_rank: 512
+  qk_head_dim: 128
+  qk_pos_emb_head_dim: 64
+  v_head_dim: 128
+  rotary_scaling_factor: 40
+  mscale: 1.0
+  mscale_all_dim: 1.0
+  qk_layernorm: true
+  num_experts: 256
+  moe_layer_freq: ([0]*3+[1]*58)
+  moe_ffn_hidden_size: 2048
+  moe_shared_expert_intermediate_size: 2048
+  moe_router_load_balancing_type: seq_aux_loss
+  moe_router_topk: 8
+  moe_router_pre_softmax: true
+  moe_aux_loss_coeff: 1e-4
+  moe_router_group_topk: 4
+  moe_router_num_groups: 8
+  moe_router_topk_scaling_factor: 2.5
+  moe_router_score_function: sigmoid
+  moe_router_enable_expert_bias: true
+  moe_router_bias_update_rate: 1e-3
+  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: 8
+  expert_model_parallel_size: 32
+  context_parallel_size: 1
+  expert_tensor_parallel_size: 1
+  pipeline_model_parallel_layout: Et*4|(tttt|)*14tmL
+  use_distributed_optimizer: true
+  sequence_parallel: true
+  delay_wgrad_compute: true
+  overlap_moe_expert_parallel_comm: true
+  moe_token_dispatcher_type: flex
+  moe_flex_dispatcher_backend: deepep
+  moe_deepep_num_sms: 24
+  moe_grouped_gemm: true
+  moe_permute_fusion: true
+  moe_router_fusion: true
+  moe_router_dtype: fp32
+  recompute_granularity: selective
+  recompute_modules:
+  - mla_up_proj
+  - mlp
+  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: te
+  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
+  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
+  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
+  log_timers_to_tensorboard: true
+  log_memory_to_tensorboard: true
+  log_validation_ppl_to_tensorboard: true
+  logging_level: 40
+  tensorboard_dir: ${OUTPUT_PATH}/tensorboard
+  wandb_exp_name: DeepSeek-V3-B200-MXFP8-TP1PP8EP32-GBS2048
+  enable_experimental: true