Merge pull request #228 from unslothai/fix/cap-num-proc-multigpu-deadlock

Cap dataset.map num_proc on multi-GPU machines to prevent fork deadlocks

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")
 

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}"
 

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"

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',
 ]

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