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

[ngu3]

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

[gemini-code-assist]

Code Review

This pull request introduces a Helm chart recipe for pretraining Llama 3.1 405B workloads on a4x-max GKE node pools using the Nvidia Megatron-Bridge framework. Feedback focuses on improving deployment flexibility and robustness: replacing hardcoded 'default' namespaces in the JobSet metadata and FQDN environment variables with release namespace templates, avoiding runtime package installations in the container startup command, and securing the 'HF_TOKEN' environment variable. Additionally, recommendations were made to fix invalid YAML syntax in 'values.yaml', use a quoted heredoc delimiter in the launcher script, ensure portable timestamping instead of relying on GNU awk extensions, and sanitize the workload name in the documentation to comply with Kubernetes naming standards.

[gemini-code-assist]

Hardcoding the namespace to default in the JobSet metadata prevents the chart from being deployed in other namespaces. If a user deploys the Helm release in a custom namespace (e.g., helm install -n custom-ns), the JobSet will still be created in the default namespace, while other resources (like ConfigMaps and Services) will be created in custom-ns. This will cause the JobSet to fail to find its referenced ConfigMaps.

  namespace: "{{ .Release.Namespace }}"

[gemini-code-assist]

The environment variables RANK_0_FQDN, DOMAIN_NAME, and MASTER_ADDR hardcode the default namespace in their FQDN values. If the Helm chart is deployed in a custom namespace, these addresses will resolve incorrectly or fail to resolve entirely, breaking the torch distributed initialization. Use {{ .Release.Namespace }} instead.

              - 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"

[gemini-code-assist]

Running apt update and installing complex packages like doca-ofed-userspace and nccl-gib-plugins at runtime inside the container's startup command is a major anti-pattern. It significantly slows down pod startup, introduces flakiness (due to network/repository downtime), and will fail completely in private GKE clusters without internet access. These dependencies should be pre-baked into a custom container image, or installed via a pre-built init container (similar to the nccl-plugin-installer pattern used in other charts).

[gemini-code-assist]

Hardcoding export HF_TOKEN=<YOUR_HF_TOKEN> in the launcher script will overwrite any HF_TOKEN environment variable passed to the container with the literal string <YOUR_HF_TOKEN>. Since the token is sensitive, it should be passed securely via Kubernetes Secrets and referenced in the environment variables, rather than being hardcoded or templated directly in the script.

Suggested change

	export HF_TOKEN=<YOUR_HF_TOKEN>
	if [[ -z "${HF_TOKEN}" ]]; then
	echo "Warning: HF_TOKEN is not set. Hugging Face downloads may fail if authentication is required."
	fi

[gemini-code-assist]

Using an unquoted heredoc delimiter cat <<- EOM causes the shell to perform parameter expansion and command substitution at definition time. This forces you to escape variables like \$RANK and \$((LOCAL_RANK/2)). Quoting the delimiter as cat <<- 'EOM' prevents this expansion, making the script cleaner and less error-prone.

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
)

[gemini-code-assist]

Using strftime in awk is a GNU awk (gawk) extension. In many minimal container images (including Ubuntu-based ones), the default awk is mawk, which does not support strftime and will crash with a syntax error. Since Python 3 is guaranteed to be installed in this environment, you can use a portable Python one-liner to prepend timestamps to stdout/stderr unbuffered.

Suggested change

	--no-python stdbuf -oL -eL bash worker_command.sh 2>&1 | awk '{ print strftime("[%Y-%m-%d %H:%M:%S]"), $0; fflush(); }'
	--no-python stdbuf -oL -eL bash worker_command.sh 2>&1 | python3 -u -c 'import sys, datetime; [sys.stdout.write(f"[{datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")}]: {line}") for line in iter(sys.stdin.readline, "")]'

[gemini-code-assist]

The key subnetworks[]: null is invalid YAML syntax and is likely a typo. It should be defined as subnetworks: [] or subnetworks: null.

  subnetworks: null

[gemini-code-assist]

Using $USER directly in the WORKLOAD_NAME can cause Helm installation failures if the user's username contains uppercase letters, underscores, or other characters that do not conform to Kubernetes RFC 1123 naming standards (which only allow lowercase alphanumeric characters, '-' or '.'). Sanitizing the username using tr ensures a valid resource name.

Suggested change

	export WORKLOAD_NAME=$USER-a4x-max-llama3-1-405b-256gpus
	export WORKLOAD_NAME=$(echo "${USER}-a4x-max-llama3-1-405b-256gpus" | tr '[:upper:]' '[:lower:]' | tr '_' '-')