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/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 new file mode 100644 index 000000000..eeec1b43f --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/README.md @@ -0,0 +1,170 @@ +# 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 | ~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. + +## 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 + +> **Disk space:** The container build requires ~50 GB of disk space in TMPDIR. +> `enroot import` needs `sudo` and TMPDIR pointing to FSx or another file system (not `/tmp`, which is too small). + +### Clone this repo and change it its directory +```bash +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-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-pretraining:latest +``` + +### H200 Cluster (2x p5en.48xlarge) + +```bash +# 1. Change to directory +cd h200 + +# 2. Submit training job +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 (2x p6-b300.48xlarge) + +```bash +# 1. Change to directory +cd b300 + +# 2. Submit training job +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 + +| Parameter | Value | +|-----------|-------| +| Layers | 36 | +| Hidden dim (d_model) | 4096 | +| Q-heads | 32 | +| KV-heads | 8 (GQA) | +| FFN dim | 12288 (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** | Selective (core_attn only) | Selective (core_attn only) | +| Distributed optimizer | Yes (sharded Adam) | Yes (sharded Adam) | +| Overlap grad reduce | Yes | Yes | +| 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 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. **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 + +| | 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 (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/ +│ ├── train.py ← Megatron-Bridge training script +│ └── slurm/ +│ └── run.sh ← Slurm submission script +└── b300/ + ├── train.py ← Megatron-Bridge training script + └── slurm/ + └── run.sh ← Slurm submission script +``` + +## License + +MIT-0 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..eadaba6be --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/slurm/run.sh @@ -0,0 +1,23 @@ +#!/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 --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 +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 +export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" + +# Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init +/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/b300/train.py b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/train.py new file mode 100644 index 000000000..f8763bedd --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/b300/train.py @@ -0,0 +1,55 @@ +#!/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. +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, + 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, + ) + + # No gradient checkpointing (B300 has 288 GB - fits MBS=4 without recompute) + + # 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/b300" + 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/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..03361b6ed --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/slurm/run.sh @@ -0,0 +1,23 @@ +#!/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 --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 +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 +export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" + +# Launch - Megatron uses SLURM env vars (SLURM_PROCID, SLURM_LOCALID) for distributed init +/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 new file mode 100755 index 000000000..bfceeef6d --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/h200/train.py @@ -0,0 +1,69 @@ +#!/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 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 +""" +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 + +# 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() + + # 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 + 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 + cfg.optimizer.overlap_grad_reduce = True + cfg.optimizer.overlap_param_gather = True + + # Logging and checkpoints + cfg.logger.log_interval = 5 + 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) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000..edf0fee33 --- /dev/null +++ b/3.test_cases/megatron/nemo/qwen3-8b-pretraining/preprocessing/preprocess.py @@ -0,0 +1,185 @@ +#!/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 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 + 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 + + +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, 4096) + 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() 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..901ac0192 --- /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/preprocessing/ +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 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