[Feat] streaming multipack

src/axolotl/utils/data.py

@@ -2,9 +2,11 @@
 import functools
 import hashlib
 import logging
+from collections import defaultdict
 from pathlib import Path
 from typing import Dict, List, Tuple, Union
 
+import numpy as np
 import torch
 from datasets import (
  Dataset,
@@ -14,6 +16,7 @@
  load_from_disk,
 )
 from huggingface_hub import hf_hub_download
+from torch.utils.data import RandomSampler
 from transformers import PreTrainedTokenizerBase
 
 from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
@@ -41,9 +44,11 @@
 )
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import is_main_process, zero_first
+from axolotl.utils.samplers.multipack import MultipackBatchSampler
 from axolotl.utils.trainer import (
  calculate_total_num_steps,
  process_datasets_for_packing,
+ process_pretraining_datasets_for_packing,
 )
 
 LOG = logging.getLogger("axolotl")
@@ -819,3 +824,68 @@ def load_pretraining_dataset(path, tokenizer, max_tokens=2048, seed=42):
  remove_columns=dataset.features.keys(),
  )
  return dataset
+
+
+def encode_packed_pretraining(
+ tokenizer: PreTrainedTokenizerBase,
+ examples: List[str],
+ max_seq_length: int = 8192,
+ sample_packing_efficiency: int = 1,
+) -> Dict[str, List]:
+ # pylint: disable=duplicate-code
+ # tokenize all the examples
+ # rows get split with stride (overlap)
+ res = tokenizer(
+ examples,
+ truncation=True,
+ max_length=max_seq_length - 1,
+ add_special_tokens=True,
+ return_overflowing_tokens=True,
+ stride=256,
+ )
+
+ input_ids = [seq + [tokenizer.eos_token_id] for seq in res["input_ids"]]
+ attention_mask = [seq + [1] for seq in res["attention_mask"]]
+
+ tokenized_examples = {
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ }
+
+ train_dataset = Dataset.from_dict(tokenized_examples)
+ train_dataset = process_pretraining_datasets_for_packing(
+ train_dataset, max_seq_length
+ )
+
+ sampler = MultipackBatchSampler(
+ RandomSampler(train_dataset),
+ batch_size=1,
+ drop_last=True,
+ batch_max_len=max_seq_length,
+ lengths=(
+ train_dataset.data.column("position_ids")
+ .to_pandas()
+ .apply(lambda x: x[-1] + 1)
+ .values
+ ),
+ packing_efficiency_estimate=sample_packing_efficiency,
+ )
+
+ chunked_data = defaultdict(list)
+
+ for data in sampler:
+ features = train_dataset[data]
+
+ for feature in features.keys():
+ if feature == "length":
+ continue
+ if feature == "attention_mask":
+ arrays = [(1) * np.array(item) for item in features[feature]]
+ chunked_data[feature].append(np.concatenate(arrays))
+ else:
+ arrays = [np.array(item) for item in features[feature]]
+ chunked_data[feature].append(np.concatenate(arrays))
+
+ chunked_data["labels"] = chunked_data["input_ids"].copy()
+
+ return chunked_data

src/axolotl/utils/trainer.py

@@ -143,6 +143,16 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
  return train_dataset, eval_dataset
 
 
+def process_pretraining_datasets_for_packing(train_dataset, sequence_len):
+ drop_long = partial(drop_long_seq, sequence_len=sequence_len)
+
+ train_dataset = train_dataset.filter(drop_long)
+ train_dataset = train_dataset.map(
+ add_position_ids,
+ )
+ return train_dataset
+
+
 def calculate_total_num_steps(cfg, train_dataset, update=True):
  if not cfg.total_num_tokens:
  total_num_tokens = np.sum(

tests/test_packed_pretraining.py

@@ -0,0 +1,87 @@
+"""Module for testing streaming dataset sequence packing"""
+import math
+import unittest
+from functools import partial
+
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+from axolotl.utils.collators import DataCollatorForSeq2Seq
+from axolotl.utils.data import encode_packed_pretraining
+
+
+class TestPacking(unittest.TestCase):
+ """
+ Test class for packing streaming dataset sequences
+ """
+
+ def setUp(self) -> None:
+ # pylint: disable=duplicate-code
+ self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
+ self.tokenizer.add_special_tokens(
+ {
+ "bos_token": "<s>",
+ "eos_token": "</s>",
+ "unk_token": "<unk>",
+ "pad_token": "[PAD]",
+ }
+ )
+ self.max_seq_length = 8192
+ self.batch_size = 6
+ self.sample_packing_efficiency = 1
+ self.data_collator_kwargs = {
+ "padding": True,
+ "pad_to_multiple_of": 64 * math.ceil(self.max_seq_length / 64),
+ }
+
+ def test_packing_stream_dataset(self):
+ # pylint: disable=duplicate-code
+ dataset = load_dataset(
+ "c4",
+ "en",
+ streaming=True,
+ )["train"]
+
+ encode = partial(
+ encode_packed_pretraining,
+ self.tokenizer,
+ max_seq_length=self.max_seq_length,
+ sample_packing_efficiency=self.sample_packing_efficiency,
+ )
+
+ dataset = dataset.map(
+ encode,
+ batched=True,
+ input_columns="text",
+ remove_columns=dataset.features.keys(),
+ )
+
+ data_collator_fn = DataCollatorForSeq2Seq(
+ self.tokenizer,
+ return_tensors="pt",
+ **self.data_collator_kwargs,
+ )
+
+ trainer_loader = DataLoader(
+ dataset,
+ batch_size=self.batch_size,
+ collate_fn=data_collator_fn,
+ drop_last=True,
+ )
+ idx = 0
+ for data in trainer_loader:
+ if idx > 10:
+ break
+ assert data["input_ids"].shape == (self.batch_size, self.max_seq_length)
+ assert data["position_ids"].shape == (self.batch_size, self.max_seq_length)
+ assert data["labels"].shape == (self.batch_size, self.max_seq_length)
+ assert data["attention_mask"].shape == (
+ self.batch_size,
+ self.max_seq_length,
+ )
+ idx += 1
+
+
+if __name__ == "__main__":
+ unittest.main()