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

Grpo compile settings update (#3927)

Browse source 
* 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>
This commit is contained in:
pluesclues 2026-01-24 20:17:55 -05:00 • committed by GitHub
commit 3fde3a91ee
1 changed files with 20 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth/models/rl.py
View file

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