Add LLaMA 3.1 405B 256-GPU FP8-MX pretraining recipe on GKE A4X (NeMo 26.06.01) #recipebot - #284
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 pretraining recipe for LLaMA 3.1 405B on A4X GKE Node Pools using NVIDIA Megatron-Bridge and NeMo. The review feedback identifies several critical issues: the NCCL tuner configuration is set for the wrong platform (A3U instead of A4X); the Helm installation instructions in the README are missing a required file parameter; the launcher script contains hardcoded GPU and node counts, unconditionally overwrites the Hugging Face token, ignores parsed configuration overrides, and lacks safety checks for empty artifact directories; the Helm templates lack guards for null GCS mount values; and cloning the repository at runtime across 64 nodes simultaneously risks GitHub rate limiting.
|
|
||
| # Overriding NCCL_SOCKET_IFNAME definition | ||
| export NCCL_SOCKET_IFNAME="eth0,eth1" | ||
| export NCCL_TUNER_CONFIG_PATH=/usr/local/gib/configs/tuner_config_a3u.txtpb |
There was a problem hiding this comment.
The NCCL tuner config path is set to tuner_config_a3u.txtpb, which is for the A3U (H100) platform. Since this recipe is specifically for A4X (GB200 NVL72), it should use the A4 tuner config (tuner_config_a4.txtpb) to ensure optimal network performance and correct topology tuning.
export NCCL_TUNER_CONFIG_PATH=/usr/local/gib/configs/tuner_config_a4.txtpb| helm install llama31-405b-256gpus-fp8mx . \ | ||
| --set queue=${KUEUE_NAME} \ | ||
| --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \ | ||
| --set volumes.gcsMounts[0].mountPath=/runtime-logs \ | ||
| --set workload.envs[0].name=ARTIFACT_DIR \ | ||
| --set workload.envs[0].value=/runtime-logs/llama31-405b-256gpus-fp8mx/artifacts | ||
| ``` |
There was a problem hiding this comment.
The helm install command is missing --set-file workload_launcher=launcher.sh. Without this, the chart will deploy the default "No workload launcher specified" script and fail immediately.
| helm install llama31-405b-256gpus-fp8mx . \ | |
| --set queue=${KUEUE_NAME} \ | |
| --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \ | |
| --set volumes.gcsMounts[0].mountPath=/runtime-logs \ | |
| --set workload.envs[0].name=ARTIFACT_DIR \ | |
| --set workload.envs[0].value=/runtime-logs/llama31-405b-256gpus-fp8mx/artifacts | |
| ``` | |
| helm install llama31-405b-256gpus-fp8mx . \ | |
| --set queue=${KUEUE_NAME} \ | |
| --set-file workload_launcher=launcher.sh \ | |
| --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \ | |
| --set volumes.gcsMounts[0].mountPath=/runtime-logs \ | |
| --set workload.envs[0].name=ARTIFACT_DIR \ | |
| --set workload.envs[0].value=/runtime-logs/llama31-405b-256gpus-fp8mx/artifacts |
| @@ -0,0 +1,169 @@ | |||
| export HF_TOKEN=<YOUR_HF_TOKEN> | |||
There was a problem hiding this comment.
This line unconditionally overwrites HF_TOKEN with the literal string <YOUR_HF_TOKEN>, which prevents users from securely passing their Hugging Face token via Kubernetes Secrets or standard environment variables. It should use a default value expansion instead.
| export HF_TOKEN=<YOUR_HF_TOKEN> | |
| export HF_TOKEN="${HF_TOKEN:-<YOUR_HF_TOKEN>}" |
| --nproc-per-node="4" \ | ||
| --nnodes="64" \ |
There was a problem hiding this comment.
The GPU and node counts are hardcoded in the torchrun command. They should use the environment variables GPUS_PER_NODE and NNODES provided by the JobSet to allow dynamic scaling and customization via Helm values.
| --nproc-per-node="4" \ | |
| --nnodes="64" \ | |
| --nproc-per-node="${GPUS_PER_NODE:-4}" \ | |
| --nnodes="${NNODES:-64}" \ |
| --num_gpus 256 \ | ||
| --gpus_per_node 4 \ |
There was a problem hiding this comment.
The GPU and node counts are hardcoded in the Python script arguments. They should use the environment variables WORLD_SIZE and GPUS_PER_NODE provided by the JobSet to allow dynamic scaling and customization via Helm values.
| --num_gpus 256 \ | |
| --gpus_per_node 4 \ | |
| --num_gpus \$WORLD_SIZE \ | |
| --gpus_per_node \$GPUS_PER_NODE \ |
| logger.log_throughput=True \ | ||
| train.manual_gc_interval=100 |
There was a problem hiding this comment.
The script parses command-line arguments into config_overrides but never actually passes them to the Python training script (run_script.py). This means any configuration overrides provided by the user will be silently ignored. You should append ${config_overrides} to the Python command.
| logger.log_throughput=True \ | |
| train.manual_gc_interval=100 | |
| logger.log_throughput=True \ | |
| train.manual_gc_interval=100 \ | |
| ${config_overrides} |
| 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 |
There was a problem hiding this comment.
If ARTIFACT_DIR is empty or null, mkdir -p will fail, and the subsequent cp command will attempt to copy log files to the root directory (/), which can cause permission errors or unintended file overwrites. Adding a check to ensure ARTIFACT_DIR is non-empty before proceeding is highly recommended.
| 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 |
| {{- range $gcs := $root.Values.volumes.gcsMounts }} | ||
| - name: "{{ $gcs.bucketName }}" | ||
| csi: | ||
| driver: gcsfuse.csi.storage.gke.io | ||
| volumeAttributes: | ||
| bucketName: "{{ $gcs.bucketName }}" | ||
| {{- if $gcs.mountOptions }} | ||
| mountOptions: "{{ $gcs.mountOptions }}" | ||
| {{- end }} | ||
| {{- end}} |
There was a problem hiding this comment.
The gcsMounts loop is evaluated even if gcsVolumes is false or if bucketName is null, which can cause Helm rendering errors or invalid Kubernetes manifests. It should be guarded with {{- if $gcs.bucketName }}.
{{- range $gcs := $root.Values.volumes.gcsMounts }}
{{- if $gcs.bucketName }}
- name: "{{ $gcs.bucketName }}"
csi:
driver: gcsfuse.csi.storage.gke.io
volumeAttributes:
bucketName: "{{ $gcs.bucketName }}"
{{- if $gcs.mountOptions }}
mountOptions: "{{ $gcs.mountOptions }}"
{{- end }}
{{- end }}
{{- end}}| {{- range $gcs := $root.Values.volumes.gcsMounts }} | ||
| - name: "{{ $gcs.bucketName }}" | ||
| mountPath: "{{ $gcs.mountPath }}" | ||
| {{- end }} |
There was a problem hiding this comment.
The gcsMounts volumeMounts loop is evaluated even if gcsVolumes is false or if bucketName is null, which can cause Helm rendering errors or invalid Kubernetes manifests. It should be guarded with {{- if $gcs.bucketName }}.
{{- range $gcs := $root.Values.volumes.gcsMounts }}
{{- if $gcs.bucketName }}
- name: "{{ $gcs.bucketName }}"
mountPath: "{{ $gcs.mountPath }}"
{{- end }}
{{- end }}| cd /opt | ||
| rm -rf Megatron-Bridge | ||
| git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git | ||
| cd Megatron-Bridge | ||
| git checkout fcbb6031103d0ca845c1a54d4fee55ecfcca17b6 | ||
| git submodule update --init --recursive |
There was a problem hiding this comment.
Cloning the Megatron-Bridge repository at runtime on all 64 nodes simultaneously can trigger GitHub rate limits, leading to connection resets and job failures. It is highly recommended to pre-bake this repository into the container image or use a shared volume. If runtime cloning is necessary, consider adding a retry loop or cloning with --depth 1 where possible.
1 participant
Adds the verified 256-GPU LLaMA 3.1 405B FP8-MX pretraining recipe for GKE A4X (NVIDIA GB200 NVL72) using NeMo 26.06.01 and Megatron-Bridge.
Benchmark & Performance Highlights