[automodel] consolidate sft peft scripts (#13634)

* move files

Signed-off-by: Alexandros Koumparoulis <[email protected]>

* remove peft

Signed-off-by: Alexandros Koumparoulis <[email protected]>

* add --lora option

Signed-off-by: Alexandros Koumparoulis <[email protected]>

* fix paths

Signed-off-by: Alexandros Koumparoulis <[email protected]>

---------

Signed-off-by: Alexandros Koumparoulis <[email protected]>

Author: akoumpa

examples/llm/sft/automodel.py renamed to examples/llm/finetune/automodel.py

@@ -27,12 +27,13 @@
 
 # Run this example with torchrun, for example:
 # torchrun --nproc-per-node=8 \
-#   examples/llm/peft/automodel.py \
+#   examples/llm/finetune/automodel.py \
 #   --strategy fsdp2 \
 #   --devices 8 \
 #   --model meta-llama/Llama-3.2-1B \
 #   --ckpt-folder "output"
 #
+# For PEFT please also pass --lora
 # Note: ensure that the --nproc-per-node and --devices values match.
 
 
@@ -201,7 +202,7 @@ def logger(ckpt_folder, save_every_n_train_steps) -> nl.NeMoLogger:
     )
 
     return nl.NeMoLogger(
-        name="nemo2_sft",
+        name="nemo2_finetune",
         log_dir=ckpt_folder,
         use_datetime_version=False,  # must be false if using auto resume
         ckpt=ckpt,
@@ -210,7 +211,7 @@ def logger(ckpt_folder, save_every_n_train_steps) -> nl.NeMoLogger:
 
 
 def main():
-    """Example script to run SFT with a HF transformers-instantiated model on squad."""
+    """Example script to run SFT/PEFT with a HF transformers-instantiated model on squad."""
     import argparse
 
     parser = argparse.ArgumentParser()
@@ -308,6 +309,7 @@ def main():
         'run with --attn-implementation=flash_attention_2',
     )
     parser.add_argument('--start-of-turn-token', default=None, help='Chat turn token')
+    parser.add_argument('--lora', action='store_true', help='Enables PEFT (LoRA) finetuning; Default: off (SFT).')
 
     args = parser.parse_args()
     if args.dp_size is None:
@@ -371,7 +373,7 @@ def main():
         trust_remote_code=args.trust_remote_code,
         use_liger_kernel=args.liger,
         enable_grad_ckpt=args.enable_grad_ckpt,
-        use_linear_ce_loss=not args.no_lce,
+        use_linear_ce_loss=args.no_lce,
     )
 
     assert (
@@ -427,7 +429,13 @@ def main():
             limit_dataset_samples=args.limit_dataset_samples,
             fp8=args.fp8,
         )
-
+    if args.peft:
+        peft = llm.peft.LoRA(
+            target_modules=['*_proj'],
+            dim=8,
+        )
+    else:
+        peft = None
     llm.api.finetune(
         model=model,
         data=dataset,
@@ -448,6 +456,7 @@ def main():
             callbacks=callbacks,
             precision="bf16-mixed",
         ),
+        peft=peft,
         optim=optimizer,
         log=logger(args.ckpt_folder, args.max_steps // 2),
         resume=resume,