Fix gpt temporary patch for grpo to happen after compile

[mmathew23]

Fixes #4175

The gpt oss grpo patch in temporary_patches runs 3 times. Since it runs before compiler, compiler doesnt match the causallm pattern so we lose the connection to unsloth fused loss. This PR proposes phases for temporary patches so patches can determine when they want to run.

Depends on unslothai/unsloth-zoo#536

Notebooks Tested:
gpt finetuning
https://colab.research.google.com/drive/11Hi8srn79s5gGoVbOhtw09LCkq6wfJnb?usp=sharing

gpt long context
https://colab.research.google.com/drive/1jcMg3xQ5c1C1Vxd-6DcZjIuzgfEBh56R?usp=sharing

gpt grpo
https://colab.research.google.com/drive/1JwlBpfg-cpnUDtlIdtziSfDxKP_YlMwW?usp=sharing

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue where certain temporary patches, particularly the GPT GRPO patch, were executing at an incorrect stage relative to the compiler, leading to a loss of connection to unsloth fused loss. The core change introduces a robust phasing mechanism for temporary patches, allowing them to be explicitly scheduled to run during initialization, before compilation, or after compilation. This refactoring provides fine-grained control over patch execution order, thereby resolving the timing conflict and ensuring proper integration and functionality of all temporary patches within the system.

Highlights

• Phased Patch Execution: Introduced a new _run_temporary_patches function that enables temporary patches to be executed in specific phases: 'init', 'pre_compile', and 'post_compile'.
• Dynamic Patch Argument Handling: Implemented logic within _run_temporary_patches to dynamically check if a temporary patch function accepts a phase argument, allowing for more flexible patch design.
• Corrected Patch Timing: Refactored all existing temporary patch application points to utilize the new phased system, ensuring patches like the GPT GRPO patch run at the correct stage relative to the compiler, resolving previous timing conflicts.

Changelog

• unsloth/models/_utils.py
  • Introduced a new helper function _run_temporary_patches to manage the execution of temporary patches based on a specified phase.
  • Modified the initial application of temporary patches to use _run_temporary_patches("init").
  • Updated the unsloth_compile_transformers function to call _run_temporary_patches("pre_compile") before compilation.
  • Updated the unsloth_compile_transformers function to call _run_temporary_patches("post_compile") after compilation.

Activity

• No specific activity has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1a93b57df6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Narrow TypeError fallback to signature introspection

Catching TypeError around the whole patch invocation can mis-handle real runtime failures inside a temporary patch: if temporary_patch(phase=phase) raises a TypeError from its own logic (not from argument binding), this code retries temporary_patch() and can execute patch side effects twice while obscuring the original failure path. This regression only appears when a patch throws TypeError, but in that case it can leave monkey-patched state inconsistent and make debugging much harder.

Useful? React with 👍 / 👎.

[gemini-code-assist]

Code Review

This pull request introduces a phased execution for temporary patches, allowing them to run at different stages like 'init', 'pre_compile', and 'post_compile'. This is achieved by adding a _run_temporary_patches function that inspects patch signatures for a phase parameter. My review focuses on improving the robustness of the new error handling logic. Overall, the change is well-structured to address the issue of patch ordering.

[gemini-code-assist]

The except (ValueError, TypeError) block is overly broad and can mask TypeError exceptions that signal a genuine issue with a patch function, such as incorrect arguments. By catching TypeError and then unconditionally calling temporary_patch(), you might suppress the original error and cause a new, potentially more confusing one. It's better to only catch ValueError, which inspect.signature raises for callables it can't inspect (like some C built-ins). A TypeError should generally be allowed to propagate to indicate a problem with the patch or its invocation, leading to faster bug detection.

Suggested change

	temporary_patch()
	except (ValueError, TypeError):
	except ValueError:
	temporary_patch()