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

Fix multiprocessing crash on Windows/macOS and unify num_proc logic (#3999)

Browse source 
On Windows and macOS (Python 3.8+), multiprocessing uses the spawn
start method. When datasets .map(num_proc=N) is called, it creates a
Pool(N) which re-imports __main__ in each worker, causing infinite
recursion and a RuntimeError during bootstrapping.

Guard the auto-computed dataset_num_proc in the generated Config
__init__ by checking multiprocessing.get_start_method() != 'fork'.
When the start method is not fork (spawn/forkserver), force
dataset_num_proc = None so datasets takes the single-process path.
Linux fork behavior is unchanged.

Also replace the fixed memory threshold logic with the simpler
adaptive approach: cap at 64, then min(num_proc, int(available_gb)),
with a safety floor of 1 when available memory is at or below 2GB.

Co-authored-by: Daniel Hanchen <danielhanchen@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-02-08 02:50:06 -08:00 • committed by GitHub
commit df720b642c
1 changed files with 6 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

11
unsloth/models/rl.py
View file

@ -932,14 +932,15 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
# Edit dataset_num_proc
if "dataset_num_proc" in call_args:
num_proc_check = (
"if dataset_num_proc is None:\n"
"import multiprocessing as _mp\n"
"if _mp.get_start_method() != 'fork':\n"
" dataset_num_proc = None\n"
"elif dataset_num_proc is None:\n"
" import psutil\n"
" dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n"
" memory_gb_left = psutil.virtual_memory().available / (1024**3)\n"
" if memory_gb_left <= 4: dataset_num_proc = 1 # Too risky, so set to 1\n"
" elif memory_gb_left <= 6: dataset_num_proc = min(2, dataset_num_proc)\n"
" elif memory_gb_left <= 10: dataset_num_proc = min(4, dataset_num_proc)\n"
" elif memory_gb_left <= 14: dataset_num_proc = min(6, dataset_num_proc)\n"
" if memory_gb_left <= 2: dataset_num_proc = 1\n"
" else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))\n"
)
extra_args += num_proc_check