Clear MRoPE after generation for GRPO (#5683)

* clear mrope state after generation

* move clear mrope to here

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [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>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

studio/backend/tests/test_cpu_threads.py

@@ -63,16 +63,18 @@ def test_cpu_thread_cap_is_opt_in(raw):
 
 
 # Anything that is not a positive integer raises a clear ValueError.
-@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"])
+@pytest.mark.parametrize(
+    "raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]
+)
 def test_cpu_thread_cap_requires_positive_integer(raw):
-    with pytest.raises(ValueError, match="must be a positive integer"):
+    with pytest.raises(ValueError, match = "must be a positive integer"):
         configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw})
 
 
 # env=None path uses real os.environ (production call from run.py / main.py).
 def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch):
     for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
-        monkeypatch.delenv(variable, raising=False)
+        monkeypatch.delenv(variable, raising = False)
     monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3")
 
     configure_cpu_threads()
@@ -84,7 +86,7 @@ def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch):
 # Calling twice must not flip any seeded value.
 def test_cpu_thread_cap_idempotent(monkeypatch):
     for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
-        monkeypatch.delenv(variable, raising=False)
+        monkeypatch.delenv(variable, raising = False)
     monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5")
 
     configure_cpu_threads()
@@ -138,9 +140,9 @@ def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point):
 
     result = subprocess.run(
         [sys.executable, str(entry_point)],
-        env=env,
-        capture_output=True,
-        text=True,
+        env = env,
+        capture_output = True,
+        text = True,
     )
 
     assert result.returncode == 1

unsloth/models/rl_replacements.py

@@ -53,6 +53,23 @@ from ._utils import _get_inference_mode_context_manager
 RL_EXTRA_ARGS = defaultdict(list)
 RL_FUNCTIONS = defaultdict(list)
 RL_PRE_ITEMS = defaultdict(list)
+
+
+def _unsloth_clear_stateful_mrope(model):
+    modules = getattr(model, "modules", None)
+    if modules is None:
+        return False
+
+    cleared = False
+    for module in modules():
+        if hasattr(module, "compute_3d_position_ids") and hasattr(
+            module, "rope_deltas"
+        ):
+            module.rope_deltas = None
+            cleared = True
+    return cleared
+
+
 RL_CONFIG_CHANGES = defaultdict(list)
 RL_METRICS_CHANGES = defaultdict(list)
 RL_ADDITIONAL_FUNCTIONS = defaultdict(list)
@@ -897,6 +914,18 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
         )
         function = function.replace(string_to_find, replacement_string)
 
+    _generate_return = """        ) = self._generate(prompts)"""
+    if _generate_return in function and "_unsloth_clear_stateful_mrope" not in function:
+        function = function.replace(
+            _generate_return,
+            _generate_return
+            + """
+
+        _unsloth_clear_stateful_mrope(
+            self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False)
+        )""",
+        )
+
     if "wake_up()" not in function:
         # Sleep functionality has been added to trl in v0.23.0. We do not want to redo this.
         # https://github.com/huggingface/trl/commit/edbe8234bc7e528f72ac76607de9d3e4753e2709
@@ -1509,6 +1538,7 @@ RL_PRE_ITEMS["grpo_trainer"].append(
 )
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids))
+RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_clear_stateful_mrope))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss))