Skip to content

Commit d7dcf9f

Browse files
ngu3ngu3
committed
Add A4XMAX gpt-oss-120b FP8mx 1024 GPUs recipe
1 parent eef14ea commit d7dcf9f

8 files changed

Lines changed: 850 additions & 0 deletions

File tree

‎training/a4x-max/gpt-oss-120b/megatron-bridge-gke/nemo2604/1024gpus-fp8mx-seq4096-gbs20480/recipe/Chart.yaml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: v2
16+
name: a4x_max_jobset_workload
17+
description: a4x_max_jobset_workload
18+
type: application
19+
version: 0.1.0
20+
appVersion: "1.16.0"

‎training/a4x-max/gpt-oss-120b/megatron-bridge-gke/nemo2604/1024gpus-fp8mx-seq4096-gbs20480/recipe/README.md‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
<!-- mdformat global-off -->
2+
# Pretrain gpt-oss-120b workloads on a4x-max GKE Node pools with Nvidia Megatron-Bridge Framework
3+
4+
This recipe outlines the steps for running a gpt-oss-120b pretraining
5+
workload on [a4x-max GKE Node pools](https://cloud.google.com/kubernetes-engine) by using the
6+
[Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge).
7+
8+
## Orchestration and deployment tools
9+
10+
For this recipe, the following setup is used:
11+
12+
- Orchestration - [Google Kubernetes Engine (GKE)](https://cloud.google.com/kubernetes-engine)
13+
- Pretraining job configuration and deployment - A Helm chart is used to
14+
configure and deploy the [Kubernetes Jobset](https://kubernetes.io/blog/2025/03/23/introducing-jobset) resource which manages the execution of the
15+
[Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge).
16+
17+
## Test environment
18+
19+
This recipe has been optimized for and tested with the following configuration:
20+
21+
- GKE cluster
22+
Please follow Cluster Toolkit [instructions](https://github.com/GoogleCloudPlatform/cluster-toolkit/)
23+
to create your a4x-max GKE cluster.
24+
25+
## Training dataset
26+
27+
This recipe uses a mock pretraining dataset provided by the Megatron-Bridge framework.
28+
29+
## Docker container image
30+
31+
This recipe uses the following docker images:
32+
33+
- `nvcr.io/nvidia/nemo:26.04.01`
34+
**Installed Plugins:**
35+
- `nccl-gib-plugins` version: 1.1.2-1
36+
37+
## Run the recipe
38+
39+
From your client workstation, complete the following steps:
40+
41+
### Configure environment settings
42+
43+
Set the environment variables to match your environment:
44+
45+
```bash
46+
export PROJECT_ID=<PROJECT_ID>
47+
export CLUSTER_REGION=<CLUSTER_REGION>
48+
export CLUSTER_NAME=<CLUSTER_NAME>
49+
export GCS_BUCKET=<GCS_BUCKET> # Note: path should not be prefixed with gs://
50+
export KUEUE_NAME=<KUEUE_NAME>
51+
export HF_TOKEN=<YOUR_HF_TOKEN>
52+
```
53+
54+
Replace the following values:
55+
56+
- `<PROJECT_ID>`: your Google Cloud project ID.
57+
- `<CLUSTER_REGION>`: the region where your cluster is located.
58+
- `<CLUSTER_NAME>`: the name of your GKE cluster.
59+
- `<GCS_BUCKET>`: the name of your Cloud Storage bucket. Don't include the `gs://` prefix.
60+
- `<KUEUE_NAME>`: the name of the Kueue local queue. The default queue created by the cluster toolkit is `a4x-max`. Make sure to verify the name of the local queue in your cluster.
61+
- `<YOUR_HF_TOKEN>`: Your HuggingFace token.
62+
63+
Set the default project:
64+
65+
```bash
66+
gcloud config set project $PROJECT_ID
67+
```
68+
69+
### Get the recipe
70+
71+
Clone the `gpu-recipes` repository and set a reference to the recipe folder.
72+
73+
```
74+
git clone https://github.com/ai-hypercomputer/gpu-recipes.git
75+
cd gpu-recipes
76+
export REPO_ROOT=`git rev-parse --show-toplevel`
77+
export RECIPE_ROOT=$REPO_ROOT/training/a4x-max/gpt-oss-120b/megatron-bridge-gke/nemo2604/1024gpus-fp8mx-seq4096-gbs20480/recipe
78+
cd $RECIPE_ROOT
79+
```
80+
81+
### Get cluster credentials
82+
83+
```
84+
gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION
85+
```
86+
87+
### Configure and submit a pretraining job
88+
89+
#### Using 256 node (1024 gpus) fp8mx precision
90+
To execute the job with the default settings, run the following command from
91+
your client:
92+
93+
```bash
94+
cd $RECIPE_ROOT
95+
export WORKLOAD_NAME=$USER-a4x-max-gpt-oss-120b-1024gpus
96+
helm install $WORKLOAD_NAME . -f values.yaml \
97+
--set-file workload_launcher=launcher.sh \
98+
--set workload.image=nvcr.io/nvidia/nemo:26.04.01 \
99+
--set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
100+
--set volumes.gcsMounts[0].mountPath=/job-logs \
101+
--set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \
102+
--set queue=${KUEUE_NAME}
103+
```
104+
105+
**Examples**
106+
107+
- To set the number of training steps to 100, run the following command from
108+
your client:
109+
110+
```bash
111+
cd $RECIPE_ROOT
112+
export WORKLOAD_NAME=$USER-a4x-max-gpt-oss-120b-1024gpus
113+
helm install $WORKLOAD_NAME . -f values.yaml \
114+
--set-file workload_launcher=launcher.sh \
115+
--set workload.image=nvcr.io/nvidia/nemo:26.04.01 \
116+
--set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
117+
--set volumes.gcsMounts[0].mountPath=/job-logs \
118+
--set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \
119+
--set queue=${KUEUE_NAME} \
120+
--set workload.arguments[0]="trainer.max_steps=100"
121+
```
122+
123+
### Monitor the job
124+
125+
To check the status of pods in your job, run the following command:
126+
127+
```
128+
kubectl get pods | grep $USER-a4x-max-gpt-oss-120b-256node
129+
```
130+
131+
Replace the following:
132+
133+
- JOB_NAME_PREFIX - your job name prefix. For example $USER-a4x-max-gpt-oss-120b-256node.
134+
135+
To get the logs for one of the pods, run the following command:
136+
137+
```
138+
kubectl logs POD_NAME
139+
```
140+
141+
Information about the training job's progress, including crucial details such as
142+
loss, step count, and step time, is generated by the rank 0 process.
143+
This process runs on the pod whose name begins with
144+
`JOB_NAME_PREFIX-workload-0-0`.
145+
For example: `$USER-a4x-max-gpt-oss-120b-256node-workload-0-0-s9zrv`.
146+
147+
### Uninstall the Helm release
148+
149+
You can delete the job and other resources created by the Helm chart. To
150+
uninstall Helm, run the following command from your client:
151+
152+
```bash
153+
helm uninstall $USER-a4x-max-gpt-oss-120b-256node
154+
```

‎training/a4x-max/gpt-oss-120b/megatron-bridge-gke/nemo2604/1024gpus-fp8mx-seq4096-gbs20480/recipe/launcher.sh‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
usage()
2+
{
3+
cat << EOF
4+
usage: bash ./launcher.sh [config-override [config-override ...]]
5+
config-override (Optional) A NeMo configuration override. E.g. trainer.max_steps=10000.
6+
EOF
7+
}
8+
9+
parse_args() {
10+
while [[ "$1" != "" ]]; do
11+
case $(grep -o "=" <<< "$1" | wc -l) in
12+
1 )
13+
config_overrides+=("$1")
14+
;;
15+
* )
16+
echo "Invalid config override: $1"
17+
usage
18+
exit 1
19+
esac
20+
shift
21+
done
22+
config_overrides="${config_overrides[*]}"
23+
}
24+
25+
config_overrides=()
26+
parse_args "$@"
27+
28+
if [[ -z "${config_overrides[*]}" ]]; then
29+
echo "No NeMo config overrides specified"
30+
else
31+
echo "NeMo config overrides:"
32+
echo " ${config_overrides}"
33+
fi
34+
35+
export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH"
36+
ldconfig "$LD_LIBRARY_PATH"
37+
echo "Added $LD_LIBRARY_PATH to ldconfig:"
38+
ldconfig -p | grep libcuda | sed 's/^/ /'
39+
echo ""
40+
41+
if [[ -n "${EXPLICIT_LOG_DIR}" ]]; then
42+
explicit_log_dir="${EXPLICIT_LOG_DIR}"
43+
else
44+
explicit_log_dir="workload_logs"
45+
fi
46+
47+
# Ensure explicit_log_dir is an absolute path before any cd commands
48+
if [[ "$explicit_log_dir" != /* ]]; then
49+
explicit_log_dir="${PWD}/${explicit_log_dir}"
50+
fi
51+
echo "Logging to ${explicit_log_dir}"
52+
53+
if [[ -n "${TOKENIZER_PATH}" ]]; then
54+
echo "Getting tokenizer files"
55+
cp "${TOKENIZER_PATH}"/* .
56+
echo ""
57+
fi
58+
59+
echo "Launching Torch distributed on the node rank $JOB_COMPLETION_INDEX out of $NNODES nodes"
60+
61+
# Create the nsys directory.
62+
mkdir -p "${explicit_log_dir}/nsys"
63+
64+
# Collect diagnostics
65+
linux_kv="$(uname --kernel-release)"
66+
cuda_driver_v=""
67+
driver_v=""
68+
vbios_v=""
69+
if command -v nvidia-smi &> /dev/null; then
70+
cuda_driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<cuda_version>).*(?=</cuda_version>)' || true)
71+
driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<driver_version>).*(?=</driver_version>)' || true)
72+
vbios_v=$(nvidia-smi -q -x | grep -Po '(?<=<vbios_version>).*(?=</vbios_version>)' | head -n1 || true)
73+
fi
74+
nccl_v=$(python3 -c "import torch; v=torch.cuda.nccl.version() if hasattr(torch.cuda, 'nccl') else 'unknown'; print('.'.join(map(str, v)) if isinstance(v, tuple) else v)" || echo "unknown")
75+
cuda_container_v=$(python3 -c "import torch; print(torch.version.cuda)" || echo "unknown")
76+
77+
kv="{\"linux_kernel_version\": \"${linux_kv}\""
78+
kv="${kv}, \"cuda_driver_version\": \"${cuda_driver_v}\""
79+
kv="${kv}, \"cuda_container_version\": \"${cuda_container_v}\""
80+
kv="${kv}, \"gpu_driver_version\": \"${driver_v}\""
81+
kv="${kv}, \"vbios_version\": \"${vbios_v}\""
82+
kv="${kv}, \"nccl_version\": \"${nccl_v}\"}"
83+
84+
echo "VERSION_DIAGNOSTICS: ${kv}"
85+
86+
87+
export HF_TOKEN=<YOUR_HF_TOKEN>
88+
export NCCL_IB_SPLIT_DATA_ON_QPS=1
89+
export NCCL_RAS_ENABLE=0
90+
91+
cd /opt
92+
rm -rf Megatron-Bridge
93+
git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
94+
cd Megatron-Bridge
95+
git checkout c810129341a84e58f4cbed3093f70668a088c028
96+
git submodule update --init --recursive && sed -i 's/timeout=60/timeout=600/g' src/megatron/bridge/models/hf_pretrained/safe_config_loader.py
97+
sed -i -e '/pretrain(config=recipe/i \ recipe.dist.distributed_timeout_minutes = 10' scripts/performance/run_script.py
98+
ls
99+
100+
101+
102+
worker_command=$(cat <<- EOM
103+
if [ "\$RANK" -eq "0" ]; then
104+
echo "Worker 0 is stalling for a few seconds.." ;
105+
sleep 3 ;
106+
echo "The detected environment within worker rank 0 is:" ;
107+
env | sed 's/^/ /' ;
108+
else
109+
echo "Worker \$RANK is running" ;
110+
fi ;
111+
112+
echo "Collect nvidia-smi telemetry"
113+
gpu_uuid=\$(nvidia-smi --id \$LOCAL_RANK --query-gpu uuid --format noheader)
114+
115+
telemetry_dir="/runtime-logs/${JOBSET_NAME}/nvidia-smi"
116+
mkdir -p \$telemetry_dir
117+
touch "\$telemetry_dir/\$RANK.csv"
118+
119+
nv_smi_pid=\$!
120+
121+
cd /opt/Megatron-Bridge ;
122+
123+
numactl \
124+
--cpunodebind=\$((LOCAL_RANK/2)) \
125+
--membind=\$((LOCAL_RANK/2)) \
126+
nice -10 \
127+
python scripts/performance/run_script.py \
128+
--model_family_name gpt_oss \
129+
--model_recipe_name gpt_oss_120b \
130+
--config_variant v2 \
131+
--gpu gb300 \
132+
--num_gpus 1024 \
133+
--gpus_per_node 4 \
134+
--compute_dtype fp8_mx \
135+
--seq_length 4096 \
136+
--global_batch_size 20480 \
137+
--micro_batch_size 4 \
138+
--tensor_model_parallel_size 1 \
139+
--pipeline_model_parallel_size 1 \
140+
--context_parallel_size 1 \
141+
--expert_model_parallel_size 64 \
142+
--expert_tensor_parallel_size 1 \
143+
--cuda_graph_impl transformer_engine \
144+
--cuda_graph_scope attn,moe_router,moe_preprocess \
145+
--max_step 50 \
146+
logger.log_throughput=True
147+
148+
149+
kill -9 \$nv_smi_pid
150+
EOM
151+
)
152+
153+
echo "$worker_command" > worker_command.sh
154+
chmod 777 worker_command.sh
155+
156+
mkdir -p "/runtime-logs/${JOBSET_NAME}/logs"
157+
158+
torchrun \
159+
--nproc-per-node="4" \
160+
--nnodes="256" \
161+
--node_rank="${JOB_COMPLETION_INDEX}" \
162+
--rdzv_id="${JOB_IDENTIFIER}" \
163+
--master_addr="${MASTER_ADDR}" \
164+
--master_port="${MASTER_PORT}" \
165+
--no-python bash worker_command.sh 2>&1 | python3 -u -c "import sys, time; [sys.stdout.write('[{}] {}'.format(time.strftime('%Y-%m-%d %H:%M:%S'), line)) for line in iter(sys.stdin.readline, '')]" | tee "/runtime-logs/${JOBSET_NAME}/logs/nemo_mb_RANK_${JOB_COMPLETION_INDEX}.log"
166+
167+
168+
if [[ "$JOB_COMPLETION_INDEX" == "0" ]]; then
169+
mkdir -p "${ARTIFACT_DIR}"
170+
cp -r "${explicit_log_dir}"/* "${ARTIFACT_DIR}/"
171+
env > "${ARTIFACT_DIR}/environ.txt"
172+
ls "${ARTIFACT_DIR}"
173+
fi
174+
echo "Training completed"
175+
echo "Pod on $(hostname --fqdn) is exiting"

‎training/a4x-max/gpt-oss-120b/megatron-bridge-gke/nemo2604/1024gpus-fp8mx-seq4096-gbs20480/recipe/templates/workload-config-configmap.yaml‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
# yamllint disable
2+
# Copyright 2026 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
{{- if .Values.workload.configFile }}
17+
apiVersion: v1
18+
kind: ConfigMap
19+
metadata:
20+
name: "{{ .Release.Name }}-config"
21+
data:
22+
workload-configuration: |-
23+
{{- if .Values.workload_config }}
24+
{{ .Values.workload_config | nindent 4 }}
25+
{{- else }}
26+
{{ "config: null" | nindent 4 }}
27+
{{- end }}
28+
{{- end }}

0 commit comments

Comments
 (0)