Copybara import of gpu-recipes:

  - ddc7b4335809a0077ddac5cba1660cff3a937bf3 Merge "Modifications on Llama-3.1-405B recipes" into main

GitOrigin-RevId: ddc7b4335809a0077ddac5cba1660cff3a937bf3

Author: Copybara

README.md

@@ -32,6 +32,7 @@ Welcome to the reproducible benchmark recipes repository for GPUs! This reposito
 | ---------------- | ---------------- | --------- | ------------------- | ------------ | ------------------ |
 | **Llama-3.1-70B**     | [A3 Ultra (NVIDIA H200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-ultra-vms)    | MaxText  | Pre-training   | GKE          | [Link](./training/a3ultra/llama-3.1-70b/maxtext-pretraining-gke/README.md)
 | **Llama-3.1-70B**     | [A3 Ultra (NVIDIA H200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-ultra-vms)    | NeMo  | Pre-training   | GKE          | [Link](./training/a3ultra/llama-3.1-70b/nemo-pretraining-gke/README.md)
+| **Llama-3.1-405B**     | [A3 Ultra (NVIDIA H200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-ultra-vms)    | NeMo  | Pre-training   | GKE          | [Link](./training/a3ultra/llama-3.1-405b/nemo-pretraining-gke/README.md)
 | **Mixtral-8-7B**     | [A3 Ultra (NVIDIA H200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-ultra-vms)    | MaxText  | Pre-training   | GKE          | [Link](./training/a3ultra/mixtral-8x7b/maxtext-pretraining-gke/README.md)
 | **Mixtral-8-7B**     | [A3 Ultra (NVIDIA H200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-ultra-vms)    | NeMo  | Pre-training   | GKE          | [Link](./training/a3ultra/mixtral-8x7b/nemo-pretraining-gke/README.md)            |
 

src/frameworks/a3ultra/nemo-configs/llama-3.1-405b-576gpus-a3ultra-bf16.yaml

@@ -1,5 +1,5 @@
 run:
-  name: llama-3.1-405b-a3u-fp8
+  name: llama-3.1-405b-a3u-bf16
   time_limit: 0-02:30:00
   dependency: singleton
 trainer:
@@ -124,7 +124,7 @@ model:
   deterministic_mode: false
   transformer_engine: true
   fp8: false
-  ub_tp_comm_overlap: false
+  ub_tp_comm_overlap: true
   use_flash_attention: true
   fsdp: false
   fsdp_sharding_strategy: full

src/frameworks/a3ultra/nemo-configs/llama-3.1-405b-576gpus-a3ultra-fp8.yaml

@@ -130,7 +130,7 @@ model:
   fp8_interval: 1
   fp8_amax_history_len: 1024
   fp8_amax_compute_algo: max
-  ub_tp_comm_overlap: false
+  ub_tp_comm_overlap: true
   use_flash_attention: true
   fsdp: false
   fsdp_sharding_strategy: full

training/a3ultra/llama-3.1-405b/nemo-pretraining-gke/values.yaml

@@ -40,4 +40,6 @@ network:
   - name: NCCL_DEBUG
     value: "VERSION"
   - name: NCCL_WORK_FIFO_DEPTH 
-    value: "4194304"
+    value: "4194304"
+  - name: NVTE_UB_SOCKET_IFNAME
+    value: "eth1"