From 834013aae5041193c9ce284c03947361aad8a6e0 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Mon, 23 Feb 2026 14:25:31 +0000 Subject: [PATCH] Cap dataset.map num_proc on multi-GPU machines to prevent fork deadlocks --- studio/backend/core/training/trainer.py | 6 +- .../backend/utils/datasets/chat_templates.py | 12 ++-- .../utils/datasets/format_conversion.py | 18 +++-- studio/backend/utils/hardware/__init__.py | 4 ++ studio/backend/utils/hardware/hardware.py | 66 +++++++++++++++++++ 5 files changed, 93 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 063d87ead6..804baa33c0 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -7,7 +7,7 @@ import os os.environ["TOKENIZERS_PARALLELISM"] = "false" import torch -from utils.hardware import clear_gpu_cache +from utils.hardware import clear_gpu_cache, safe_num_proc torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported from unsloth.chat_templates import get_chat_template @@ -711,7 +711,7 @@ class UnslothTrainer: "output_dir": output_dir, "report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none", "include_num_input_tokens_seen": True, # Enable token counting - "dataset_num_proc": max(1, os.cpu_count() // 4), + "dataset_num_proc": safe_num_proc(max(1, os.cpu_count() // 4)), } # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps @@ -871,7 +871,7 @@ class UnslothTrainer: self.trainer, instruction_part=instruction_part, response_part=response_part, - num_proc=config_args.get("dataset_num_proc", max(1, os.cpu_count() // 4)), + num_proc=config_args.get("dataset_num_proc", safe_num_proc(max(1, os.cpu_count() // 4))), ) print("Train on responses only configured successfully\n") diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 6420aa899d..52c876291d 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -283,9 +283,11 @@ def apply_chat_template_to_dataset( } if not isinstance(dataset, IterableDataset): - from multiprocessing import cpu_count + from utils.hardware import safe_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = max(1, cpu_count() // 3) + num_proc = safe_num_proc() + else: + num_proc = safe_num_proc(num_proc) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Applying template to Alpaca format" @@ -347,9 +349,11 @@ def apply_chat_template_to_dataset( } if not isinstance(dataset, IterableDataset): - from multiprocessing import cpu_count + from utils.hardware import safe_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = max(1, cpu_count() // 3) + num_proc = safe_num_proc() + else: + num_proc = safe_num_proc(num_proc) 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 a97db20ba4..6436e7a82a 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -107,10 +107,12 @@ def standardize_chat_format( } if not isinstance(dataset, IterableDataset): - from multiprocessing import cpu_count + from utils.hardware import safe_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = max(1, cpu_count() // 3) + num_proc = safe_num_proc() + else: + num_proc = safe_num_proc(num_proc) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Standardizing chat format" @@ -173,10 +175,12 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None): } if not isinstance(dataset, IterableDataset): - from multiprocessing import cpu_count + from utils.hardware import safe_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = max(1, cpu_count() // 3) + num_proc = safe_num_proc() + else: + num_proc = safe_num_proc(num_proc) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format" @@ -221,10 +225,12 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None): } if not isinstance(dataset, IterableDataset): - from multiprocessing import cpu_count + from utils.hardware import safe_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = max(1, cpu_count() // 3) + num_proc = safe_num_proc() + else: + num_proc = safe_num_proc(num_proc) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format" diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 667466efea..e8256eb9aa 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -13,6 +13,8 @@ from .hardware import ( get_gpu_summary, get_package_versions, get_gpu_utilization, + get_physical_gpu_count, + safe_num_proc, ) __all__ = [ @@ -27,4 +29,6 @@ __all__ = [ 'get_gpu_summary', 'get_package_versions', 'get_gpu_utilization', + 'get_physical_gpu_count', + 'safe_num_proc', ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 0d7cc97cfb..b885e130d5 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -385,3 +385,69 @@ def get_gpu_utilization() -> Dict[str, Any]: "power_limit_w": power_limit, "power_utilization_pct": power_pct, } + + +# ========== Multi-GPU Detection & Safe num_proc ========== + +_physical_gpu_count: Optional[int] = None + +def get_physical_gpu_count() -> int: + """ + Return the number of physical NVIDIA GPUs on the machine. + + Uses ``nvidia-smi -L`` which is NOT affected by CUDA_VISIBLE_DEVICES, + so it always reflects the true hardware count. + Result is cached after the first call. + """ + global _physical_gpu_count + if _physical_gpu_count is not None: + return _physical_gpu_count + + try: + import subprocess + result = subprocess.run( + ["nvidia-smi", "-L"], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + _physical_gpu_count = len(result.stdout.strip().splitlines()) + else: + _physical_gpu_count = 1 + except Exception: + _physical_gpu_count = 1 + + return _physical_gpu_count + + +def safe_num_proc(desired: Optional[int] = None) -> int: + """ + Return a safe ``num_proc`` for ``dataset.map()`` calls. + + On multi-GPU machines the NVIDIA driver spawns extra background threads, + making ``os.fork()`` prone to deadlocks when many workers are created. + This helper caps ``num_proc`` to 4 on such machines. + + On single-GPU (or CPU-only) machines the original value is returned + unchanged. + + Args: + desired: The num_proc you *want*. If None, auto-computes from + ``os.cpu_count()``. + + Returns: + A safe integer ≥ 1. + """ + import os + + if desired is None or not isinstance(desired, int): + desired = max(1, os.cpu_count() // 3) + + if get_physical_gpu_count() > 1: + capped = min(4, desired) + print( + f"⚙️ Multi-GPU detected ({get_physical_gpu_count()} GPUs) — " + f"capping num_proc {desired} → {capped} to avoid fork deadlocks" + ) + return capped + + return desired