From 2763449f6400d95c30a228ed578e8f639efd386f Mon Sep 17 00:00:00 2001
From: wassname <git@wassname.org>
Date: Mon, 20 Nov 2023 10:01:56 +0800
Subject: [PATCH] bug fix

---
 config/trainer.yaml       |  2 +-
 src/models/transformer.py | 11 +++++------
 src/utils.py              |  4 +++-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/config/trainer.yaml b/config/trainer.yaml
index b4fe7c5..c52d41c 100644
--- a/config/trainer.yaml
+++ b/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
diff --git a/src/models/transformer.py b/src/models/transformer.py
index 6237c57..cc91255 100644
--- a/src/models/transformer.py
+++ b/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
diff --git a/src/utils.py b/src/utils.py
index d64273d..07ba101 100644
--- a/src/utils.py
+++ b/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!"