Add A4X MAX Llama-3.1-405B NVFP4 256 GPUs recipe by ngu3 · Pull Request #275 · AI-Hypercomputer/gpu-recipes

ngu3 · 2026-08-13T00:32:15Z

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

gemini-code-assist

Code Review

This pull request introduces a Helm chart and a launcher script to run Llama 3.1 405B pretraining workloads on GKE using the Nvidia Megatron-Bridge framework. Key feedback on these changes includes enabling set -eo pipefail in the launcher script to prevent silent failures, avoiding hardcoded token placeholders, and checking if ARTIFACT_DIR is set before copying logs to prevent writing to the root directory. Additionally, installing heavy system packages like DOCA-OFED and NCCL plugins at container startup should be avoided in favor of pre-building them into the Docker image. Other issues identified include a priority inversion with the nice command, a potential memory leak from using a list comprehension for log streaming, and Kubernetes API validation errors caused by null values in values.yaml.

gemini-code-assist · 2026-08-13T00:34:07Z

+usage()
+{


The script does not enable set -eo pipefail at the beginning. Since it is executed as a subshell via bash /workload/launcher/launch-workload.sh, it does not inherit the set -e from the parent shell. Any failure in critical setup steps (such as git clone, git checkout, git submodule update, or ldconfig) will not stop the script, leading to silent failures or hard-to-debug errors later.

Suggested change

usage()

{

set -eo pipefail

usage()

{

gemini-code-assist · 2026-08-13T00:34:08Z

+echo "VERSION_DIAGNOSTICS: ${kv}"
+
+
+export HF_TOKEN=<YOUR_HF_TOKEN>


Hardcoding export HF_TOKEN=<YOUR_HF_TOKEN> will overwrite any valid Hugging Face token passed via Kubernetes environment variables or secrets. It is better to validate that a token is provided and only use the existing environment variable.

Suggested change

export HF_TOKEN=<YOUR_HF_TOKEN>

if [[ -z "${HF_TOKEN}" || "${HF_TOKEN}" == "<YOUR_HF_TOKEN>" ]]; then

echo "Error: HF_TOKEN is not set or is set to the placeholder. Please provide a valid Hugging Face token."

exit 1

fi

gemini-code-assist · 2026-08-13T00:34:08Z

+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


If ARTIFACT_DIR is empty or unset (which is the default in values.yaml), the script will attempt to run mkdir -p "" (which fails) and then cp -r .../* / which will copy all logs directly to the root directory / and write environ.txt to /environ.txt. We should check that ARTIFACT_DIR is non-empty before attempting these operations.

Suggested change

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

if [[ "$JOB_COMPLETION_INDEX" == "0" && -n "${ARTIFACT_DIR}" ]]; then

mkdir -p "${ARTIFACT_DIR}"

cp -r "${explicit_log_dir}"/* "${ARTIFACT_DIR}/"

env > "${ARTIFACT_DIR}/environ.txt"

ls "${ARTIFACT_DIR}"

fi

gemini-code-assist · 2026-08-13T00:34:08Z

+  gcsMounts:
+    - bucketName: null
+      mountPath: null


Defining gcsMounts with a single element containing null values means that the Helm chart will always attempt to render and mount a volume named "null" to path "null" when no GCS mounts are specified. It is much cleaner and safer to default gcsMounts to an empty list [].

gcsMounts: []

gemini-code-assist · 2026-08-13T00:34:08Z

+    - name: ARTIFACT_DIR
+      value: null


In Kubernetes, the value field of an environment variable must be a string. Setting value: null in values.yaml will render as value: null in the JobSet manifest, which will be rejected by the Kubernetes API server with a validation error. It should be set to "" (empty string) instead.

- name: ARTIFACT_DIR value: ""

gemini-code-assist · 2026-08-13T00:34:08Z

+                apt update -y
+                apt install -y curl
+                export DOCA_URL="https://linux.mellanox.com/public/repo/doca/3.1.0/ubuntu22.04/arm64-sbsa/"
+                BASE_URL=$([ "${DOCA_PREPUBLISH:-false}" = "true" ] && echo https://doca-repo-prod.nvidia.com/public/repo/doca || echo https://linux.mellanox.com/public/repo/doca)
+                DOCA_SUFFIX=${DOCA_URL#*public/repo/doca/}; DOCA_URL="$BASE_URL/$DOCA_SUFFIX"
+                curl $BASE_URL/GPG-KEY-Mellanox.pub | gpg --dearmor > /etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub
+                echo "deb [signed-by=/etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub] $DOCA_URL ./" > /etc/apt/sources.list.d/doca.list
+                apt update
+                apt install -y --allow-downgrades --allow-change-held-packages -o Dpkg::Options::="--force-overwrite" doca-ofed-userspace || apt --fix-broken install -y
+
+                # Install NCCL and nccl-gib-plugins package
+                apt install --only-upgrade --allow-change-held-packages -y libnccl2 libnccl-dev
+
+                # If image not from Google, trust the GCP signing key
+                curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /etc/apt/trusted.gpg.d/cloud.google.gpg
+
+                # Add gpudirect-gib-apt repo
+                echo 'deb https://packages.cloud.google.com/apt gpudirect-gib-apt main' | tee /etc/apt/sources.list.d/nccl-gib.list
+
+                apt update
+                apt install -y nccl-gib-plugins


Running apt update and installing heavy packages like doca-ofed-userspace and nccl-gib-plugins at container startup on 64 nodes (256 GPUs) simultaneously is highly inefficient and extremely unreliable. It introduces a critical dependency on external package repositories and network connectivity at job launch time. If any single node fails to download or install these packages, the entire 256-GPU training job will fail. These dependencies should be pre-built into the Docker image to ensure fast, reliable, and offline-capable startup.

gemini-code-assist · 2026-08-13T00:34:08Z

+  numactl \
+    --cpunodebind=\$((LOCAL_RANK/2)) \
+    --membind=\$((LOCAL_RANK/2)) \
+  nice -10 \


In the nice command, nice -10 actually sets the niceness to +10 (which lowers the process priority). If the intention was to run the training process with higher priority (negative niceness), you should use nice -n -10 or nice --10.

Suggested change

nice -10 \

nice -n -10 \

gemini-code-assist · 2026-08-13T00:34:08Z

+--rdzv_id="${JOB_IDENTIFIER}" \
+--master_addr="${MASTER_ADDR}" \
+--master_port="${MASTER_PORT}" \
+--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, '')]"


The Python one-liner uses a list comprehension [sys.stdout.write(...) for line in ...] solely for its side effects. In Python, list comprehensions evaluate and store the entire list of results in memory. For long-running training jobs with millions of log lines, this will cause a memory leak and eventually an Out-Of-Memory (OOM) crash. Using a standard for loop avoids this issue and is much more readable.

Suggested change

--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, '')]"

--no-python bash worker_command.sh 2>&1 | python3 -u -c "

import sys, time

for line in iter(sys.stdin.readline, ''):

sys.stdout.write('[{}] {}'.format(time.strftime('%Y-%m-%d %H:%M:%S'), line))

"

Add A4XMAX llama3-1-405b NVFP4 256 GPUs recipe

7cc041f

gemini-code-assist Bot reviewed Aug 13, 2026

View reviewed changes

ngu3 requested a review from Alina-PANG August 13, 2026 04:34

Navigation Menu

Uh oh!

Add A4X MAX Llama-3.1-405B NVFP4 256 GPUs recipe - #275

Add A4X MAX Llama-3.1-405B NVFP4 256 GPUs recipe#275
ngu3 wants to merge 1 commit into
mainfrom
publish-ninggu-ubench-sgmszgah

ngu3 commented Aug 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

-usage()
-{
+set -eo pipefail
+usage()
+{

		echo "VERSION_DIAGNOSTICS: ${kv}"


		export HF_TOKEN=<YOUR_HF_TOKEN>

-export HF_TOKEN=<YOUR_HF_TOKEN>
+if [[ -z "${HF_TOKEN}" || "${HF_TOKEN}" == "<YOUR_HF_TOKEN>" ]]; then
+  echo "Error: HF_TOKEN is not set or is set to the placeholder. Please provide a valid Hugging Face token."
+  exit 1
+fi

Uh oh!

Conversation

ngu3 commented Aug 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant