[#recipebot] Add Llama 3.1 405B nvfp4 Megatron-Bridge A4X recipe

[wenqinI]

Adding a recipe for Llama 3.1 405B Megatron-Bridge on A4X.

run_id: wenqinyang-ubench-m9hm78qy

Final Diff: https://paste.googleplex.com/6704000242679808?raw

[gemini-code-assist]

Code Review

This pull request adds a Helm chart and launcher recipe for pretraining llama31-405b workloads on GKE using the Megatron-Bridge framework. Feedback highlights several critical improvements: passing parsed configuration overrides to the training script, using environment variables instead of hardcoded node/GPU counts in torchrun, avoiding memory accumulation in the log-streaming script, adding retries for runtime git clones, defensively checking privileges before executing dmesg, and correcting a version mismatch in the documentation.

[gemini-code-assist]

The parsed config_overrides are never passed to the Python execution command. This means any custom arguments passed via Helm (such as trainer.max_steps=100 as shown in the README examples) will be completely ignored, and the job will run with default parameters. Append ${config_overrides} to the python command to ensure overrides are applied.

  python scripts/performance/run_script.py \
    --model_family_name llama \
    --model_recipe_name llama31_405b \
    --config_variant v2 \
    --gpu gb200 \
    --num_gpus \${WORLD_SIZE} \
    --gpus_per_node \${GPUS_PER_NODE} \
    --compute_dtype nvfp4 \
    --seq_length 8192 \
    --global_batch_size 1536 \
    --micro_batch_size 1 \
    --tensor_model_parallel_size 4 \
    --pipeline_model_parallel_size 16 \
    --virtual_pipeline_model_parallel_size 8 \
    --context_parallel_size 1 \
    --expert_model_parallel_size 1 \
    --expert_tensor_parallel_size 1 \
    --cuda_graph_impl none \
    --max_steps 30 \
    logger.log_throughput=True \
    dataset.num_workers=0 \
    train.manual_gc_interval=100 \
    ${config_overrides}