ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: subprocess crash during map operation on Windows (#4507)

Browse source 
* fix: handle Windows subprocess crash during dataset.map()

Windows uses spawn (not fork) for multiprocessing. Spawned workers
cannot resolve Unsloth's dynamically compiled cache modules from
unsloth_compiled_cache/, causing ModuleNotFoundError and RuntimeError
during dataset.map() tokenization.

Add two platform-guarded patches for sys.platform == "win32":
1. Force HF_DATASETS_MULTITHREADING_MAX_WORKERS=1 and set spawn method
2. Monkey-patch Dataset.map() to force num_proc=None

Fixes #4490

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review: extend spawn fix to macOS, add multiprocess fallback

- Change platform checks from sys.platform == "win32" to
  sys.platform != "linux" so macOS (also spawn-based) is covered
- Wrap multiprocess import in try/except falling back to stdlib
  multiprocessing when the multiprocess package isn't installed
- Rename _win32_safe_map to _spawn_safe_map to reflect broader scope

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: replace global Dataset.map monkey-patch with targeted num_proc routing

The previous approach had issues: Patch 1 set HF_DATASETS_MULTITHREADING_MAX_WORKERS
and forced set_start_method (dead code on platforms already using spawn), and Patch 2
globally monkey-patched Dataset.map() (too broad, missed Dataset.filter()).

Replace with a two-layer fix:

1. Studio layer: Add dataset_map_num_proc() that returns None on spawn platforms
   (Windows, macOS). Unlike num_proc=1 which still creates Pool(1) and spawns a
   worker, num_proc=None runs Dataset.map()/filter() truly in-process.
   Update all dataset.map() callsites to use it. ThreadPoolExecutor callers
   (format_conversion.py) keep using safe_num_proc() since threads are unaffected.

2. Root-cause layer: Propagate UNSLOTH_COMPILE_LOCATION via PYTHONPATH on spawn
   platforms so spawned workers can import compiled modules. Mirrors the .venv_t5
   pattern in worker.py. Does not import unsloth_zoo.compiler (heavy torch/triton
   imports). Completely skipped on Linux.

Also extend safe_num_proc() to return 1 on macOS (was only guarding Windows),
and narrow the transformers 5.x dataloader guard from != "linux" to explicit
("win32", "darwin").

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: add safe_thread_num_proc() for ThreadPoolExecutor callsites

safe_num_proc() correctly caps to 1 on macOS/Windows for process-based
multiprocessing, but format_conversion.py reuses it for ThreadPoolExecutor
workers. Threads share address space and are unaffected by spawn, so
capping to 1 makes image URL downloads sequential -- a real regression.

Add safe_thread_num_proc() that skips the platform guard but keeps the
cpu_count heuristic, and switch both ThreadPoolExecutor callsites in
format_conversion.py to use it.

* fix: remove double-wrap in dataset_num_proc + fix num_proc=1 in datasets route

- trainer.py:3009: Replace safe_num_proc(max(1, os.cpu_count() // 4))
  with max(1, (os.cpu_count() or 1) // 4) to avoid double-wrapping
  inside dataset_map_num_proc which already calls safe_num_proc
- trainer.py:15-20: Clarify comment on PYTHONPATH propagation
- datasets.py:445: Change num_proc=1 to num_proc=None for 10-row
  preview slice (avoids unnecessary multiprocessing overhead)

* fix: guard os.cpu_count() against None in worker-count helpers

os.cpu_count() can return None on some platforms. Use (os.cpu_count() or 1)
to prevent TypeError in safe_num_proc() and safe_thread_num_proc().

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Andrew Barnes 2026-03-22 08:21:09 -04:00 • committed by GitHub
commit 2c5d3c48ec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 101 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

42
studio/backend/core/training/trainer.py
View file

@ -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."):

2
studio/backend/routes/datasets.py
View file

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

12
studio/backend/utils/datasets/chat_templates.py
View file

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

26
studio/backend/utils/datasets/format_conversion.py
View file

@ -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()

4
studio/backend/utils/hardware/__init__.py
View file

@ -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",
]

47
studio/backend/utils/hardware/hardware.py
View file

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