|
| 1 | +DEPENDENCIES: |
| 2 | + pytorch_base_image: nvcr.io/nvidia/pytorch:26.03-py3 |
| 3 | + dockerfile: | |
| 4 | + # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 5 | + # IMAGE_NAME: gb200-torch2603 |
| 6 | +
|
| 7 | + ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:26.03-py3 |
| 8 | + FROM ${FROM_IMAGE_NAME} AS base |
| 9 | +
|
| 10 | + ENV SHELL=/bin/bash |
| 11 | + ENV DEBIAN_FRONTEND=noninteractive |
| 12 | +
|
| 13 | + ARG YQ_VERSION=4.27.5 |
| 14 | + ARG TE_TAG=release_v2.14 |
| 15 | + ARG NVTE_CUDA_ARCHS="100a" |
| 16 | +
|
| 17 | + # Install system dependencies |
| 18 | + RUN bash -ex <<"EOF" |
| 19 | + rm -rf /opt/megatron-lm |
| 20 | + apt-get update |
| 21 | + apt-get install -y --no-install-recommends git curl gettext sudo |
| 22 | + ARCH=$(uname -m) |
| 23 | + case "${ARCH}" in |
| 24 | + "x86_64") YQ_ARCH=amd64 ;; |
| 25 | + "aarch64") YQ_ARCH=arm64 ;; |
| 26 | + *) echo "Unsupported architecture: ${ARCH}" && exit 1 ;; |
| 27 | + esac |
| 28 | + wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/bin/yq |
| 29 | + chmod +x /usr/bin/yq |
| 30 | + apt-get clean |
| 31 | + rm -rf /var/lib/apt/lists/* |
| 32 | + EOF |
| 33 | +
|
| 34 | + # Install Megatron-Core dependencies (mlm + dev + test groups from pyproject.toml) |
| 35 | + RUN bash -ex <<"EOF" |
| 36 | + unset PIP_CONSTRAINT |
| 37 | + pip install \ |
| 38 | + sentencepiece tiktoken transformers accelerate \ |
| 39 | + wandb einops tqdm datasets nvtx \ |
| 40 | + flask-restful flask[async] fastapi hypercorn \ |
| 41 | + nltk wrapt pydantic pyyaml omegaconf \ |
| 42 | + pytest==8.3.5 pytest-mock pytest-cov pytest-random-order pytest-asyncio \ |
| 43 | + coverage tensorboard |
| 44 | + EOF |
| 45 | +
|
| 46 | + # Install TransformerEngine |
| 47 | + RUN unset PIP_CONSTRAINT && \ |
| 48 | + NVTE_CUDA_ARCHS="${NVTE_CUDA_ARCHS}" NVTE_FRAMEWORK=pytorch \ |
| 49 | + pip install --no-build-isolation --no-cache-dir \ |
| 50 | + "git+https://github.com/nvidia/TransformerEngine.git@${TE_TAG}" |
| 51 | +
|
| 52 | + # Install Flash Attention 4, CUTLASS DSL, and cuDNN frontend with CuTe DSL support. |
| 53 | + # Keep cudnn-frontend at 1.22.1: 1.23.0 removes the c_dtype kwarg from |
| 54 | + # grouped_gemm_quant_wrapper_sm100, which TE 2.14 still passes. |
| 55 | + RUN pip install --no-cache-dir flash-attn-4==4.0.0b4 "nvidia-cutlass-dsl[cu13]==4.4.2" && \ |
| 56 | + pip install --no-cache-dir --no-deps "nvidia-cudnn-frontend==1.22.1" |
| 57 | +
|
| 58 | + # Install Flash Linear Attention for gated-delta-net Qwen3.5-VL recipes. |
| 59 | + RUN pip install --no-cache-dir --no-deps \ |
| 60 | + fla-core==0.4.2 \ |
| 61 | + flash-linear-attention==0.4.2 |
| 62 | +
|
| 63 | + # ========================= |
| 64 | + # Option 1: HybridEP |
| 65 | + # ========================= |
| 66 | + FROM base AS hybridep |
| 67 | +
|
| 68 | + RUN bash -ex <<"EOF" |
| 69 | + cd /workspace |
| 70 | + git clone https://github.com/linux-rdma/rdma-core.git |
| 71 | + cd rdma-core && git checkout tags/v60.0 && sh build.sh |
| 72 | + apt-get update |
| 73 | + apt-get install -y --no-install-recommends libnvidia-ml-dev |
| 74 | + git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git |
| 75 | + cd DeepEP && git checkout 7febc6e25660af0f54d95dd781ecdcd62265ecca |
| 76 | + RDMA_CORE_HOME=/workspace/rdma-core/build HYBRID_EP_MULTINODE=1 \ |
| 77 | + TORCH_CUDA_ARCH_LIST="10.0" MAX_JOBS=8 \ |
| 78 | + pip install --no-build-isolation . |
| 79 | + apt-get purge -y libnvidia-ml-dev |
| 80 | + apt-get autoremove -y |
| 81 | + rm -rf /root/.cache /tmp/* /var/lib/apt/lists/* |
| 82 | + EOF |
| 83 | +
|
| 84 | + WORKDIR /workspace/ |
| 85 | +ENV_VARS: |
| 86 | + NVTE_ALLOW_NONDETERMINISTIC_ALGO: '1' |
| 87 | + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True |
| 88 | + NCCL_NVLS_ENABLE: '0' |
| 89 | + HF_HUB_OFFLINE: '1' |
| 90 | + NVTE_FUSED_ATTN: '1' |
| 91 | + NVTE_NORM_FWD_USE_CUDNN: '1' |
| 92 | + NVTE_NORM_BWD_USE_CUDNN: '1' |
| 93 | + CUDA_DEVICE_MAX_CONNECTIONS: '1' |
| 94 | +ARGS: |
| 95 | + tokenizer_type: HuggingFaceTokenizer |
| 96 | + tokenizer_model: Qwen/Qwen3-30B-A3B |
| 97 | + num_layers: 48 |
| 98 | + hidden_size: 2048 |
| 99 | + ffn_hidden_size: 6144 |
| 100 | + num_attention_heads: 32 |
| 101 | + kv_channels: 128 |
| 102 | + max_position_embeddings: 40960 |
| 103 | + group_query_attention: true |
| 104 | + num_query_groups: 4 |
| 105 | + qk_layernorm: true |
| 106 | + normalization: RMSNorm |
| 107 | + norm_epsilon: 1e-6 |
| 108 | + swiglu: true |
| 109 | + disable_bias_linear: true |
| 110 | + untie_embeddings_and_output_weights: true |
| 111 | + position_embedding_type: rope |
| 112 | + rotary_percent: 1.0 |
| 113 | + rotary_base: 1000000 |
| 114 | + make_vocab_size_divisible_by: 1187 |
| 115 | + num_experts: 128 |
| 116 | + moe_ffn_hidden_size: 768 |
| 117 | + moe_router_load_balancing_type: aux_loss |
| 118 | + moe_router_topk: 8 |
| 119 | + moe_router_pre_softmax: false |
| 120 | + moe_aux_loss_coeff: 1e-3 |
| 121 | + attention_dropout: 0.0 |
| 122 | + hidden_dropout: 0.0 |
| 123 | + mock_data: true |
| 124 | + seq_length: 4096 |
| 125 | + moe_router_force_load_balancing: true |
| 126 | + tensor_model_parallel_size: 1 |
| 127 | + pipeline_model_parallel_size: 1 |
| 128 | + expert_model_parallel_size: 16 |
| 129 | + context_parallel_size: 1 |
| 130 | + expert_tensor_parallel_size: 1 |
| 131 | + use_distributed_optimizer: true |
| 132 | + sequence_parallel: true |
| 133 | + moe_token_dispatcher_type: flex |
| 134 | + moe_flex_dispatcher_backend: hybridep |
| 135 | + moe_grouped_gemm: true |
| 136 | + moe_permute_fusion: true |
| 137 | + moe_router_fusion: true |
| 138 | + moe_router_dtype: fp32 |
| 139 | + use_mcore_models: true |
| 140 | + use_flash_attn: true |
| 141 | + transformer_impl: transformer_engine |
| 142 | + micro_batch_size: 4 |
| 143 | + global_batch_size: 512 |
| 144 | + train_samples: 268554688 |
| 145 | + exit_duration_in_mins: 230 |
| 146 | + no_create_attention_mask_in_dataloader: true |
| 147 | + cross_entropy_loss_fusion: true |
| 148 | + cross_entropy_fusion_impl: te |
| 149 | + manual_gc: true |
| 150 | + manual_gc_interval: 5 |
| 151 | + lr: 0.00012 |
| 152 | + min_lr: 1.2e-05 |
| 153 | + lr_decay_style: cosine |
| 154 | + lr_decay_samples: 255126953 |
| 155 | + lr_warmup_samples: 162761 |
| 156 | + weight_decay: 0.1 |
| 157 | + clip_grad: 1.0 |
| 158 | + adam_beta1: 0.9 |
| 159 | + adam_beta2: 0.95 |
| 160 | + bf16: true |
| 161 | + init_method_std: 0.02 |
| 162 | + eval_iters: 32 |
| 163 | + eval_interval: 500 |
| 164 | + finetune: true |
| 165 | + auto_detect_ckpt_format: true |
| 166 | + no_load_rng: true |
| 167 | + no_load_optim: true |
| 168 | + load: ${LOAD_PATH} |
| 169 | + save_interval: 500 |
| 170 | + dist_ckpt_strictness: log_all |
| 171 | + log_throughput: true |
| 172 | + log_interval: 1 |
| 173 | + log_timers_to_tensorboard: true |
| 174 | + log_memory_to_tensorboard: true |
| 175 | + log_num_zeros_in_grad: true |
| 176 | + log_params_norm: true |
| 177 | + log_validation_ppl_to_tensorboard: true |
| 178 | + logging_level: 40 |
| 179 | + tensorboard_dir: ${OUTPUT_PATH}/tensorboard |
| 180 | + wandb_exp_name: Qwen3-30B-GB200-BF16-TP1PP1EP16-GBS512 |
| 181 | + enable_experimental: true |
0 commit comments