From 3fde3a91eecaa85f14fe930122654ed3b49d905a Mon Sep 17 00:00:00 2001 From: pluesclues <136766175+pluesclues@users.noreply.github.com> Date: Sat, 24 Jan 2026 20:17:55 -0500 Subject: [PATCH] Grpo compile settings update (#3927) * Add torch compile options for GRPOTrainer * Update CUDA settings based on device capability * Add triton persistent TMA matmul condition * Fix syntax for triton.enable_persistent_tma_matmul * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update rl.py * Update rl.py --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/rl.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 9788207c99..803153e608 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1090,6 +1090,26 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): align_logprobs_with_mask_code = align_logprobs_with_mask_code, ) + if RLTrainer_name == "GRPOTrainer": + new_options = """torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + #"combo_kernels" : torch.cuda.get_device_capability()[0] >= 10, + "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] >= 9, + "cuda.cutlass_epilogue_fusion_enabled": torch.cuda.get_device_capability()[0] >= 9, + "cuda.cutlass_tma_only": torch.cuda.get_device_capability()[0] >= 9, + "cuda.compile_opt_level" : "-O2", + "cuda.enable_cuda_lto" : True, + }""" + + pattern = r"torch_compile_options\s*=\s*\{[^}]*\}" + + RLTrainer_source = re.sub( + pattern, new_options, RLTrainer_source, flags = re.DOTALL + ) + if RLTrainer_name == "SFTTrainer": original_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask"]' new_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask","labels"]'