diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index ede213647d..9b320c93a8 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -12,8 +12,28 @@ import sys # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork os.environ["TOKENIZERS_PARALLELISM"] = "false" +# Ensure compiled cache modules are importable by any subprocess. +# On spawn-based platforms (Windows, macOS), spawned dataset.map() workers must +# re-import all top-level modules. The compiled cache's trainer files import +# torch and unsloth_zoo (which initializes CUDA), making spawn impractical. +# Propagating UNSLOTH_COMPILE_LOCATION via PYTHONPATH ensures any subprocess +# (not just Pool workers) can find compiled modules. +# NOTE: Do NOT import unsloth_zoo.compiler here -- it triggers heavy torch/triton imports. +if sys.platform in ("win32", "darwin"): + _compile_cache = os.environ.get( + "UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache" + ) + if not os.path.isabs(_compile_cache): + _compile_cache = os.path.abspath(_compile_cache) + os.environ["UNSLOTH_COMPILE_LOCATION"] = _compile_cache + _pp = os.environ.get("PYTHONPATH", "") + if _compile_cache not in _pp.split(os.pathsep): + os.environ["PYTHONPATH"] = _compile_cache + (os.pathsep + _pp if _pp else "") + if _compile_cache not in sys.path: + sys.path.insert(0, _compile_cache) + import torch -from utils.hardware import clear_gpu_cache, safe_num_proc +from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported @@ -1465,7 +1485,10 @@ class UnslothTrainer: self._update_progress(status_message = "Formatting audio VLM dataset...") dataset = dataset.map( - format_messages, batched = True, batch_size = 4, num_proc = safe_num_proc(4) + format_messages, + batched = True, + batch_size = 4, + num_proc = dataset_map_num_proc(4), ) logger.info(f"Audio VLM dataset formatted: {len(dataset)} examples\n") return dataset @@ -2980,9 +3003,11 @@ class UnslothTrainer: "output_dir": output_dir, "report_to": _build_report_targets(training_args), "include_num_input_tokens_seen": True, # Enable token counting - "dataset_num_proc": 1 - if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) - else safe_num_proc(max(1, os.cpu_count() // 4)), + "dataset_num_proc": dataset_map_num_proc( + 1 + if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) + else max(1, (os.cpu_count() or 1) // 4) + ), "max_seq_length": training_args.get("max_seq_length", 2048), } if training_args.get("enable_tensorboard", False): @@ -2993,9 +3018,10 @@ class UnslothTrainer: f"[DEBUG] dataset_num_proc={config_args['dataset_num_proc']} (is_audio={self.is_audio}, is_audio_vlm={self.is_audio_vlm}, _cuda_audio_used={self._cuda_audio_used})" ) - # On Windows with transformers 5.x, disable DataLoader multiprocessing - # to avoid issues with modified sys.path (.venv_t5) in spawned workers. - if sys.platform == "win32": + # On spawn-based platforms (Windows, macOS) with transformers 5.x, + # disable DataLoader multiprocessing to avoid issues with modified + # sys.path (.venv_t5) in spawned workers. + if sys.platform in ("win32", "darwin"): import transformers as _tf if _tf.__version__.startswith("5."): diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 611a34e39d..8333009626 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -445,7 +445,7 @@ def check_format( format_result = format_dataset( preview_slice, format_type = "auto", - num_proc = 1, # Only 10 preview rows — no need for multiprocessing + num_proc = None, # Only 10 preview rows -- no need for multiprocessing ) processed = format_result["dataset"] preview_samples = _serialize_preview_rows(processed) diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index ee68c6bc0c..0c498281a4 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -291,11 +291,11 @@ def apply_chat_template_to_dataset( } if not isinstance(dataset, IterableDataset): - from utils.hardware import safe_num_proc + from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = safe_num_proc() + num_proc = dataset_map_num_proc() else: - num_proc = safe_num_proc(num_proc) + num_proc = dataset_map_num_proc(num_proc) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Applying template to Alpaca format" @@ -357,11 +357,11 @@ def apply_chat_template_to_dataset( } if not isinstance(dataset, IterableDataset): - from utils.hardware import safe_num_proc + from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = safe_num_proc() + num_proc = dataset_map_num_proc() else: - num_proc = safe_num_proc(num_proc) + num_proc = dataset_map_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 264789c41c..f29ea0e559 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -127,12 +127,12 @@ def standardize_chat_format( } if not isinstance(dataset, IterableDataset): - from utils.hardware import safe_num_proc + from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = safe_num_proc() + num_proc = dataset_map_num_proc() else: - num_proc = safe_num_proc(num_proc) + num_proc = dataset_map_num_proc(num_proc) dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Standardizing chat format" @@ -197,12 +197,12 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None): } if not isinstance(dataset, IterableDataset): - from utils.hardware import safe_num_proc + from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = safe_num_proc() + num_proc = dataset_map_num_proc() else: - num_proc = safe_num_proc(num_proc) + num_proc = dataset_map_num_proc(num_proc) dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Converting ChatML to Alpaca format" @@ -247,12 +247,12 @@ def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None): } if not isinstance(dataset, IterableDataset): - from utils.hardware import safe_num_proc + from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: - num_proc = safe_num_proc() + num_proc = dataset_map_num_proc() else: - num_proc = safe_num_proc(num_proc) + num_proc = dataset_map_num_proc(num_proc) dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Converting Alpaca to ChatML format" @@ -435,9 +435,9 @@ def convert_to_vlm_format( if has_urls and total > PROBE_SIZE: import time from concurrent.futures import ThreadPoolExecutor, as_completed - from utils.hardware import safe_num_proc + from utils.hardware import safe_thread_num_proc - num_workers = safe_num_proc() + num_workers = safe_thread_num_proc() _notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...") logger.info( f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..." @@ -521,9 +521,9 @@ def convert_to_vlm_format( # Parallel conversion for URL-based datasets import time from concurrent.futures import ThreadPoolExecutor, as_completed - from utils.hardware import safe_num_proc + from utils.hardware import safe_thread_num_proc - num_workers = safe_num_proc() + num_workers = safe_thread_num_proc() batch_size = 500 start_time = time.time() diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index f2a34ecde9..f86a56d186 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -21,6 +21,8 @@ from .hardware import ( get_physical_gpu_count, get_visible_gpu_count, safe_num_proc, + safe_thread_num_proc, + dataset_map_num_proc, ) __all__ = [ @@ -39,4 +41,6 @@ __all__ = [ "get_physical_gpu_count", "get_visible_gpu_count", "safe_num_proc", + "safe_thread_num_proc", + "dataset_map_num_proc", ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 7d1bc528c6..b055c4e52d 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -510,13 +510,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int: import os import sys - # Windows uses 'spawn' for multiprocessing -- the overhead of re-importing - # torch/transformers/unsloth per worker is typically slower than single-process. - if sys.platform == "win32": + # Windows and macOS use 'spawn' for multiprocessing -- the overhead of + # re-importing torch/transformers/unsloth per worker is typically slower + # than single-process. + if sys.platform in ("win32", "darwin"): return 1 if desired is None or not isinstance(desired, int): - desired = max(1, os.cpu_count() // 3) + desired = max(1, (os.cpu_count() or 1) // 3) visible = get_visible_gpu_count() if visible > 1: @@ -528,3 +529,41 @@ def safe_num_proc(desired: Optional[int] = None) -> int: return capped return desired + + +def safe_thread_num_proc(desired: Optional[int] = None) -> int: + """ + Return a safe worker count for ``ThreadPoolExecutor`` calls. + + Unlike ``safe_num_proc()``, this does NOT cap to 1 on macOS/Windows. + Threads share the parent process address space and are unaffected by + the ``spawn`` vs ``fork`` distinction. + + Args: + desired: The thread count 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() or 1) // 3) + + return desired + + +def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: + """ + Return a safe ``num_proc`` for ``Dataset.map()`` and ``Dataset.filter()``. + + Returns ``None`` on spawn-based platforms (Windows, macOS) because + ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``). + Only ``num_proc=None`` guarantees in-process execution. + """ + import sys + + if sys.platform in ("win32", "darwin"): + return None + return safe_num_proc(desired)