Fix GRPO logit scaling when model is wrapped by DDP (#5955)

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>

tests/python/test_grpo_ddp_model_config.py

@@ -0,0 +1,31 @@
+"""GRPO logit-scaling helpers must read config through DDP wrappers."""
+
+from __future__ import annotations
+
+import os
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
+SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py")
+
+
+def _read_source() -> str:
+    with open(SOURCE_PATH, "r") as fh:
+        return fh.read()
+
+
+def test_grpo_logit_scaling_uses_model_config_helper():
+    src = _read_source()
+    # Helper exists and unwraps DDP/Accelerate wrappers via `.module`.
+    assert "def _unsloth_get_model_config(model):" in src
+    assert 'getattr(model.module, "config", None)' in src
+    # Softcapping takes the model and tolerates a missing config.
+    assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src
+    assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1]
+    # Logit scale/divide read through the unwrapped config, not bare model.config.
+    assert 'getattr(model_config, "logit_scale", 0)' in src
+    assert 'getattr(model_config, "logits_scaling", 0)' in src
+    assert src.count("model_config = _unsloth_get_model_config(model)") >= 2
+    # Helper source is injected into the compiled GRPO trainer.
+    assert "inspect.getsource(_unsloth_get_model_config)" in src
+    # No direct model.config access remains in the RL logit path.
+    assert "model.config" not in src

unsloth/models/rl_replacements.py

@@ -1337,11 +1337,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                     image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end))
 
             temperature = self.temperature
-            logit_softcapping = _unsloth_get_final_logit_softcapping(model.config)
-            logit_scale_multiply = getattr(model.config, "logit_scale", 0)
+            model_config = _unsloth_get_model_config(model)
+            logit_softcapping = _unsloth_get_final_logit_softcapping(model)
+            logit_scale_multiply = getattr(model_config, "logit_scale", 0)
             if logit_scale_multiply is None:
                 logit_scale_multiply = 0
-            logit_scale_divide = getattr(model.config, "logits_scaling", 0)
+            logit_scale_divide = getattr(model_config, "logits_scaling", 0)
             if logit_scale_divide is None:
                 logit_scale_divide = 0
 
@@ -1471,7 +1472,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
 RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies)
 
 
-def _unsloth_get_final_logit_softcapping(config):
+def _unsloth_get_model_config(model):
+    """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers."""
+    config = getattr(model, "config", None)
+    if config is None and hasattr(model, "module"):
+        config = getattr(model.module, "config", None)
+    return config
+
+
+def _unsloth_get_final_logit_softcapping(model):
     """Return final_logit_softcapping for a model config, falling back to the
     nested text sub-config for composite models. Handles both:
       - Gemma-4-style configs where the attribute lives on ``config.text_config``
@@ -1479,6 +1488,9 @@ def _unsloth_get_final_logit_softcapping(config):
         reachable via ``config.get_text_config()``
     Returns 0 if unset, matching the previous behaviour.
     """
+    config = _unsloth_get_model_config(model)
+    if config is None:
+        return 0
     softcap = getattr(config, "final_logit_softcapping", None)
     if softcap is None:
         text_cfg = getattr(config, "text_config", None)
@@ -1499,6 +1511,7 @@ grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"]
 UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"]
 grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"]
 grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"]
+RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping))
 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))
@@ -1616,11 +1629,12 @@ def grpo_trainer_compute_loss(function_name, function):
         input_ids = input_ids[:, -logits_to_keep:]
 
         # Get logit softcapping and logit scale
-        logit_softcapping = _unsloth_get_final_logit_softcapping(model.config)  # Gemma
-        logit_scale_multiply = getattr(model.config, "logit_scale", 0)  # Cohere
+        model_config = _unsloth_get_model_config(model)
+        logit_softcapping = _unsloth_get_final_logit_softcapping(model)  # Gemma
+        logit_scale_multiply = getattr(model_config, "logit_scale", 0)  # Cohere
         if logit_scale_multiply is None:
             logit_scale_multiply = 0
-        logit_scale_divide = getattr(model.config, "logits_scaling", 0)  # Granite
+        logit_scale_divide = getattr(model_config, "logits_scaling", 0)  # Granite
         if logit_scale_divide is None:
             logit_scale_divide = 0