add optimization for group-by-len (#563)

axolotl-ai-cloud · Sep 13, 2023 · 60c32b0 · 60c32b0
1 parent e625338
commit 60c32b0
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
@@ -358,7 +358,14 @@ def create_scheduler(
 
 
 def add_position_ids(sample):
+    sample_len = len(sample["input_ids"])
     sample["position_ids"] = torch.arange(len(sample["input_ids"]))
+    sample["length"] = sample_len
+    return sample
+
+
+def add_length(sample):
+    sample["length"] = len(sample["input_ids"])
     return sample
 
 
@@ -382,6 +389,9 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
         if eval_dataset:
             eval_dataset = eval_dataset.filter(drop_long, num_proc=os.cpu_count())
 
+        if cfg.group_by_length:
+            train_dataset = train_dataset.map(add_length, num_proc=os.cpu_count())
+
         if cfg.sample_packing:
             train_dataset = train_dataset.map(add_position_ids, num_proc=os.cpu_count())
             if eval_dataset: