Fix UMT5 decoder self-attention not being causal by the-cross-art · Pull Request #49135 · huggingface/transformers

the-cross-art · 2026-09-27T21:16:41Z

What does this PR do?

UMT5LayerSelfAttention never passes is_causal to UMT5Attention, so the decoder's
self-attention is built with is_causal=False. Under SDPA the causal mask is skipped when there
is no padding (_ignore_causal_mask_sdpa → create_causal_mask returns None), leaving the
flag as the only thing applying causality. The decoder can therefore attend to future positions.

T5LayerSelfAttention and MT5LayerSelfAttention both pass is_causal=config.is_decoder
(t5,
mt5).
UMT5 is the only member of the family missing it, so this PR applies the same approach.

First bad commit is 560f36c (#47014, which moved the family to SDPA): it added the
is_causal parameter to UMT5Attention but not the call sites. v5.14.1 is correct (eager,
mask materialised); v5.15.0 onward is not.

The fix is one line, plus is_causal=False stated explicitly on the cross-attention to match T5.

Reproducer

A causal decoder's logits at position 0 must not change when tokens are appended after it.

import torch
from transformers import AutoTokenizer, UMT5ForConditionalGeneration
torch.set_grad_enabled(False)

# google/umt5-small has no "model_type" in its config, hence the explicit class
tok = AutoTokenizer.from_pretrained("google/umt5-small", use_fast=False)
m = UMT5ForConditionalGeneration.from_pretrained("google/umt5-small", dtype=torch.float32).eval()

ids = tok("Translate English to French: The house is wonderful.", return_tensors="pt").input_ids
tgt = tok("La maison est merveilleuse et grande.", return_tensors="pt").input_ids[0].tolist()
start = m.config.decoder_start_token_id

for n in (1, 2, 6):
    lg = m(input_ids=ids, decoder_input_ids=torch.tensor([[start] + tgt[:n - 1]])).logits[0]
    print(n, float(torch.log_softmax(lg[0].float(), -1)[tgt[0]]))

before            after (this PR)     v5.14.1
1  -12.547372     -12.547372          -12.547372
2  -12.444402     -12.547371          -12.547372
6  -12.522055     -12.547370          -12.547372
   drift 1.03e-01    drift 1.91e-06      drift 1.0e-06

The patched values match v5.14.1, so this restores previous behaviour rather than changing it.

generate() is unaffected — incremental decoding feeds one query position at a time, so there
is no future to leak from. Only full-sequence teacher forcing is affected, i.e. training loss,
evaluation perplexity and offline scoring.

The behaviour was also batch-shape dependent, since the mask is only skipped without padding:

decoder_attention_mask=None       drift 1.03e-01   leaking
decoder_attention_mask=all ones   drift 1.03e-01   leaking
decoder_attention_mask=padded     drift 1.91e-06   causal

DataCollatorForSeq2Seq emits no decoder_attention_mask, so the default Seq2SeqTrainer path
was in the leaking case.

test_decoder_self_attention_is_causal asserts the position-0 logits are unchanged when the
decoder input is extended, with no decoder_attention_mask (the case where SDPA skips the mask).
test_decoder_self_attention_is_causal_flag asserts the decoder self-attention is causal while
the encoder and cross-attention are not. Both fail on main (greatest absolute difference 3.40
against a 1e-4 tolerance) and pass here.

The behavioural test raises initializer_factor from the common tester's 0.002 to 1.0. At
0.002 the violation measures ~6e-11, so a test written against the default config passes even
with the bug present — worth knowing if you add similar tests elsewhere.

tests/models/umt5/test_modeling_umt5.py: 231 passed, 279 skipped, 2995 subtests passed.

Note on `test_eager_matches_sdpa_inference`

Not addressed here, but that test cannot catch this for the T5 family: it builds the eager model
as deepcopy(model_sdpa) then set_attn_implementation("eager"), while T5Stack/UMT5Stack
deepcopy the config in __init__ — and it is that copy which create_causal_mask and the
attention-interface lookup read. The stacks keep sdpa, so it compares sdpa against sdpa.
from_pretrained(attn_implementation="eager") does propagate, which is why the dispatch tests
pass. Happy to open a separate issue.

Fixes #49134

Code Agent Policy

(First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent

I can't tick that box honestly: I used an AI assistant for this work, and I'd rather disclose it
than leave it ambiguous.

For context on how I got here, since I appreciate the policy is aimed at drive-by agent PRs: I
wasn't looking at Transformers at all. I was adding UMT5 support to CTranslate2
(OpenNMT/CTranslate2#2103), and its maintainer asked me to explain a per-token log-probability
mismatch against Transformers rather than rely on matching greedy output. Tracking that down over
a few hours is what surfaced this — the reference was the thing that was wrong, not my port. So
this is a byproduct of work I was already doing, not an issue I went looking for.

The parts your guide asks for were done and verified: the first bad commit is identified, the
reproducer was run on real weights, the behaviour is compared against T5/MT5 with links, and the
diff is one line plus tests.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline and the Pull Request checks?
Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case. → [Regression v5.15.0+] UMT5 decoder self-attention is not causal: is_causal never reaches UMT5Attention #49134
Did you make sure to update the documentation with your changes according to the guidelines? (no documented behaviour changes)
Did you write any new necessary tests?

Who can review?

@vasqu @jiqing-feng — from git blame: #47014 introduced this, authored by @jiqing-feng and
merged by @vasqu, who is also listed for attention.

vasqu

Yep that was indeed missed but also note that we will probably deprecate / remove this feature soon

vasqu · 2026-09-28T13:20:20Z

+    def test_decoder_self_attention_is_causal(self):
+        # The logits at a decoder position must not depend on the tokens that follow it.
+        # Regression test for the decoder self-attention losing `is_causal`: under SDPA the
+        # causal mask is skipped when there is no padding, so causality comes from the flag
+        # alone. See https://github.com/huggingface/transformers/issues/49134
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        # The common tester initializes weights with initializer_factor=0.002, which makes the
+        # attention output too small for a causality violation to be measurable.
+        config = copy.deepcopy(config)
+        config.initializer_factor = 1.0
+
+        model = UMT5ForConditionalGeneration(config).to(torch_device).eval()
+        input_ids = inputs_dict["input_ids"]
+        decoder_input_ids = inputs_dict["decoder_input_ids"]
+        self.assertGreater(decoder_input_ids.shape[1], 1)
+
+        # No decoder_attention_mask on purpose: that is the case where SDPA skips the mask.
+        with torch.no_grad():
+            truncated = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids[:, :1])
+            full = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)
+
+        torch.testing.assert_close(truncated.logits[:, 0], full.logits[:, 0], rtol=1e-4, atol=1e-4)
+
+    def test_decoder_self_attention_is_causal_flag(self):
+        # The decoder self-attention must be built causal, the encoder and the cross-attention
+        # bidirectional.
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        model = UMT5ForConditionalGeneration(config)
+
+        self.assertTrue(all(block.layer[0].SelfAttention.is_causal for block in model.decoder.block))
+        self.assertFalse(any(block.layer[0].SelfAttention.is_causal for block in model.encoder.block))
+        self.assertFalse(any(block.layer[1].EncDecAttention.is_causal for block in model.decoder.block))


Can we add a test as in #48760? I think that sole test would be enough then

…lidation and improve test reliability

github-actions · 2026-09-28T13:38:08Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: umt5

vasqu · 2026-09-28T13:40:00Z

        pass

+    def test_decoder_causal_mask(self):
+        # Regression test for #49134


Suggested change

# Regression test for #49134

# Regression test for #49135

no?

happy to switch it if you'd rather point at the PR

Ah it's the issue sorry i didnt notice, no should be fine

vasqu · 2026-09-28T13:49:07Z

run-slow: umt5

github-actions · 2026-09-28T13:50:21Z

AMD CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on AMD:

models: ["models/umt5"]

github-actions · 2026-09-28T13:50:32Z

Nvidia CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on Nvidia:

models: ["models/umt5"]
quantizations: []

github-actions · 2026-09-28T13:54:54Z

CI Results (AMD)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	6ef1ac0b	workflow commit (merge commit)
PR	72041646	branch commit (from PR)
main	885320ea	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

github-actions · 2026-09-28T14:02:05Z

CI recap

Dashboard: View test results in Grafana
Latest run: 36431341880:1
Result: success | Jobs: 1 | Tests: 30 | Failures: 0 | Duration: 53s

github-actions · 2026-09-28T14:13:07Z

CI Results (Nvidia)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	6ef1ac0b	workflow commit (merge commit)
PR	72041646	branch commit (from PR)
main	885320ea	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

HuggingFaceDocBuilderDev · 2026-09-28T14:16:38Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Adds a contribution table for the merged huggingface/transformers#49135 and OpenNMT/CTranslate2#2103 PRs (a UMT5 correctness fix found while adding UMT5 support to CTranslate2), plus a merged-PRs badge, inspired by how github.com/Bruce-Yii surfaces upstream OSS work. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Add causal attention flag to UMT5 decoder and implement regression tests

b7efeb7

vasqu approved these changes Sep 28, 2026

View reviewed changes

Refactor decoder self-attention tests for UMT5: update causal mask va…

7204164

…lidation and improve test reliability

vasqu reviewed Sep 28, 2026

View reviewed changes

Merge branch 'main' into fix-umt5-decoder-causal

2191dd5

vasqu enabled auto-merge September 28, 2026 14:11

vasqu added this pull request to the merge queue Sep 28, 2026

Merged via the queue into huggingface:main with commit d7ba937 Sep 28, 2026
39 checks passed

the-cross-art mentioned this pull request Sep 28, 2026

Add UMT5 support (per-layer relative attention bias) OpenNMT/CTranslate2#2103

Merged

tarekziade mentioned this pull request Oct 6, 2026

set_attn_implementation() does not reach T5 encoder/decoder: model reports SDPA but still runs eager #49372

Open

Conversation

the-cross-art commented Sep 27, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Reproducer

Note on test_eager_matches_sdpa_inference

Code Agent Policy

Before submitting

Who can review?

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

vasqu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 28, 2026

Uh oh!

vasqu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

the-cross-art Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu commented Sep 28, 2026

Uh oh!

github-actions Bot commented Sep 28, 2026

AMD CI

Uh oh!

github-actions Bot commented Sep 28, 2026

Nvidia CI

Uh oh!

github-actions Bot commented Sep 28, 2026

CI Results (AMD)

Commit Info

Uh oh!

github-actions Bot commented Sep 28, 2026

CI recap

Uh oh!

github-actions Bot commented Sep 28, 2026

CI Results (Nvidia)

Commit Info

Uh oh!

HuggingFaceDocBuilderDev commented Sep 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

the-cross-art commented Sep 27, 2026 •

edited by github-actions Bot

Loading

Note on `test_eager_matches_sdpa_inference`