From 6d50004da20dd64ebb602a1df531771826604274 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Dec 2025 05:35:06 -0800 Subject: [PATCH] Update rl.py --- unsloth/models/rl.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 31316e45b7..4ffd790d1e 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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