Repository navigation
Always tie embeddings for LongT5 and Pop2Piano - #47620
Conversation
`tie_word_embeddings` only controls whether the output embeddings are tied to the input embeddings, but `get_expanded_tied_weights_keys` was returning an empty mapping whenever it was False, dropping the structural entries as well. For encoder-decoder models, `encoder.embed_tokens` and `decoder.embed_tokens` are aliases of `shared` and are absent from the checkpoints, so they ended up randomly initialized (e.g. longt5, which has `tie_word_embeddings=False`). Only drop the entries pointing at `get_output_embeddings()` instead.
`tie_word_embeddings` only controls whether the output embeddings are tied to the input embeddings, but `get_expanded_tied_weights_keys` was returning an empty mapping whenever it was False, dropping the structural entries as well. For encoder-decoder models, `encoder.embed_tokens` and `decoder.embed_tokens` are aliases of `shared` and are absent from the checkpoints, so they ended up randomly initialized (e.g. longt5, which has `tie_word_embeddings=False`). Keep the entries tying an input embedding to another input embedding, and only drop the ones targeting the output head.
bd66db0 to
84193d1
Compare
`tie_word_embeddings=False` only unties the output head; the encoder/decoder embedding aliases are structural and stay tied. These two counts were 2 from huggingface#26292/huggingface#26422 until huggingface#41580 bumped them to 3/4 to match the all-or-nothing gate.
zucchini-nlp
left a comment
zucchini-nlp
left a comment
There was a problem hiding this comment.
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
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
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.
|
Thanks for the detailed check — the predicted behaviour matches what I measured. Added the T5 test T5 test. Added One correction on the motivation though: FSMT does exercise a partially-filtered fan-out. Its two Regex-form keys. You're right that they don't resolve to a module and get dropped. I scanned all So there's no behaviour change today, and dropping it is semantically correct since
|
|
Hey! Sorry there's a bit too much to read.
Is wrong. these are not structural. |
Reverts the framework-level change: `_tied_weights_keys` is an all-or-nothing
mapping gated by `tie_word_embeddings`, so `get_expanded_tied_weights_keys`
should not split it into "structural" and "head" entries. `modeling_utils` and
the fsmt/mbart/t5 test expectations are back to their state on main.
The actual bug is per-model. LongT5 and Pop2Piano are derived from T5 and
inherit its `tie_word_embeddings` quirk: the flag historically encoded "do not
scale the decoder outputs" (T5 1.1), not "do not share the embeddings". Every
official checkpoint sets it to `False` (`google/long-t5-local-base`,
`google/long-t5-tglobal-base`,
`Stancld/longt5-tglobal-large-16384-pubmed-3k_steps`, `sweetcocoa/pop2piano`)
while only storing `shared.weight`, so the aliased embeddings end up randomly
initialized.
Apply the same fix T5/MT5/UMT5 already have: force `tie_word_embeddings=True`
and move the scaling to its own `scale_decoder_outputs` flag.
Before, on `google/long-t5-local-base`:
encoder.embed_tokens.weight | MISSING
decoder.embed_tokens.weight | MISSING
lm_head.weight | MISSING
shared std 10.09 vs embed std 1.00
Pop2Piano is affected too, though only latently: the hub ships a `.bin`, which
preserves shared storage, but a safetensors round-trip stores `shared.weight`
alone and both `embed_tokens` come back randomly initialized.
|
Thanks for the review @ArthurZucker — reverted the Digging further, the bug turned out to be a regression from #42815, not a framework issue. Details below. The suggested fix doesn't apply
Root cause: #42815 missed two models
Two models were missed:
ImpactAll official checkpoints set the flag to Also affects Pop2Piano is latent: the Hub ships a The fixSame two lines t5/mt5/umt5 already have, applied to both configs, plus a Checks
Open questionThis leaves the underlying wart: one bool gating the entire Should I retitle the PR to "Always tie embeddings for LongT5 and Pop2Piano"? |
| # Same quirk as T5, since LongT5 is derived from it: the original T5 always scaled the decoder outputs, | ||
| # but the 1.1v does not. The model code was relying on saved configs where `tie_word_embeddings` is set | ||
| # to `False` in 1.1v and using it as indicator of whether to scale or not. But in fact we tie weights | ||
| # always, as the checkpoints only store `shared.weight`, so we force it to be `True`. |
|
run-slow: longt5, pop2piano |
|
This comment contains models: ["models/longt5", "models/pop2piano"] |
jiqing-feng
changed the title
tie_word_embeddings
|
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. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: longt5, pop2piano |
CI recapDashboard: View test results in Grafana |
| # Same quirk as T5: `tie_word_embeddings=False` indicates the 1.1v (no decoder output scaling), but the | ||
| # checkpoints only store `shared.weight` so weights are always tied. | ||
| self.scale_decoder_outputs = self.tie_word_embeddings | ||
| self.tie_word_embeddings = True | ||
|
|
There was a problem hiding this comment.
yes, this is exactly what I meant!
|
This comment contains models: ["models/longt5", "models/pop2piano"] |
Problem
LongT5andPop2Pianoload with randomly initialized embeddings from their official checkpoints.Loading
google/long-t5-local-base:Generation is garbage:
Stancld/longt5-tglobal-large-16384-pubmed-3k_stepsproduces'muralvoenseenseensevo(0)(0)(0)...'.Cause
tie_word_embeddings=Falsehistorically meant "don't scale the decoder outputs" (T5 1.1), not"don't share the embeddings". The checkpoints store only
shared.weight, so the flag must not beallowed to disable tying. Every official checkpoint sets it to
False:google/long-t5-local-base,google/long-t5-tglobal-base,Stancld/longt5-tglobal-large-16384-pubmed-3k_steps,sweetcocoa/pop2piano.#42815 handled exactly this for
t5,mt5,umt5,udopandvilt— force the flag toTrueandmove the scaling to a dedicated
scale_decoder_outputsattribute. Two models did not get the guard:self.tie_encoder_decoder = True # forcing itfrom the config without adding a replacement.
It became load-breaking once
get_expanded_tied_weights_keysstarted returning{}for the whole_tied_weights_keysmapping when the flag isFalse.Fix
Apply the same two lines the other T5-family configs already have:
and read the new attribute where the scaling happens:
Plus a
test_tie_word_embeddingsfor each model, matching the onemt5already has.Net diff: 6 files, +73/-4.
modeling_utils.pyis untouched.Notes
transformers, which still readstie_word_embeddingsto decide whether to scale — 1.1 checkpoints would silently get the T5 1.0scaling.
.binand pickle preserves shared storage,but a safetensors round-trip loses it.
lm_head. Ongoogle/flan-t5-small, which stores itseparately,
tie_weightssees both keys present with different values, warns, and leaves themuntied.
tie=True, scale=False.Tests
tests/models/{longt5,pop2piano,t5}— 738 passed, 889 skipped, 5668 subtests passed.Remaining wart (out of scope)
One bool still gates the entire
_tied_weights_keysdict, and in the T5 family its meaning isoverloaded. That is why #42815 missed two models. Renaming
tie_word_embeddings→tie_weights,as suggested in review, would be a good follow-up.