[Bugfix] Apply logits scaling in the Transformers backend

[tdoublep]

Purpose

Transformers scales the logits in the ForCausalLM wrapper. This backend only loads the decoder and provides its own head and logits processor, so whatever the wrapper does after the decoder has to be reproduced here. Only logit_scale was, which means:

• Granite-family and MiniCPM3 models returned logprobs logits_scalingx too confident: 9x for ibm/PowerLM-3b, 10x for ibm-granite/granite-swash-2b.
• HyperCLOVAX silently dropped its muP multiplier.

Scaling does not reorder the logits, so the generated tokens are unaffected. That is why this went unnoticed, and it means the impact is on consumers of logprob values: sampling with a temperature or penalties, logprob-based evals, speculative decoding acceptance, and anything ranking sequences by score.

There is no config field or API to read the scale from, and no single convention for it either — logit_scale multiplies (Cohere), logits_scaling divides (Granite) except under muP where it multiplies (HyperCLOVAX), logits_mup_width_multiplier divides the hidden states (Inkling), dim_model_base implies a divisor (MiniCPM3), lm_head_multiplier multiplies (FalconH1). Rather than enumerate that, this PR measures it: at model init, build the model's own ForCausalLM class with its decoder and head replaced by stubs (no weights, no checkpoint), run its forward over known hidden states, and take the ratio. Whatever the wrapper does to the logits is then the only thing left in the output, and the value is inherited from Transformers instead of transcribed from it.

If the ratio is not constant across the hidden states, the transform is not a scale (final logit softcapping) and cannot be folded into a logits processor, so we warn and leave the logits alone — as today. Same if the probe cannot be run at all. Neither is a regression, but both are now visible in the log rather than silent.

Found while fixing attention sinks on the same backend (#52156), but independent of it.

Test plan

Three new tests in tests/models/transformers/test_backend.py:

• test_get_logit_scale pins the measured scale for one config per convention (Granite, HyperCLOVAX, Cohere, MiniCPM3, Inkling), plus Llama for no scaling and Gemma2 to assert softcapping is not folded in as a scale. Pure CPU, no weights, runs in seconds.
• test_probe_logit_scale_runs_for_scaled_models is the canary for the approach: it sweeps MODEL_FOR_CAUSAL_LM_MAPPING, sets every float config field that looks like a logit multiplier to a non-default value, and asserts the probe can still be run for each. If a Transformers refactor breaks the probe for a model that does scale its logits, we would silently fall back to 1.0; this fails CI instead.
• test_logits_scaling compares the spread of the top-5 logprobs between the native and Transformers implementations of ibm/PowerLM-3b. Logprob values have to be compared because check_logprobs_close only compares token ids, which a missing division cannot change.

Commands and results (1x GB200, transformers 5.15.0):

pytest tests/models/transformers/test_backend.py -k "logit_scale"      # 8 passed
pytest tests/models/transformers/test_backend.py -k "logits_scaling"   # 1 passed

Mutation checks that the new tests bite:

• Reverting the fix entirely: test_logits_scaling fails on 8/8 positions, max difference 48.06 against a 0.2 tolerance.
• Dropping the config scalars from the decoder stub: the sweep fails on FalconH1, which reads its multiplier off the decoder.
• Dropping weight from the head stub: the sweep fails on xLSTM, which reads the head dtype off it.

What the probe finds across the 173 entries of MODEL_FOR_CAUSAL_LM_MAPPING, with default configs: 135 apply no scaling, 5 scale (Cohere/Cohere2/Cohere2Moe 0.0625, Inkling 1/24, MiniCPM3 0.1), 7 are nonlinear (Gemma2, Gemma3n, VaultGemma, RecurrentGemma, NanoChat, xLSTM — all final logit softcapping), 2 cannot have a default config built (Musicgen), and 24 wrappers cannot be run by the probe. Those 24 are encoder, seq2seq and legacy families (BERT variants, Whisper, TrOCR, XLNet, ...) whose head is not lm_head, plus Llama4/Mllama; none of them declares a logit scale, and they behave exactly as they do today.

Model evaluation

Prompt logprobs for a 55-token paragraph, ibm-granite/granite-swash-2b, bf16, against HF eager as the reference (mean logprob -1.5943). Measured on top of #52156, since that model also needs the sink fix to produce sane output at all:

build	mean logprob	mean abs delta	max abs delta
#52156 alone	-7.8748	6.8942	81.0792
#52156 + this PR	-1.5990	0.0165	0.0863

The remaining 0.0165 is bf16/kernel-level noise.

Notes

• Not a duplicate: no open PR touches logits scaling in the Transformers backend.
• Final logit softcapping is still dropped by this backend, as it was before. It is not a scale, so it does not belong in the logits processor; Attention(..., soft_cap=) handles the attention-logit variant but the final one needs handling in compute_logits. Left for a separate PR, and the probe now warns when it is hit.
• HyperCLOVAX's muP direction is covered by the unit test only. The smallest registered checkpoint is naver-hyperclovax/HyperCLOVAX-SEED-Think-14B, too big for an e2e test here.
• AI assistance was used for this change (Claude Code).

[tdoublep]

Closing this. I think it is better addressed at the transformers level by making the behaviour consistent.