Fix hf dataset hang on small dataset (#1370)

mosaicml · Jul 18, 2024 · 006f251 · 006f251
1 parent 59b9c2a
commit 006f251
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py
@@ -913,6 +913,8 @@ def dataset_mapper(example: Dict):
             detected_cpu_count = os.cpu_count() or 1
             detected_cpus_with_margin = detected_cpu_count - 8
             num_cpus_to_use = max(1, detected_cpus_with_margin)
+            if len(dataset) < num_cpus_to_use:
+                num_cpus_to_use = 1
 
             columns_to_remove = list(dataset[0].keys())
             tokenized_dataset = dataset.map(