Add A4X MAX Llama-3.1-405B FP8cs 256 GPUs recipe - #273
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 Llama 3.1 405B workloads on GKE using the Megatron-Bridge framework with 256 GPUs. The recipe includes orchestration templates, a launcher script, and configuration files. The review feedback highlights several critical improvements: adding robust error handling (set -eo pipefail) and safety checks for environment variables (like HF_TOKEN and ARTIFACT_DIR) in the launcher script, correcting an invalid ldconfig call, fixing a YAML syntax typo in values.yaml, and strongly recommending baking system dependencies (such as DOCA-OFED and NCCL plugins) directly into the container image rather than installing them at runtime across 64 nodes.
| @@ -0,0 +1,166 @@ | |||
| usage() | |||
There was a problem hiding this comment.
The script lacks set -eo pipefail. Without it, if any command (such as git clone or torchrun) fails, the script will continue executing. Furthermore, because torchrun is piped to python3 for logging, any failure in torchrun will be masked and the Kubernetes Job will incorrectly report success. Adding set -eo pipefail at the top ensures the script exits immediately on any command or pipeline failure.
| usage() | |
| set -eo pipefail | |
| usage() |
| echo "VERSION_DIAGNOSTICS: ${kv}" | ||
|
|
||
|
|
||
| export HF_TOKEN=<YOUR_HF_TOKEN> |
There was a problem hiding this comment.
Overwriting HF_TOKEN with the literal placeholder <YOUR_HF_TOKEN> will discard any valid Hugging Face token passed via Kubernetes environment variables. Use a default fallback pattern instead so that if HF_TOKEN is already set in the environment, it is preserved.
| export HF_TOKEN=<YOUR_HF_TOKEN> | |
| export HF_TOKEN="${HF_TOKEN:-<YOUR_HF_TOKEN>}" |
| export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH" | ||
| ldconfig "$LD_LIBRARY_PATH" |
There was a problem hiding this comment.
ldconfig does not accept colon-separated paths (like LD_LIBRARY_PATH) and will fail with an error. Since the dynamic linker automatically searches directories in LD_LIBRARY_PATH at runtime, running ldconfig on it is redundant and incorrect. You can safely remove the ldconfig call.
| export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH" | |
| ldconfig "$LD_LIBRARY_PATH" | |
| export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH" |
| 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 not set or is empty, mkdir -p "${ARTIFACT_DIR}" will fail and crash the script. Add a check to ensure ARTIFACT_DIR is non-empty before attempting to create the directory and copy logs.
| 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 |
| # 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/" | ||
| 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 |
There was a problem hiding this comment.
Installing heavy system packages, drivers (DOCA-OFED), and NCCL plugins at runtime via apt is a major anti-pattern for large-scale Kubernetes jobs. With 64 nodes (256 GPUs) starting simultaneously, this will cause significant startup latency, risk hitting external repository rate limits, and make the job highly vulnerable to network/repository downtime.
Recommendation: Bake these dependencies (DOCA-OFED, NCCL, and nccl-gib-plugins) directly into a custom container image instead of installing them at runtime.
| network: | ||
| hostNetwork: false | ||
| hostIPC: true | ||
| subnetworks[]: null |
There was a problem hiding this comment.
subnetworks[]: null is invalid/atypical YAML syntax and appears to be a typo. It should be defined as a standard list key subnetworks: [].
subnetworks: []
1 participant
Add A4X MAX Llama-3.1-405B 256 GPUs FP8cs recipe