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

fix: preprocess yielding whole dataset to each worker (#2503) [skip ci]

Browse Source
This commit is contained in:
NanoCode012
2025-04-17 05:02:35 +07:00
committed by GitHub
parent f776f889a1
commit 32637fad00
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/axolotl/utils/data/sft.py
View File

@@ -332,16 +332,23 @@ def load_tokenized_prepared_datasets(
if cfg.local_rank == 0 and not cfg.skip_prepare_dataset: if cfg.local_rank == 0 and not cfg.skip_prepare_dataset:
LOG.info(f"Saving merged prepared dataset to disk... {prepared_ds_path}") LOG.info(f"Saving merged prepared dataset to disk... {prepared_ds_path}")
if isinstance(dataset, IterableDataset): if isinstance(dataset, IterableDataset):
num_workers = cfg.dataset_processes
def gen_from_iter_ds(_ds, _=None): def gen_from_iter_ds(_ds, worker_id: List[int], num_workers: List[int]):
yield from _ds """Generator function to correctly splice the dataset for each worker"""
for i, item in enumerate(_ds):
if i % num_workers[0] == worker_id[0]:
yield item
ds_from_iter = Dataset.from_generator( ds_from_iter = Dataset.from_generator(
functools.partial(gen_from_iter_ds, dataset), functools.partial(gen_from_iter_ds, dataset),
features=dataset.features, features=dataset.features,
num_proc=cfg.dataset_processes, num_proc=num_workers,
split=split, split=split,
gen_kwargs={"_": list(range(cfg.dataset_processes))}, gen_kwargs={
"worker_id": list(range(num_workers)),
"num_workers": [num_workers] * num_workers,
},
) )
ds_from_iter.save_to_disk(str(prepared_ds_path)) ds_from_iter.save_to_disk(str(prepared_ds_path))
else: else: