Always tie structural weights, regardless of tie_word_embeddings

[jiqing-feng]

Problem

get_expanded_tied_weights_keys returns {} as soon as config.tie_word_embeddings is False, dropping
every entry of _tied_weights_keys. But that dict mixes two kinds of entries:

• output head — {"lm_head.weight": "model.shared.weight"}; this is what the config flag is meant to gate.
• structural — {"encoder.embed_tokens.weight": "shared.weight", ...}; module aliasing the architecture
  relies on, not configurable.

So for encoder-decoder models with tie_word_embeddings=False, the encoder/decoder embeddings end up
randomly initialized. longt5 (Stancld/longt5-tglobal-large-16384-pubmed-3k_steps) generates
'muralvoenseenseensevo(0)(0)(0)...'; loading reports encoder.embed_tokens.weight | MISSING and their
std is 1.00 instead of 8.92. The Hub checkpoint only stores shared.weight + lm_head.weight, so this
is not recoverable. Regressed in #41541 / #41580.

Fix

Keep the entries whose both ends are nn.Embedding, drop the rest:

if not getattr(self.config, "tie_word_embeddings", False):
    modules = dict(self.named_modules(remove_duplicate=False))
    tied_mapping = {
        target: source
        for target, source in tied_mapping.items()
        if isinstance(modules.get(target.rsplit(".", 1)[0]), nn.Embedding)
        and isinstance(modules.get(source.rsplit(".", 1)[0]), nn.Embedding)
    }

Output heads are nn.Linear, input embeddings are nn.Embedding, so this isolates the structural entries
without relying on naming or get_output_embeddings(). Of the 362 classes with a dict
_tied_weights_keys, 46 return None from get_output_embeddings() — including models whose head is not
named lm_head (BlipForConditionalGeneration, LxmertForPreTraining, FlavaForPreTraining, ...), which
a name-based rule would have wrongly kept tied.

Test updates

mbart / fsmt test_ensure_weights_are_shared count distinct data_ptr()s under
tie_word_embeddings=False. Restored to their pre-#41580 values:

commit	mbart	fsmt
#26422 (2023) … 6f6095e0cf^	2	2
6f6095e0cf (#41580)	4	3
this PR	2	2

#41580 is the commit that introduced the gate, and it bumped these numbers to match the new behaviour.
Consistent with the checkpoints: facebook/mbart-large-50-many-to-many-mmt stores only
model.shared.weight, facebook/wmt19-en-de only model.encoder.embed_tokens.weight.

Verification

• RUN_SLOW=1 longt5 test_summarization → PASSED (fails on main); lm_head stays untied.
• tests/models/{fsmt,mbart,bart,longt5,t5,blip,lxmert,flava} → green.

[zucchini-nlp]

Not 100% sure since I didn't check all official LongT5 ckpt, prob it is same as other T5 models that tie weights always? We had issue with T6 in teh past and hardcoded the value like this

transformers/src/transformers/models/t5/configuration_t5.py

Line 83 in 7dc6341

self.tie_word_embeddings = True

[ErenAta16]

Ran the filter's predicate against the real _tied_weights_keys on 2ef79f8 and it picks exactly the intended split, but T5 shows a case the FSMT test doesn't cover.

T5ForConditionalGeneration with tie_word_embeddings=False, applying the new predicate by hand:

entry	parent modules	kept
lm_head.weight → shared.weight	(Linear, Embedding)	no
encoder.embed_tokens.weight → shared.weight	(Embedding, Embedding)	yes
decoder.embed_tokens.weight → shared.weight	(Embedding, Embedding)	yes

That's the behaviour the PR argues for: the nn.Linear head stays gated on the flag, the two input-embedding aliases survive. FSMT splits the same way — decoder.output_projection is nn.Linear at modeling_fsmt.py:534 so it drops, and the encoder.embed_tokens → decoder.embed_tokens pair is (Embedding, Embedding) so it stays.

Worth adding a T5 case next to the FSMT one, because T5 exercises something FSMT can't: two targets sharing one source, where the dropped entry (lm_head) and the kept entries all point at shared.weight. FSMT's mapping has the source as decoder.embed_tokens.weight for both, so it never tests a partially-filtered fan-out. A regression that dropped the whole mapping when any entry fails the predicate would pass the FSMT test and fail on T5.

Two things about the predicate itself:

modules.get(...) returns None for a name that isn't a module, and isinstance(None, nn.Embedding) is False, so a regex-pattern key — which _tied_weights_keys explicitly supports, that's what the expansion below exists for — gets silently dropped when tie_word_embeddings=False. d_fine is called out in the comments a few lines down as having "complicated regex patterns". Those patterns won't rsplit into a real module name, so if any model has a regex-form entry that is structurally an embedding alias, this filter removes it rather than expanding it. Might be worth resolving the regex first, or at least a comment saying regex-form entries are intentionally treated as head-tying.

self.named_modules(remove_duplicate=False) is built on every call, and get_expanded_tied_weights_keys is called per submodule from the all_submodels=True branch. Only on the tie_word_embeddings=False path, so it's not the common one — just noting it's inside a loop.

Measured with real model classes (T5ForConditionalGeneration, FSMTForConditionalGeneration) on tiny configs, transformers 5.13.1, applying the predicate to _tied_weights_keys as read from main @ 2ef79f8. I did not run the modified get_expanded_tied_weights_keys end to end — the installed 5.13.1 predates the if not tie_word_embeddings: return {} early return this PR is rewriting, so its output isn't comparable. The parent-module types and the mappings above are from the current source.

[jiqing-feng]

Thanks for the detailed check — the predicted behaviour matches what I measured. Added the T5 test
and a comment on the regex case in 4e3835a.

T5 test. Added test_ensure_weights_are_shared to T5ModelTest: 1 distinct storage with
tie_word_embeddings=True, 2 with False. It fails on the base commit with 4 != 2, so it does
guard the fix.

One correction on the motivation though: FSMT does exercise a partially-filtered fan-out. Its two
entries both point at decoder.embed_tokens.weight, so a regression that dropped the whole mapping
when any entry fails the predicate would give 3 there, not 2, and the FSMT test would catch it. T5 is
still worth adding for a different reason — its source (shared) is a third, standalone module
rather than one of the tied embeddings, and it's the example the PR comment names.

Regex-form keys. You're right that they don't resolve to a module and get dropped. I scanned all
74 model classes with regex-form entries in _tied_weights_keys; exactly one has
tie_word_embeddings=False:

Qwen3OmniMoeTalkerForConditionalGeneration
  _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"}
  type(codec_head)   = nn.Linear
  expanded           = {}      # same as main

So there's no behaviour change today, and dropping it is semantically correct since codec_head is a
head. But it's only accidentally correct — "codec_head" has no dot, so rsplit(".", 1)[0] returns
the key itself rather than a parent module. Added a comment recording that regex/module-form entries
are intentionally treated as head-tying, and that no model currently pairs them with
tie_word_embeddings=False. Resolving the regex first would be dead code right now, so I left it.

named_modules in a loop. Agreed, but it only runs on the tie_word_embeddings=False path
during __init__/from_pretrained, never in the forward path, so I left it as is.

[github-actions]

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

run-slow: fsmt, mbart, t5