more gpt-neox long ctx fixes

src/axolotl/utils/callbacks.py

@@ -61,6 +61,7 @@ class SaveBetterTransformerModelCallback(
 
             model = BetterTransformer.reverse(kwargs["model"])
             model.save_pretrained(checkpoint_folder)
+            # FIXME - need to cleanup old checkpoints
 
             # since we're saving here, we don't need the trainer loop to attempt to save too b/c
             # the trainer will raise an exception since it can't save a BetterTransformer wrapped model

src/axolotl/utils/data.py

@@ -381,9 +381,13 @@ def load_prepare_datasets(
             index=cfg.dataset_shard_idx,
         )
 
-    dataset = dataset.train_test_split(test_size=cfg.val_set_size, shuffle=False)
+    if cfg.val_set_size:
-    train_dataset = dataset["train"]
+        dataset = dataset.train_test_split(test_size=cfg.val_set_size, shuffle=False)
-    eval_dataset = dataset["test"]
+        train_dataset = dataset["train"]
+        eval_dataset = dataset["test"]
+    else:
+        train_dataset = dataset
+        eval_dataset = None
 
     return train_dataset, eval_dataset
 

src/axolotl/utils/models.py

@@ -262,6 +262,12 @@ def load_model(
     embeddings_len = math.ceil(len(tokenizer) / 32) * 32
     model.resize_token_embeddings(embeddings_len)
 
+    if cfg.sequence_len >= model.config.max_position_embeddings:
+        logging.warning(
+            f"increasing model.config.max_position_embeddings to {cfg.sequence_len}"
+        )
+        model.config.max_position_embeddings = cfg.sequence_len
+
     if not cfg.gptq and (
         (cfg.adapter == "lora" and load_in_8bit)
         or (cfg.adapter == "qlora" and cfg.load_in_4bit)

src/axolotl/utils/validation.py

@@ -70,4 +70,11 @@ def validate_config(cfg):
     # TODO
     # MPT 7b
     # https://github.com/facebookresearch/bitsandbytes/issues/25
-    # no 8bit adamw w bf16
+    # no 8bit adaAmw w bf16
+
+    # GPT-NeoX
+    # evals broken when extending context len
+    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 162, in forward                        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
+    # File "/root/miniconda3/envs/py3.9/lib/python3.9/site-packages/optimum/bettertransformer/models/attention.py", line 74, in gpt2_wrapped_scaled_dot_product
+    # attention_mask = causal_mask + attention_mask
+    # RuntimeError: The size of tensor a (2048) must match the size of tensor b (8132) at non-singleton dimension 3