diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py index f2bd0239c8..edbfcc28c7 100644 --- a/llmfoundry/data/finetuning/tasks.py +++ b/llmfoundry/data/finetuning/tasks.py @@ -339,14 +339,20 @@ def dataset_mapper(example: Dict): example = preprocessing_fn(example) return _tokenize_formatted_example(example, tokenizer) + detected_cpu_count = os.cpu_count() or 1 + num_cpus_to_use = max(1, detected_cpu_count - 4) + columns_to_remove = list(dataset[0].keys()) tokenized_dataset = dataset.map( dataset_mapper, batched=False, remove_columns=columns_to_remove, + num_proc=num_cpus_to_use, ) prompt_length_filtered_dataset = tokenized_dataset.filter( - lambda example: len(example['input_ids']) < max_seq_len) + lambda example: len(example['input_ids']) < max_seq_len, + num_proc=num_cpus_to_use, + ) examples_removed = len(tokenized_dataset) - len( prompt_length_filtered_dataset) diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py index d2f203d3a0..d2c2a9e1c9 100644 --- a/tests/test_hf_conversion_script.py +++ b/tests/test_hf_conversion_script.py @@ -5,7 +5,7 @@ import os import pathlib import sys -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from composer import Trainer from composer.loggers import MLFlowLogger @@ -254,6 +254,7 @@ def test_callback_inits_with_defaults(): @pytest.mark.parametrize( 'hf_save_interval,save_interval,max_duration,expected_hf_checkpoints,expected_normal_checkpoints', [('3ba', '2ba', '7ba', 3, 4), ('1dur', '2ba', '1ep', 1, 4)]) +@patch('os.cpu_count', MagicMock(return_value=None)) def test_huggingface_conversion_callback(model: str, tmp_path: pathlib.Path, fsdp_state_dict_type: Optional[str], log_to_mlflow: bool,