Updating Configs for LLAMA3 70B LoRa (NVIDIA-NeMo#2292)

Signed-off-by: Raghav Hrishikeshan Mukundan <rmukundan@nvidia.com>
Signed-off-by: sowmen <sowmendipta@gmail.com>

Author: rhmukundan

scripts/performance/configs/llama/llama3_llm_finetune.py

@@ -205,7 +205,7 @@ def llama3_70b_lora_config_gb300(precision: str = "bf16", config_variant: str =
         peft="lora",
         precision_config=precision_config,
         packed_sequence=True,
-        seq_length=2048,
+        seq_length=4096,
     )
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
@@ -234,11 +234,14 @@ def llama3_70b_lora_config_gb200(precision: str = "bf16", config_variant: str =
     )
     precision_config = get_precision_config(precision)
 
+    # BF16 uses seq_length=2048, FP8 variants use seq_length=4096
+    seq_length = 2048 if precision.lower() == "bf16" else 4096
+
     cfg = llama3_70b_finetune_config(
         peft="lora",
         precision_config=precision_config,
         packed_sequence=True,
-        seq_length=2048,
+        seq_length=seq_length,
     )
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)

scripts/performance/configs/llama/llama3_workload_base_configs.py

@@ -554,55 +554,72 @@
     BASE_LLAMA3_70B_CONFIG,
     num_gpus=8,
     peft="lora",
-    # pipeline_model_parallel_size=4,
-    # virtual_pipeline_model_parallel_size=20,
+    tensor_model_parallel_size=1,
+    pipeline_model_parallel_size=1,
+    context_parallel_size=1,
     micro_batch_size=1,
-    global_batch_size=64,
+    global_batch_size=32,
     cuda_graph_impl="transformer_engine",
     cuda_graph_scope="mlp",
 )
 
 LLAMA3_70B_LORA_CONFIG_GB300_BF16_V1 = _LLAMA3_70B_LORA_CONFIG_GB300
 LLAMA3_70B_LORA_CONFIG_GB300_FP8_CS_V1 = _LLAMA3_70B_LORA_CONFIG_GB300
 LLAMA3_70B_LORA_CONFIG_GB300_FP8_MX_V1 = replace(
-    LLAMA3_70B_LORA_CONFIG_GB300_FP8_CS_V1,
-    pipeline_model_parallel_size=2,  # PP=1 is OOM
+    _LLAMA3_70B_LORA_CONFIG_GB300,
+    pipeline_model_parallel_size=2,
 )
 
 
 _LLAMA3_70B_LORA_CONFIG_GB200 = replace(
     BASE_LLAMA3_70B_CONFIG,
     num_gpus=8,
     peft="lora",
-    pipeline_model_parallel_size=4,
-    virtual_pipeline_model_parallel_size=20,
+    tensor_model_parallel_size=1,
+    pipeline_model_parallel_size=1,
+    context_parallel_size=1,
     micro_batch_size=1,
     global_batch_size=64,
     cuda_graph_impl="transformer_engine",
     cuda_graph_scope="mlp",
 )
 
 LLAMA3_70B_LORA_CONFIG_GB200_BF16_V1 = _LLAMA3_70B_LORA_CONFIG_GB200
-LLAMA3_70B_LORA_CONFIG_GB200_FP8_CS_V1 = _LLAMA3_70B_LORA_CONFIG_GB200
+LLAMA3_70B_LORA_CONFIG_GB200_FP8_CS_V1 = replace(
+    BASE_LLAMA3_70B_CONFIG,
+    num_gpus=8,
+    peft="lora",
+    tensor_model_parallel_size=1,
+    pipeline_model_parallel_size=2,
+    context_parallel_size=1,
+    micro_batch_size=1,
+    global_batch_size=32,
+    cuda_graph_impl="transformer_engine",
+    cuda_graph_scope="mlp",
+)
 LLAMA3_70B_LORA_CONFIG_GB200_FP8_MX_V1 = LLAMA3_70B_LORA_CONFIG_GB200_FP8_CS_V1
 
 
 _LLAMA3_70B_LORA_CONFIG_H100 = replace(
     BASE_LLAMA3_70B_CONFIG,
     num_gpus=8,
     peft="lora",
-    tensor_model_parallel_size=2,
+    tensor_model_parallel_size=1,
     pipeline_model_parallel_size=4,
+    context_parallel_size=1,
     virtual_pipeline_model_parallel_size=20,
     micro_batch_size=1,
     global_batch_size=32,
 )
 
 LLAMA3_70B_LORA_CONFIG_H100_BF16_V1 = replace(
     _LLAMA3_70B_LORA_CONFIG_H100,
-    recompute_num_layers=2,
+    recompute_num_layers=1,
+)
+LLAMA3_70B_LORA_CONFIG_H100_FP8_CS_V1 = replace(
+    _LLAMA3_70B_LORA_CONFIG_H100,
+    tensor_model_parallel_size=2,
 )
-LLAMA3_70B_LORA_CONFIG_H100_FP8_CS_V1 = _LLAMA3_70B_LORA_CONFIG_H100
 
 
 __all__ = [