Add A4X MAX Qwen3-235B-A22B FP8mx 256 GPUs recipe - #278
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Helm chart recipe to pretrain the qwen3-235b-a22b model on a4x-max GKE node pools using the Nvidia Megatron-Bridge framework. The review feedback focuses on improving script robustness and deployment flexibility. Key recommendations include adding error handling and capturing the exit status of torchrun in launcher.sh, avoiding hardcoded default namespaces in Kubernetes manifests to support multi-namespace deployments, dynamically detecting the OS version for DOCA-OFED installation, and properly handling the HF_TOKEN environment variable across the configuration files, Helm templates, and documentation.
| @@ -0,0 +1,184 @@ | |||
| usage() | |||
There was a problem hiding this comment.
Add set -eo pipefail at the beginning of the script to ensure that any failures during setup (such as cloning the repository or checking out the commit) are not silently ignored.
| usage() | |
| set -eo pipefail | |
| usage() |
| 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 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, '')]" | ||
|
|
||
|
|
||
|
|
||
| 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" |
There was a problem hiding this comment.
Capture the exit status of torchrun and exit with it at the end of the script. This ensures that if training fails, the Kubernetes job is correctly marked as failed. Additionally, this guarantees that the rank 0 log/artifact copying still runs even if the training fails, which is crucial for debugging.
| 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 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, '')]" | |
| 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" | |
| rc=0 | |
| 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 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, '')]" || rc=$? | |
| 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" | |
| exit $rc |
| echo "VERSION_DIAGNOSTICS: ${kv}" | ||
|
|
||
|
|
||
| export HF_TOKEN=<YOUR_HF_TOKEN> |
There was a problem hiding this comment.
Do not overwrite HF_TOKEN with a hardcoded placeholder if it is already set in the environment. Instead, preserve the existing environment variable and warn the user if it is missing or still set to the placeholder.
| export HF_TOKEN=<YOUR_HF_TOKEN> | |
| if [[ -z "${HF_TOKEN}" || "${HF_TOKEN}" == "<YOUR_HF_TOKEN>" ]]; then | |
| echo "WARNING: HF_TOKEN is not set or is a placeholder. Hugging Face downloads may fail." | |
| fi |
| kind: JobSet | ||
| metadata: | ||
| name: "{{ .Release.Name }}" | ||
| namespace: default |
There was a problem hiding this comment.
Avoid hardcoding the default namespace in the JobSet metadata. Use {{ .Release.Namespace }} to allow deploying the Helm chart to other namespaces.
namespace: "{{ .Release.Namespace }}"| - name: RANK_0_FQDN | ||
| value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.default.svc.cluster.local" | ||
| - name: HOSTNAME_PREFIX | ||
| value: "{{.Release.Name}}-workload-" | ||
| - name: DOMAIN_NAME | ||
| value: "{{.Release.Name}}.default.svc.cluster.local" | ||
| - name: MASTER_ADDR | ||
| value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.default.svc.cluster.local" | ||
| - name: MASTER_PORT | ||
| value: "6002" |
There was a problem hiding this comment.
Avoid hardcoding the default namespace in the FQDN and master address environment variables. Use {{ .Release.Namespace }} to ensure correct DNS resolution when deploying to non-default namespaces.
- name: RANK_0_FQDN
value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
- name: HOSTNAME_PREFIX
value: "{{.Release.Name}}-workload-"
- name: DOMAIN_NAME
value: "{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
- name: MASTER_ADDR
value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
- name: MASTER_PORT
value: "6002"| # Install DOCA-OFED | ||
| apt update -y | ||
| apt install -y curl | ||
| export DOCA_URL="https://linux.mellanox.com/public/repo/doca/3.1.0/ubuntu22.04/arm64-sbsa/" |
There was a problem hiding this comment.
Instead of hardcoding ubuntu22.04 in the DOCA-OFED repository URL, dynamically detect the OS version from /etc/os-release. This ensures compatibility with newer container images (such as those based on Ubuntu 24.04), which is highly likely for a NeMo 26.06 container.
. /etc/os-release
export DOCA_URL="https://linux.mellanox.com/public/repo/doca/3.1.0/ubuntu${VERSION_ID}/arm64-sbsa/"| gcsMounts: | ||
| - bucketName: null | ||
| mountPath: null |
There was a problem hiding this comment.
Setting a default element with null values in gcsMounts will cause Helm to render invalid Kubernetes volume definitions when using default values. Change the default to an empty list []. Users can still easily override it via --set or custom values files.
gcsMounts: []| envs: | ||
| - name: ARTIFACT_DIR | ||
| value: null | ||
| - name: GLOO_SOCKET_IFNAME | ||
| value: eth0 |
There was a problem hiding this comment.
Add HF_TOKEN to the default environment variables list in values.yaml so that it can be easily passed to the container via Helm.
envs:
- name: ARTIFACT_DIR
value: null
- name: GLOO_SOCKET_IFNAME
value: eth0
- name: HF_TOKEN
value: null| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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} |
There was a problem hiding this comment.
Update the helm install command to pass the HF_TOKEN environment variable to the container, since it is requested to be exported but never actually passed to the Helm release.
| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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} | |
| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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 workload.envs[2].value=${HF_TOKEN} \ | |
| --set queue=${KUEUE_NAME} |
| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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" |
There was a problem hiding this comment.
Update the example helm install command to pass the HF_TOKEN environment variable to the container.
| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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" | |
| export WORKLOAD_NAME=$USER-a4x-max-qwen3-235b-a22b-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 workload.envs[2].value=${HF_TOKEN} \ | |
| --set queue=${KUEUE_NAME} \ | |
| --set workload.arguments[0]="trainer.max_steps=100" |
1 participant
Add A4X MAX Qwen3-235B-A22B 256 GPUs FP8mx recipe