tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Wing Lian
2023-09-13 10:57:12 -04:00
committed by GitHub
parent fdb777bc06
commit e5bb22a56b
1 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/axolotl/utils/trainer.py
View File

@@ -358,7 +358,14 @@ class ReLoRATrainer(AxolotlTrainer):
def add_position_ids(sample): def add_position_ids(sample):
sample_len = len(sample["input_ids"])
sample["position_ids"] = torch.arange(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 return sample
@@ -382,6 +389,9 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
if eval_dataset: if eval_dataset:
eval_dataset = eval_dataset.filter(drop_long, num_proc=os.cpu_count()) 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: if cfg.sample_packing:
train_dataset = train_dataset.map(add_position_ids, num_proc=os.cpu_count()) train_dataset = train_dataset.map(add_position_ids, num_proc=os.cpu_count())
if eval_dataset: if eval_dataset: