fix attetion mask with packing

update pip install command for apex

README.md

@@ -195,6 +195,10 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"message_1": "...", "message_2": "..."}
   ```
+- `alpaca_w_system.load_open_orca`: support for open orca datasets with included system prompts, instruct
+  ```json
+  {"system_prompt": "...", "question": "...", "response": "..."}
+  ```
 - `context_qa`: in context question answering from an article
   ```json
   {"article": "...", "question": "...", "answer": "..."}
@@ -336,6 +340,8 @@ datasets:
 dataset_prepared_path: data/last_run_prepared
 # push prepared dataset to hub
 push_dataset_to_hub: # repo path
+# push checkpoints to hub
+push_to_hub_model_id: # repo path
 # whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
 # required to be true when used in combination with `push_dataset_to_hub`
 hf_use_auth_token: # boolean

src/axolotl/datasets.py

@@ -79,11 +79,13 @@ class ConstantLengthDataset(IterableDataset):
         buffer = {"input_ids": [], "attention_mask": [], "labels": []}
         buffer_len = 0
         for dataset in self.datasets:
+            idx = 0
             iterator = iter(dataset)
             more_examples = True
             while more_examples:
                 try:
                     example = next(iterator)
+                    idx += 1
                 except StopIteration:
                     more_examples = False
                     example = None
@@ -124,6 +126,7 @@ class ConstantLengthDataset(IterableDataset):
                         "labels": [],
                     }
                     buffer_len = 0
+                    idx = 1
 
                 if example:
                     # FIXME
@@ -132,11 +135,6 @@ class ConstantLengthDataset(IterableDataset):
                         input_ids = example["input_ids"]
                         attention_mask = example["attention_mask"]
                         labels = example["labels"]
-                        if (
-                            buffer["input_ids"]
-                            and input_ids[0] == self.tokenizer.bos_token_id
-                        ):
-                            attention_mask[0] = 0
 
                         if add_concat_token:
                             input_ids.append(self.concat_token_id)
@@ -147,7 +145,7 @@ class ConstantLengthDataset(IterableDataset):
                             input_ids, dtype=self.tokens_dtype
                         )
                         attention_mask_with_concat = torch.tensor(
-                            attention_mask, dtype=self.tokens_dtype
+                            [idx * m for m in attention_mask], dtype=torch.int16
                         )
                         labels_with_concat = torch.tensor(
                             labels, dtype=self.tokens_dtype

src/axolotl/prompt_strategies/alpaca_w_system.py

@@ -75,10 +75,46 @@ class SystemDataPrompter(AlpacaPrompter):
         yield res
 
 
+class OpenOrcaPromptTokenizingStrategy(InstructionWSystemPromptTokenizingStrategy):
+    """
+    Tokenizing strategy for OpenOrca datasets
+    """
+
+    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str]:
+        return (
+            prompt["question"],
+            "",
+            prompt["response"],
+            prompt["system_prompt"],
+        )
+
+
 def load(tokenizer, cfg):
+    return load_chat(tokenizer, cfg)
+
+
+def load_instruct(tokenizer, cfg):
+    return InstructionWSystemPromptTokenizingStrategy(
+        SystemDataPrompter(PromptStyle.INSTRUCT.value),
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )
+
+
+def load_chat(tokenizer, cfg):
     return InstructionWSystemPromptTokenizingStrategy(
         SystemDataPrompter(PromptStyle.CHAT.value),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,
     )
+
+
+def load_open_orca(tokenizer, cfg):
+    return OpenOrcaPromptTokenizingStrategy(
+        SystemDataPrompter(PromptStyle.INSTRUCT.value),
+        tokenizer,
+        cfg.train_on_inputs,
+        cfg.sequence_len,
+    )

src/axolotl/prompt_tokenizers.py

@@ -440,7 +440,7 @@ def parse_tokenized_to_result(
     result: Dict[str, List[int]],
     current_len: int,
     res: Dict[str, List[int]],
-    labels: list[int],
+    labels: List[int],
     pad_token_id: Union[int, None] = None,
 ) -> Tuple[Dict[str, List[int]], int]:
     """

src/axolotl/utils/data.py

@@ -37,7 +37,7 @@ from axolotl.prompters import (
 
 
 def load_tokenized_prepared_datasets(
-    tokenizer, cfg, default_dataset_prepared_path
+    split, tokenizer, cfg, default_dataset_prepared_path
 ) -> DatasetDict:
     tokenizer_name = tokenizer.__class__.__name__
     ds_hash = str(
@@ -49,6 +49,8 @@ def load_tokenized_prepared_datasets(
                     sorted([f"{d.path}:{d.type}:{d.shards}" for d in cfg.datasets])
                 )
                 + "|"
+                + split
+                + "|"
                 + tokenizer_name
             ).encode("utf-8")
         ).hexdigest()
@@ -66,7 +68,7 @@ def load_tokenized_prepared_datasets(
                 f"{cfg.push_dataset_to_hub}/{ds_hash}",
                 use_auth_token=use_auth_token,
             )
-            dataset = dataset["train"]
+            dataset = dataset[split]
     except Exception:  # pylint: disable=broad-except # nosec
         pass
 
@@ -134,8 +136,8 @@ def load_tokenized_prepared_datasets(
                 raise ValueError("unhandled dataset load")
             # support for using a subset of the data
             if d.shards:
-                if "train" in ds:
+                if split in ds:
-                    ds = ds.shuffle(seed=seed)["train"].shard(
+                    ds = ds.shuffle(seed=seed)[split].shard(
                         num_shards=d.shards, index=0
                     )
                 else:
@@ -144,8 +146,8 @@ def load_tokenized_prepared_datasets(
             d_type_split = d_type.split(":")
             d_base_type = d_type_split[0]
             d_prompt_style = d_type_split[1] if len(d_type_split) > 1 else None
-            if "train" in ds:
+            if split in ds:
-                ds = ds["train"]
+                ds = ds[split]
             if ds_strategy := load(d.type, tokenizer, cfg):
                 ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
                 datasets.append(ds_wrapper)
@@ -319,7 +321,6 @@ def load_prepare_datasets(
                     f"{cfg.push_dataset_to_hub}/{ds_hash}",
                     use_auth_token=use_auth_token,
                 )
-                dataset = dataset["train"]
         except Exception:  # pylint: disable=broad-except # nosec
             pass
 
@@ -339,28 +340,37 @@ def load_prepare_datasets(
                     f"{cfg.push_dataset_to_hub}/{ds_hash}", private=True
                 )
         else:
-            dataset = load_tokenized_prepared_datasets(
+            dataset_train = load_tokenized_prepared_datasets(
-                tokenizer, cfg, default_dataset_prepared_path
+                "train", tokenizer, cfg, default_dataset_prepared_path
             )
-
+            dataset_test = load_tokenized_prepared_datasets(
+                "test", tokenizer, cfg, default_dataset_prepared_path
+            )
+            dataset = DatasetDict({"train": dataset_train, "test": dataset_test})
             if cfg.seed:
                 dataset = dataset.shuffle(seed=cfg.seed)
 
-            constant_len_dataset = ConstantLengthDataset(
+            constant_len_dataset_train = ConstantLengthDataset(
                 tokenizer,
-                [dataset],
+                [dataset["train"]],
+                seq_length=max_packed_sequence_len,
+            )
+            constant_len_dataset_test = ConstantLengthDataset(
+                tokenizer,
+                [dataset["test"]],
                 seq_length=max_packed_sequence_len,
             )
             logging.info(
                 f"packing master dataset to len: {cfg.max_packed_sequence_len}"
             )
-            dataset = Dataset.from_list(list(constant_len_dataset))
+            dataset_train = Dataset.from_list(list(constant_len_dataset_train))
+            dataset_test = Dataset.from_list(list(constant_len_dataset_test))
 
             # filter out bad data
-            dataset = Dataset.from_list(
+            dataset_train = Dataset.from_list(
                 [
                     d
-                    for d in dataset
+                    for d in dataset_train
                     if len(d["input_ids"]) < cfg.sequence_len
                     and len(d["input_ids"]) > 0
                     and len(d["input_ids"]) == len(d["attention_mask"])
@@ -368,6 +378,19 @@ def load_prepare_datasets(
                 ]
             )
 
+            # filter out bad data
+            dataset_test = Dataset.from_list(
+                [
+                    d
+                    for d in dataset_test
+                    if len(d["input_ids"]) < cfg.sequence_len
+                    and len(d["input_ids"]) > 0
+                    and len(d["input_ids"]) == len(d["attention_mask"])
+                    and len(d["input_ids"]) == len(d["labels"])
+                ]
+            )
+            dataset = DatasetDict({"train": dataset_train, "test": dataset_test})
+
             if cfg.local_rank == 0:
                 logging.info(
                     f"Saving packed prepared dataset to disk... {prepared_ds_path}"
@@ -382,9 +405,13 @@ def load_prepare_datasets(
                         private=True,
                     )
     else:
-        dataset = load_tokenized_prepared_datasets(
+        dataset_train = load_tokenized_prepared_datasets(
-            tokenizer, cfg, default_dataset_prepared_path
+            "train", tokenizer, cfg, default_dataset_prepared_path
         )
+        dataset_test = load_tokenized_prepared_datasets(
+            "test", tokenizer, cfg, default_dataset_prepared_path
+        )
+        dataset = DatasetDict({"train": dataset_train, "test": dataset_test})
 
     if cfg.dataset_shard_num and cfg.dataset_shard_idx is not None:
         logging.info(
@@ -399,6 +426,9 @@ def load_prepare_datasets(
         dataset = dataset.train_test_split(test_size=cfg.val_set_size, shuffle=False)
         train_dataset = dataset["train"]
         eval_dataset = dataset["test"]
+    elif "train" in dataset:
+        train_dataset = dataset["train"]
+        eval_dataset = dataset["test"]
     else:
         train_dataset = dataset
         eval_dataset = None

src/axolotl/utils/trainer.py

@@ -137,9 +137,9 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
         eval_accumulation_steps=cfg.gradient_accumulation_steps,
         num_train_epochs=cfg.num_epochs,
         learning_rate=cfg.learning_rate,
-        evaluation_strategy="steps" if cfg.val_set_size > 0 else "no",
+        evaluation_strategy="steps",
         save_strategy="steps" if cfg.save_steps else "epoch",
-        eval_steps=cfg.eval_steps if cfg.val_set_size > 0 else None,
+        eval_steps=cfg.eval_steps,
         save_steps=cfg.save_steps,
         output_dir=cfg.output_dir,
         save_total_limit=3,

tests/test_packed_dataset.py

@@ -27,7 +27,7 @@ class TestPacking(unittest.TestCase):
             }
         )
 
-    def test_resets_attention(self):
+    def test_increments_attention(self):
         prompter = AlpacaPrompter("chat")
         strat = AlpacaPromptTokenizingStrategy(
             prompter,
@@ -58,7 +58,7 @@ class TestPacking(unittest.TestCase):
 
         # but subsequent one does
         assert example["input_ids"][next_bos_index] == self.tokenizer.bos_token_id
-        assert example["attention_mask"][next_bos_index] == 0
+        assert example["attention_mask"][next_bos_index] == 2
 
 
 if __name__ == "__main__":