Add refactored recipes for finetuning

scripts/performance/configs/llama/llama3_llm_finetune.py

@@ -18,7 +18,11 @@
 from utils.precision import get_precision_config
 from utils.utils import get_workload_base_config
 
-from megatron.bridge.recipes.llama import llama3_8b_finetune_config, llama3_70b_finetune_config
+from megatron.bridge.recipes.llama import (
+    llama3_8b_sft_config,
+    llama3_70b_peft_config,
+    llama3_70b_sft_config,
+)
 from megatron.bridge.training.comm_overlap import (
     CommOverlapConfig,
 )
@@ -57,12 +61,11 @@ def llama3_8b_sft_config_gb200(precision: str = "bf16", config_variant: str = "v
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_8b_finetune_config(
-        peft="none",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=16384,
-    )
+    cfg = llama3_8b_sft_config()
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 16384
+    cfg.dataset.seq_length = 16384
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 16384
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
 
@@ -81,12 +84,11 @@ def llama3_8b_sft_config_h100(precision: str = "bf16", config_variant: str = "v1
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_8b_finetune_config(
-        peft="none",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_8b_sft_config()
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
 
@@ -105,12 +107,11 @@ def llama3_70b_sft_config_gb300(precision: str = "bf16", config_variant: str = "
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_70b_finetune_config(
-        peft="none",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_70b_sft_config()
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
 
@@ -141,12 +142,11 @@ def llama3_70b_sft_config_gb200(precision: str = "bf16", config_variant: str = "
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_70b_finetune_config(
-        peft="none",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_70b_sft_config()
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
 
@@ -171,12 +171,11 @@ def llama3_70b_sft_config_h100(precision: str = "bf16", config_variant: str = "v
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_70b_finetune_config(
-        peft="none",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_70b_sft_config()
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
 
@@ -201,12 +200,11 @@ def llama3_70b_lora_config_gb300(precision: str = "bf16", config_variant: str =
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_70b_finetune_config(
-        peft="lora",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_70b_peft_config(peft_scheme="lora")
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     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))
@@ -238,12 +236,11 @@ def llama3_70b_lora_config_gb200(precision: str = "bf16", config_variant: str =
     # 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=seq_length,
-    )
+    cfg = llama3_70b_peft_config(peft_scheme="lora")
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = seq_length
+    cfg.dataset.seq_length = seq_length
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = seq_length
     set_llama3_common_peft_configs(cfg)
     set_workload_base_configs(cfg, base_cfg)
     # Enable pad_cu_seqlens for CUDA graphs compatibility with packed sequences.
@@ -271,12 +268,11 @@ def llama3_70b_lora_config_h100(precision: str = "bf16", config_variant: str = "
     )
     precision_config = get_precision_config(precision)
 
-    cfg = llama3_70b_finetune_config(
-        peft="lora",
-        precision_config=precision_config,
-        packed_sequence=True,
-        seq_length=4096,
-    )
+    cfg = llama3_70b_peft_config(peft_scheme="lora")
+    cfg.mixed_precision = precision_config
+    cfg.model.seq_length = 4096
+    cfg.dataset.seq_length = 4096
+    cfg.dataset.packed_sequence_specs.packed_sequence_size = 4096
     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))

src/megatron/bridge/recipes/common.py

@@ -16,6 +16,8 @@
 
 from megatron.core.distributed import DistributedDataParallelConfig
 
+from megatron.bridge.peft.lora import LoRA
+from megatron.bridge.recipes.utils.finetune_utils import default_squad_config
 from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing
 from megatron.bridge.training.config import (
     CheckpointConfig,
@@ -126,3 +128,210 @@ def _pretrain_common() -> ConfigContainer:
     )
 
     return cfg
+
+
+def _sft_common() -> ConfigContainer:
+    """Create a base SFT (Supervised Fine-Tuning) ConfigContainer with common defaults.
+
+    This function returns a ConfigContainer template with sensible defaults for full SFT
+    (not LoRA/DoRA). The caller MUST set `cfg.model` and `cfg.tokenizer.tokenizer_model`
+    before use.
+
+    Key differences from pre-training:
+    - Uses HFDatasetConfig with SQuAD as default dataset
+    - Lower learning rate (5e-6) suitable for full fine-tuning
+    - Fewer training iterations (1000)
+    - Smaller batch sizes
+    - Supports pretrained_checkpoint loading
+    - No PEFT (full parameter training)
+
+    Returns:
+        ConfigContainer: Base configuration template for full SFT.
+    """
+    # Default output directories
+    base_output_dir = os.path.join(os.getcwd(), "nemo_experiments")
+    run_output_dir = os.path.join(base_output_dir, "default")
+    checkpoint_dir = os.path.join(run_output_dir, "checkpoints")
+    tensorboard_dir = os.path.join(run_output_dir, "tb_logs")
+
+    # Default sequence length for SFT
+    seq_length = 2048
+
+    # Packed sequence is enabled by default for training efficiency
+    # pad_seq_to_mult should be set to context_parallel_size * 2 if CP > 1
+    packed_sequence = True
+    pad_seq_to_mult = 1  # Override in model config if context_parallel_size > 1
+
+    # Optimizer and scheduler with lower LR for full SFT
+    opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing(
+        lr_warmup_iters=50,
+        lr_decay_iters=None,  # Defaults to train_iters during validation
+        max_lr=5e-6,  # Lower LR for full fine-tuning
+        min_lr=0.0,
+        adam_beta2=0.98,  # Common for fine-tuning
+    )
+
+    cfg = ConfigContainer(
+        # Model - MUST be set by each recipe before use
+        model=None,  # type: ignore[arg-type]
+        # Training config - shorter training for SFT
+        train=TrainingConfig(
+            train_iters=1000,
+            global_batch_size=128,
+            micro_batch_size=1,
+        ),
+        validation=ValidationConfig(
+            eval_interval=100,
+            eval_iters=32,
+        ),
+        # Optimizer and scheduler
+        optimizer=opt_cfg,
+        scheduler=scheduler_cfg,
+        # DDP config - minimal settings, model-specific configs can override
+        ddp=DistributedDataParallelConfig(
+            check_for_nan_in_grad=True,
+            grad_reduce_in_fp32=True,
+        ),
+        # Dataset config - uses SQuAD with packed sequences by default
+        dataset=default_squad_config(
+            seq_length=seq_length, packed_sequence=packed_sequence, pad_seq_to_mult=pad_seq_to_mult
+        ),
+        # Logger config
+        logger=LoggerConfig(
+            log_interval=1,
+            tensorboard_dir=tensorboard_dir,
+            log_timers_to_tensorboard=True,
+        ),
+        # Tokenizer - placeholder, each recipe should set tokenizer_model
+        tokenizer=TokenizerConfig(
+            tokenizer_type="HuggingFaceTokenizer",
+            tokenizer_model=None,  # Must be set by each recipe
+        ),
+        # Checkpoint config with pretrained_checkpoint support
+        checkpoint=CheckpointConfig(
+            save_interval=100,
+            save=checkpoint_dir,
+            load=checkpoint_dir,
+            pretrained_checkpoint=None,  # Set to load from pretrained weights
+            ckpt_format="torch_dist",
+            fully_parallel_save=True,
+        ),
+        # RNG config - different seed from pretrain
+        rng=RNGConfig(seed=5678),
+        # Distributed init config
+        dist=DistributedInitConfig(),
+        comm_overlap=None,
+        # Mixed precision - bf16 by default
+        mixed_precision="bf16_mixed",
+        # No PEFT for full SFT
+        peft=None,
+    )
+
+    return cfg
+
+
+def _peft_common() -> ConfigContainer:
+    """Create a base PEFT (Parameter-Efficient Fine-Tuning) ConfigContainer with LoRA defaults.
+
+    This function returns a ConfigContainer template with sensible defaults for PEFT
+    using LoRA. The caller MUST set `cfg.model` and `cfg.tokenizer.tokenizer_model`
+    before use.
+
+    Key differences from full SFT:
+    - Higher learning rate (1e-4) suitable for adapter training
+    - LoRA enabled by default with standard settings (dim=32, alpha=32)
+    - Targets all linear layers: linear_qkv, linear_proj, linear_fc1, linear_fc2
+
+    Returns:
+        ConfigContainer: Base configuration template for PEFT with LoRA.
+    """
+    # Default output directories
+    base_output_dir = os.path.join(os.getcwd(), "nemo_experiments")
+    run_output_dir = os.path.join(base_output_dir, "default")
+    checkpoint_dir = os.path.join(run_output_dir, "checkpoints")
+    tensorboard_dir = os.path.join(run_output_dir, "tb_logs")
+
+    # Default sequence length for PEFT
+    seq_length = 2048
+
+    # Packed sequence is enabled by default for training efficiency
+    # pad_seq_to_mult should be set to context_parallel_size * 2 if CP > 1
+    packed_sequence = True
+    pad_seq_to_mult = 1  # Override in model config if context_parallel_size > 1
+
+    # Optimizer and scheduler with higher LR for PEFT (only training adapters)
+    opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing(
+        lr_warmup_iters=50,
+        lr_decay_iters=None,  # Defaults to train_iters during validation
+        max_lr=1e-4,  # Higher LR for adapter training
+        min_lr=0.0,
+        adam_beta2=0.98,  # Common for fine-tuning
+    )
+
+    cfg = ConfigContainer(
+        # Model - MUST be set by each recipe before use
+        model=None,  # type: ignore[arg-type]
+        # Training config - shorter training for PEFT
+        train=TrainingConfig(
+            train_iters=1000,
+            global_batch_size=128,
+            micro_batch_size=1,
+        ),
+        validation=ValidationConfig(
+            eval_interval=100,
+            eval_iters=32,
+        ),
+        # Optimizer and scheduler
+        optimizer=opt_cfg,
+        scheduler=scheduler_cfg,
+        # DDP config - minimal settings for PEFT
+        ddp=DistributedDataParallelConfig(
+            check_for_nan_in_grad=True,
+            grad_reduce_in_fp32=True,
+        ),
+        # Dataset config - uses SQuAD with packed sequences by default
+        dataset=default_squad_config(
+            seq_length=seq_length, packed_sequence=packed_sequence, pad_seq_to_mult=pad_seq_to_mult
+        ),
+        # Logger config
+        logger=LoggerConfig(
+            log_interval=1,
+            tensorboard_dir=tensorboard_dir,
+            log_timers_to_tensorboard=True,
+        ),
+        # Tokenizer - placeholder, each recipe should set tokenizer_model
+        tokenizer=TokenizerConfig(
+            tokenizer_type="HuggingFaceTokenizer",
+            tokenizer_model=None,  # Must be set by each recipe
+        ),
+        # Checkpoint config with pretrained_checkpoint support
+        checkpoint=CheckpointConfig(
+            save_interval=100,
+            save=checkpoint_dir,
+            load=checkpoint_dir,
+            pretrained_checkpoint=None,  # Set to load from pretrained weights
+            ckpt_format="torch_dist",
+            fully_parallel_save=True,
+        ),
+        # RNG config - different seed from pretrain
+        rng=RNGConfig(seed=5678),
+        # Distributed init config
+        dist=DistributedInitConfig(),
+        comm_overlap=None,
+        # Mixed precision - bf16 by default
+        mixed_precision="bf16_mixed",
+        # LoRA config with standard defaults
+        peft=LoRA(
+            target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
+            dim=32,
+            alpha=32,
+            dropout=0.0,
+            dropout_position="pre",
+            lora_A_init_method="xavier",
+            lora_B_init_method="zero",
+            a2a_experimental=False,
+            lora_dtype=None,  # Uses model's dtype
+        ),
+    )
+
+    return cfg