Fix CohereASR training-loss double-shift (same as Moonshine fix #46784)#46895
Conversation
…ngface#46784) CohereAsrForConditionalGeneration shifts labels right via shift_tokens_right() to create decoder_input_ids, then calls self.loss_function() which maps to ForCausalLMLoss — that function shifts labels again internally. This causes the model to train against labels[..., 1:] instead of labels. Replace self.loss_function() with plain CrossEntropyLoss, matching the fix applied to Moonshine in huggingface#46784.
| 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)) |
There was a problem hiding this comment.
Not sure if I agree with this whole approach tbh.
Trying to reconstruct from what I understood.
- We fall into the case where labels is not None but the input is None
- This triggers the right shift of tokens as input (with the first token being the bos token)
- Now we apply the loss but this unnecessarily shifts the labels
- This leads to the double shift you mentioned
However, wouldn't it make more sense to detect when we have this special case and pass it as shifted labels directly? Or am I missing something?
| loss = None | ||
| if labels is not None: | ||
| loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size) | ||
| loss = self.loss_function( |
There was a problem hiding this comment.
I meant it more like when we enter the case not in general, so something along
do_shift_labels = False
if labels is not None:
if decoder_input_ids is None and decoder_inputs_embeds is None:
decoder_input_ids = shift_tokens_right(
labels, self.config.pad_token_id, self.config.decoder_start_token_id
)
do_shift_labels = True
...
if labels is not None:
shift_labels = kwargs.pop("shift_labels", None)
if do_shift_labels and shift_labels is None:
shift_labels = labels
loss = self.loss_function(logits=logits, labels=labels, shift_labels=shift_labels, vocab_size=self.config.vocab_size, **kwargs)There was a problem hiding this comment.
Thanks for the suggestion! Updated the fix to use do_shift_labels flag - now shift_labels is only passed when shift_tokens_right() was actually called. Also added a test (test_training_loss_no_double_shift). Please take a look!
…as called (per reviewer)
| loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size) | ||
| shift_labels = kwargs.pop("shift_labels", None) | ||
| if do_shift_labels and shift_labels is None: | ||
| shift_labels = labels |
There was a problem hiding this comment.
I feel like the labels should never be shifted no and only decoder_input_ids should be. So maybe we can even simplify to this + remove do_shift_labels.
| shift_labels = labels | |
| shift_labels = kwargs.pop("shift_labels", labels) |
eustlb
left a comment
eustlb
left a comment
There was a problem hiding this comment.
#47090 is the correct fix IMO, see #47090 (review) for explaination
cc @vasqu @SunMarc if you agree
|
[For maintainers] Suggested jobs to run (before merge) run-slow: cohere_asr |
CI recapDashboard: View test results in Grafana |
|
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. |
* upstream/main: (39 commits) Remove deprecated training args and `is_fast` property (huggingface#46917) Consistent output shape from `get_image_features` (huggingface#46405) Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423) fix failed test cases for qwen3_omni_moe model (huggingface#47449) Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233) Move `value` padding into the attention interfaces that need it (huggingface#47451) Simplify function dispatch for linear attention (huggingface#47450) [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447) Fix double-shifted training loss in GitForCausalLM (huggingface#47395) Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895) Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379) Update bug report list (huggingface#46607) Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952) Optimize flash attention max seqlen computation in vision attention (huggingface#47170) fix: remove unreachable return in special token builder (huggingface#47420) Add Harry to slow CI (huggingface#47454) BLT: vectorize patch length processing (huggingface#47385) Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935) [Kimi] add integration tests (huggingface#47383) Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974) ...
5 participants
What does this PR do?
CohereAsrForConditionalGeneration.forward()has a training-loss double-shift bug — the same pattern fixed in Moonshine via #46784.The problem:
shift_tokens_right()to createdecoder_input_idsself.loss_function()is called, which maps toForCausalLMLossForCausalLMLossshifts labels again internally (labels[..., 1:])This means the model trains against
labels[..., 1:]instead oflabels.The fix: Replace
self.loss_function()with plainCrossEntropyLoss(), matching Whisper, Bart, and the now-fixed Moonshine model.Fixes #46894
Code Agent Policy
Before submitting
@eustlb @ebezzam @vasqu