diff --git a/tests/test_chunked_ce_partial_forward.py b/tests/test_chunked_ce_partial_forward.py new file mode 100644 index 00000000000..b857b7d02bb --- /dev/null +++ b/tests/test_chunked_ce_partial_forward.py @@ -0,0 +1,47 @@ +# Copyright 2020-2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 + +import functools +import inspect +import types +from unittest.mock import MagicMock + +import torch +import torch.nn as nn + +from trl.trainer.sft_trainer import _patch_chunked_ce_lm_head + + +class _TinyLM(nn.Module): + def __init__(self): + super().__init__() + self.config = MagicMock() + self.config.text_config = None + # top-level config attrs used when is_vlm=False + self.config.final_logit_softcapping = None + self.config.logit_scale = 1.0 + self.config.output_router_logits = False + self.lm_head = nn.Linear(4, 8, bias=False) + self.base_model = MagicMock() + + def get_output_embeddings(self): + return self.lm_head + + def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs): + return MagicMock(loss=torch.tensor(0.0), logits=None) + + +def test_patch_chunked_ce_accepts_partial_forward(): + """Qwen3.5-style models expose forward as functools.partial (#6483).""" + model = _TinyLM() + real = model.forward + model.forward = functools.partial(real) # no __func__ + + # Should not raise AttributeError: 'partial' object has no attribute '__func__' + _patch_chunked_ce_lm_head(model, chunk_size=2, is_vlm=False) + + assert isinstance(model.forward, types.MethodType) + sig = inspect.signature(model.forward) + # Bound method signature should include original kwargs surface + assert "input_ids" in sig.parameters or len(sig.parameters) >= 1 diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py index c88601cebc6..6a0cc710671 100644 --- a/trl/trainer/sft_trainer.py +++ b/trl/trainer/sft_trainer.py @@ -13,6 +13,7 @@ # limitations under the License. import contextlib +import functools import inspect import json import os @@ -373,8 +374,16 @@ def _chunked_ce_forward( # Keep the original forward signature so `generate`'s `_validate_model_kwargs` still sees the # model's real inputs (e.g. VLM `pixel_values`, `spatial_shapes`) and doesn't reject them. The - # unbound `__func__` signature makes `MethodType`'s `self`-stripping land correctly. - _chunked_ce_forward.__signature__ = inspect.signature(original_forward.__func__) + # unbound function signature makes `MethodType`'s `self`-stripping land correctly. + # + # Some multimodal models (e.g. Qwen3.5) expose `forward` as a `functools.partial` rather than a + # bound method, so `original_forward.__func__` is unavailable — unwrap partials / fall back to + # the callable itself (see #6483). + forward_for_sig = original_forward + while isinstance(forward_for_sig, functools.partial): + forward_for_sig = forward_for_sig.func + unbound = getattr(forward_for_sig, "__func__", forward_for_sig) + _chunked_ce_forward.__signature__ = inspect.signature(unbound) model.forward = types.MethodType(_chunked_ce_forward, model)