support resume_from_checkpoint

llama_lora/lib/finetune.py

@@ -33,7 +33,7 @@ def train(
     num_train_epochs: int = 3,
     learning_rate: float = 3e-4,
     cutoff_len: int = 256,
-    val_set_size: int = 2000,  # TODO: use percentage
+    val_set_size: int = 2000,
     # lora hyperparams
     lora_r: int = 8,
     lora_alpha: int = 16,
@@ -46,7 +46,7 @@ def train(
     train_on_inputs: bool = True,  # if False, masks out inputs in loss
     group_by_length: bool = False,  # faster, but produces an odd training loss curve
     # either training checkpoint or final adapter
-    resume_from_checkpoint: str = None,
+    resume_from_checkpoint = None,
     save_steps: int = 200,
     save_total_limit: int = 3,
     logging_steps: int = 10,
@@ -68,6 +68,7 @@ def train(
         'num_train_epochs': num_train_epochs,
         'learning_rate': learning_rate,
         'cutoff_len': cutoff_len,
+        'val_set_size': val_set_size,
         'lora_r': lora_r,
         'lora_alpha': lora_alpha,
         'lora_dropout': lora_dropout,
@@ -78,7 +79,12 @@ def train(
         'save_total_limit': save_total_limit,
         'logging_steps': logging_steps,
     }
+    if val_set_size and val_set_size > 0:
+        finetune_args['val_set_size'] = val_set_size
+    if resume_from_checkpoint:
+        finetune_args['resume_from_checkpoint'] = resume_from_checkpoint
 
+    wandb = None
     if wandb_api_key:
         os.environ["WANDB_API_KEY"] = wandb_api_key
 
@@ -220,7 +226,7 @@ def train(
             adapters_weights = torch.load(checkpoint_name)
             model = set_peft_model_state_dict(model, adapters_weights)
         else:
-            print(f"Checkpoint {checkpoint_name} not found")
+            raise ValueError(f"Checkpoint {checkpoint_name} not found")
 
     # Be more transparent about the % of trainable params.
     model.print_trainable_parameters()
@@ -315,4 +321,7 @@ def train(
     with open(os.path.join(output_dir, "train_output.json"), 'w') as train_output_json_file:
         json.dump(train_output, train_output_json_file, indent=2)
 
+    if use_wandb and wandb:
+        wandb.finish()
+
     return train_output

llama_lora/ui/finetune_ui.py

@@ -306,6 +306,17 @@ def do_train(
 ):
     try:
         base_model_name = Global.base_model_name
+
+        resume_from_checkpoint = None
+        if continue_from_model == "-" or continue_from_model == "None":
+            continue_from_model = None
+        if continue_from_checkpoint == "-" or continue_from_checkpoint == "None":
+            continue_from_checkpoint = None
+        if continue_from_model:
+            resume_from_checkpoint = os.path.join(Global.data_dir, "lora_models", continue_from_model)
+            if continue_from_checkpoint:
+                resume_from_checkpoint = os.path.join(resume_from_checkpoint, continue_from_checkpoint)
+
         output_dir = os.path.join(Global.data_dir, "lora_models", model_name)
         if os.path.exists(output_dir):
             if (not os.path.isdir(output_dir)) or os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
@@ -376,6 +387,8 @@ Train options: {json.dumps({
     'lora_dropout': lora_dropout,
     'lora_target_modules': lora_target_modules,
     'model_name': model_name,
+    'continue_from_model': continue_from_model,
+    'continue_from_checkpoint': continue_from_checkpoint,
 }, indent=2)}
 
 Train data (first 10):
@@ -386,7 +399,7 @@ Train data (first 10):
             return message
 
         if not should_training_progress_track_tqdm:
-            progress(0, desc="Preparing model for training...")
+            progress(0, desc=f"Preparing model {base_model_name} for training...")
 
         log_history = []
 
@@ -461,6 +474,10 @@ Train data (first 10):
                 # 'lora_dropout': lora_dropout,
                 # 'lora_target_modules': lora_target_modules,
             }
+            if continue_from_model:
+                info['continued_from_model'] = continue_from_model
+                if continue_from_checkpoint:
+                    info['continued_from_checkpoint'] = continue_from_checkpoint
             json.dump(info, info_json_file, indent=2)
 
         if not should_training_progress_track_tqdm:
@@ -490,7 +507,7 @@ Train data (first 10):
             lora_target_modules,  # lora_target_modules
             train_on_inputs,  # train_on_inputs
             False,  # group_by_length
-            None,  # resume_from_checkpoint
+            resume_from_checkpoint,  # resume_from_checkpoint
             save_steps,  # save_steps
             save_total_limit,  # save_total_limit
             logging_steps,  # logging_steps
@@ -582,6 +599,8 @@ def handle_load_params_from_model(
                 cutoff_len = value
             elif key == "evaluate_data_count":
                 evaluate_data_count = value
+            elif key == "val_set_size":
+                evaluate_data_count = value
             elif key == "micro_batch_size":
                 micro_batch_size = value
             elif key == "gradient_accumulation_steps":
@@ -610,6 +629,8 @@ def handle_load_params_from_model(
                 logging_steps = value
             elif key == "group_by_length":
                 pass
+            elif key == "resume_from_checkpoint":
+                pass
             else:
                 unknown_keys.append(key)
     except Exception as e: