fixes for smoke tests

axolotl-ai-cloud · Jan 22, 2024 · f77a465 · f77a465
1 parent 5824228
commit f77a465
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
@@ -157,6 +157,12 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
             if eval_dataset:
                 eval_dataset = eval_dataset.remove_columns("attention_mask")
 
+        if cfg.model_config_type == "falcon":
+            LOG.info("dropping token_type_ids column")
+            train_dataset = train_dataset.remove_columns("token_type_ids")
+            if eval_dataset:
+                eval_dataset = eval_dataset.remove_columns("token_type_ids")
+
     return train_dataset, eval_dataset
 
 

diff --git a/tests/e2e/patched/test_falcon_samplepack.py b/tests/e2e/patched/test_falcon_samplepack.py
@@ -39,6 +39,7 @@ def test_qlora(self, temp_dir):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
+                "lora_modules_to_save": ["word_embeddings", "lm_head"],
                 "val_set_size": 0.1,
                 "special_tokens": {
                     "bos_token": "<|endoftext|>",