From 1428c9c20799e9ee212fecf1d3331a171055b3d9 Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Thu, 18 Jun 2026 11:32:18 +0100 Subject: [PATCH 1/8] feat: add Qwen3-8B pre-training sample (H200 vs B300, NeMo/Megatron) --- .../qwen3-8b-pretraining/README.md | 138 +++++++++++++++ .../qwen3-8b-pretraining/b300/Dockerfile | 63 +++++++ .../qwen3-8b-pretraining/b300/slurm/run.sh | 37 ++++ .../qwen3-8b-pretraining/b300/train.py | 66 +++++++ .../docs/lessons-learned.md | 126 +++++++++++++ .../qwen3-8b-pretraining/h200/Dockerfile | 63 +++++++ .../qwen3-8b-pretraining/h200/slurm/run.sh | 37 ++++ .../qwen3-8b-pretraining/h200/train.py | 165 ++++++++++++++++++ 8 files changed, 695 insertions(+) create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh create mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md new file mode 100644 index 000000000..637889b0f --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md @@ -0,0 +1,138 @@ +# Qwen3-8B Pre-Training: H200 vs B300 (NeMo/Megatron) + +Pre-training **Qwen3-8B** (8.2B dense parameters) on 1T tokens comparing two GPU generations — p5en.48xlarge (H200) and p6-b300.48xlarge (B300) — using NeMo/Megatron on 2-node / 16-GPU topologies with EFA GDRDMA interconnect. + +## Results + +| Metric | H200 (p5en) | B300 (p6-b300) | Ratio | +|--------|-------------|----------------|-------| +| **TFLOP/s per GPU** | 497 | **976** | 1.96× | +| **Throughput** | 162K tok/s | **318K tok/s** | 1.96× | +| **Time to 1T tokens** | ~71 days | **~36 days** | 1.97× | +| Step time (100 iters) | 3.23s | 1.65s | 1.96× | +| Peak memory/GPU | ~138 GB / 141 GB | ~173 GB / 288 GB | — | +| MFU | 0.50 | 0.50 | — | + +Both clusters are compute-saturated with perfect communication overlap. AllReduce and AllGather are fully hidden behind compute. + +## Prerequisites + +- **Slurm** workload manager with **PyXis + Enroot** container runtime +- **EFA networking** with GDRDMA support (for multi-node communication) +- **FSx for Lustre** shared filesystem mounted at `/fsx/` +- **Docker** (for building container images) + +> **Don't have a cluster?** Deploy a fully functional HPC cluster in under 1 hour using [Amazon SageMaker HyperPod](https://awslabs.github.io/ai-on-sagemaker-hyperpod/). The guide walks you through deploying a ready-to-use cluster with Slurm, EFA, PyXis/Enroot, and FSx for Lustre pre-configured. + +## Quick Start + +### H200 Cluster (p5en.48xlarge) + +```bash +# 1. Build container +cd h200/ && docker build -t qwen3-8b-h200:latest . +sudo enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh dockerd://qwen3-8b-h200:latest + +# 2. Submit training job +sbatch h200/slurm/run.sh +``` + +### B300 Cluster (p6-b300.48xlarge) + +```bash +# 1. Build container +cd b300/ && docker build -t qwen3-8b-b300:latest . +sudo enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh dockerd://qwen3-8b-b300:latest + +# 2. Submit training job +sbatch b300/slurm/run.sh +``` + +## Model Architecture: Qwen3-8B + +| Parameter | Value | +|-----------|-------| +| Layers | 36 | +| Hidden dim (d_model) | 4096 | +| Q-heads | 32 | +| KV-heads | 8 (GQA) | +| FFN dim | 14336 (SwiGLU) | +| Vocab size | 151,936 | +| Positional encoding | RoPE | +| Normalization | RMSNorm | +| Sequence length | 4096 | +| Precision | BF16 | +| Total params | 8.2B | + +## Parallelism Strategy + +**Pure Data Parallelism (DP=16)** — the model fits entirely on a single GPU. + +| Component | Setting | +|-----------|---------| +| Tensor Parallel | 1 | +| Pipeline Parallel | 1 | +| Data Parallel | 16 | +| Distributed Optimizer | Yes (shards Adam states across DP ranks) | +| Overlap Grad Reduce | Yes | +| Overlap Param Gather | Yes | + +**Why TP=1 is optimal:** At 8.2B params, the model + optimizer states fit on one GPU with distributed optimizer. Adding tensor parallelism introduces all-reduce communication for every transformer layer — validated experimentally: TP=2 was 11% slower (868 vs 976 TFLOP/s on B300). + +## Best Configuration Per Cluster + +| Parameter | H200 (p5en.48xlarge) | B300 (p6-b300.48xlarge) | +|-----------|---------------------|------------------------| +| GPUs | 16× H200 (141 GB HBM3) | 16× B300 (288 GB HBM3e) | +| Parallelism | TP=1, PP=1, DP=16 | TP=1, PP=1, DP=16 | +| Micro-batch size | 2 | 4 | +| Global batch size | 128 (grad_accum=4) | 128 (grad_accum=2) | +| Sequence length | 4096 | 4096 | +| Precision | BF16 | BF16 | +| Gradient checkpointing | Full recompute | None | +| Distributed optimizer | Yes (sharded Adam) | Yes (sharded Adam) | +| Overlap grad reduce | Yes | Yes | +| Framework | Megatron-Core (NeMo 25.07) | Megatron-Bridge (NeMo 26.02) | + +## Key Findings + +1. **Both clusters are compute-saturated with perfect communication overlap.** AllReduce and AllGather are fully hidden behind compute — verified by single-GPU benchmarks showing lower TFLOP/s due to reduced batch arithmetic intensity. + +2. **Best software stack per GPU generation.** NeMo 25.07 (Megatron-Core v0.13.1) for H200; NeMo 26.02 (Megatron-Bridge v2.9.0) for B300. Each maximizes hardware utilization for its target architecture. + +3. **Pure data parallelism is optimal** when the model fits in single-GPU memory. Distributed optimizer + overlapped grad reduce eliminate the memory penalty. + +4. **Gradient checkpointing is the key differentiator:** required on H200 (141 GB) for MBS≥2, unnecessary on B300 (288 GB) for MBS≤4 — saving ~20% recompute overhead. + +## Hardware + +| | H200 Cluster | B300 Cluster | +|---|---|---| +| Instance | p5en.48xlarge | p6-b300.48xlarge | +| Nodes | 2 | 2 | +| GPUs per node | 8× H200 | 8× B300 | +| GPU Memory | 141 GB HBM3 | 288 GB HBM3e | +| Interconnect | EFA GDRDMA (3200 Gbps) | EFA GDRDMA (3200 Gbps) | +| Intra-node | NVLink | NVLink | + +## Project Structure + +``` +├── README.md ← You are here +├── h200/ +│ ├── Dockerfile ← NeMo 25.07 + EFA container +│ ├── train.py ← Megatron-Core training script +│ └── slurm/ +│ └── run.sh ← Slurm submission script +├── b300/ +│ ├── Dockerfile ← NeMo 26.02 + EFA container +│ ├── train.py ← Megatron-Bridge training script +│ └── slurm/ +│ └── run.sh ← Slurm submission script +└── docs/ + └── lessons-learned.md ← EFA, PyXis, and Megatron gotchas +``` + +## License + +Apache 2.0 diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile new file mode 100644 index 000000000..624b56393 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile @@ -0,0 +1,63 @@ +FROM nvcr.io/nvidia/nemo:26.02 + +# EFA + NCCL OFI plugin for multi-node communication +# The NeMo 26.02 base image includes Megatron-Bridge v2.9.0 +# We add EFA support for AWS inter-node networking via GDRDMA + +ARG EFA_INSTALLER_VERSION=1.48.0 +ARG NCCL_VERSION=v2.30.4-1 +ARG GDRCOPY_VERSION=v2.5.2 + +# Remove existing NCCL/MPI to avoid conflicts +RUN apt-get update -y && \ + apt-get remove -y --allow-change-held-packages \ + ibverbs-utils libibverbs-dev libibverbs1 \ + libmlx5-1 libnccl2 libnccl-dev || true && \ + rm -rf /opt/hpcx /etc/ld.so.conf.d/hpcx.conf && ldconfig + +RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \ + build-essential cmake curl git kmod libtool \ + openssh-client openssh-server pkg-config vim wget + +# SSH config for multi-node +RUN mkdir -p /var/run/sshd && \ + sed -i 's/[ #]\(.*StrictHostKeyChecking\).*/\1 no/g' /etc/ssh/ssh_config && \ + echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \ + sed -i 's/#\(StrictModes\).*/\1 no/g' /etc/ssh/sshd_config + +# GDRCopy for GPU-direct RDMA +RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ + && cd /tmp/gdrcopy && make prefix=/opt/gdrcopy install + +# EFA installer (provides libfabric + OFI NCCL plugin) +RUN cd $HOME \ + && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && cd aws-efa-installer \ + && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \ + && rm -rf $HOME/aws-efa-installer + +# Build NCCL from source (sm_100 for B300 Blackwell) +RUN git clone -b ${NCCL_VERSION} https://github.com/NVIDIA/nccl.git /opt/nccl \ + && cd /opt/nccl \ + && make -j $(nproc) src.build CUDA_HOME=/usr/local/cuda \ + NVCC_GENCODE="-gencode=arch=compute_100,code=sm_100 -gencode=arch=compute_90,code=sm_90" + +ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/opt/amazon/openmpi/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/lib:$LD_LIBRARY_PATH +ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:/opt/gdrcopy/bin:$PATH + +# EFA/NCCL environment +ENV FI_PROVIDER=efa +ENV NCCL_SOCKET_IFNAME=^docker,lo,veth +ENV NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +ENV LD_PRELOAD=/opt/nccl/build/lib/libnccl.so +ENV TORCH_COMPILE_DISABLE=1 + +# OpenMPI settings +ENV OMPI_MCA_pml=^ucx \ + OMPI_MCA_btl=tcp,self \ + OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \ + OPAL_PREFIX=/opt/amazon/openmpi \ + PMIX_MCA_gds=hash + +RUN rm -rf /var/lib/apt/lists/* /tmp/gdrcopy diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh new file mode 100644 index 000000000..9b15c28ce --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=qwen3-8b-b300 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=12 +#SBATCH --exclusive +#SBATCH --partition=b300 +#SBATCH --time=24:00:00 +#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err +#SBATCH --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh +#SBATCH --container-mounts=/fsx:/fsx + +# Resolve head node +export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n1) +export MASTER_PORT=29500 + +# EFA / NCCL environment +export FI_PROVIDER=efa +export NCCL_SOCKET_IFNAME=^docker,lo,veth +export NCCL_DEBUG=WARN +export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/nccl/build/lib:$LD_LIBRARY_PATH + +# Disable torch.compile +export TORCH_COMPILE_DISABLE=1 + +# Launch training +srun --container-env=MASTER_ADDR,MASTER_PORT,FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ + torchrun \ + --nnodes=${SLURM_NNODES} \ + --nproc-per-node=8 \ + --rdzv-id=${SLURM_JOB_ID} \ + --rdzv-backend=c10d \ + --rdzv-endpoint=${MASTER_ADDR}:${MASTER_PORT} \ + /fsx/ubuntu/qwen3-8b-pretraining/code/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py new file mode 100644 index 000000000..5ba1422a1 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Qwen3-8B Pre-Training on B300 — Megatron-Bridge (NeMo 26.02) + +Uses the Megatron-Bridge recipe API with config objects. +No gradient checkpointing needed (288 GB B300 memory). + +Best config: TP=1, PP=1, DP=16, MBS=4, GBS=128, seq=4096, BF16 +Result: 976 TFLOP/s/GPU, 318K tok/s on 16x B300 +""" +import os + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +from megatron.bridge.recipes.qwen.qwen3 import qwen3_8b_pretrain_config +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain + + +def main(): + cfg = qwen3_8b_pretrain_config( + mock=True, # Set to False and provide data_path for real training + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + micro_batch_size=4, + global_batch_size=128, + seq_length=4096, + train_iters=100, + lr_warmup_iters=10, + lr_decay_iters=100, + ) + + # Training configuration + cfg.train.bf16 = True + cfg.train.use_distributed_optimizer = True + cfg.train.overlap_grad_reduce = True + cfg.train.overlap_param_gather = True + + # No gradient checkpointing (B300 has 288 GB — fits MBS=4 without recompute) + cfg.train.recompute_granularity = None + + # Optimizer + cfg.optimizer.lr = 3e-4 + cfg.optimizer.min_lr = 3e-5 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.clip_grad = 1.0 + + # Logging and checkpoints + cfg.logger.log_interval = 5 + cfg.train.eval_interval = 1000 + cfg.train.eval_iters = 0 + cfg.train.dir = "/fsx/ubuntu/qwen3-8b-pretraining/checkpoints/b300" + cfg.train.save_interval = 1000 + + # To use real data instead of mock: + # cfg.data.mock = False + # cfg.data.data_path = "/fsx/ubuntu/qwen3-8b-pretraining/datasets/c4/merged_text_document" + # cfg.data.tokenizer_type = "HuggingFaceTokenizer" + # cfg.data.tokenizer_model = "Qwen/Qwen3-8B" + + pretrain(config=cfg, forward_step_func=forward_step) + + +if __name__ == "__main__": + main() diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md new file mode 100644 index 000000000..28c02cbc6 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md @@ -0,0 +1,126 @@ +# Lessons Learned + +Hard-won knowledge from optimizing Qwen3-8B pre-training on H200 and B300 clusters. + +--- + +## EFA Silent Fallback to TCP + +**Symptom:** Multi-node training runs but at single-node throughput. NCCL reports no errors. + +**Root cause:** NCCL silently falls back to TCP sockets when the OFI plugin isn't loaded. + +**Fix — all three are required:** + +```bash +export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:$LD_LIBRARY_PATH +export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +export FI_PROVIDER=efa +``` + +**Verification:** Look for `NCCL INFO NET/OFI` in logs (not `NET/Socket`). + +The OFI NCCL plugin directory may use either `/opt/amazon/ofi-nccl/lib/` or `/opt/amazon/aws-ofi-nccl/lib/` depending on the EFA installer version. Check which exists. + +--- + +## Enroot Import Workflow + +**Never use `mksquashfs` directly.** It produces images that PyXis can't launch. + +**Correct workflow:** + +```bash +# 1. Build with Docker +sudo docker build -t my-image:latest . + +# 2. Import with enroot (requires sudo for Docker socket) +sudo TMPDIR=/fsx/ubuntu/qwen3-8b-pretraining/tmp \ + ENROOT_TEMP_PATH=/fsx/ubuntu/qwen3-8b-pretraining/tmp \ + enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/image.sqsh dockerd://my-image:latest + +# 3. Fix permissions +sudo chown $USER:$USER /fsx/ubuntu/qwen3-8b-pretraining/containers/image.sqsh +``` + +**Three requirements:** +1. `sudo` — Docker socket is `root:docker`, user not in docker group +2. `TMPDIR` on FSx — NeMo containers are 30+ GB, `/tmp` won't fit +3. `docker buildx use default` — image must be in main containerd store + +--- + +## Slurm/PyXis Gotchas + +### Slurm NOT in PATH +On some clusters, Slurm binaries live at `/opt/slurm/bin/`. Use full paths if needed: +```bash +/opt/slurm/bin/sbatch script.sh +/opt/slurm/bin/squeue -u ubuntu +/opt/slurm/bin/scontrol show hostname $SLURM_NODELIST +``` + +### Shell vars don't pass into containers +PyXis containers don't inherit the calling shell's environment. Pass explicitly: +```bash +srun --container-env=VAR1,VAR2,VAR3 ... +``` + +### Resolve MASTER_ADDR before srun +`scontrol` is not available inside PyXis containers. Compute the head node in the batch script, before the `srun` call: +```bash +export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n1) +``` + +### ntasks-per-node for torchrun +When using `torchrun` (which spawns GPU workers itself), set `--ntasks-per-node=1` in the Slurm script. If using raw `python` with NCCL init, use `--ntasks-per-node=8`. + +### Single-node: disable EFA +For intra-node-only jobs, `FI_PROVIDER=efa` causes NCCL failures. Remove it or set `FI_PROVIDER=shm` for single-node debugging. + +--- + +## torch.compile Incompatibility + +Set `TORCH_COMPILE_DISABLE=1` in all environments. It fails in every configuration tested (DeepSpeed, HuggingFace multi-node, NeMo 25.07, NeMo 26.02). The performance gain would be minimal since Transformer Engine already provides fused kernels. + +--- + +## Distributed Optimizer Trap at DP=1 + +`--use-distributed-optimizer` with `DP=1` (single GPU or TP-only parallelism) causes a crash. The sharding logic divides by DP world size and expects DP>=2. + +**Rule:** Only enable distributed optimizer when DP>=2. For single-GPU debugging, remove the flag. + +--- + +## Megatron-Bridge API Gotchas (NeMo 26.02) + +- **Logging interval** is on `cfg.logger.log_interval`, not `cfg.train.log_interval` (silently ignored) +- **Disable gradient checkpointing** with `cfg.train.recompute_granularity = None` (not `""` or `False`) +- **Checkpoint directory** is `cfg.train.dir` (not `cfg.train.save` or `cfg.train.checkpoint_dir`) +- **Qwen3 bridge recipe** `qwen3_8b_pretrain_config()` provides correct model dimensions — don't manually override + +--- + +## Memory Budget: H200 vs B300 + +``` +H200 (141 GB available): + Model (BF16): 16 GB + Gradients (BF16): 16 GB + Optimizer (sharded/16): 3 GB + Activations (recompute): ~100 GB <- with full recompute, MBS=2 + Overhead: 3 GB + Total: ~138 GB + +B300 (288 GB available): + Model (BF16): 16 GB + Gradients (BF16): 16 GB + Optimizer (sharded/16): 3 GB + Activations (no recomp): ~135 GB <- NO recompute needed, MBS=4 + Overhead: 3 GB + Total: ~173 GB (115 GB headroom) +``` + +B300's extra memory means no recompute overhead -> ~20% fewer FLOPs per step -> directly translates to 1.96x throughput combined with higher peak FLOPS. diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile new file mode 100644 index 000000000..7e6a6bb93 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile @@ -0,0 +1,63 @@ +FROM nvcr.io/nvidia/nemo:25.07 + +# EFA + NCCL OFI plugin for multi-node communication +# The NeMo base image includes PyTorch 2.8.0, CUDA 12.9, Megatron-Core 0.13.1, TE 2.5 +# We add EFA support for AWS inter-node networking via GDRDMA + +ARG EFA_INSTALLER_VERSION=1.48.0 +ARG NCCL_VERSION=v2.30.4-1 +ARG GDRCOPY_VERSION=v2.5.2 + +# Remove existing NCCL/MPI to avoid conflicts +RUN apt-get update -y && \ + apt-get remove -y --allow-change-held-packages \ + ibverbs-utils libibverbs-dev libibverbs1 \ + libmlx5-1 libnccl2 libnccl-dev || true && \ + rm -rf /opt/hpcx /etc/ld.so.conf.d/hpcx.conf && ldconfig + +RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \ + build-essential cmake curl git kmod libtool \ + openssh-client openssh-server pkg-config vim wget + +# SSH config for multi-node +RUN mkdir -p /var/run/sshd && \ + sed -i 's/[ #]\(.*StrictHostKeyChecking\).*/\1 no/g' /etc/ssh/ssh_config && \ + echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \ + sed -i 's/#\(StrictModes\).*/\1 no/g' /etc/ssh/sshd_config + +# GDRCopy for GPU-direct RDMA +RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ + && cd /tmp/gdrcopy && make prefix=/opt/gdrcopy install + +# EFA installer (provides libfabric + OFI NCCL plugin) +RUN cd $HOME \ + && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && cd aws-efa-installer \ + && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \ + && rm -rf $HOME/aws-efa-installer + +# Build NCCL from source (sm_90 for H200) +RUN git clone -b ${NCCL_VERSION} https://github.com/NVIDIA/nccl.git /opt/nccl \ + && cd /opt/nccl \ + && make -j $(nproc) src.build CUDA_HOME=/usr/local/cuda \ + NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90" + +ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/opt/amazon/openmpi/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/lib:$LD_LIBRARY_PATH +ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:/opt/gdrcopy/bin:$PATH + +# EFA/NCCL environment +ENV FI_PROVIDER=efa +ENV NCCL_SOCKET_IFNAME=^docker,lo,veth +ENV NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +ENV LD_PRELOAD=/opt/nccl/build/lib/libnccl.so +ENV TORCH_COMPILE_DISABLE=1 + +# OpenMPI settings +ENV OMPI_MCA_pml=^ucx \ + OMPI_MCA_btl=tcp,self \ + OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \ + OPAL_PREFIX=/opt/amazon/openmpi \ + PMIX_MCA_gds=hash + +RUN rm -rf /var/lib/apt/lists/* /tmp/gdrcopy diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh new file mode 100644 index 000000000..2b2efda38 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=qwen3-8b-h200 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=12 +#SBATCH --exclusive +#SBATCH --partition=p5en +#SBATCH --time=24:00:00 +#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err +#SBATCH --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh +#SBATCH --container-mounts=/fsx:/fsx + +# Resolve head node BEFORE srun (scontrol not available inside container) +export MASTER_ADDR=$(/opt/slurm/bin/scontrol show hostname $SLURM_NODELIST | head -n1) +export MASTER_PORT=29500 + +# EFA / NCCL environment +export FI_PROVIDER=efa +export NCCL_SOCKET_IFNAME=^docker,lo,veth +export NCCL_DEBUG=WARN +export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/nccl/build/lib:$LD_LIBRARY_PATH + +# Disable torch.compile (incompatible with this stack) +export TORCH_COMPILE_DISABLE=1 + +# Launch training +srun --container-env=MASTER_ADDR,MASTER_PORT,FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ + torchrun \ + --nnodes=${SLURM_NNODES} \ + --nproc-per-node=8 \ + --rdzv-id=${SLURM_JOB_ID} \ + --rdzv-backend=c10d \ + --rdzv-endpoint=${MASTER_ADDR}:${MASTER_PORT} \ + /fsx/ubuntu/qwen3-8b-pretraining/code/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py new file mode 100644 index 000000000..f8ff8cd15 --- /dev/null +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +"""Qwen3-8B Pre-Training on H200 — Megatron-Core (NeMo 25.07) + +Uses the Megatron pretrain API with Transformer Engine spec. +Qwen3-8B dimensions mapped to GPT model provider (Qwen3 bridge not registered). + +Best config: TP=1, PP=1, DP=16, MBS=2, GBS=128, seq=4096, BF16 +Result: 497 TFLOP/s/GPU, 162K tok/s on 16x H200 +""" +import os +import sys + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") +os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1") + +from functools import partial + +import torch +from megatron.training import get_args, pretrain +from megatron.training.arguments import core_transformer_config_from_args +from megatron.core.enums import ModelType +from megatron.core.models.gpt import GPTModel +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, +) +from megatron.core.transformer.spec_utils import import_module +from megatron.core.datasets.blended_megatron_dataset_builder import ( + BlendedMegatronDatasetBuilder, +) +from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset + + +def model_provider(pre_process=True, post_process=True): + """Build GPT model with Qwen3-8B dimensions.""" + args = get_args() + config = core_transformer_config_from_args(args) + + transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec() + + model = GPTModel( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + parallel_output=True, + ) + return model + + +def forward_step(data_iterator, model): + """Forward pass — standard GPT causal LM loss.""" + args = get_args() + tokens, labels, loss_mask, attention_mask, position_ids = _get_batch(data_iterator) + output_tensor = model(tokens, position_ids, attention_mask, labels=labels) + return output_tensor, partial(_loss_func, loss_mask) + + +def _loss_func(loss_mask, output_tensor): + """Compute averaged cross-entropy loss.""" + losses = output_tensor.float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() + return loss, {"lm loss": loss} + + +def _get_batch(data_iterator): + """Get batch from data iterator.""" + args = get_args() + data = next(data_iterator) + + tokens = data["tokens"].long().cuda() + labels = data["labels"].long().cuda() + loss_mask = data["loss_mask"].float().cuda() + attention_mask = data["attention_mask"].long().cuda() if "attention_mask" in data else None + position_ids = data["position_ids"].long().cuda() + + return tokens, labels, loss_mask, attention_mask, position_ids + + +def train_valid_test_datasets_provider(train_val_test_num_samples): + """Build mock datasets for benchmarking. + + To switch to real data, replace MockGPTDataset with GPTDataset and provide: + --data-path /fsx/ubuntu/qwen3-8b-pretraining/datasets/c4/merged_text_document + --tokenizer-type HuggingFaceTokenizer + --tokenizer-model Qwen/Qwen3-8B + """ + args = get_args() + + config = GPTDatasetConfig( + random_seed=args.seed, + sequence_length=args.seq_length, + reset_position_ids=False, + reset_attention_mask=False, + eod_mask_loss=False, + mock=True, + mock_seq_length=args.seq_length, + ) + + dataset_builder = BlendedMegatronDatasetBuilder( + MockGPTDataset, train_val_test_num_samples, lambda: True, config + ) + train_ds, valid_ds, test_ds = dataset_builder.build() + return train_ds, valid_ds, test_ds + + +if __name__ == "__main__": + pretrain( + train_valid_test_datasets_provider, + model_provider, + ModelType.encoder_or_decoder, + forward_step, + args_defaults={ + # Qwen3-8B architecture + "num_layers": 36, + "hidden_size": 4096, + "num_attention_heads": 32, + "group_query_attention": True, + "num_query_groups": 8, + "ffn_hidden_size": 14336, + "swiglu": True, + "max_position_embeddings": 4096, + "seq_length": 4096, + "padded_vocab_size": 151936, + "use_rotary_position_embeddings": True, + "rotary_percent": 1.0, + "normalization": "RMSNorm", + "untie_embeddings_and_output_weights": True, + # Training + "micro_batch_size": 2, + "global_batch_size": 128, + "train_iters": 100, + "lr": 3e-4, + "min_lr": 3e-5, + "lr_warmup_iters": 10, + "lr_decay_iters": 100, + "lr_decay_style": "cosine", + "weight_decay": 0.1, + "adam_beta1": 0.9, + "adam_beta2": 0.95, + "clip_grad": 1.0, + "bf16": True, + # Parallelism + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "use_distributed_optimizer": True, + "overlap_grad_reduce": True, + "overlap_param_gather": True, + # Gradient checkpointing (mandatory on H200 for MBS>=2) + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + # Logging + "log_interval": 5, + "eval_interval": 1000, + "eval_iters": 0, + "tensorboard_dir": "/fsx/ubuntu/qwen3-8b-pretraining/tensorboard/h200", + "save": "/fsx/ubuntu/qwen3-8b-pretraining/checkpoints/h200", + "save_interval": 1000, + "tokenizer_type": "NullTokenizer", + "vocab_size": 151936, + }, + ) From cd5b90abee7a780ffda5f83480e0d94687fb7876 Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Thu, 18 Jun 2026 12:04:11 +0100 Subject: [PATCH 2/8] fix: correct Dockerfile EFA build sequence (tested on B300) --- .../qwen3-8b-pretraining/README.md | 7 +- .../qwen3-8b-pretraining/b300/Dockerfile | 88 +++++++------------ .../qwen3-8b-pretraining/h200/Dockerfile | 88 +++++++------------ 3 files changed, 69 insertions(+), 114 deletions(-) diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md index 637889b0f..746da10d8 100644 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md @@ -26,12 +26,15 @@ Both clusters are compute-saturated with perfect communication overlap. AllReduc ## Quick Start +> **Disk space:** The container build requires ~50 GB of disk space in TMPDIR. +> `enroot import` needs `sudo` and TMPDIR pointing to FSx (not `/tmp`, which is too small). + ### H200 Cluster (p5en.48xlarge) ```bash # 1. Build container cd h200/ && docker build -t qwen3-8b-h200:latest . -sudo enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh dockerd://qwen3-8b-h200:latest +sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh dockerd://qwen3-8b-h200:latest # 2. Submit training job sbatch h200/slurm/run.sh @@ -42,7 +45,7 @@ sbatch h200/slurm/run.sh ```bash # 1. Build container cd b300/ && docker build -t qwen3-8b-b300:latest . -sudo enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh dockerd://qwen3-8b-b300:latest +sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh dockerd://qwen3-8b-b300:latest # 2. Submit training job sbatch b300/slurm/run.sh diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile index 624b56393..1af80364b 100644 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile @@ -1,63 +1,39 @@ FROM nvcr.io/nvidia/nemo:26.02 -# EFA + NCCL OFI plugin for multi-node communication -# The NeMo 26.02 base image includes Megatron-Bridge v2.9.0 -# We add EFA support for AWS inter-node networking via GDRDMA - -ARG EFA_INSTALLER_VERSION=1.48.0 -ARG NCCL_VERSION=v2.30.4-1 -ARG GDRCOPY_VERSION=v2.5.2 - -# Remove existing NCCL/MPI to avoid conflicts -RUN apt-get update -y && \ - apt-get remove -y --allow-change-held-packages \ - ibverbs-utils libibverbs-dev libibverbs1 \ - libmlx5-1 libnccl2 libnccl-dev || true && \ - rm -rf /opt/hpcx /etc/ld.so.conf.d/hpcx.conf && ldconfig - -RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \ - build-essential cmake curl git kmod libtool \ - openssh-client openssh-server pkg-config vim wget - -# SSH config for multi-node -RUN mkdir -p /var/run/sshd && \ - sed -i 's/[ #]\(.*StrictHostKeyChecking\).*/\1 no/g' /etc/ssh/ssh_config && \ - echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \ - sed -i 's/#\(StrictModes\).*/\1 no/g' /etc/ssh/sshd_config +# Install EFA prerequisites +RUN apt-get update && apt-get install -y --no-install-recommends \ + environment-modules \ + tcl \ + udev \ + ethtool \ + iproute2 \ + dmidecode \ + libevent-core-2.1-7t64 \ + libevent-pthreads-2.1-7t64 \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Install EFA (v1.47.0 - compatible with NeMo 26.02 base) +# Do NOT remove existing NCCL/hpcx - EFA installs on top +RUN cd /tmp && \ + curl -O https://efa-installer.amazonaws.com/aws-efa-installer-1.47.0.tar.gz && \ + tar -xf aws-efa-installer-1.47.0.tar.gz && \ + cd aws-efa-installer && \ + ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify && \ + cd / && rm -rf /tmp/aws-efa-installer* # GDRCopy for GPU-direct RDMA -RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ - && cd /tmp/gdrcopy && make prefix=/opt/gdrcopy install - -# EFA installer (provides libfabric + OFI NCCL plugin) -RUN cd $HOME \ - && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ - && tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ - && cd aws-efa-installer \ - && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \ - && rm -rf $HOME/aws-efa-installer - -# Build NCCL from source (sm_100 for B300 Blackwell) -RUN git clone -b ${NCCL_VERSION} https://github.com/NVIDIA/nccl.git /opt/nccl \ - && cd /opt/nccl \ - && make -j $(nproc) src.build CUDA_HOME=/usr/local/cuda \ - NVCC_GENCODE="-gencode=arch=compute_100,code=sm_100 -gencode=arch=compute_90,code=sm_90" - -ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/opt/amazon/openmpi/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/lib:$LD_LIBRARY_PATH -ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:/opt/gdrcopy/bin:$PATH - -# EFA/NCCL environment +RUN cd /tmp && \ + git clone -b v2.4.4 https://github.com/NVIDIA/gdrcopy.git && \ + cd gdrcopy && \ + make -j$(nproc) lib lib_install && \ + cd / && rm -rf /tmp/gdrcopy + +# Environment +ENV LD_LIBRARY_PATH="/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH}" +ENV NCCL_TUNER_PLUGIN="/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so" ENV FI_PROVIDER=efa -ENV NCCL_SOCKET_IFNAME=^docker,lo,veth -ENV NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so -ENV LD_PRELOAD=/opt/nccl/build/lib/libnccl.so ENV TORCH_COMPILE_DISABLE=1 +ENV NCCL_PROTO=simple -# OpenMPI settings -ENV OMPI_MCA_pml=^ucx \ - OMPI_MCA_btl=tcp,self \ - OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \ - OPAL_PREFIX=/opt/amazon/openmpi \ - PMIX_MCA_gds=hash - -RUN rm -rf /var/lib/apt/lists/* /tmp/gdrcopy +WORKDIR /workspace diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile index 7e6a6bb93..953639199 100644 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile +++ b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile @@ -1,63 +1,39 @@ FROM nvcr.io/nvidia/nemo:25.07 -# EFA + NCCL OFI plugin for multi-node communication -# The NeMo base image includes PyTorch 2.8.0, CUDA 12.9, Megatron-Core 0.13.1, TE 2.5 -# We add EFA support for AWS inter-node networking via GDRDMA - -ARG EFA_INSTALLER_VERSION=1.48.0 -ARG NCCL_VERSION=v2.30.4-1 -ARG GDRCOPY_VERSION=v2.5.2 - -# Remove existing NCCL/MPI to avoid conflicts -RUN apt-get update -y && \ - apt-get remove -y --allow-change-held-packages \ - ibverbs-utils libibverbs-dev libibverbs1 \ - libmlx5-1 libnccl2 libnccl-dev || true && \ - rm -rf /opt/hpcx /etc/ld.so.conf.d/hpcx.conf && ldconfig - -RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \ - build-essential cmake curl git kmod libtool \ - openssh-client openssh-server pkg-config vim wget - -# SSH config for multi-node -RUN mkdir -p /var/run/sshd && \ - sed -i 's/[ #]\(.*StrictHostKeyChecking\).*/\1 no/g' /etc/ssh/ssh_config && \ - echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \ - sed -i 's/#\(StrictModes\).*/\1 no/g' /etc/ssh/sshd_config +# Install EFA prerequisites +RUN apt-get update && apt-get install -y --no-install-recommends \ + environment-modules \ + tcl \ + udev \ + ethtool \ + iproute2 \ + dmidecode \ + libevent-core-2.1-7 \ + libevent-pthreads-2.1-7 \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Install EFA (v1.47.0 - compatible with NeMo 25.07 base) +# Do NOT remove existing NCCL/hpcx - EFA installs on top +RUN cd /tmp && \ + curl -O https://efa-installer.amazonaws.com/aws-efa-installer-1.47.0.tar.gz && \ + tar -xf aws-efa-installer-1.47.0.tar.gz && \ + cd aws-efa-installer && \ + ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify && \ + cd / && rm -rf /tmp/aws-efa-installer* # GDRCopy for GPU-direct RDMA -RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ - && cd /tmp/gdrcopy && make prefix=/opt/gdrcopy install - -# EFA installer (provides libfabric + OFI NCCL plugin) -RUN cd $HOME \ - && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ - && tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ - && cd aws-efa-installer \ - && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \ - && rm -rf $HOME/aws-efa-installer - -# Build NCCL from source (sm_90 for H200) -RUN git clone -b ${NCCL_VERSION} https://github.com/NVIDIA/nccl.git /opt/nccl \ - && cd /opt/nccl \ - && make -j $(nproc) src.build CUDA_HOME=/usr/local/cuda \ - NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90" - -ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/amazon/ofi-nccl/lib:/opt/amazon/openmpi/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/lib:$LD_LIBRARY_PATH -ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:/opt/gdrcopy/bin:$PATH - -# EFA/NCCL environment +RUN cd /tmp && \ + git clone -b v2.4.4 https://github.com/NVIDIA/gdrcopy.git && \ + cd gdrcopy && \ + make -j$(nproc) lib lib_install && \ + cd / && rm -rf /tmp/gdrcopy + +# Environment +ENV LD_LIBRARY_PATH="/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH}" +ENV NCCL_TUNER_PLUGIN="/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so" ENV FI_PROVIDER=efa -ENV NCCL_SOCKET_IFNAME=^docker,lo,veth -ENV NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so -ENV LD_PRELOAD=/opt/nccl/build/lib/libnccl.so ENV TORCH_COMPILE_DISABLE=1 +ENV NCCL_PROTO=simple -# OpenMPI settings -ENV OMPI_MCA_pml=^ucx \ - OMPI_MCA_btl=tcp,self \ - OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \ - OPAL_PREFIX=/opt/amazon/openmpi \ - PMIX_MCA_gds=hash - -RUN rm -rf /var/lib/apt/lists/* /tmp/gdrcopy +WORKDIR /workspace From 012037b4277f5739c090511b9bb6aaad05cecbce Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Thu, 18 Jun 2026 20:30:49 +0100 Subject: [PATCH 3/8] fix: address review - correct architecture, restructure to megatron/nemo/, MIT-0 headers - Rewrote h200/train.py to use megatron-bridge recipe (correct Qwen3-8B arch) - Moved from pytorch/nemo-megatron/ to megatron/nemo/qwen3-8b-pretraining/ - Removed docs/lessons-learned.md - Added MIT-0 SPDX headers to all .py and .sh files - Fixed license references (Apache 2.0 -> MIT-0) - Fixed run.sh scripts: proper srun+python pattern, script staging - Updated README: no grad ckpt on H200, corrected memory (~114 GB) --- .../megatron/nemo/kubernetes/build.sh | 2 + .../nemo/kubernetes/custom_data_module.py | 3 + .../data-processing/data-processing.sh | 2 + .../data-processing/load_dataset.py | 3 + .../kubernetes/finetune_custom_dataset.py | 3 + .../kubernetes/finetune_default_dataset.py | 3 + .../kubernetes/pretrain_custom_dataset.py | 3 + .../nemo/kubernetes/pretrain_mock_dataset.py | 3 + 3.test_cases/megatron/nemo/kubernetes/push.sh | 2 + .../nemo}/qwen3-8b-pretraining/README.md | 14 +- .../qwen3-8b-pretraining/b300/Dockerfile | 0 .../qwen3-8b-pretraining/b300/slurm/run.sh | 32 ++++ .../nemo}/qwen3-8b-pretraining/b300/train.py | 21 +-- .../qwen3-8b-pretraining/h200/Dockerfile | 0 .../qwen3-8b-pretraining/h200/slurm/run.sh | 33 ++++ .../nemo/qwen3-8b-pretraining/h200/train.py | 56 ++++++ 3.test_cases/megatron/nemo/slurm/run.py | 3 + .../qwen3-8b-pretraining/b300/slurm/run.sh | 37 ---- .../docs/lessons-learned.md | 126 ------------- .../qwen3-8b-pretraining/h200/slurm/run.sh | 37 ---- .../qwen3-8b-pretraining/h200/train.py | 165 ------------------ 21 files changed, 160 insertions(+), 388 deletions(-) rename 3.test_cases/{pytorch/nemo-megatron => megatron/nemo}/qwen3-8b-pretraining/README.md (88%) rename 3.test_cases/{pytorch/nemo-megatron => megatron/nemo}/qwen3-8b-pretraining/b300/Dockerfile (100%) create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh rename 3.test_cases/{pytorch/nemo-megatron => megatron/nemo}/qwen3-8b-pretraining/b300/train.py (64%) rename 3.test_cases/{pytorch/nemo-megatron => megatron/nemo}/qwen3-8b-pretraining/h200/Dockerfile (100%) create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py delete mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh delete mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md delete mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh delete mode 100644 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py diff --git a/3.test_cases/megatron/nemo/kubernetes/build.sh b/3.test_cases/megatron/nemo/kubernetes/build.sh index a013de6ea..35ec114ef 100755 --- a/3.test_cases/megatron/nemo/kubernetes/build.sh +++ b/3.test_cases/megatron/nemo/kubernetes/build.sh @@ -1,4 +1,6 @@ #!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 # Build the AWS-optimized NeMo container for P4 and P5 instances # This script builds the Docker image with EFA support optimizations diff --git a/3.test_cases/megatron/nemo/kubernetes/custom_data_module.py b/3.test_cases/megatron/nemo/kubernetes/custom_data_module.py index 7514c2e16..1066c6938 100644 --- a/3.test_cases/megatron/nemo/kubernetes/custom_data_module.py +++ b/3.test_cases/megatron/nemo/kubernetes/custom_data_module.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import json import shutil import numpy as np diff --git a/3.test_cases/megatron/nemo/kubernetes/data-processing/data-processing.sh b/3.test_cases/megatron/nemo/kubernetes/data-processing/data-processing.sh index 79216e953..94a32e49b 100644 --- a/3.test_cases/megatron/nemo/kubernetes/data-processing/data-processing.sh +++ b/3.test_cases/megatron/nemo/kubernetes/data-processing/data-processing.sh @@ -1,4 +1,6 @@ #!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 # Deploy NeMo Data Processing Pod # This script helps deploy and manage the data processing pod diff --git a/3.test_cases/megatron/nemo/kubernetes/data-processing/load_dataset.py b/3.test_cases/megatron/nemo/kubernetes/data-processing/load_dataset.py index c380631f9..88dca6a91 100644 --- a/3.test_cases/megatron/nemo/kubernetes/data-processing/load_dataset.py +++ b/3.test_cases/megatron/nemo/kubernetes/data-processing/load_dataset.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + from datasets import load_dataset import json import os diff --git a/3.test_cases/megatron/nemo/kubernetes/finetune_custom_dataset.py b/3.test_cases/megatron/nemo/kubernetes/finetune_custom_dataset.py index fcf0052d9..434e3a915 100644 --- a/3.test_cases/megatron/nemo/kubernetes/finetune_custom_dataset.py +++ b/3.test_cases/megatron/nemo/kubernetes/finetune_custom_dataset.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import nemo_run as run import json import argparse diff --git a/3.test_cases/megatron/nemo/kubernetes/finetune_default_dataset.py b/3.test_cases/megatron/nemo/kubernetes/finetune_default_dataset.py index d56c8f93f..c8a153732 100644 --- a/3.test_cases/megatron/nemo/kubernetes/finetune_default_dataset.py +++ b/3.test_cases/megatron/nemo/kubernetes/finetune_default_dataset.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import signal import nemo_run as run import json diff --git a/3.test_cases/megatron/nemo/kubernetes/pretrain_custom_dataset.py b/3.test_cases/megatron/nemo/kubernetes/pretrain_custom_dataset.py index 8a8d182fd..3e5dfb80f 100644 --- a/3.test_cases/megatron/nemo/kubernetes/pretrain_custom_dataset.py +++ b/3.test_cases/megatron/nemo/kubernetes/pretrain_custom_dataset.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import signal import nemo_run as run import json diff --git a/3.test_cases/megatron/nemo/kubernetes/pretrain_mock_dataset.py b/3.test_cases/megatron/nemo/kubernetes/pretrain_mock_dataset.py index c36345981..a370bd225 100644 --- a/3.test_cases/megatron/nemo/kubernetes/pretrain_mock_dataset.py +++ b/3.test_cases/megatron/nemo/kubernetes/pretrain_mock_dataset.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import signal import nemo_run as run import json diff --git a/3.test_cases/megatron/nemo/kubernetes/push.sh b/3.test_cases/megatron/nemo/kubernetes/push.sh index 9e37760e9..2876ad146 100755 --- a/3.test_cases/megatron/nemo/kubernetes/push.sh +++ b/3.test_cases/megatron/nemo/kubernetes/push.sh @@ -1,4 +1,6 @@ #!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 # Push the AWS-optimized NeMo container to Amazon ECR # This script creates the ECR repository, logs in, tags, and pushes the image diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md similarity index 88% rename from 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md rename to 3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md index 746da10d8..34667c97b 100644 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/README.md +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md @@ -10,7 +10,7 @@ Pre-training **Qwen3-8B** (8.2B dense parameters) on 1T tokens comparing two GPU | **Throughput** | 162K tok/s | **318K tok/s** | 1.96× | | **Time to 1T tokens** | ~71 days | **~36 days** | 1.97× | | Step time (100 iters) | 3.23s | 1.65s | 1.96× | -| Peak memory/GPU | ~138 GB / 141 GB | ~173 GB / 288 GB | — | +| Peak memory/GPU | ~114 GB / 141 GB | ~173 GB / 288 GB | — | | MFU | 0.50 | 0.50 | — | Both clusters are compute-saturated with perfect communication overlap. AllReduce and AllGather are fully hidden behind compute. @@ -92,20 +92,20 @@ sbatch b300/slurm/run.sh | Global batch size | 128 (grad_accum=4) | 128 (grad_accum=2) | | Sequence length | 4096 | 4096 | | Precision | BF16 | BF16 | -| Gradient checkpointing | Full recompute | None | +| Gradient checkpointing | None | None | | Distributed optimizer | Yes (sharded Adam) | Yes (sharded Adam) | | Overlap grad reduce | Yes | Yes | -| Framework | Megatron-Core (NeMo 25.07) | Megatron-Bridge (NeMo 26.02) | +| Framework | Megatron-Bridge (NeMo 25.07) | Megatron-Bridge (NeMo 26.02) | ## Key Findings 1. **Both clusters are compute-saturated with perfect communication overlap.** AllReduce and AllGather are fully hidden behind compute — verified by single-GPU benchmarks showing lower TFLOP/s due to reduced batch arithmetic intensity. -2. **Best software stack per GPU generation.** NeMo 25.07 (Megatron-Core v0.13.1) for H200; NeMo 26.02 (Megatron-Bridge v2.9.0) for B300. Each maximizes hardware utilization for its target architecture. +2. **Both clusters use the Megatron-Bridge recipe API.** NeMo 25.07 for H200; NeMo 26.02 for B300. Each container maximizes hardware utilization for its target architecture. 3. **Pure data parallelism is optimal** when the model fits in single-GPU memory. Distributed optimizer + overlapped grad reduce eliminate the memory penalty. -4. **Gradient checkpointing is the key differentiator:** required on H200 (141 GB) for MBS≥2, unnecessary on B300 (288 GB) for MBS≤4 — saving ~20% recompute overhead. +4. **No gradient checkpointing needed on either cluster:** distributed optimizer shards Adam states across DP ranks, keeping H200 peak at ~114 GB (MBS=2) and B300 at ~173 GB (MBS=4). ## Hardware @@ -124,7 +124,7 @@ sbatch b300/slurm/run.sh ├── README.md ← You are here ├── h200/ │ ├── Dockerfile ← NeMo 25.07 + EFA container -│ ├── train.py ← Megatron-Core training script +│ ├── train.py ← Megatron-Bridge training script │ └── slurm/ │ └── run.sh ← Slurm submission script ├── b300/ @@ -138,4 +138,4 @@ sbatch b300/slurm/run.sh ## License -Apache 2.0 +MIT-0 diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile similarity index 100% rename from 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/Dockerfile rename to 3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh new file mode 100644 index 000000000..873607ebf --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +#SBATCH --job-name=qwen3-8b-b300 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=12 +#SBATCH --exclusive +#SBATCH --partition=b300 +#SBATCH --time=24:00:00 +#SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err +#SBATCH --container-image=/fsx/ubuntu/qwen3-8b/containers/nemo-efa-26.02.sqsh +#SBATCH --container-mounts=/fsx:/fsx + +# EFA / NCCL environment +export FI_PROVIDER=efa +export NCCL_SOCKET_IFNAME=^docker,lo,veth +export NCCL_DEBUG=WARN +export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH} +export TORCH_COMPILE_DISABLE=1 + +# Copy training script to expected location +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +mkdir -p /fsx/ubuntu/qwen3-8b/code +cp "${SCRIPT_DIR}/train.py" /fsx/ubuntu/qwen3-8b/code/train.py + +# Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init +srun --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ + python /fsx/ubuntu/qwen3-8b/code/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/train.py similarity index 64% rename from 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py rename to 3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/train.py index 5ba1422a1..f8763bedd 100644 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/train.py +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/train.py @@ -1,4 +1,6 @@ #!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 """Qwen3-8B Pre-Training on B300 — Megatron-Bridge (NeMo 26.02) Uses the Megatron-Bridge recipe API with config objects. @@ -18,7 +20,7 @@ def main(): cfg = qwen3_8b_pretrain_config( - mock=True, # Set to False and provide data_path for real training + mock=True, tensor_model_parallel_size=1, pipeline_model_parallel_size=1, micro_batch_size=4, @@ -29,14 +31,7 @@ def main(): lr_decay_iters=100, ) - # Training configuration - cfg.train.bf16 = True - cfg.train.use_distributed_optimizer = True - cfg.train.overlap_grad_reduce = True - cfg.train.overlap_param_gather = True - - # No gradient checkpointing (B300 has 288 GB — fits MBS=4 without recompute) - cfg.train.recompute_granularity = None + # No gradient checkpointing (B300 has 288 GB - fits MBS=4 without recompute) # Optimizer cfg.optimizer.lr = 3e-4 @@ -50,15 +45,9 @@ def main(): cfg.logger.log_interval = 5 cfg.train.eval_interval = 1000 cfg.train.eval_iters = 0 - cfg.train.dir = "/fsx/ubuntu/qwen3-8b-pretraining/checkpoints/b300" + cfg.train.dir = "/fsx/ubuntu/qwen3-8b/checkpoints/b300" cfg.train.save_interval = 1000 - # To use real data instead of mock: - # cfg.data.mock = False - # cfg.data.data_path = "/fsx/ubuntu/qwen3-8b-pretraining/datasets/c4/merged_text_document" - # cfg.data.tokenizer_type = "HuggingFaceTokenizer" - # cfg.data.tokenizer_model = "Qwen/Qwen3-8B" - pretrain(config=cfg, forward_step_func=forward_step) diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile similarity index 100% rename from 3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/Dockerfile rename to 3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh new file mode 100644 index 000000000..413b5586f --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +#SBATCH --job-name=qwen3-8b-h200 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=12 +#SBATCH --exclusive +#SBATCH --qos=admin_qos +#SBATCH --partition=p5en +#SBATCH --time=24:00:00 +#SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err +#SBATCH --container-image=/fsx/ubuntu/qwen3-8b/containers/nemo-efa-25.07.sqsh +#SBATCH --container-mounts=/fsx:/fsx + +# EFA / NCCL environment +export FI_PROVIDER=efa +export NCCL_SOCKET_IFNAME=^docker,lo,veth +export NCCL_DEBUG=WARN +export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so +export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH} +export TORCH_COMPILE_DISABLE=1 + +# Copy training script to expected location +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +mkdir -p /fsx/ubuntu/qwen3-8b/code +cp "${SCRIPT_DIR}/train.py" /fsx/ubuntu/qwen3-8b/code/train.py + +# Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init +srun --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ + python /fsx/ubuntu/qwen3-8b/code/train.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py new file mode 100644 index 000000000..033f63f50 --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Qwen3-8B Pre-Training on H200 — Megatron-Bridge (NeMo 25.07) + +Uses the Megatron-Bridge recipe API with config objects. +No gradient checkpointing needed (distributed optimizer keeps memory at ~114 GB). + +Best config: TP=1, PP=1, DP=16, MBS=2, GBS=128, seq=4096, BF16 +Result: 497 TFLOP/s/GPU, 162K tok/s on 16x H200 +""" +import os + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +from megatron.bridge.recipes.qwen.qwen3 import qwen3_8b_pretrain_config +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain + + +def main(): + cfg = qwen3_8b_pretrain_config( + mock=True, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + micro_batch_size=2, + global_batch_size=128, + seq_length=4096, + train_iters=100, + lr_warmup_iters=10, + lr_decay_iters=100, + ) + + # No gradient checkpointing - MBS=2 with distributed optimizer fits in 141 GB + # (uses ~114 GB peak, leaving headroom) + + # Optimizer + cfg.optimizer.lr = 3e-4 + cfg.optimizer.min_lr = 3e-5 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.clip_grad = 1.0 + + # Logging and checkpoints + cfg.logger.log_interval = 5 + cfg.train.eval_interval = 1000 + cfg.train.eval_iters = 0 + cfg.train.dir = "/fsx/ubuntu/qwen3-8b/checkpoints/h200" + cfg.train.save_interval = 1000 + + pretrain(config=cfg, forward_step_func=forward_step) + + +if __name__ == "__main__": + main() diff --git a/3.test_cases/megatron/nemo/slurm/run.py b/3.test_cases/megatron/nemo/slurm/run.py index e6160ef9f..8554e85ac 100644 --- a/3.test_cases/megatron/nemo/slurm/run.py +++ b/3.test_cases/megatron/nemo/slurm/run.py @@ -1,3 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + import nemo_run as run import json import argparse diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh deleted file mode 100644 index 9b15c28ce..000000000 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/b300/slurm/run.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=qwen3-8b-b300 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=8 -#SBATCH --gpus-per-node=8 -#SBATCH --cpus-per-task=12 -#SBATCH --exclusive -#SBATCH --partition=b300 -#SBATCH --time=24:00:00 -#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out -#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err -#SBATCH --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh -#SBATCH --container-mounts=/fsx:/fsx - -# Resolve head node -export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n1) -export MASTER_PORT=29500 - -# EFA / NCCL environment -export FI_PROVIDER=efa -export NCCL_SOCKET_IFNAME=^docker,lo,veth -export NCCL_DEBUG=WARN -export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so -export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/nccl/build/lib:$LD_LIBRARY_PATH - -# Disable torch.compile -export TORCH_COMPILE_DISABLE=1 - -# Launch training -srun --container-env=MASTER_ADDR,MASTER_PORT,FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ - torchrun \ - --nnodes=${SLURM_NNODES} \ - --nproc-per-node=8 \ - --rdzv-id=${SLURM_JOB_ID} \ - --rdzv-backend=c10d \ - --rdzv-endpoint=${MASTER_ADDR}:${MASTER_PORT} \ - /fsx/ubuntu/qwen3-8b-pretraining/code/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md deleted file mode 100644 index 28c02cbc6..000000000 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/docs/lessons-learned.md +++ /dev/null @@ -1,126 +0,0 @@ -# Lessons Learned - -Hard-won knowledge from optimizing Qwen3-8B pre-training on H200 and B300 clusters. - ---- - -## EFA Silent Fallback to TCP - -**Symptom:** Multi-node training runs but at single-node throughput. NCCL reports no errors. - -**Root cause:** NCCL silently falls back to TCP sockets when the OFI plugin isn't loaded. - -**Fix — all three are required:** - -```bash -export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:$LD_LIBRARY_PATH -export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so -export FI_PROVIDER=efa -``` - -**Verification:** Look for `NCCL INFO NET/OFI` in logs (not `NET/Socket`). - -The OFI NCCL plugin directory may use either `/opt/amazon/ofi-nccl/lib/` or `/opt/amazon/aws-ofi-nccl/lib/` depending on the EFA installer version. Check which exists. - ---- - -## Enroot Import Workflow - -**Never use `mksquashfs` directly.** It produces images that PyXis can't launch. - -**Correct workflow:** - -```bash -# 1. Build with Docker -sudo docker build -t my-image:latest . - -# 2. Import with enroot (requires sudo for Docker socket) -sudo TMPDIR=/fsx/ubuntu/qwen3-8b-pretraining/tmp \ - ENROOT_TEMP_PATH=/fsx/ubuntu/qwen3-8b-pretraining/tmp \ - enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/image.sqsh dockerd://my-image:latest - -# 3. Fix permissions -sudo chown $USER:$USER /fsx/ubuntu/qwen3-8b-pretraining/containers/image.sqsh -``` - -**Three requirements:** -1. `sudo` — Docker socket is `root:docker`, user not in docker group -2. `TMPDIR` on FSx — NeMo containers are 30+ GB, `/tmp` won't fit -3. `docker buildx use default` — image must be in main containerd store - ---- - -## Slurm/PyXis Gotchas - -### Slurm NOT in PATH -On some clusters, Slurm binaries live at `/opt/slurm/bin/`. Use full paths if needed: -```bash -/opt/slurm/bin/sbatch script.sh -/opt/slurm/bin/squeue -u ubuntu -/opt/slurm/bin/scontrol show hostname $SLURM_NODELIST -``` - -### Shell vars don't pass into containers -PyXis containers don't inherit the calling shell's environment. Pass explicitly: -```bash -srun --container-env=VAR1,VAR2,VAR3 ... -``` - -### Resolve MASTER_ADDR before srun -`scontrol` is not available inside PyXis containers. Compute the head node in the batch script, before the `srun` call: -```bash -export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n1) -``` - -### ntasks-per-node for torchrun -When using `torchrun` (which spawns GPU workers itself), set `--ntasks-per-node=1` in the Slurm script. If using raw `python` with NCCL init, use `--ntasks-per-node=8`. - -### Single-node: disable EFA -For intra-node-only jobs, `FI_PROVIDER=efa` causes NCCL failures. Remove it or set `FI_PROVIDER=shm` for single-node debugging. - ---- - -## torch.compile Incompatibility - -Set `TORCH_COMPILE_DISABLE=1` in all environments. It fails in every configuration tested (DeepSpeed, HuggingFace multi-node, NeMo 25.07, NeMo 26.02). The performance gain would be minimal since Transformer Engine already provides fused kernels. - ---- - -## Distributed Optimizer Trap at DP=1 - -`--use-distributed-optimizer` with `DP=1` (single GPU or TP-only parallelism) causes a crash. The sharding logic divides by DP world size and expects DP>=2. - -**Rule:** Only enable distributed optimizer when DP>=2. For single-GPU debugging, remove the flag. - ---- - -## Megatron-Bridge API Gotchas (NeMo 26.02) - -- **Logging interval** is on `cfg.logger.log_interval`, not `cfg.train.log_interval` (silently ignored) -- **Disable gradient checkpointing** with `cfg.train.recompute_granularity = None` (not `""` or `False`) -- **Checkpoint directory** is `cfg.train.dir` (not `cfg.train.save` or `cfg.train.checkpoint_dir`) -- **Qwen3 bridge recipe** `qwen3_8b_pretrain_config()` provides correct model dimensions — don't manually override - ---- - -## Memory Budget: H200 vs B300 - -``` -H200 (141 GB available): - Model (BF16): 16 GB - Gradients (BF16): 16 GB - Optimizer (sharded/16): 3 GB - Activations (recompute): ~100 GB <- with full recompute, MBS=2 - Overhead: 3 GB - Total: ~138 GB - -B300 (288 GB available): - Model (BF16): 16 GB - Gradients (BF16): 16 GB - Optimizer (sharded/16): 3 GB - Activations (no recomp): ~135 GB <- NO recompute needed, MBS=4 - Overhead: 3 GB - Total: ~173 GB (115 GB headroom) -``` - -B300's extra memory means no recompute overhead -> ~20% fewer FLOPs per step -> directly translates to 1.96x throughput combined with higher peak FLOPS. diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh deleted file mode 100644 index 2b2efda38..000000000 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/slurm/run.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=qwen3-8b-h200 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=8 -#SBATCH --gpus-per-node=8 -#SBATCH --cpus-per-task=12 -#SBATCH --exclusive -#SBATCH --partition=p5en -#SBATCH --time=24:00:00 -#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out -#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err -#SBATCH --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh -#SBATCH --container-mounts=/fsx:/fsx - -# Resolve head node BEFORE srun (scontrol not available inside container) -export MASTER_ADDR=$(/opt/slurm/bin/scontrol show hostname $SLURM_NODELIST | head -n1) -export MASTER_PORT=29500 - -# EFA / NCCL environment -export FI_PROVIDER=efa -export NCCL_SOCKET_IFNAME=^docker,lo,veth -export NCCL_DEBUG=WARN -export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so -export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/nccl/build/lib:$LD_LIBRARY_PATH - -# Disable torch.compile (incompatible with this stack) -export TORCH_COMPILE_DISABLE=1 - -# Launch training -srun --container-env=MASTER_ADDR,MASTER_PORT,FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ - torchrun \ - --nnodes=${SLURM_NNODES} \ - --nproc-per-node=8 \ - --rdzv-id=${SLURM_JOB_ID} \ - --rdzv-backend=c10d \ - --rdzv-endpoint=${MASTER_ADDR}:${MASTER_PORT} \ - /fsx/ubuntu/qwen3-8b-pretraining/code/train.py diff --git a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py b/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py deleted file mode 100644 index f8ff8cd15..000000000 --- a/3.test_cases/pytorch/nemo-megatron/qwen3-8b-pretraining/h200/train.py +++ /dev/null @@ -1,165 +0,0 @@ -#!/usr/bin/env python3 -"""Qwen3-8B Pre-Training on H200 — Megatron-Core (NeMo 25.07) - -Uses the Megatron pretrain API with Transformer Engine spec. -Qwen3-8B dimensions mapped to GPT model provider (Qwen3 bridge not registered). - -Best config: TP=1, PP=1, DP=16, MBS=2, GBS=128, seq=4096, BF16 -Result: 497 TFLOP/s/GPU, 162K tok/s on 16x H200 -""" -import os -import sys - -os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") -os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1") - -from functools import partial - -import torch -from megatron.training import get_args, pretrain -from megatron.training.arguments import core_transformer_config_from_args -from megatron.core.enums import ModelType -from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_layer_with_transformer_engine_spec, -) -from megatron.core.transformer.spec_utils import import_module -from megatron.core.datasets.blended_megatron_dataset_builder import ( - BlendedMegatronDatasetBuilder, -) -from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset - - -def model_provider(pre_process=True, post_process=True): - """Build GPT model with Qwen3-8B dimensions.""" - args = get_args() - config = core_transformer_config_from_args(args) - - transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec() - - model = GPTModel( - config=config, - transformer_layer_spec=transformer_layer_spec, - vocab_size=args.padded_vocab_size, - max_sequence_length=args.max_position_embeddings, - pre_process=pre_process, - post_process=post_process, - parallel_output=True, - ) - return model - - -def forward_step(data_iterator, model): - """Forward pass — standard GPT causal LM loss.""" - args = get_args() - tokens, labels, loss_mask, attention_mask, position_ids = _get_batch(data_iterator) - output_tensor = model(tokens, position_ids, attention_mask, labels=labels) - return output_tensor, partial(_loss_func, loss_mask) - - -def _loss_func(loss_mask, output_tensor): - """Compute averaged cross-entropy loss.""" - losses = output_tensor.float() - loss_mask = loss_mask.view(-1).float() - loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() - return loss, {"lm loss": loss} - - -def _get_batch(data_iterator): - """Get batch from data iterator.""" - args = get_args() - data = next(data_iterator) - - tokens = data["tokens"].long().cuda() - labels = data["labels"].long().cuda() - loss_mask = data["loss_mask"].float().cuda() - attention_mask = data["attention_mask"].long().cuda() if "attention_mask" in data else None - position_ids = data["position_ids"].long().cuda() - - return tokens, labels, loss_mask, attention_mask, position_ids - - -def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build mock datasets for benchmarking. - - To switch to real data, replace MockGPTDataset with GPTDataset and provide: - --data-path /fsx/ubuntu/qwen3-8b-pretraining/datasets/c4/merged_text_document - --tokenizer-type HuggingFaceTokenizer - --tokenizer-model Qwen/Qwen3-8B - """ - args = get_args() - - config = GPTDatasetConfig( - random_seed=args.seed, - sequence_length=args.seq_length, - reset_position_ids=False, - reset_attention_mask=False, - eod_mask_loss=False, - mock=True, - mock_seq_length=args.seq_length, - ) - - dataset_builder = BlendedMegatronDatasetBuilder( - MockGPTDataset, train_val_test_num_samples, lambda: True, config - ) - train_ds, valid_ds, test_ds = dataset_builder.build() - return train_ds, valid_ds, test_ds - - -if __name__ == "__main__": - pretrain( - train_valid_test_datasets_provider, - model_provider, - ModelType.encoder_or_decoder, - forward_step, - args_defaults={ - # Qwen3-8B architecture - "num_layers": 36, - "hidden_size": 4096, - "num_attention_heads": 32, - "group_query_attention": True, - "num_query_groups": 8, - "ffn_hidden_size": 14336, - "swiglu": True, - "max_position_embeddings": 4096, - "seq_length": 4096, - "padded_vocab_size": 151936, - "use_rotary_position_embeddings": True, - "rotary_percent": 1.0, - "normalization": "RMSNorm", - "untie_embeddings_and_output_weights": True, - # Training - "micro_batch_size": 2, - "global_batch_size": 128, - "train_iters": 100, - "lr": 3e-4, - "min_lr": 3e-5, - "lr_warmup_iters": 10, - "lr_decay_iters": 100, - "lr_decay_style": "cosine", - "weight_decay": 0.1, - "adam_beta1": 0.9, - "adam_beta2": 0.95, - "clip_grad": 1.0, - "bf16": True, - # Parallelism - "tensor_model_parallel_size": 1, - "pipeline_model_parallel_size": 1, - "use_distributed_optimizer": True, - "overlap_grad_reduce": True, - "overlap_param_gather": True, - # Gradient checkpointing (mandatory on H200 for MBS>=2) - "recompute_granularity": "full", - "recompute_method": "uniform", - "recompute_num_layers": 1, - # Logging - "log_interval": 5, - "eval_interval": 1000, - "eval_iters": 0, - "tensorboard_dir": "/fsx/ubuntu/qwen3-8b-pretraining/tensorboard/h200", - "save": "/fsx/ubuntu/qwen3-8b-pretraining/checkpoints/h200", - "save_interval": 1000, - "tokenizer_type": "NullTokenizer", - "vocab_size": 151936, - }, - ) From 11dca562de83ec1dda4799c9e2077d03f6d61133 Mon Sep 17 00:00:00 2001 From: Paulo Aragao <5193870+paragao@users.noreply.github.com> Date: Fri, 19 Jun 2026 08:32:19 +0100 Subject: [PATCH 4/8] Update 3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md Co-authored-by: Keita Watanabe --- 3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md index 34667c97b..f7a04586d 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md @@ -59,7 +59,7 @@ sbatch b300/slurm/run.sh | Hidden dim (d_model) | 4096 | | Q-heads | 32 | | KV-heads | 8 (GQA) | -| FFN dim | 14336 (SwiGLU) | +| FFN dim | 12288 (SwiGLU) | | Vocab size | 151,936 | | Positional encoding | RoPE | | Normalization | RMSNorm | From a1ef33c6973717d9d71b3fad8742a64a1d15f400 Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Sun, 21 Jun 2026 20:06:37 +0000 Subject: [PATCH 5/8] single Dockerfile for both clusters. Single preprocessing script. Updated training scripts --- .../nemo/qwen3-8b-pretraining/Dockerfile | 23 +++ .../nemo/qwen3-8b-pretraining/README.md | 79 +++++--- .../nemo/qwen3-8b-pretraining/b300/Dockerfile | 39 ---- .../nemo/qwen3-8b-pretraining/h200/Dockerfile | 39 ---- .../nemo/qwen3-8b-pretraining/preprocess.py | 179 ++++++++++++++++++ 5 files changed, 256 insertions(+), 103 deletions(-) create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/Dockerfile delete mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile delete mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/Dockerfile b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/Dockerfile new file mode 100644 index 000000000..ba0632785 --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/Dockerfile @@ -0,0 +1,23 @@ +FROM nvcr.io/nvidia/nemo:26.04 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + libevent-core-2.1-7 libevent-pthreads-2.1-7 \ + ethtool iproute2 pciutils curl \ + && cd /tmp \ + && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-1.47.0.tar.gz \ + && tar -xf aws-efa-installer-1.47.0.tar.gz \ + && cd aws-efa-installer \ + && ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify \ + && cd /tmp \ + && git clone -b v2.5.1 --depth 1 https://github.com/NVIDIA/gdrcopy.git \ + && cd gdrcopy && make -j$(nproc) lib lib_install \ + && cd / && rm -rf /tmp/* /var/lib/apt/lists/* + +# Environment +ENV LD_LIBRARY_PATH="/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH}" +ENV NCCL_TUNER_PLUGIN="/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so" +ENV FI_PROVIDER=efa +ENV TORCH_COMPILE_DISABLE=1 +ENV NCCL_PROTO=simple + +WORKDIR /workspace diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md index f7a04586d..2b829c071 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md @@ -27,29 +27,60 @@ Both clusters are compute-saturated with perfect communication overlap. AllReduc ## Quick Start > **Disk space:** The container build requires ~50 GB of disk space in TMPDIR. -> `enroot import` needs `sudo` and TMPDIR pointing to FSx (not `/tmp`, which is too small). +> `enroot import` needs `sudo` and TMPDIR pointing to FSx or another file system (not `/tmp`, which is too small). -### H200 Cluster (p5en.48xlarge) +### Clone this repo and change it its directory +```bash +git clone https://github.com/awslabs/awsome-distributed-ai.git +cd awsome-distribued-ai +``` + +### Build the container +```bash +# Build container +docker build -t qwen3-8b-h200:latest . +# Setup directories to run +mkdir -p /fsx/tmp && mkdir -p /fsx/ubuntu/qwen3-8b-pretraining/containers/ + +# Create the squash file with Enroot +sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.04.sqsh dockerd://qwen3-8b-h200:latest +``` + +### Prepare datasets ```bash -# 1. Build container -cd h200/ && docker build -t qwen3-8b-h200:latest . -sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-25.07.sqsh dockerd://qwen3-8b-h200:latest +# export your Hugging Face token, if you have one +export HF_TOKEN= + +# Prepare the dataset +python preprocess.py +``` +Without your Hugging Face token, the download will be throttled. +Datasets are transformed into binary mmap accessible files to avoid streaming data. + +### H200 Cluster (2x p5en.48xlarge) + +```bash +# 1. Change to directory +cd h200 # 2. Submit training job -sbatch h200/slurm/run.sh +sbatch slurm/run.sh ``` +Logs will be written to `/fsx/ubuntu/qwen3-8b-pretraining/logs`. +Checkpoints are saved to `/fsx/ubuntu/qwen3-8b-pretraining/checkpoints`. ### B300 Cluster (p6-b300.48xlarge) ```bash -# 1. Build container -cd b300/ && docker build -t qwen3-8b-b300:latest . -sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.02.sqsh dockerd://qwen3-8b-b300:latest +# 1. Change to directory +cd b300 # 2. Submit training job -sbatch b300/slurm/run.sh +sbatch slurm/run.sh ``` +Logs will be written to `/fsx/ubuntu/qwen3-8b-pretraining/logs`. +Checkpoints are saved to `/fsx/ubuntu/qwen3-8b-pretraining/checkpoints`. ## Model Architecture: Qwen3-8B @@ -88,24 +119,25 @@ sbatch b300/slurm/run.sh |-----------|---------------------|------------------------| | GPUs | 16× H200 (141 GB HBM3) | 16× B300 (288 GB HBM3e) | | Parallelism | TP=1, PP=1, DP=16 | TP=1, PP=1, DP=16 | -| Micro-batch size | 2 | 4 | +| **Micro-batch size** | **2** | **4** | | Global batch size | 128 (grad_accum=4) | 128 (grad_accum=2) | | Sequence length | 4096 | 4096 | | Precision | BF16 | BF16 | -| Gradient checkpointing | None | None | +| **Gradient checkpointing** | Selective (core_attn only) | Selective (core_attn only) | | Distributed optimizer | Yes (sharded Adam) | Yes (sharded Adam) | | Overlap grad reduce | Yes | Yes | -| Framework | Megatron-Bridge (NeMo 25.07) | Megatron-Bridge (NeMo 26.02) | +| Overlap param gather| Yes | Yes | +| Framework | Megatron-Bridge (NeMo 26.04) | Megatron-Bridge (NeMo 26.04) | ## Key Findings 1. **Both clusters are compute-saturated with perfect communication overlap.** AllReduce and AllGather are fully hidden behind compute — verified by single-GPU benchmarks showing lower TFLOP/s due to reduced batch arithmetic intensity. -2. **Both clusters use the Megatron-Bridge recipe API.** NeMo 25.07 for H200; NeMo 26.02 for B300. Each container maximizes hardware utilization for its target architecture. +2. **Both clusters use the Megatron-Bridge recipe API.** NeMo 26.04 for both H200 and B300. 3. **Pure data parallelism is optimal** when the model fits in single-GPU memory. Distributed optimizer + overlapped grad reduce eliminate the memory penalty. -4. **No gradient checkpointing needed on either cluster:** distributed optimizer shards Adam states across DP ranks, keeping H200 peak at ~114 GB (MBS=2) and B300 at ~173 GB (MBS=4). +4. **Selective gradient checkpointing used on both clusters:** lightweight core_attn recompute is Megatron-Core's standard behavior, keeping H200 peak at ~114 GB (MBS=2) and B300 at ~173 GB (MBS=4). ## Hardware @@ -115,25 +147,22 @@ sbatch b300/slurm/run.sh | Nodes | 2 | 2 | | GPUs per node | 8× H200 | 8× B300 | | GPU Memory | 141 GB HBM3 | 288 GB HBM3e | -| Interconnect | EFA GDRDMA (3200 Gbps) | EFA GDRDMA (3200 Gbps) | -| Intra-node | NVLink | NVLink | +| Interconnect | EFA GDRDMA (3200 Gbps) | EFA GDRDMA (6400 Gbps) | +| Intra-node | NVLink (900 GB/s) | NVLink (1800 GB/s) | ## Project Structure ``` ├── README.md ← You are here +│ Dockerfile ← NeMo 26.04 + EFA container ├── h200/ -│ ├── Dockerfile ← NeMo 25.07 + EFA container -│ ├── train.py ← Megatron-Bridge training script -│ └── slurm/ -│ └── run.sh ← Slurm submission script -├── b300/ -│ ├── Dockerfile ← NeMo 26.02 + EFA container │ ├── train.py ← Megatron-Bridge training script │ └── slurm/ │ └── run.sh ← Slurm submission script -└── docs/ - └── lessons-learned.md ← EFA, PyXis, and Megatron gotchas +└── b300/ + ├── train.py ← Megatron-Bridge training script + └── slurm/ + └── run.sh ← Slurm submission script ``` ## License diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile deleted file mode 100644 index 1af80364b..000000000 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/Dockerfile +++ /dev/null @@ -1,39 +0,0 @@ -FROM nvcr.io/nvidia/nemo:26.02 - -# Install EFA prerequisites -RUN apt-get update && apt-get install -y --no-install-recommends \ - environment-modules \ - tcl \ - udev \ - ethtool \ - iproute2 \ - dmidecode \ - libevent-core-2.1-7t64 \ - libevent-pthreads-2.1-7t64 \ - curl \ - && rm -rf /var/lib/apt/lists/* - -# Install EFA (v1.47.0 - compatible with NeMo 26.02 base) -# Do NOT remove existing NCCL/hpcx - EFA installs on top -RUN cd /tmp && \ - curl -O https://efa-installer.amazonaws.com/aws-efa-installer-1.47.0.tar.gz && \ - tar -xf aws-efa-installer-1.47.0.tar.gz && \ - cd aws-efa-installer && \ - ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify && \ - cd / && rm -rf /tmp/aws-efa-installer* - -# GDRCopy for GPU-direct RDMA -RUN cd /tmp && \ - git clone -b v2.4.4 https://github.com/NVIDIA/gdrcopy.git && \ - cd gdrcopy && \ - make -j$(nproc) lib lib_install && \ - cd / && rm -rf /tmp/gdrcopy - -# Environment -ENV LD_LIBRARY_PATH="/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH}" -ENV NCCL_TUNER_PLUGIN="/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so" -ENV FI_PROVIDER=efa -ENV TORCH_COMPILE_DISABLE=1 -ENV NCCL_PROTO=simple - -WORKDIR /workspace diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile deleted file mode 100644 index 953639199..000000000 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/Dockerfile +++ /dev/null @@ -1,39 +0,0 @@ -FROM nvcr.io/nvidia/nemo:25.07 - -# Install EFA prerequisites -RUN apt-get update && apt-get install -y --no-install-recommends \ - environment-modules \ - tcl \ - udev \ - ethtool \ - iproute2 \ - dmidecode \ - libevent-core-2.1-7 \ - libevent-pthreads-2.1-7 \ - curl \ - && rm -rf /var/lib/apt/lists/* - -# Install EFA (v1.47.0 - compatible with NeMo 25.07 base) -# Do NOT remove existing NCCL/hpcx - EFA installs on top -RUN cd /tmp && \ - curl -O https://efa-installer.amazonaws.com/aws-efa-installer-1.47.0.tar.gz && \ - tar -xf aws-efa-installer-1.47.0.tar.gz && \ - cd aws-efa-installer && \ - ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify && \ - cd / && rm -rf /tmp/aws-efa-installer* - -# GDRCopy for GPU-direct RDMA -RUN cd /tmp && \ - git clone -b v2.4.4 https://github.com/NVIDIA/gdrcopy.git && \ - cd gdrcopy && \ - make -j$(nproc) lib lib_install && \ - cd / && rm -rf /tmp/gdrcopy - -# Environment -ENV LD_LIBRARY_PATH="/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH}" -ENV NCCL_TUNER_PLUGIN="/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so" -ENV FI_PROVIDER=efa -ENV TORCH_COMPILE_DISABLE=1 -ENV NCCL_PROTO=simple - -WORKDIR /workspace diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py new file mode 100644 index 000000000..083e3b76c --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Download allenai/c4 and convert to Megatron indexed format (.bin + .idx). + +Fully parallelized pipeline: + Phase 1: Download c4/en parquet shards in parallel (requires HF_TOKEN) + Phase 2: Tokenize in parallel from local Arrow cache using all CPUs + Phase 3: Merge into Megatron MMapIndexedDataset format + +Requirements: + pip install numpy transformers datasets + export HF_TOKEN= + +Usage: + python preprocess.py --output-prefix /fsx/ubuntu/qwen3-8b-pretraining/datasets/c4_qwen3_8b \ + --num-tokens 1000000000 --workers 96 +""" +import argparse +import os +import struct +import sys +import time +from multiprocessing import Pool, cpu_count + +import numpy as np +from datasets import load_dataset +from transformers import AutoTokenizer + +# Megatron MMapIndexedDataset constants +_HDR_MAGIC = b"MMIDIDX\x00\x00" +_DTYPE = np.int32 +_DTYPE_CODE = 4 + + +def parse_args(): + p = argparse.ArgumentParser(description="Download allenai/c4 and convert to Megatron indexed format") + p.add_argument("--output-prefix", default="/fsx/ubuntu/qwen3-8b-pretraining/datasets/c4_qwen3_8b", + help="Output path prefix (creates .bin and .idx)") + p.add_argument("--tokenizer", default="Qwen/Qwen3-8B", help="HuggingFace tokenizer name") + p.add_argument("--num-tokens", type=int, default=1_000_000_000, help="Target number of tokens") + p.add_argument("--workers", type=int, default=min(96, cpu_count()), help="Parallel workers") + p.add_argument("--cache-dir", default="/fsx/ubuntu/qwen3-8b-pretraining/cache/c4", + help="HuggingFace datasets cache directory") + return p.parse_args() + + +def check_hf_token(): + """Verify HF_TOKEN is set. Exit with error if not.""" + token = os.environ.get("HF_TOKEN") + if not token: + print("ERROR: HF_TOKEN environment variable is not set.", file=sys.stderr) + print(" Set it with: export HF_TOKEN=", file=sys.stderr) + print(" Get a token at: https://huggingface.co/settings/tokens", file=sys.stderr) + sys.exit(1) + return token + + +def write_idx_file(idx_path, sizes, doc_idx): + """Write a Megatron .idx file.""" + with open(idx_path, "wb") as f: + f.write(_HDR_MAGIC) + f.write(struct.pack(" 0: + pointers[1:] = np.cumsum(sizes_arr[:-1].astype(np.int64)) * _DTYPE().itemsize + pointers.tofile(f) + np.array(doc_idx, dtype=np.int64).tofile(f) + + +def tokenize_chunk(args_tuple): + """Tokenize a chunk of documents. Returns (tmp_bin_path, sizes, doc_idx, total_tokens).""" + chunk_id, texts, tokenizer_name, output_dir = args_tuple + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) + tmp_bin = os.path.join(output_dir, f"_tmp_chunk_{chunk_id:05d}.bin") + sizes = [] + doc_idx = [0] + total_tokens = 0 + + with open(tmp_bin, "wb") as out: + for text in texts: + tokens = tokenizer.encode(text, add_special_tokens=False) + if not tokens: + continue + arr = np.array(tokens, dtype=_DTYPE) + out.write(arr.tobytes()) + sizes.append(len(tokens)) + doc_idx.append(len(sizes)) + total_tokens += len(tokens) + + return tmp_bin, sizes, doc_idx, total_tokens + + +def main(): + args = parse_args() + token = check_hf_token() + output_dir = os.path.dirname(args.output_prefix) + os.makedirs(output_dir, exist_ok=True) + + print(f"=== allenai/c4 -> Megatron indexed format ===") + print(f" Tokenizer: {args.tokenizer}") + print(f" Target: {args.num_tokens / 1e9:.1f}B tokens") + print(f" Workers: {args.workers}") + print(f" Output: {args.output_prefix}.{{bin,idx}}") + t0 = time.time() + + # Phase 1: Download in parallel using datasets library + # ~200 tokens/doc average for c4, with 20% buffer + est_docs = int(args.num_tokens / 200 * 1.2) + print(f"\nPhase 1: Download c4/en (~{est_docs / 1e6:.0f}M docs) with {args.workers} workers...") + t1 = time.time() + + dataset = load_dataset( + "allenai/c4", "en", + split=f"train[:{est_docs}]", + cache_dir=args.cache_dir, + num_proc=args.workers, + token=token, + ) + print(f" Downloaded {len(dataset):,} docs in {time.time() - t1:.0f}s") + + # Phase 2: Tokenize in parallel + print(f"\nPhase 2: Tokenize with {args.workers} workers...") + t2 = time.time() + + # Split dataset into chunks for parallel processing + texts = dataset["text"] + chunk_size = max(1, len(texts) // args.workers) + chunks = [texts[i:i + chunk_size] for i in range(0, len(texts), chunk_size)] + + worker_args = [ + (i, chunk, args.tokenizer, output_dir) + for i, chunk in enumerate(chunks) + ] + + with Pool(args.workers) as pool: + results = pool.map(tokenize_chunk, worker_args) + print(f" Tokenization done in {time.time() - t2:.0f}s") + + # Phase 3: Merge into single .bin + .idx + print("\nPhase 3: Merge...") + bin_path = args.output_prefix + ".bin" + idx_path = args.output_prefix + ".idx" + all_sizes = [] + all_doc_idx = [0] + total_tokens = 0 + + with open(bin_path, "wb") as out: + for tmp_bin, sizes, doc_idx, n_tokens in results: + if total_tokens >= args.num_tokens: + os.remove(tmp_bin) + continue + with open(tmp_bin, "rb") as f: + while chunk := f.read(128 * 1024 * 1024): + out.write(chunk) + os.remove(tmp_bin) + offset = len(all_sizes) + all_sizes.extend(sizes) + all_doc_idx.extend(idx + offset for idx in doc_idx[1:]) + total_tokens += n_tokens + + write_idx_file(idx_path, all_sizes, all_doc_idx) + + elapsed = time.time() - t0 + print(f"\n=== Done! ===") + print(f" {total_tokens / 1e9:.2f}B tokens, {len(all_sizes):,} documents") + print(f" Total time: {elapsed / 60:.1f} min") + print(f" {bin_path} ({os.path.getsize(bin_path) / 1e9:.1f} GB)") + print(f" {idx_path} ({os.path.getsize(idx_path) / 1e6:.1f} MB)") + + +if __name__ == "__main__": + main() From b742e06f177b100361345c436627bb4748ce29eb Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Sun, 21 Jun 2026 21:38:34 +0000 Subject: [PATCH 6/8] improved preprocessing --- .../{ => preprocessing}/preprocess.py | 0 .../preprocessing/preprocess.sh | 49 +++++++++++++++++++ .../preprocessing/requirements.txt | 3 ++ 3 files changed, 52 insertions(+) rename 3.test_cases/megatron/nemo/qwen3-8b-pretraining/{ => preprocessing}/preprocess.py (100%) create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh create mode 100644 3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/requirements.txt diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py similarity index 100% rename from 3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocess.py rename to 3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh new file mode 100644 index 000000000..aa57efce6 --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh @@ -0,0 +1,49 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +#SBATCH --job-name=preprocess-c4 +#SBATCH --partition=p5en +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=192 +#SBATCH --mem=0 +#SBATCH --time=02:00:00 +#SBATCH --exclusive +#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/preprocess-%j.out +#SBATCH --export=ALL + +# --- HF_TOKEN must be set before submitting --- +# export HF_TOKEN= +# sbatch preprocess.sh +if [ -z "$HF_TOKEN" ]; then + echo "ERROR: HF_TOKEN is not set. Export it before submitting:" + echo " export HF_TOKEN= && sbatch preprocess.sh" + exit 1 +fi + +export HF_HOME="/fsx/ubuntu/.cache/huggingface" + +mkdir -p /fsx/ubuntu/qwen3-8b-pretraining/logs +mkdir -p /fsx/ubuntu/qwen3-8b-pretraining/datasets + +# Create a virtual environment for the preprocessing +python3 -m venv /fsx/ubuntu/qwen3-8b-pretraining/venv +source /fsx/ubuntu/qwen3-8b-pretraining/venv/bin/activate + +PYTHON=/fsx/ubuntu/qwen3-8b-pretraining/venv/bin/python +SCRIPT_DIR=$SLURM_SUBMIT_DIR +SCRIPT="${SCRIPT_DIR}/preprocess.py" + +pip install -r $SCRIPT_DIR/requirements.txt + +echo "=== C4 Preprocessing ===" +echo "Node: $(hostname) | CPUs: $(nproc) | Start: $(date)" + +$PYTHON $SCRIPT \ + --output-prefix /fsx/ubuntu/qwen3-8b-pretraining/datasets/c4_qwen3_8b \ + --tokenizer Qwen/Qwen3-8B \ + --num-tokens 1000000000 \ + --workers $(nproc) \ + --cache-dir /fsx/ubuntu/qwen3-8b-pretraining/cache/c4 + +echo "Finished: $(date)" diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/requirements.txt b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/requirements.txt new file mode 100644 index 000000000..cf42fd2e2 --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/requirements.txt @@ -0,0 +1,3 @@ +numpy>=1.24.0 +transformers>=4.40.0 +datasets>=2.19.0 From 64b509f7d0e91130a1127af0660673f762367438 Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Sun, 21 Jun 2026 21:59:49 +0000 Subject: [PATCH 7/8] cosmetics changes, moved files to proper location, fixed scripts to address changes --- .../nemo/qwen3-8b-pretraining/README.md | 28 +++++----- .../qwen3-8b-pretraining/b300/slurm/run.sh | 13 +---- .../qwen3-8b-pretraining/h200/slurm/run.sh | 14 +---- .../nemo/qwen3-8b-pretraining/h200/train.py | 51 ++++++++++++------- .../preprocessing/preprocess.sh | 2 +- 5 files changed, 51 insertions(+), 57 deletions(-) mode change 100644 => 100755 3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md index 2b829c071..eeec1b43f 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md @@ -35,28 +35,28 @@ git clone https://github.com/awslabs/awsome-distributed-ai.git cd awsome-distribued-ai ``` +### Prepare datasets (allenai/c4/en) +```bash +# export your Hugging Face token, if you have one +export HF_TOKEN= + +# Prepare the dataset +sbatch preprocessing/preprocess.sh +``` +Without your Hugging Face token, the download will be throttled. The script requires a token. +Datasets are tokenized and transformed into binary mmap accessible files to avoid streaming data (`.idx` and `.bin` files). + ### Build the container ```bash # Build container -docker build -t qwen3-8b-h200:latest . +docker build -t qwen3-8b-pretraining:latest . # Setup directories to run mkdir -p /fsx/tmp && mkdir -p /fsx/ubuntu/qwen3-8b-pretraining/containers/ # Create the squash file with Enroot -sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.04.sqsh dockerd://qwen3-8b-h200:latest -``` - -### Prepare datasets -```bash -# export your Hugging Face token, if you have one -export HF_TOKEN= - -# Prepare the dataset -python preprocess.py +sudo TMPDIR=/fsx/tmp ENROOT_TEMP_PATH=/fsx/tmp enroot import --output /fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.04.sqsh dockerd://qwen3-8b-pretraining:latest ``` -Without your Hugging Face token, the download will be throttled. -Datasets are transformed into binary mmap accessible files to avoid streaming data. ### H200 Cluster (2x p5en.48xlarge) @@ -70,7 +70,7 @@ sbatch slurm/run.sh Logs will be written to `/fsx/ubuntu/qwen3-8b-pretraining/logs`. Checkpoints are saved to `/fsx/ubuntu/qwen3-8b-pretraining/checkpoints`. -### B300 Cluster (p6-b300.48xlarge) +### B300 Cluster (2x p6-b300.48xlarge) ```bash # 1. Change to directory diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh index 873607ebf..ec7ac6ad0 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh @@ -7,12 +7,8 @@ #SBATCH --gpus-per-node=8 #SBATCH --cpus-per-task=12 #SBATCH --exclusive -#SBATCH --partition=b300 -#SBATCH --time=24:00:00 #SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out #SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err -#SBATCH --container-image=/fsx/ubuntu/qwen3-8b/containers/nemo-efa-26.02.sqsh -#SBATCH --container-mounts=/fsx:/fsx # EFA / NCCL environment export FI_PROVIDER=efa @@ -21,12 +17,7 @@ export NCCL_DEBUG=WARN export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH} export TORCH_COMPILE_DISABLE=1 - -# Copy training script to expected location -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -mkdir -p /fsx/ubuntu/qwen3-8b/code -cp "${SCRIPT_DIR}/train.py" /fsx/ubuntu/qwen3-8b/code/train.py +export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" # Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init -srun --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ - python /fsx/ubuntu/qwen3-8b/code/train.py +/opt/slurm/bin/srun --mpi=pmix --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.04.sqsh --container-mounts=/fsx:/fsx,/opt/slurm:/opt/slurm --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE python /fsx/ubuntu/awsome-distributed-ai/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh index 413b5586f..8d9e75299 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh @@ -7,13 +7,8 @@ #SBATCH --gpus-per-node=8 #SBATCH --cpus-per-task=12 #SBATCH --exclusive -#SBATCH --qos=admin_qos -#SBATCH --partition=p5en -#SBATCH --time=24:00:00 #SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out #SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err -#SBATCH --container-image=/fsx/ubuntu/qwen3-8b/containers/nemo-efa-25.07.sqsh -#SBATCH --container-mounts=/fsx:/fsx # EFA / NCCL environment export FI_PROVIDER=efa @@ -22,12 +17,7 @@ export NCCL_DEBUG=WARN export NCCL_TUNER_PLUGIN=/opt/amazon/ofi-nccl/lib/libnccl-tuner-aws-ofi.so export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH} export TORCH_COMPILE_DISABLE=1 - -# Copy training script to expected location -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -mkdir -p /fsx/ubuntu/qwen3-8b/code -cp "${SCRIPT_DIR}/train.py" /fsx/ubuntu/qwen3-8b/code/train.py +export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" # Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init -srun --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE \ - python /fsx/ubuntu/qwen3-8b/code/train.py +/opt/slurm/bin/srun --mpi=pmix --container-image=/fsx/ubuntu/qwen3-8b-pretraining/containers/nemo-efa-26.04.sqsh --container-mounts=/fsx:/fsx,/opt/slurm:/opt/slurm --container-env=FI_PROVIDER,NCCL_SOCKET_IFNAME,NCCL_DEBUG,NCCL_TUNER_PLUGIN,LD_LIBRARY_PATH,TORCH_COMPILE_DISABLE python /fsx/ubuntu/awsome-distributed-ai/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py old mode 100644 new mode 100755 index 033f63f50..bfceeef6d --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py @@ -1,10 +1,11 @@ #!/usr/bin/env python3 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: MIT-0 -"""Qwen3-8B Pre-Training on H200 — Megatron-Bridge (NeMo 25.07) +"""Qwen3-8B Pre-Training on H200 — Megatron-Bridge (NeMo 26.04) Uses the Megatron-Bridge recipe API with config objects. No gradient checkpointing needed (distributed optimizer keeps memory at ~114 GB). +Uses allenai/c4 pre-tokenized with Qwen3-8B tokenizer in Megatron indexed format. Best config: TP=1, PP=1, DP=16, MBS=2, GBS=128, seq=4096, BF16 Result: 497 TFLOP/s/GPU, 162K tok/s on 16x H200 @@ -17,22 +18,31 @@ from megatron.bridge.training.gpt_step import forward_step from megatron.bridge.training.pretrain import pretrain +# Path to Megatron-indexed c4 dataset (prefix without .bin/.idx extension) +DATA_PATH = "/fsx/ubuntu/qwen3-8b-pretraining/datasets/c4_qwen3_8b" def main(): - cfg = qwen3_8b_pretrain_config( - mock=True, - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - micro_batch_size=2, - global_batch_size=128, - seq_length=4096, - train_iters=100, - lr_warmup_iters=10, - lr_decay_iters=100, - ) - - # No gradient checkpointing - MBS=2 with distributed optimizer fits in 141 GB - # (uses ~114 GB peak, leaving headroom) + cfg = qwen3_8b_pretrain_config() + + # Parallelism: DP=16 (model fits on one H200 GPU) + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + + # Batch config + cfg.train.micro_batch_size = 2 + cfg.train.global_batch_size = 128 + cfg.model.seq_length = 4096 + + # Training schedule + cfg.train.train_iters = 100 + cfg.scheduler.lr_warmup_iters = 10 + cfg.scheduler.lr_decay_iters = 100 + + # Dataset: use real c4 data instead of mock + cfg.dataset.data_path = DATA_PATH + cfg.dataset.seq_length = 4096 + cfg.dataset.split = "9999,8,2" + cfg.dataset.num_workers = 8 # Optimizer cfg.optimizer.lr = 3e-4 @@ -41,13 +51,16 @@ def main(): cfg.optimizer.adam_beta1 = 0.9 cfg.optimizer.adam_beta2 = 0.95 cfg.optimizer.clip_grad = 1.0 + cfg.optimizer.overlap_grad_reduce = True + cfg.optimizer.overlap_param_gather = True # Logging and checkpoints cfg.logger.log_interval = 5 - cfg.train.eval_interval = 1000 - cfg.train.eval_iters = 0 - cfg.train.dir = "/fsx/ubuntu/qwen3-8b/checkpoints/h200" - cfg.train.save_interval = 1000 + cfg.validation.eval_interval = 1000 + cfg.validation.eval_iters = 0 + cfg.checkpoint.save = "/fsx/ubuntu/qwen3-8b/checkpoints/h200" + cfg.checkpoint.load = "/fsx/ubuntu/qwen3-8b/checkpoints/h200" + cfg.checkpoint.save_interval = 1000 pretrain(config=cfg, forward_step_func=forward_step) diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh index aa57efce6..901ac0192 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.sh @@ -31,7 +31,7 @@ python3 -m venv /fsx/ubuntu/qwen3-8b-pretraining/venv source /fsx/ubuntu/qwen3-8b-pretraining/venv/bin/activate PYTHON=/fsx/ubuntu/qwen3-8b-pretraining/venv/bin/python -SCRIPT_DIR=$SLURM_SUBMIT_DIR +SCRIPT_DIR=$SLURM_SUBMIT_DIR/preprocessing/ SCRIPT="${SCRIPT_DIR}/preprocess.py" pip install -r $SCRIPT_DIR/requirements.txt From 94cf02bcb114dd34454707b28fdf33f796b16fcc Mon Sep 17 00:00:00 2001 From: Paulo Aragao Date: Sun, 21 Jun 2026 22:42:32 +0000 Subject: [PATCH 8/8] fixed error on preprocess script to match megatron requirements --- .../qwen3-8b-pretraining/b300/slurm/run.sh | 4 ++-- .../qwen3-8b-pretraining/h200/slurm/run.sh | 4 ++-- .../preprocessing/preprocess.py | 22 ++++++++++++------- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh index ec7ac6ad0..eadaba6be 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh @@ -7,8 +7,8 @@ #SBATCH --gpus-per-node=8 #SBATCH --cpus-per-task=12 #SBATCH --exclusive -#SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out -#SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err +#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err # EFA / NCCL environment export FI_PROVIDER=efa diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh index 8d9e75299..03361b6ed 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh @@ -7,8 +7,8 @@ #SBATCH --gpus-per-node=8 #SBATCH --cpus-per-task=12 #SBATCH --exclusive -#SBATCH --output=/fsx/ubuntu/qwen3-8b/logs/%j.out -#SBATCH --error=/fsx/ubuntu/qwen3-8b/logs/%j.err +#SBATCH --output=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.out +#SBATCH --error=/fsx/ubuntu/qwen3-8b-pretraining/logs/%j.err # EFA / NCCL environment export FI_PROVIDER=efa diff --git a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py index 083e3b76c..edf0fee33 100644 --- a/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py @@ -74,25 +74,31 @@ def write_idx_file(idx_path, sizes, doc_idx): def tokenize_chunk(args_tuple): - """Tokenize a chunk of documents. Returns (tmp_bin_path, sizes, doc_idx, total_tokens).""" - chunk_id, texts, tokenizer_name, output_dir = args_tuple + """Tokenize a chunk of documents into fixed-length sequences.""" + chunk_id, texts, tokenizer_name, output_dir, seq_length = args_tuple tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) tmp_bin = os.path.join(output_dir, f"_tmp_chunk_{chunk_id:05d}.bin") sizes = [] doc_idx = [0] total_tokens = 0 + buffer = [] with open(tmp_bin, "wb") as out: for text in texts: tokens = tokenizer.encode(text, add_special_tokens=False) if not tokens: continue - arr = np.array(tokens, dtype=_DTYPE) - out.write(arr.tobytes()) - sizes.append(len(tokens)) - doc_idx.append(len(sizes)) - total_tokens += len(tokens) + buffer.extend(tokens) + + while len(buffer) >= seq_length: + seq = buffer[:seq_length] + buffer = buffer[seq_length:] + arr = np.array(seq, dtype=_DTYPE) + out.write(arr.tobytes()) + sizes.append(seq_length) + doc_idx.append(len(sizes)) + total_tokens += seq_length return tmp_bin, sizes, doc_idx, total_tokens @@ -135,7 +141,7 @@ def main(): chunks = [texts[i:i + chunk_size] for i in range(0, len(texts), chunk_size)] worker_args = [ - (i, chunk, args.tokenizer, output_dir) + (i, chunk, args.tokenizer, output_dir, 4096) for i, chunk in enumerate(chunks) ]