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"]'