Add resume for adapter_v2, enable continued finetuning for adapter (#1354)

Co-authored-by: alal <[email protected]>

Author: altria-zewei-wang

litgpt/finetune/adapter_v2.py

@@ -37,6 +37,7 @@
     instantiate_bnb_optimizer,
     instantiate_torch_optimizer,
     load_checkpoint,
+    load_checkpoint_update,
     num_parameters,
     parse_devices,
     save_hyperparameters,
@@ -51,6 +52,7 @@ def setup(
     quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8-training"]] = None,
     devices: Union[int, str] = 1,
     num_nodes: int = 1,
+    resume: Optional[bool] = False,
     data: Optional[DataModule] = None,
     train: TrainArgs = TrainArgs(
         save_interval=1000,
@@ -137,7 +139,7 @@ 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, resume, checkpoint_dir, out_dir, train, eval, optimizer, num_nodes)
 
 
 def main(
@@ -146,6 +148,7 @@ def main(
     seed: int,
     config: Config,
     data: DataModule,
+    resume: bool,
     checkpoint_dir: Path,
     out_dir: Path,
     train: TrainArgs,
@@ -191,9 +194,22 @@ def main(
 
     optimizer = fabric.setup_optimizers(optimizer)
     scheduler = get_lr_scheduler(optimizer, warmup_steps=train.lr_warmup_steps, max_steps=lr_max_steps)
+    if resume:
+        # Finding last trace of adapter training
+        try:
+            resume = max(out_dir.rglob("step-*/*.pth.adapter_v2"), key=(lambda p: int(p.parent.name.split("-")[1])))
+            fabric.print(f"Resuming training from {resume}")
+            load_checkpoint_update(fabric, resume, model, checkpoint_path, strict=False)
+            resume = True
+        except ValueError:
+            fabric.print("No previous adapter found. Finetune from start.")
+            resume = False
+            load_checkpoint(fabric, model, checkpoint_path, strict=False)
+    else:
+        # strict=False because missing keys due to Adapter weights not contained in state dict
+        load_checkpoint(fabric, model, checkpoint_path, strict=False)
 
-    # strict=False because missing keys due to Adapter weights not contained in state dict
-    load_checkpoint(fabric, model, checkpoint_path, strict=False)
+    mark_only_adapter_v2_as_trainable(model)
 
     train_time = time.perf_counter()
     token_counts = fit(
@@ -204,6 +220,7 @@ def main(
         train_dataloader=train_dataloader,
         val_dataloader=val_dataloader,
         devices=devices,
+        resume=resume,
         num_nodes=num_nodes,
         checkpoint_dir=checkpoint_dir,
         out_dir=out_dir,
@@ -241,6 +258,7 @@ def fit(
     train_dataloader: DataLoader,
     val_dataloader: DataLoader,
     devices: int,
+    resume: bool,
     checkpoint_dir: Path,
     out_dir: Path,
     train: TrainArgs,
@@ -283,7 +301,15 @@ def fit(
         "raw_tokens_plus_prompt_template_and_padding": torch.tensor(0, device=fabric.device, dtype=torch.long),
     }
 
-    while step_count < max_steps:
+    if not resume:
+        try:
+            iter_match = max(out_dir.rglob("step-*/*.pth.adapter_v2"), key=lambda p: int(p.parent.name.split("-")[1]))
+            step_count = int(iter_match.parent.name.split("-")[1]) if iter_match else 0
+        except ValueError:
+            step_count = 0
+
+    fabric.print(f"Starting at step count {step_count}")
+    while step_count < max_steps and train_iterator.epoch < train.epochs:
         iter_num += 1
         iter_t0 = time.perf_counter()
         batch = next(train_iterator)

litgpt/utils.py

@@ -391,6 +391,19 @@ def load_checkpoint(fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, s
         model.load_state_dict(state_dict, strict=strict)
 
 
+def load_checkpoint_update(
+    fabric: L.Fabric, adapter_path: Path, model: nn.Module, checkpoint_path: Path, strict: bool = True
+) -> None:
+    if isinstance(fabric.strategy, FSDPStrategy):
+        fabric.load_raw(checkpoint_path, model, strict=strict)
+    else:
+        state_dict = lazy_load(checkpoint_path)
+        state_dict = state_dict.get("model", state_dict)
+        adapter_cp = lazy_load(adapter_path)
+        state_dict.update(adapter_cp)
+        model.load_state_dict(state_dict, strict=strict)
+
+
 def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
     flops_per_token = 2 * n_params  # each parameter is used for a MAC (2 FLOPS) per network operation
     # this assumes that all samples have a fixed length equal to the block size