Add UMT5 support (per-layer relative attention bias)

[the-cross-art]

Add UMT5 support (per-layer relative attention bias)

---

Closes #2102. Follows up on #1478, which requested UMT5 support and was closed without a fix.

What this does

UMT5 is structurally identical to mT5 except for one thing: every self-attention layer owns its own relative attention bias table, where T5 and mT5 compute the bias once in layer 0 and share it. CTranslate2 currently cannot convert UMT5 at all, and its runtime assumes the shared-table layout.

Two commits, 96 lines:

1. UMT5Loader — registers UMT5Config and overrides set_stack to keep each layer's own bias table. T5Loader.set_stack reads each layer's bias and then overwrites layers 1..N with layer 0's, which for UMT5 discards 14 of 16 real tables, converts without error, and produces fluent but incorrect output. It also overrides get_vocabulary: UMT5 tokenizers already contain the <extra_id_*> sentinels, so the inherited padding would append duplicates.

2. Per-layer position bias in the Transformer stacks — TransformerEncoder::operator() and TransformerDecoder::operator() create one position_bias buffer per forward pass and thread it through every layer; attention.cc fills it only when empty, so layer 0's bias is used everywhere.

Why not just remove the cache

The comment on #1478 proposed removing the position_bias->empty() guard, and noted it "may lead to performance degradation in T5 and MT5 models". That is what stalled it — it would make every existing T5/mT5 model recompute the bias in every layer.

Instead, this detects which layout a model has and only takes the per-layer path when needed. The per-layer path already exists in MultiHeadAttention: when the caller passes nullptr, each layer falls back to a local buffer and computes from its own table. This change makes that branch reachable rather than adding new computation.

Detection uses pointer identity and needs no new config flag, no format change, and no spec_revision bump. _alias_variables already serializes byte-identical tensors as aliases; register_variable_alias resolves an alias to the same shared_ptr<StorageView>, and Model::copy_to preserves that across device copies. So for T5 every layer's bias resolves to one pointer, and for UMT5 they differ.

Models with no relative attention bias hit the early return, and models with fewer than two layers take the previous path, so both are byte-for-byte unchanged.

No regressions

• t5-small model.bin is byte-identical (same sha256) before and after the converter change.
• Detection: umt5-small → per-layer on both stacks; t5-small and mt5-small → shared on both stacks, i.e. the existing code path exactly.
• t5-small numerical parity vs Hugging Face unchanged at max |diff| = 1.07e-06.
• C++ test suite, macOS/CPU: 197 ran, 196 passed, 1 skipped — identical to baseline.
• C++ test suite, Linux + CUDA: 344/345 passed, 1 pre-existing skip, 0 failures, covering Attention, Transformer, Model, Translator and LayerDevice on both CPU and GPU.

Verification

Same model and harness, only the runtime differs:

	greedy 20 tokens vs Hugging Face
converter commit only	diverges at token 1
both commits	exact match, all 20 tokens

Also validated on a 2.9B-parameter UMT5 checkpoint (24 encoder + 24 decoder layers) on an NVIDIA L4: correct conversion, and 139/140 prompts valid on a production benchmark at float32, replacing a path that previously failed 100% of the time.

Known limitations

• Teacher-forced per-token log-probs for UMT5 differ from Hugging Face by up to 0.605 (mean 0.206), against a 1.07e-06 noise floor measured on t5-small with the same harness in float32. Greedy and beam output match exactly, so generation is correct and only scores/sampling are affected. This is not introduced by these commits — a 1-layer UMT5 on the unchanged shared path still differs (0.022), zeroing every bias table still differs (0.025), and the error is flat in sequence length. I ruled out the activation, head_dim, and a scalar logit factor. Detail and reasoning are in UMT5 support: per-layer relative attention bias without regressing T5/mT5 #2102; I would appreciate a pointer if this is a known characteristic of the T5-family path.
• float16 on GPU is unusable for large T5-family models, but this predates the change: google/mt5-small, which takes the untouched shared-bias path, fails identically on the stock unpatched 4.8.2 PyPI wheel. Instrumentation showed clean bias values with progressive activation overflow from layer 4 onward. float32 is fine; bfloat16 is correct at beam=1.
• Stock google/umt5-* repositories omit model_type from config.json, so AutoConfig.from_pretrained raises before the loader registry is consulted and the key has to be added locally. A fallback dispatching on config.json's architectures entry would fix that for any pre-4.31 repository, but it touches shared converter code so I left it out of this PR. Happy to add it here or separately.

[jordimas]

The per-layer bias approach looks reasonable. Before merging, could you add regression tests for UMT5’s distinct bias tables and T5/mT5 sharing, investigate the reported log-probability mismatch (matching greedy output doesn’t establish sampling or beam-search parity), and document the model_type workaround with a working conversion example?

[the-cross-art]

@jordimas Thanks for the review. All three items are addressed now — tests and docs are pushed, and the
log-probability question turned out not to be a CTranslate2 problem at all: it's a regression in
transformers.

The log-probability mismatch

You were right that greedy proves nothing — argmax doesn't care about scaling. Chasing it
properly, I isolated it to transformers' UMT5 decoder, which isn't causal.

The test needs no reference implementation: a causal decoder's score for position 0 can't
change when you append tokens after it. Extending the target and watching position 0:

• CTranslate2: -12.547371 at every target length 1-6. Drift 0.000000.
• transformers: -12.547372 at length 1, then -12.444402, -12.541924, ... Drift 0.102970.
• transformers with t5-small: drift 0.000000.

They agree exactly at length 1 — the only length where there's no future to leak from.

UMT5Stack does call create_causal_mask, which initially looked like it ruled this out. But
under SDPA that returns None on purpose and hands causality to is_causal on the attention
module. T5 and MT5 pass is_causal=config.is_decoder; UMT5 doesn't:

# t5/modeling_t5.py:375-380  and  mt5/modeling_mt5.py:382-386
T5Attention(config, ..., layer_idx=layer_idx, is_causal=config.is_decoder)

# umt5/modeling_umt5.py:375   <- no is_causal, defaults to False even in the decoder
UMT5Attention(config, has_relative_attention_bias=True, layer_idx=layer_idx)

So no mask and is_causal=False means no causal masking at all.

I verified the cause rather than just reading it — setting is_causal=True on the decoder's
self-attention at runtime, changing nothing else, takes the drift from 0.102970 to 0.000002.
And with that one-line patch, CTranslate2 vs transformers across target lengths 1-7 goes from
6.09e-01 to 4.77e-06. So the whole discrepancy was the reference.

It's also a regression, not longstanding: transformers 5.14.1 has UMT5 correctly causal
(eager attention, mask materialised). #47014 moved the T5 family to SDPA on 2026-07-30 and added
the is_causal parameter to UMT5 but never the call sites. Affected: v5.15.0 through v5.17.0
and main. Filed upstream as huggingface/transformers#49134.

I'd like to drop the "known limitation" from the PR description on this basis — happy to
replace it with a short note that score_batch will differ from transformers until that's
fixed, if you'd prefer something stay.

Sampling and beam search

google/umt5-small, 3 prompts, against HF generate:

• beam 1, 2, 4, 8 — all match
• length_penalty 0.6, 1.0, 2.0 at beam 4 — all match

length_penalty re-ranks on accumulated scores rather than per-step argmax, so that does
exercise the distribution. And the step-1 distribution itself matches: over 40 candidate tokens
(HF's top-12 plus 28 random), std(CT2 - HF) = 0.00000, best fit a=1.00000, b=+0.00002,
identical ranking — which is what sampling draws from.

One more isolation result worth having on record: ctranslate2.Encoder rejects a seq2seq spec,
so I built a standalone encoder via TransformerEncoderModelSpec populated by the same
UMT5Loader.set_stack, and compared against HF's encoder_last_hidden_state —
max |diff| 3.73e-06 (t5-small floor is 1.07e-06). That's a direct check of the per-layer
bias path, and it also clears gated-GELU, 6-head/384-inner attention and RMSNorm, since the
encoder uses all three.

For what it's worth, before finding the cause I'd confirmed the residual survived making the
decoder's bias tables identical (forcing the shared path) and zeroing them entirely, so the
per-layer mechanism was never implicated.

Tests

Added in python/tests/test_transformers.py:

• test_umt5_loader_keeps_one_relative_attention_bias_per_layer — builds a tiny UMT5 in memory,
  fills layer i's bias table with i + 1, runs UMT5Loader.set_stack, and asserts the tables
  come out as [1.0, 2.0, 3.0] for both the encoder and the decoder.
• test_t5_loader_shares_the_first_relative_attention_bias — the same shape for T5, asserting
  every layer's bias is layer 0's object.

Both construct the model from UMT5Config / T5Config with d_model=16 and 3 layers, so there
is no checkpoint download and they run in a few seconds.

I went with testing the loader rather than _alias_variables in test_spec.py as I first
suggested, because this exercises the code the PR actually changes. And I checked the tests can
fail: removing UMT5Loader.set_stack so it falls back to T5Loader's gives

FAILED test_umt5_loader_keeps_one_relative_attention_bias_per_layer
At index 1 diff: 1.0 != 2.0

which is exactly the failure mode — layer 1 receiving layer 0's table. The T5 test keeps passing,
as it should.

That also settles the first of the two fixture questions I asked: the synthetic in-memory model
won, and it's what's in the diff.

The second stands only as an offer. A C++-level test of has_shared_position_bias itself would
need either its declaration moved to a header or two small model fixtures under
tests/data/models/; the same applies to turning the causality check above into a test, since it
needs a converted model. The Python tests pin the converter behaviour and the existing C++ suite
covers the runtime, so I left both out — happy to add either if you'd prefer the detection covered
at that level.

Docs

• docs/guides/transformers.md: a ## UMT5 section following the ## T5 pattern, plus UMT5 in
  the supported-model list. The section carries the model_type workaround as a runnable snippet
  — stock google/umt5-* configs predate that key, so AutoConfig raises before the loader
  registry is reached — then the converter command and a translation example. I ran all three
  steps verbatim in a clean directory before committing.
• One caveat I added while testing that: google/umt5-* are pretrained-only checkpoints (span
  corruption, no downstream fine-tuning), so they answer that prompt with sentinel tokens rather
  than a translation. CTranslate2 matches Hugging Face either way, but the example would read as
  broken without the note.
• Also noted preferring compute_type="float32" for large T5-family models on GPU, per the
  float16 dynamic-range behaviour described earlier.
• README.md: UMT5 added to the encoder-decoder model list.
• CHANGELOG.md: one entry under new features, one under fixes and improvements.

[the-cross-art]

@jordimas
Update: the root cause is now fixed upstream.

huggingface/transformers#49135 (commit d7ba937) landed yesterday. UMT5LayerSelfAttention now passes is_causal=config.is_decoder, which fully explains the log-probability mismatch reported earlier.

After re-running validation against transformers main (5.18.0.dev0, no patches), CT2 and HF now match within float32 noise (max abs diff = 4.77e-06). Position-0 drift in transformers is effectively gone, and greedy/beam decoding matches exactly across all test cases.

I also re-converted google/umt5-small against transformers main; the generated model.bin is byte-identical to the one produced with 5.17.0, so the converter itself is unaffected. CI on 232edee is green (21/21 jobs), and the branch still merges cleanly onto current master.

For the PR description, I'd like to remove the first "Known limitations" bullet since it was caused by this upstream issue. I'll replace it with a brief note that teacher-forced scores may differ from transformers versions older than the release containing #49135. The other limitations remain valid.

We're already running this branch on a 2.9B UMT5 model on an L4 with good results, so I'd prefer to move away from a patched build if possible.

Let me know if there's anything else you'd like addressed before merging. Happy to rebase, reorganize commits, or update the docs.

[jordimas]

The converter tests look good. Could you also add a runtime regression test? The current tests check that each layer’s bias table is preserved during conversion, but they would still pass if the C++ runtime reused layer 0’s bias for every layer.

A small synthetic UMT5 model with distinct, nonconstant bias tables would work well. Convert it, run score_batch, and compare the per-token log probabilities against Hugging Face using eager attention to avoid the upstream SDPA masking issue. Ideally, cover distinct tables in the encoder only, decoder only, and both.

The tables should be nonconstant because a uniform bias cancels during softmax and could hide the bug. The test should pass with this PR and fail when the runtime is forced back to sharing layer 0’s bias.

[the-cross-art]

Added in python/tests/test_transformers.py:

• test_umt5_runtime_uses_per_layer_bias_in_encoder
• test_umt5_runtime_uses_per_layer_bias_in_decoder
• test_umt5_runtime_matches_transformers_with_distinct_biases_in_both_stacks

Each builds a tiny UMT5 (3+3 layers, d_model=32, vocab 64) with distinct, non-constant
per-layer tables, converts it in-process, and compares against Transformers with
_attn_implementation = "eager".

Verification

I reproduced the pre-PR behaviour by forcing has_shared_position_bias() to return true
unconditionally and rebuilding, which is exactly "the runtime reuses layer 0's bias":

assertion	this PR	runtime forced to share layer 0	threshold
encoder output vs Transformers	1.19e-07	1.07e-01	< 1e-5
encoder: distinct vs shared tables	0.1068	0.000000	> 1e-2
decoder log-probs vs Transformers	9.54e-07	2.90e-03	< 1e-4
decoder: distinct vs shared tables	0.002903	0.000000	> 1e-4
both stacks distinct, vs Transformers	9.54e-07	2.90e-03	< 1e-4

All three fail with the forced-shared runtime and pass with the PR. The 0.000000 rows are
the clearest signal: with the bug, a model with distinct tables and a model with every table
set to layer 0's produce bit-identical output, i.e. layers 1..n-1 are ignored outright.

Worth noting that both existing converter tests still pass against the broken runtime —
your diagnosis was exactly right, and that gap is now closed.

Also: black --check, isort --check-only and flake8 clean; the rest of the Python suite
is unaffected (59 passed, 125 skipped, 0 failures).

Two things worth flagging

1. Non-constant tables matter even more than it first appears. You were right that a
uniform table cancels in the softmax. There's a second trap in the same family: the tables
have to be non-constant per bucket, not merely different per layer. My first attempt gave
layer i the table t_i + 2i — different per layer, but the per-layer part was a constant
offset, which cancels. Sensitivity measured ~1e-05, so the test was nearly vacuous. The
committed version draws each table independently from U(-5, 5).

2. The synthetic model has to use Transformers' own initialiser. This cost me the most
time, so it's probably worth recording. T5 omits the 1/sqrt(head_dim) attention scaling
(UMT5Attention.scaling = 1.0) and compensates in the q initialiser, which uses
(d_model * d_kv)**-0.5. I had hand-rolled a uniform 1/sqrt(fan_in) init, which makes q
too large by sqrt(d_kv), saturates the attention softmax, and then CTranslate2 and
Transformers disagree by ~1.5e-01 on a 1-layer model with every bias table zeroed —
nothing to do with this PR. T5 and UMT5 produced bit-identical errors under that init,
which is what confirmed it wasn't the per-layer path. Using UMT5ForConditionalGeneration's
own init with a small initializer_factor brings it to 1e-06. Comments in the test record
this so the next person doesn't repeat it.

One deviation from your suggestion, and why

For the encoder I assert on the encoder output (via ctranslate2.Encoder on a
TransformerEncoderModelSpec populated by the same UMT5Loader.set_stack) rather than on
score_batch log-probs.

The reason is measurement, not convenience: on a randomly initialised tiny model, changing
every encoder bias table moves the final token scores by exactly 0.000000. The weights
are small enough that the decoder's cross-attention contribution is swamped by the residual
stream, so a score-level encoder assertion would pass whatever the encoder did — the same
failure mode as a constant bias table, reached by a different route. I tried to recover the
signal by amplifying the cross-attention value path; that lifts encoder sensitivity to
~8e-04 but degrades parity to ~1e-03 in lockstep, so the margin never exceeds ~1x.

Checked at the encoder output instead, the same comparison separates the two cases by about
six orders of magnitude (1.19e-07 parity vs 0.1068 sensitivity), and it is still a
genuine runtime check through ctranslate2.Encoder. The decoder and both-stacks cases do use
score_batch as you asked.

If you'd rather have a score_batch assertion for the encoder too, the way to get real
signal is a trained checkpoint rather than a random one — google/umt5-small gives
max |CT2 - HF| = 4.77e-06 across target lengths 1-7 end-to-end — but that needs a 1.2 GB
download, so I kept it out of the unit tests. Happy to add it behind a marker if useful.