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

Make dataset_processes configurable (#651)

Browse Source 
I'm using the Axolotl script to train models on https://modal.com serverless GPUs. Unfortunately, their environment seems to have some kind of bug where if I try to run `datasets.filter` with too high a `num_proc`, it throws an error and dies.

This PR adds a new configuration option `dataset_processes`, which lets you explicitly set the number of processes used to map/filter the dataset. If not included, this defaults to the current behavior of setting that to `os.cpu_count()`.
This commit is contained in:
Kyle Corbitt
2023-09-28 21:22:22 -07:00
committed by GitHub
parent 590d6032fd
commit 9ec20777ba
3 changed files with 16 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
README.md
View File

@@ -487,6 +487,9 @@ datasets:
dataset_prepared_path: data/last_run_prepared dataset_prepared_path: data/last_run_prepared
# push prepared dataset to hub # push prepared dataset to hub
push_dataset_to_hub: # repo path push_dataset_to_hub: # repo path
# The maximum number of processes to use while preprocessing your input dataset. This defaults to `os.cpu_count()`
# if not set.
dataset_processes: # defaults to os.cpu_count() if not set
# push checkpoints to hub # push checkpoints to hub
hub_model_id: # repo path to push finetuned model hub_model_id: # repo path to push finetuned model
# how to push checkpoints to hub # how to push checkpoints to hub

2
src/axolotl/utils/config.py
View File

@@ -75,6 +75,8 @@ def normalize_config(cfg):
else: else:
cfg.torch_dtype = torch.float32 cfg.torch_dtype = torch.float32
cfg.dataset_processes = cfg.dataset_processes or os.cpu_count()
model_config = load_model_config(cfg) model_config = load_model_config(cfg)
cfg.model_config_type = model_config.model_type cfg.model_config_type = model_config.model_type

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

@@ -400,19 +400,25 @@ def disable_datasets_caching():
def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer): def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
drop_long = partial(drop_long_seq, sequence_len=cfg.sequence_len) drop_long = partial(drop_long_seq, sequence_len=cfg.sequence_len)
with zero_first(is_main_process()): with zero_first(is_main_process()):
train_dataset = train_dataset.filter(drop_long, num_proc=os.cpu_count()) train_dataset = train_dataset.filter(drop_long, num_proc=cfg.dataset_processes)
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=cfg.dataset_processes
)
if cfg.group_by_length: if cfg.group_by_length:
train_dataset = train_dataset.map(add_length, num_proc=os.cpu_count()) train_dataset = train_dataset.map(
add_length, num_proc=cfg.dataset_processes
)
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=cfg.dataset_processes
)
if cfg.eval_sample_packing is not False: if cfg.eval_sample_packing is not False:
if eval_dataset: if eval_dataset:
eval_dataset = eval_dataset.map( eval_dataset = eval_dataset.map(
add_position_ids, num_proc=os.cpu_count() add_position_ids, num_proc=cfg.dataset_processes
) )
# Phi doesn't want the attention_mask feature when training # Phi doesn't want the attention_mask feature when training