Add A4X MAX Llama-3.1-405B FP8mx 256 GPUs recipe

training/a4x-max/llama3-1-405b/megatron-bridge-gke/nemo2606/256gpus-fp8mx-seq8192-gbs1536/recipe/Chart.yaml

@@ -0,0 +1,20 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: v2
+name: a4x_max_jobset_workload
+description: a4x_max_jobset_workload
+type: application
+version: 0.1.0
+appVersion: "1.16.0"

training/a4x-max/llama3-1-405b/megatron-bridge-gke/nemo2606/256gpus-fp8mx-seq8192-gbs1536/recipe/README.md

@@ -0,0 +1,154 @@
+<!-- mdformat global-off -->
+# Pretrain llama3-1-405b workloads on a4x-max GKE Node pools with Nvidia Megatron-Bridge Framework
+
+This recipe outlines the steps for running a llama3-1-405b pretraining
+workload on [a4x-max GKE Node pools](https://cloud.google.com/kubernetes-engine) by using the
+[Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
+## Orchestration and deployment tools
+
+For this recipe, the following setup is used:
+
+- Orchestration - [Google Kubernetes Engine (GKE)](https://cloud.google.com/kubernetes-engine)
+- Pretraining job configuration and deployment - A Helm chart is used to
+  configure and deploy the [Kubernetes Jobset](https://kubernetes.io/blog/2025/03/23/introducing-jobset) resource which manages the execution of the
+  [Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge).
+
+## Test environment
+
+This recipe has been optimized for and tested with the following configuration:
+
+- GKE cluster
+Please follow Cluster Toolkit [instructions](https://github.com/GoogleCloudPlatform/cluster-toolkit/)
+to create your a4x-max GKE cluster.
+
+## Training dataset
+
+This recipe uses a mock pretraining dataset provided by the Megatron-Bridge framework.
+
+## Docker container image
+
+This recipe uses the following docker images:
+
+- `nvcr.io/nvidia/nemo:26.06.01`
+**Installed Plugins:**
+- `nccl-gib-plugins` version: 1.1.2-1
+
+## Run the recipe
+
+From your client workstation, complete the following steps:
+
+### Configure environment settings
+
+Set the environment variables to match your environment:
+
+ ```bash
+ export PROJECT_ID=<PROJECT_ID>
+ export CLUSTER_REGION=<CLUSTER_REGION>
+ export CLUSTER_NAME=<CLUSTER_NAME>
+ export GCS_BUCKET=<GCS_BUCKET> # Note: path should not be prefixed with gs://
+ export KUEUE_NAME=<KUEUE_NAME>
+ export HF_TOKEN=<YOUR_HF_TOKEN>
+ ```
+
+Replace the following values:
+
+ - `<PROJECT_ID>`: your Google Cloud project ID.
+ - `<CLUSTER_REGION>`: the region where your cluster is located.
+ - `<CLUSTER_NAME>`: the name of your GKE cluster.
+ - `<GCS_BUCKET>`: the name of your Cloud Storage bucket. Don't include the `gs://` prefix.
+ - `<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.
+ - `<YOUR_HF_TOKEN>`: Your HuggingFace token.
+
+Set the default project:
+
+ ```bash
+ gcloud config set project $PROJECT_ID
+ ```
+
+### Get the recipe
+
+Clone the `gpu-recipes` repository and set a reference to the recipe folder.
+
+```
+git clone https://github.com/ai-hypercomputer/gpu-recipes.git
+cd gpu-recipes
+export REPO_ROOT=`git rev-parse --show-toplevel`
+export RECIPE_ROOT=$REPO_ROOT/training/a4x-max/llama3-1-405b/megatron-bridge-gke/nemo2606/256gpus-fp8mx-seq8192-gbs1536/recipe
+cd $RECIPE_ROOT
+```
+
+### Get cluster credentials
+
+```
+gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION
+```
+
+### Configure and submit a pretraining job
+
+#### Using 64 node (256 gpus) fp8mx precision
+To execute the job with the default settings, run the following command from
+your client:
+
+```bash
+cd $RECIPE_ROOT
+export WORKLOAD_NAME=$USER-a4x-max-llama3-1-405b-256gpus
+helm install $WORKLOAD_NAME . -f values.yaml \
+--set-file workload_launcher=launcher.sh \
+--set workload.image=nvcr.io/nvidia/nemo:26.06.01 \
+--set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
+--set volumes.gcsMounts[0].mountPath=/job-logs \
+--set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \
+--set queue=${KUEUE_NAME}
+```
+
+**Examples**
+
+-   To set the number of training steps to 100, run the following command from
+    your client:
+
+    ```bash
+    cd $RECIPE_ROOT
+    export WORKLOAD_NAME=$USER-a4x-max-llama3-1-405b-256gpus
+    helm install $WORKLOAD_NAME . -f values.yaml \
+    --set-file workload_launcher=launcher.sh \
+    --set workload.image=nvcr.io/nvidia/nemo:26.06.01 \
+    --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
+    --set volumes.gcsMounts[0].mountPath=/job-logs \
+    --set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \
+    --set queue=${KUEUE_NAME} \
+    --set workload.arguments[0]="trainer.max_steps=100"
+    ```
+
+### Monitor the job
+
+To check the status of pods in your job, run the following command:
+
+```
+kubectl get pods | grep $USER-a4x-max-llama3-1-405b-256gpus
+```
+
+Replace the following:
+
+- JOB_NAME_PREFIX - your job name prefix. For example $USER-a4x-max-llama3-1-405b-256gpus.
+
+To get the logs for one of the pods, run the following command:
+
+```
+kubectl logs POD_NAME
+```
+
+Information about the training job's progress, including crucial details such as
+loss, step count, and step time, is generated by the rank 0 process.
+This process runs on the pod whose name begins with
+`JOB_NAME_PREFIX-workload-0-0`.
+For example: `$USER-a4x-max-llama3-1-405b-256gpus-workload-0-0-s9zrv`.
+
+### Uninstall the Helm release
+
+You can delete the job and other resources created by the Helm chart. To
+uninstall Helm, run the following command from your client:
+
+```bash
+helm uninstall $USER-a4x-max-llama3-1-405b-256gpus
+```

training/a4x-max/llama3-1-405b/megatron-bridge-gke/nemo2606/256gpus-fp8mx-seq8192-gbs1536/recipe/launcher.sh

@@ -0,0 +1,166 @@
+usage()
+{
+cat << EOF
+usage: bash ./launcher.sh [config-override  [config-override ...]]
+config-override  (Optional) A  NeMo configuration override. E.g. trainer.max_steps=10000.
+EOF
+}
+
+parse_args() {
+  while [[ "$1" != "" ]]; do
+    case $(grep -o "=" <<< "$1" | wc -l) in
+        1  )
+        config_overrides+=("$1")
+        ;;
+        * )
+            echo "Invalid config override: $1"
+            usage
+            exit 1
+    esac
+    shift
+  done
+  config_overrides="${config_overrides[*]}"
+}
+
+config_overrides=()
+parse_args "$@"
+
+if [[ -z "${config_overrides[*]}" ]]; then
+  echo "No NeMo config overrides specified"
+else
+  echo "NeMo config overrides:"
+  echo "  ${config_overrides}"
+fi
+
+export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH"
+ldconfig "$LD_LIBRARY_PATH"
+echo "Added $LD_LIBRARY_PATH to ldconfig:"
+ldconfig -p | grep libcuda | sed 's/^/  /'
+echo ""
+
+if [[ -n "${EXPLICIT_LOG_DIR}" ]]; then
+  explicit_log_dir="${EXPLICIT_LOG_DIR}"
+else
+  explicit_log_dir="workload_logs"
+fi
+
+# Ensure explicit_log_dir is an absolute path before any cd commands
+if [[ "$explicit_log_dir" != /* ]]; then
+  explicit_log_dir="${PWD}/${explicit_log_dir}"
+fi
+echo "Logging to ${explicit_log_dir}"
+
+if [[ -n "${TOKENIZER_PATH}" ]]; then
+  echo "Getting tokenizer files"
+  cp "${TOKENIZER_PATH}"/* .
+  echo ""
+fi
+
+echo "Launching Torch distributed on the node rank $JOB_COMPLETION_INDEX out of $NNODES nodes"
+
+# Create the nsys directory.
+mkdir -p "${explicit_log_dir}/nsys"
+
+# Collect diagnostics
+linux_kv="$(uname --kernel-release)"
+cuda_driver_v=""
+driver_v=""
+vbios_v=""
+if command -v nvidia-smi &> /dev/null; then
+  cuda_driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<cuda_version>).*(?=</cuda_version>)' || true)
+  driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<driver_version>).*(?=</driver_version>)' || true)
+  vbios_v=$(nvidia-smi -q -x | grep -Po '(?<=<vbios_version>).*(?=</vbios_version>)' | head -n1 || true)
+fi
+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")
+cuda_container_v=$(python3 -c "import torch; print(torch.version.cuda)" || echo "unknown")
+
+kv="{\"linux_kernel_version\": \"${linux_kv}\""
+kv="${kv}, \"cuda_driver_version\": \"${cuda_driver_v}\""
+kv="${kv}, \"cuda_container_version\": \"${cuda_container_v}\""
+kv="${kv}, \"gpu_driver_version\": \"${driver_v}\""
+kv="${kv}, \"vbios_version\": \"${vbios_v}\""
+kv="${kv}, \"nccl_version\": \"${nccl_v}\"}"
+
+echo "VERSION_DIAGNOSTICS: ${kv}"
+
+
+export HF_TOKEN=<YOUR_HF_TOKEN>
+export PYTHONUNBUFFERED=1
+export CUDA_DEVICE_MAX_CONNECTIONS=32
+export NCCL_P2P_NET_CHUNKSIZE=2097152
+export NVTE_BWD_LAYERNORM_SM_MARGIN=16
+export NVTE_FWD_LAYERNORM_SM_MARGIN=16
+export NCCL_RAS_ENABLE=0
+
+cd /opt
+rm -rf Megatron-Bridge
+git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
+cd Megatron-Bridge
+git checkout 5cb3444c43f7499cf3872b2d46870cf8bc2e00ce
+git submodule update --init --recursive && sed -i 's/timeout=60/timeout=600/g' src/megatron/bridge/models/hf_pretrained/safe_config_loader.py
+sed -i -e '/pretrain(config=recipe/i \    recipe.dist.distributed_timeout_minutes = 10' scripts/performance/run_script.py
+ls
+
+
+
+worker_command=$(cat <<- EOM
+  if [ "\$RANK" -eq "0" ]; then
+    echo "Worker 0 is stalling for a few seconds.." ;
+    sleep 3 ;
+    echo "The detected environment within worker rank 0 is:" ;
+    env | sed 's/^/  /' ;
+  else
+    echo "Worker \$RANK is running" ;
+  fi ;
+
+  cd /opt/Megatron-Bridge ;
+
+  numactl \
+    --cpunodebind=\$((LOCAL_RANK/2)) \
+    --membind=\$((LOCAL_RANK/2)) \
+  nice -10 \
+  python scripts/performance/run_script.py \
+    --model_family_name llama \
+    --model_recipe_name llama31_405b \
+    --config_variant v2 \
+    --gpu gb300 \
+    --num_gpus 256 \
+    --gpus_per_node 4 \
+    --compute_dtype fp8_mx \
+    --seq_length 8192 \
+    --global_batch_size 1536 \
+    --micro_batch_size 1 \
+    --tensor_model_parallel_size 2 \
+    --pipeline_model_parallel_size 8 \
+    --virtual_pipeline_model_parallel_size 4 \
+    --context_parallel_size 2 \
+    --expert_model_parallel_size 1 \
+    --expert_tensor_parallel_size 1 \
+    --max_step 50 \
+    logger.log_throughput=True \
+    train.manual_gc_interval=100
+EOM
+)
+
+echo "$worker_command" > worker_command.sh
+chmod 777 worker_command.sh
+
+torchrun \
+--nproc-per-node="4" \
+--nnodes="64" \
+--node_rank="${JOB_COMPLETION_INDEX}" \
+--rdzv_id="${JOB_IDENTIFIER}" \
+--master_addr="${MASTER_ADDR}" \
+--master_port="${MASTER_PORT}" \
+--no-python stdbuf -oL -eL bash worker_command.sh 2>&1 | awk '{ print strftime("[%Y-%m-%d %H:%M:%S]"), $0; fflush(); }'
+
+
+
+if [[ "$JOB_COMPLETION_INDEX" == "0" ]]; then
+  mkdir -p "${ARTIFACT_DIR}"
+  cp -r "${explicit_log_dir}"/* "${ARTIFACT_DIR}/"
+  env > "${ARTIFACT_DIR}/environ.txt"
+  ls "${ARTIFACT_DIR}"
+fi
+echo "Training completed"
+echo "Pod on $(hostname --fqdn) is exiting"

training/a4x-max/llama3-1-405b/megatron-bridge-gke/nemo2606/256gpus-fp8mx-seq8192-gbs1536/recipe/templates/workload-config-configmap.yaml

@@ -0,0 +1,28 @@
+# yamllint disable
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+{{- if .Values.workload.configFile }}
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: "{{ .Release.Name }}-config"
+data:
+  workload-configuration: |-
+{{- if .Values.workload_config }}
+{{ .Values.workload_config | nindent 4 }}
+{{- else }}
+{{ "config: null" | nindent 4 }}
+{{- end }}
+{{- end }}