diff --git a/tests/python/test_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py new file mode 100644 index 0000000000..5af31f65b8 --- /dev/null +++ b/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 diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 61a07b686d..d3ada23cf9 100644 --- a/unsloth/models/rl_replacements.py +++ b/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