GraniteSpeech: logits_scaling not applied since v5.10 (logits are 8x too large) #49937

System Info

• transformers version: 5.20.2
• Platform: macOS-26.7-arm64-arm-64bit
• Python version: 3.12.14
• Huggingface_hub version: 2.2.0
• Safetensors version: 0.8.0
• Accelerate version: 1.14.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.13.0 (NA)
• Using distributed or parallel set-up in script?: no

Same result on main (944e22e). The repro below runs on CPU.

Who can help?

@eustlb

Reproduction

Hi! I was comparing ibm-granite/granite-speech-4.1-2b-plus outputs between transformers and the MLX port in mlx-audio, and the transformers logits came out exactly 8x larger. All the Granite Speech checkpoints have text_config.logits_scaling = 8, and it looks like the composite model doesn't apply it anymore.

• background: I added supported for the model to mlx-audio last week-ish, and was looking at compatibility as IBM updated the weights/repo with a new v2610 improved version just a few days ago
• my journey down the rabbit hole/comparisons on inference in MLX brought me to dig through the modeling code, and here we are

GraniteForCausalLM divides by config.logits_scaling at the end of forward (here). Up to 5.9 the Granite Speech model wrapped a full GraniteForCausalLM, so the division happened there. Since #45534 (v5.10.0) it builds the base GraniteModel with its own lm_head, and the division didn't come along (here):

# modeling_granite_speech.py (and the generated modeling_granite_speech_plus.py)
hidden_states = slice_logits_to_keep(hidden_states, logits_to_keep)
logits = self.lm_head(hidden_states)  # no / self.config.text_config.logits_scaling

Minimal repro, no checkpoint needed:

import torch
from transformers import (
    GraniteConfig,
    GraniteForCausalLM,
    GraniteSpeechConfig,
    GraniteSpeechForConditionalGeneration,
)

# tiny random model, same logits_scaling as the real checkpoints (8.0)
text_config = GraniteConfig(
    hidden_size=64,
    intermediate_size=128,
    num_hidden_layers=2,
    num_attention_heads=4,
    num_key_value_heads=2,
    vocab_size=128,
    logits_scaling=8.0,
)
config = GraniteSpeechConfig(text_config=text_config, has_lora_adapter=False)
model = GraniteSpeechForConditionalGeneration(config).eval()

# same text weights in a plain GraniteForCausalLM
lm = GraniteForCausalLM(text_config).eval()
lm.model.load_state_dict(model.model.language_model.state_dict())
lm.lm_head.load_state_dict(model.lm_head.state_dict())

input_ids = torch.randint(0, 100, (1, 8))
with torch.no_grad():
    speech_logits = model(input_ids=input_ids).logits
    lm_logits = lm(input_ids=input_ids).logits

print(torch.allclose(speech_logits, lm_logits, atol=1e-5))      # False
print(torch.allclose(speech_logits, 8 * lm_logits, atol=1e-5))  # True

Greedy decoding doesn't care about a positive scale, so transcripts look fine, which I assume is also why the integration tests didn't catch it. Anything that uses the actual values is off though:

• Sampling with temperature or top_p. temperature=1.0 effectively becomes 1/8.
• Beam search scores.
• output_scores and compute_transition_scores.
• The loss when passing labels for fine-tuning.
• repetition_penalty with num_beams > 1, since beam search applies it to the log-softmax of the 8x logits. IBM's model card recommends num_beams=4 with repetition_penalty=2.0.

As far as I can tell, the same thing happens in other models that put their own lm_head on top of a text backbone with output scaling:

• LlavaNextForConditionalGeneration with a Granite text config (e.g. ibm-granite/granite-vision-3.2-2b) has been 8x off since 🔴 [VLM] Add base model without head  #37033 (v4.52).
• AyaVisionForConditionalGeneration drops Cohere's logit_scale the same way.

Happy to open separate issues for those if that's better.

Expected behavior

The logits should match GraniteForCausalLM with the same weights, i.e. be divided by text_config.logits_scaling like before 5.10. granite4_vision already does this (here):

logits = logits / self.config.text_config.logits_scaling

I'm happy to open a PR with that one-line change (Plus picks it up through its modular file), plus a small test that compares against GraniteForCausalLM. Let me know if you'd prefer a different approach!