From c3b7614bd5879edbf184a058f589ab21f754d2f3 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sun, 8 Mar 2026 08:14:38 -0500 Subject: [PATCH] Fix gpt temporary patch for grpo to happen after compile (#4180) * Fix gpt temporary patch for grpo to happen after compile * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 03cdcd92fa..89ced4f7c4 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -254,8 +254,21 @@ def prefer_flex_attn_if_supported(model_class, config): return None -for temporary_patch in TEMPORARY_PATCHES: - temporary_patch() +def _run_temporary_patches(phase): + import inspect + + for temporary_patch in TEMPORARY_PATCHES: + try: + sig = inspect.signature(temporary_patch) + if "phase" in sig.parameters: + temporary_patch(phase = phase) + else: + temporary_patch() + except (ValueError, TypeError): + temporary_patch() + + +_run_temporary_patches("init") # ============================================= # Disable some warnings which can get annoying @@ -2095,8 +2108,7 @@ def unsloth_compile_transformers( # Run patches BEFORE compiler so class replacements (e.g. GptOssTopKRouter, # GptOssExperts) are in place before the compiler caches references to them. - for temporary_patch in TEMPORARY_PATCHES: - temporary_patch() + _run_temporary_patches("pre_compile") for model_type in model_types: _unsloth_compile_transformers( @@ -2128,8 +2140,7 @@ def unsloth_compile_transformers( supports_sdpa = supports_sdpa, ) # Redo patches which override compiler - for temporary_patch in TEMPORARY_PATCHES: - temporary_patch() + _run_temporary_patches("post_compile") return model_types, supports_sdpa[0]