🪛 Add resize token embeddings option

turbo-llm · Nov 26, 2024 · 4ef4f87 · 4ef4f87
2 parents 809eb0e + cde2fea
commit 4ef4f87
Show file tree

Hide file tree

Showing 13 changed files with 15 additions and 1 deletion.
diff --git a/tests/fixtures/configs/inference/classification/base.json b/tests/fixtures/configs/inference/classification/base.json
@@ -4,6 +4,7 @@
       "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "seq_cls",
+        "resize_token_embeddings": true,
         "transformers_settings": {
         },
         "model_kwargs": {

diff --git a/tests/fixtures/configs/inference/rag/base.json b/tests/fixtures/configs/inference/rag/base.json
@@ -6,6 +6,7 @@
           "model_path": "tests/fixtures/models/llama2_tiny",
           "model_type": "causal",
           "transformers_settings": {},
+          "resize_token_embeddings": true,
           "embeddings_initialization_strategy": {
             "<RS>": "<s>",
             "</RS>": "</s>",

diff --git a/tests/fixtures/configs/inference/sft/base.json b/tests/fixtures/configs/inference/sft/base.json
@@ -4,6 +4,7 @@
       "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
+        "resize_token_embeddings": true,
         "transformers_settings": {},
         "adapter_path": "tests/fixtures/models/llama2_tiny_fine_tuned_with_adapters/trainer/adapter_model"
       },

diff --git a/tests/fixtures/configs/sampling/rm.json b/tests/fixtures/configs/sampling/rm.json
@@ -22,6 +22,7 @@
   "rm": {
     "model_path": "tests/fixtures/models/gptj_tiny_for_seq_cls",
     "model_type": "seq_cls",
+    "resize_token_embeddings": true,
     "transformers_settings": {},
     "model_kwargs": {
       "num_labels": 1

diff --git a/tests/fixtures/configs/sampling/rso.json b/tests/fixtures/configs/sampling/rso.json
@@ -22,6 +22,7 @@
   "rm": {
     "model_path": "tests/fixtures/models/gptj_tiny_for_seq_cls",
     "model_type": "seq_cls",
+    "resize_token_embeddings": true,
     "transformers_settings": {},
     "model_kwargs": {
       "num_labels": 1

diff --git a/tests/fixtures/configs/train/rag/end2end.json b/tests/fixtures/configs/train/rag/end2end.json
@@ -45,6 +45,7 @@
     "generator_settings": {
       "model_path": "tests/fixtures/models/llama2_tiny",
       "model_type": "causal",
+      "resize_token_embeddings": true,
       "transformers_settings": {},
       "embeddings_initialization_strategy": {
         "<RS>": "<s>",

diff --git a/tests/fixtures/configs/train/sft/base.json b/tests/fixtures/configs/train/sft/base.json
@@ -44,6 +44,7 @@
     "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
+        "resize_token_embeddings": true,
         "transformers_settings": {
         },
         "peft_settings": {

diff --git a/tests/fixtures/configs/train/sft/prompt_tuning.json b/tests/fixtures/configs/train/sft/prompt_tuning.json
@@ -45,6 +45,7 @@
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
     "transformers_settings": {},
+    "resize_token_embeddings": true,
     "peft_settings": {
         "task_type": "CAUSAL_LM",
         "name": "PROMPT_TUNING",

diff --git a/tests/fixtures/configs/train/sft/resume_from_checkpoint.json b/tests/fixtures/configs/train/sft/resume_from_checkpoint.json
@@ -44,6 +44,7 @@
     "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
+        "resize_token_embeddings": true,
         "transformers_settings": {},
         "adapter_path": "tests/fixtures/checkpoints/llama2_tiny_debug_checkpoint",
         "is_trainable": true

diff --git a/tests/fixtures/configs/train/sft/sft_retrieval_utility.json b/tests/fixtures/configs/train/sft/sft_retrieval_utility.json
@@ -44,6 +44,7 @@
     "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
+        "resize_token_embeddings": true,
         "transformers_settings": {
         },
         "peft_settings": {

diff --git a/tests/fixtures/configs/train/sft/sft_with_rm_metric.json b/tests/fixtures/configs/train/sft/sft_with_rm_metric.json
@@ -44,6 +44,7 @@
     "model_settings": {
         "model_path": "tests/fixtures/models/llama2_tiny",
         "model_type": "causal",
+        "resize_token_embeddings": true,
         "transformers_settings": {
         },
         "peft_settings": {

diff --git a/turbo_alignment/common/tf/loaders/model/model.py b/turbo_alignment/common/tf/loaders/model/model.py
@@ -85,7 +85,8 @@ def load_model(
     if model_settings.transformers_settings.load_in_8bit:
         model = prepare_model_for_int8_training(model)
 
-    model.resize_token_embeddings(len(tokenizer))
+    if model_settings.resize_token_embeddings:
+        model.resize_token_embeddings(len(tokenizer))
 
     if model_settings.embeddings_initialization_strategy is not None:
         with torch.no_grad():

diff --git a/turbo_alignment/settings/model.py b/turbo_alignment/settings/model.py
@@ -41,6 +41,8 @@ class PreTrainedModelSettings(ExtraFieldsNotAllowedBaseModel):
 
     transformers_settings: ModelTransformersSettings
 
+    resize_token_embeddings: bool = False
+
     embeddings_initialization_strategy: dict[str, str] | None = None
 
     liger_kernels_settings: LigerKernelSettings | None = None