finetune_lora upgrades (#2086)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ysjprojects

litgpt/__main__.py

@@ -12,6 +12,7 @@
 from litgpt.finetune.adapter_v2 import setup as finetune_adapter_v2_fn
 from litgpt.finetune.full import setup as finetune_full_fn
 from litgpt.finetune.lora import setup as finetune_lora_fn
+from litgpt.finetune.lora_legacy import setup as finetune_lora_legacy_fn
 from litgpt.generate.adapter import main as generate_adapter_fn
 from litgpt.generate.adapter_v2 import main as generate_adapter_v2_fn
 from litgpt.generate.base import main as generate_base_fn
@@ -35,6 +36,7 @@ def main() -> None:
         "chat": chat_fn,
         "finetune": finetune_lora_fn,
         "finetune_lora": finetune_lora_fn,
+        "finetune_lora_legacy": finetune_lora_legacy_fn,
         "finetune_full": finetune_full_fn,
         "finetune_adapter": finetune_adapter_fn,
         "finetune_adapter_v2": finetune_adapter_v2_fn,

litgpt/args.py

@@ -28,6 +28,8 @@ class TrainArgs:
     """Total number of tokens to train on"""
     max_steps: Optional[int] = None
     """Limits the number of optimizer steps to run"""
+    max_time: Optional[float] = None
+    """Limits the number of seconds to train for"""
     max_seq_length: Optional[int] = None
     """Limits the length of samples"""
     tie_embeddings: Optional[bool] = None

litgpt/finetune/lora.py

@@ -11,7 +11,7 @@
 import lightning as L
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
-from lightning.fabric.strategies import FSDPStrategy
+from lightning.fabric.strategies import ModelParallelStrategy
 from lightning.fabric.utilities import ThroughputMonitor
 from lightning_utilities.core.imports import RequirementCache
 from torch.utils.data import ConcatDataset, DataLoader
@@ -20,7 +20,7 @@
 from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import Alpaca, DataModule
 from litgpt.generate.base import generate
-from litgpt.lora import GPT, Block, Config, lora_filter, mark_only_lora_as_trainable
+from litgpt.lora import GPT, Block, Config, mark_only_lora_as_trainable
 from litgpt.prompts import save_prompt_style
 from litgpt.scripts.merge_lora import merge_lora
 from litgpt.tokenizer import Tokenizer
@@ -70,6 +70,7 @@ def setup(
         lr_warmup_steps=100,
         epochs=5,
         max_seq_length=None,
+        max_time=None,
     ),
     log: LogArgs = LogArgs(),
     eval: EvalArgs = EvalArgs(interval=100, max_new_tokens=100, max_iters=100),
@@ -105,6 +106,7 @@ def setup(
         seed: The random seed to use for reproducibility.
         access_token: Optional API token to access models with restrictions.
     """
+
     checkpoint_dir = auto_download_checkpoint(model_name=checkpoint_dir, access_token=access_token)
     pprint(locals())
     data = Alpaca() if data is None else data
@@ -152,12 +154,10 @@ def setup(
                 "Quantization is currently not supported for multi-GPU training. Please set devices=1 and num_nodes=1"
                 " when using the --quantize flag."
             )
-        strategy = FSDPStrategy(
-            auto_wrap_policy={torch.nn.Linear},
-            activation_checkpointing_policy={Block},
-            state_dict_type="full",
-            limit_all_gathers=True,
-            cpu_offload=False,
+        strategy = ModelParallelStrategy(
+            parallelize_fn=parallelize_fn,
+            data_parallel_size=devices * num_nodes,
+            tensor_parallel_size=1,
         )
     else:
         strategy = "auto"
@@ -174,7 +174,9 @@ def setup(
     if torch.cuda.is_available() and devices > 1:
         check_nvlink_connectivity(fabric)
 
-    fabric.launch(main, devices, seed, config, data, checkpoint_dir, out_dir, train, eval, optimizer, num_nodes)
+    fabric.launch(
+        main, devices, seed, config, data, checkpoint_dir, out_dir, train, eval, optimizer, num_nodes, precision
+    )
 
 
 def main(
@@ -189,6 +191,7 @@ def main(
     eval: EvalArgs,
     optimizer: Union[str, Dict],
     num_nodes: int = 1,
+    precision: Optional[str] = None,
 ) -> None:
     validate_args(train, eval)
 
@@ -229,7 +232,6 @@ def main(
     optimizer = fabric.setup_optimizers(optimizer)
     scheduler = get_lr_scheduler(optimizer, warmup_steps=train.lr_warmup_steps, max_steps=lr_max_steps)
 
-    # strict=False because missing keys due to LoRA weights not contained in state dict
     load_checkpoint(fabric, model, checkpoint_path, strict=False)
 
     train_time = time.perf_counter()
@@ -264,12 +266,19 @@ def main(
     save_path = out_dir / "final" / "lit_model.pth.lora"
     save_path.parent.mkdir(parents=True, exist_ok=True)
     save_lora_checkpoint(fabric, model, save_path)
+
+    fabric.barrier()
     if fabric.global_rank == 0:
         # Copy checkpoint files from original checkpoint dir
         copy_config_files(checkpoint_dir, save_path.parent)
         save_hyperparameters(setup, save_path.parent)
         save_prompt_style(data.prompt_style, save_path.parent)
-        merge_lora(checkpoint_dir=save_path.parent)
+        merge_lora(
+            checkpoint_dir=save_path.parent,
+            pretrained_checkpoint_dir=checkpoint_dir,
+            precision=precision,
+        )
+    fabric.barrier()
 
 
 def fit(
@@ -316,6 +325,8 @@ def fit(
     total_lengths = 0
     total_t0 = time.perf_counter()
 
+    max_time = train.max_time or float("inf")
+
     token_counts = {
         "raw_tokens": torch.tensor(0, device=fabric.device, dtype=torch.long),
         "raw_tokens_plus_prompt_template": torch.tensor(0, device=fabric.device, dtype=torch.long),
@@ -327,6 +338,12 @@ def fit(
         iter_t0 = time.perf_counter()
         batch = next(train_iterator)
         if train_iterator.epoch >= train.epochs:
+            generate_example(fabric, model, tokenizer, eval, data)
+            fabric.print(f"Number of epochs {train.epochs} reached, stopping training...")
+            break
+        if iter_t0 - total_t0 > max_time:
+            generate_example(fabric, model, tokenizer, eval, data)
+            fabric.print(f"Max time ({max_time / 60.0:.2f}m) reached, stopping training...")
             break
         input_ids, targets = batch["input_ids"], batch["labels"]
 
@@ -497,9 +514,45 @@ def get_longest_seq_length(data: List[Dict]) -> Tuple[int, int]:
     return longest_seq_length, longest_seq_ix
 
 
+def parallelize_fn(model, device_mesh, activation_checkpointing=True):
+    from torch.distributed._composable.fsdp.fully_shard import fully_shard
+    from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper, checkpoint_wrapper
+
+    if activation_checkpointing:
+        model.transformer.h = torch.nn.ModuleList(
+            [checkpoint_wrapper(el, preserve_rng_state=False) for el in model.transformer.h]
+        )
+
+    dp_mesh = device_mesh["data_parallel"]
+
+    for m in reversed(list(model.modules())):
+        if (
+            (isinstance(m, torch.nn.Linear) and m.weight.requires_grad)
+            or isinstance(m, CheckpointWrapper)
+            or isinstance(m, Block)
+        ):
+            fully_shard(m, mesh=dp_mesh)
+
+    fully_shard(model, mesh=dp_mesh)
+
+    return model
+
+
 def save_lora_checkpoint(fabric: L.Fabric, model: torch.nn.Module, file_path: Path) -> None:
-    fabric.print(f"Saving LoRA weights to {str(file_path)!r}")
-    fabric.save(file_path, {"model": model}, filter={"model": lora_filter})
+    cpu_state_dict = {}
+    sharded_sd = model.state_dict()
+    for param_name, param in sharded_sd.items():
+        if "lora_" not in param_name:
+            continue
+        if param.is_cpu:
+            param = param.to(fabric.device)
+        if hasattr(param, "_local_tensor"):
+            param = param.full_tensor()
+        if fabric.is_global_zero:
+            cpu_state_dict[param_name] = param.cpu()
+        fabric.barrier()
+    if fabric.is_global_zero:
+        torch.save({"model": cpu_state_dict}, file_path)
 
 
 def validate_args(train: TrainArgs, eval: EvalArgs) -> None: