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

Update rl.py

Browse source
This commit is contained in:
Daniel Han 2025-12-23 05:35:06 -08:00
commit 6d50004da2
1 changed files with 7 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/rl.py
View file

@ -812,8 +812,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
if "dataset_num_proc" in call_args:
num_proc_check = (
"if dataset_num_proc is None:\n"
" from multiprocessing import cpu_count\n"
" dataset_num_proc = min(max(cpu_count()+4, 2), 64)\n"
" import psutil\n"
" dataset_num_proc = min(max(psutil.cpu_count()+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 <= 8: dataset_num_proc = min(4, dataset_num_proc)\n"
" elif memory_gb_left <= 12: dataset_num_proc = min(6, dataset_num_proc)\n"
)
extra_args += num_proc_check