Enabling TP Comm Overlap and Packed Sequencing Configs for LLAMA3 70B… (#2247)

Signed-off-by: Raghav Hrishikeshan Mukundan <rmukundan@nvidia.com>
Signed-off-by: Raghav Hrishikeshan Mukundan <102543536+rhmukundan@users.noreply.github.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>

Author: rhmukundan

scripts/performance/configs/llama/llama3_llm_finetune.py

@@ -209,6 +209,7 @@ def llama3_70b_lora_config_gb300(precision: str = "bf16", config_variant: str =
     )
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
+    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=bool(cfg.model.tensor_model_parallel_size > 1))
 
     # Enable pad_cu_seqlens for CUDA graphs compatibility with packed sequences.
     # This ensures consistent cu_seqlens tensor shapes across batches, which is required
@@ -239,6 +240,12 @@ def llama3_70b_lora_config_gb200(precision: str = "bf16", config_variant: str =
     )
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
+    # Enable pad_cu_seqlens for CUDA graphs compatibility with packed sequences.
+    # This ensures consistent cu_seqlens tensor shapes across batches, which is required
+    # for CUDA graphs and avoids NaN issues in attention kernels.
+    cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
+    cfg.dataset.dataset_kwargs["pad_to_max_length"] = True
+    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=bool(cfg.model.tensor_model_parallel_size > 1))
 
     return cfg
 
@@ -263,5 +270,6 @@ def llama3_70b_lora_config_h100(precision: str = "bf16", config_variant: str = "
     )
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
+    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=bool(cfg.model.tensor_model_parallel_size > 1))
 
     return cfg