Fix CohereAsr training-loss double-shift (train against labels, not labels[..., 1:]) by eliaghazal · Pull Request #47090 · huggingface/transformers

eliaghazal · 2026-07-05T22:36:46Z

What does this PR do?

CohereAsrForConditionalGeneration.forward right-shifts labels into decoder_input_ids via shift_tokens_right, then computes the loss with self.loss_function (ForCausalLMLoss), which shifts the labels a second time internally. The model therefore trains against labels[..., 1:] instead of labels.

Verified on main (b70d02fc72) with a tiny random-init model:

model loss:            4.990426
aligned CE (correct):  4.985212
double-shift CE (bug): 4.990426

This is the same bug fixed for Moonshine in #46784. CohereAsr inherits from Moonshine but overrides forward, so that fix didn't propagate.

The fix keeps self.loss_function and passes the (already aligned) labels as shift_labels, so ForCausalLMLoss skips its internal shift while preserving the num_items_in_batch handling needed for correct loss reduction under gradient accumulation — same pattern as Florence2 after #46898. Changed in modular_cohere_asr.py and regenerated modeling_cohere_asr.py with the modular converter.

Also adds the same regression test as Moonshine's (test_training_loss_no_double_shift, with and without -100 padding), extended with an assertion for the num_items_in_batch (summed) reduction path. It fails on current main and passes with this fix; the full tests/models/cohere_asr file passes (124 passed, 138 skipped).

Fixes #46894
Closes #46895, closes #46958 (earlier PRs for the same issue)

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, Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case. → CohereASR training-loss double-shift bug (same pattern as Moonshine #46784) #46894
Did you make sure to update the documentation with your changes?
Did you write any new necessary tests?

Who can review?

@eustlb @vasqu

…abels[..., 1:]) CohereAsrForConditionalGeneration right-shifts labels into decoder_input_ids, then computes the loss via ForCausalLMLoss, which shifts again, so the model trains against labels[..., 1:]. Use a plain CrossEntropyLoss instead, matching Whisper/Bart and the Moonshine fix in huggingface#46784 (CohereAsr overrides forward, so that fix did not propagate). Adds the same regression test (with and without -100 padding). Fixes huggingface#46894

eustlb

Nice catch! See the comment (84e2058) I added to make this explicit, but the issue here is that we were mixing two different conventions:

whisper/ legacy encoder-decoder convention: we pass labels, and forward builds decoder_input_ids from them by right-shifting: [the, black, cat] → [bos, the, black, cat]
llama-like / causal LM convention: labels == input_ids. input_ids and labels are [bos, the, black, cat] and ForCausalLMLoss shifts labels into [the, black, cat], ie we do [bos, the, black, cat] → [the, black, cat]

Ideally we would modernize Whisper and similar models to follow the causal LM convention, but that would require too much BC handling, so this fix LGTM 👌

vasqu

So we don't want to use the loss function in any case? Note that this happens on several models, not only cohere asr - would be nice to collectively fix them in one PR?

Also cc @ebezzam because we talked about it

vasqu · 2026-07-07T10:57:38Z

        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.reshape(-1, self.config.vocab_size), labels.reshape(-1))


My only concern with this is that we don't really support proper gradient accumulation with this anymore which is the reason behind the whole loss function tbh --> gradient accumulation sums and divides only at the end (before it divided on each step)

Fixed in 4629fa3 — kept self.loss_function and passed the already-aligned labels as shift_labels, so the internal shift is skipped but the num_items_in_batch reduction still applies (same pattern as Florence2 after #46898). Also extended the regression test with an assertion for the num_items_in_batch path.

vasqu · 2026-07-07T10:59:00Z

Can we add closes #... for the other PRs so we focus on only this one?

eliaghazal · 2026-07-07T16:37:48Z

@vasqu good catch on the gradient accumulation — dropping self.loss_function did lose the num_items_in_batch handling. I've updated the PR (4629fa3) to keep self.loss_function and pass the labels as shift_labels instead: forward already right-shifts them into decoder_input_ids, so they're aligned with the logits and ForCausalLMLoss can skip its internal shift while still summing and dividing by num_items_in_batch. This is the same pattern Florence2 uses since #46898. The regression test now also covers the num_items_in_batch reduction path.

On fixing collectively — from a sweep of models mixing the whisper-style label shift with ForCausalLMLoss:

Florence2 — already fixed with the shift_labels pattern (Fix Florence2 training-loss double-shift (same pattern as Moonshine #… #46898)
Moonshine — fixed in Fix Moonshine training-loss double-shift (train against labels, not labels[..., 1:]) #46784, but with a plain CrossEntropyLoss, so it has the same grad-accum gap you raised here; happy to open a follow-up PR moving it to the shift_labels pattern
PPFormulaNet — builds decoder_input_ids by shifting input_ids while the loss shifts labels internally, so it looks affected too (depends on the intended labels convention); can check it in the same follow-up

I'd keep this PR scoped to CohereAsr and do the Moonshine/PPFormulaNet pass in a follow-up, unless you'd prefer it all in one.

Also added the closes-references for the earlier duplicate PRs (#46895, #46958) to the description.

vasqu · 2026-07-10T13:00:21Z

Welp now we are getting close to #46895 😅 cc @SunMarc @eustlb to coordinate a bit

eliaghazal · 2026-07-10T13:08:48Z

Fair point 😅 — reading through #46895, the review there converged on this same shift_labels approach before my update did, so credit to @sharmax-vikas. No attachment to which PR lands — happy to close this one in favor of #46895 if that's the cleaner path.

Whichever one you pick, two bits from here may be worth carrying over: the regression test also asserts the num_items_in_batch (summed) reduction path, and the change is made in modular_cohere_asr.py with modeling_cohere_asr.py regenerated by the converter.

Either way I'll follow up on the Moonshine/PPFormulaNet sweep discussed above once this settles.

vasqu · 2026-07-10T13:12:58Z

Yep, not on you. Imo it's just about coordination now and imo probably best to sweep all affected models at once

eliaghazal · 2026-07-10T13:15:40Z

Happy to do the sweep in this PR then. I'd extend it to move Moonshine off the plain CrossEntropyLoss (it has the same num_items_in_batch gap) and check/fix PPFormulaNet, each with the same regression test including the grad accumulation reduction path. If you'd rather do it in #46895 with @sharmax-vikas driving that's fine too, just say the word and I'll share what I have.

vasqu · 2026-07-10T13:17:36Z

Let's wait on @SunMarc and @eustlb first, also happy with either 🤗 expecting sometime next week to handle this properly then.

Moonshine's huggingface#46784 fix used a plain CrossEntropyLoss, dropping num_items_in_batch handling; PPFormulaNet overrides Florence2's forward and missed the huggingface#46898 fix, double-shifting labels. Both now pass labels as shift_labels through self.loss_function, keeping gradient accumulation reduction. Regression tests use a global token count that differs from the per-step count so the summed path is distinguishable.

SunMarc · 2026-07-10T15:08:37Z

Then maybe we can merge this #46895 and @eliaghazal, if you are willing, maybe you can propagate this change to other models so that the have the same behavior across all models that have this isssue ?

eliaghazal · 2026-07-17T20:23:39Z

@SunMarc done — this PR now carries the propagation:

Moonshine: Fix Moonshine training-loss double-shift (train against labels, not labels[..., 1:]) #46784 fixed the double-shift but left a plain CrossEntropyLoss, which drops num_items_in_batch — so loss reduction under gradient accumulation is wrong. Moved it to self.loss_function(..., shift_labels=labels) like the rest.
PPFormulaNet: inherits Florence2 but overrides forward, so it missed Fix Florence2 training-loss double-shift (same pattern as Moonshine #… #46898 and still double-shifts (also reported in PPFormulaNet training-loss bug #47317). Same shift_labels fix.
Regression tests for both exercise the grad-accumulation path with a global token count that differs from the per-step count, so the summed path is distinguishable from a plain mean.
Audited the rest of the codebase for models that right-shift labels into decoder_input_ids and call self.loss_function without shift_labels: the only other hit is T5Gemma, which is fine — it uses loss_type="ForMaskedLM", which doesn't shift.

Coordination so we don't trip over each other: if #46895 lands first I'll rebase and drop the CohereASR commits here; and #46903 (@OmkumarSolanki) also covers PPFormulaNet with a config-gated common test — happy to drop my PPFormulaNet half in deference, whichever way you prefer to slice it.

modular_moonshine_streaming.py inherits MoonshineForConditionalGeneration, so the generated modeling file must be regenerated after the Moonshine change (repo-consistency CI failure). Extend its regression test with the same num_items_in_batch grad-accumulation check.

github-actions · 2026-07-17T20:35:46Z

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

run-slow: cohere_asr, moonshine, moonshine_streaming, pp_formulanet

github-actions · 2026-07-17T20:39:49Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29611022127:2
Result: success | Jobs: 5 | Tests: 867 | Failures: 1 | Duration: 4m 58s

eliaghazal and others added 2 commits July 6, 2026 01:36

better comment

84e2058

eustlb reviewed Jul 6, 2026

View reviewed changes

eustlb approved these changes Jul 6, 2026

View reviewed changes

This was referenced Jul 6, 2026

Fix CohereASR training-loss double-shift (same as Moonshine fix #46784) #46895

Merged

Fix CohereASR training loss double shift #46958

Open

vasqu reviewed Jul 7, 2026

View reviewed changes

Use loss_function with shift_labels to keep num_items_in_batch support

4629fa3

vasqu reviewed Jul 10, 2026

View reviewed changes

sharmax-vikas mentioned this pull request Jul 14, 2026

PPFormulaNet training-loss bug #47317

Open

4 tasks

ShamSaleem mentioned this pull request Jul 23, 2026

Fix PPFormulaNet training-loss double-shift in encoder-decoder forward #47495

Closed

5 tasks

Uh oh!

Conversation

eliaghazal commented Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

eustlb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

vasqu Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

eliaghazal Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eliaghazal commented Jul 7, 2026

Uh oh!

vasqu Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

eliaghazal commented Jul 10, 2026

Uh oh!

vasqu commented Jul 10, 2026

Uh oh!

eliaghazal commented Jul 10, 2026

Uh oh!

vasqu commented Jul 10, 2026

Uh oh!

SunMarc commented Jul 10, 2026

Uh oh!

eliaghazal commented Jul 17, 2026

Uh oh!

github-actions Bot commented Jul 17, 2026

Uh oh!

github-actions Bot commented Jul 17, 2026

CI recap

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

eliaghazal commented Jul 5, 2026 •

edited

Loading

eustlb left a comment •

edited

Loading

vasqu commented Jul 7, 2026 •

edited

Loading