From 3d0d1c7020147f351b5ee6e1da3d8c83c7de4d0d Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 17 Feb 2026 23:12:45 +0000 Subject: [PATCH] fix: cap dataset.map() num_proc to 8 to prevent CUDA fork deadlocks --- studio/backend/utils/datasets/chat_templates.py | 4 ++-- studio/backend/utils/datasets/format_conversion.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index e3b1bef315..3f6d73377e 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -285,7 +285,7 @@ def apply_chat_template_to_dataset( if not isinstance(dataset, IterableDataset): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Applying template to Alpaca format" @@ -349,7 +349,7 @@ def apply_chat_template_to_dataset( if not isinstance(dataset, IterableDataset): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}" diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 9367741e8e..df2ff95fc8 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -110,7 +110,7 @@ def standardize_chat_format( from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Standardizing chat format" @@ -176,7 +176,7 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format" @@ -224,7 +224,7 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"