bug fix

config/trainer.yaml

@@ -61,7 +61,7 @@ training:
     steps_per_epoch: 200
   world_model:
     batch_num_samples: 8  # pretrained models use lots of ram
-    grad_acc_steps: 1
+    grad_acc_steps: 2
     max_grad_norm: 10.0
     weight_decay: 0.01
     start_after_epochs: 25

src/models/transformer.py

@@ -111,18 +111,17 @@ def load_pretrained_model(config, device="cuda:0"):
             "mlp.up_proj",
             "mlp.down_proj",
             # "wte", "embed_tokens",
-            "lm_head",
+            # "lm_head",
         ],
-        bias=True,
+        # bias="lora_only",
         # tune the embedding layer and prediction head
-        modules_to_save = ["lm_head", "embed_tokens"],
+        modules_to_save = ["lm_head",], # we want the classifier parameters to be trained too when fine-tuning the base model on our custom dataset. To ensure that the classifier parameters are also trained, we specify modules_to_save. 
     )
     base_model_peft = peft.get_peft_model(base_model, peft_config)
-    base_model_peft.add_adapter(adapter_name="dynamics", peft_config=peft_config) # make an adapter
-    base_model_peft.set_adapter("dynamics") # use an adapter
+    base_model_peft.add_adapter(adapter_name="dynamics", peft_config=peft_config) # make and set an adapter
     disable_causal_mask_always()
     print(base_model_peft.print_trainable_parameters())
-    logger.info(f"loaded model {base_model_peft}")
+    logger.debug(f"loaded model {base_model_peft}")
     return base_model_peft
 
 @contextmanager

src/utils.py

@@ -3,7 +3,7 @@ import cv2
 from pathlib import Path
 import random
 import shutil
-
+from loguru import logger
 import numpy as np
 import torch
 import torch.nn as nn
@@ -15,6 +15,7 @@ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 
 def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
     """Credits to https://github.com/karpathy/minGPT"""
+    # FIXME: check this is still good for LoRA
     # separate out all parameters to those that will and won't experience regularizing weight decay
     decay = set()
     no_decay = set()
@@ -39,6 +40,7 @@ def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_na
     param_dict = {pn: p for pn, p in model.named_parameters()}
     inter_params = decay & no_decay
     union_params = decay | no_decay
+    logger.debug(f"decay {decay} no_decay {no_decay}")
     assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
     assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"