Adapting Whisper to the new loss_function attribute

@ArthurZucker

@ArthurZucker @muellerzr Following up on #35838, #34191, #34198, #34283

I would like to help bring in Whisper into this. I see it was not included in the last #35875 round of fixes related to the loss function bug fix (grad acc.) nor the new global "loss_function" attr. Being an encodec model derived from Bart code in many places around loss and decoder input token handling - I suspect Bart would also benefit from such attention.

So - Would like to help with the following missing support:

It does not accept kwargs

In 'forward' (For Conditional Gen) - Seems straight-forward to follow @muellerzr work and implement considering test passing. Anything special to consider there?

I does not use the global "loss_function" attr (introduced with #34191)

I find that the closest Loss implementation would be ForMaskedLMLoss since seems like the shifted labels are expected from how the existing loss is calc'd

transformers/src/transformers/models/whisper/modeling_whisper.py

Lines 1787 to 1791 in d4a6b40

    
           if labels is not None: 
        
               loss_fct = CrossEntropyLoss() 
        
               # move labels to correct device to enable PP 
        
               labels = labels.to(lm_logits.device) 
        
               loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1))

Some background on the above

I find this was derived from the Bart implementation, which forced the user to either provide decoder_input_ids or derived them from labels by shifting them to the right as part of the denoising pre-training task type - this lead to a situation where labels are expected to be left shifted compared to the logits which is properly served by the above loss calculation.

Whisper, inherited that, but has a more involved input id prefixing scheme. the model is hardly the place to grab the "decoder start token id" which is required to accomplish the "shift right" of labels and get the decoder_input_ids, and anyway - for Whisper this prefix during inference is critical to determining the task and control over that is properly reflected in other args. (language, task, notimestamps etc)

Thus, proper collators suggested by @sanchit-gandhi in his great guidance and the work on Distill-Whisper have explicitly specified both labels and decoder_input_ids that worked around the auto (now unusable) "label shift righting". (See code here)

Or otherwise "cooked" the labels to contain all but the first "decode start token id" as a hack. (More at #27384) and even the Collator code in the popular blog post about Whisper FT does:

       # if bos token is appended in previous tokenization step,
        # cut bos token here as it's append later anyways
        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

Which of course is a workaround to mitigate that Bart heritage.

WDYT @sanchit-gandhi, Did I get this right?

Anyway - this is why the ForCausalLMLoss probably won't be a fit - it will shift the labels left to match against logits positions.

Would like to know if that proper loss then to use is ForMaskedLMLoss or maybe a new ForConditionalGenerationLMLoss actually. Personally, I think a new one should exist, that does exactly what ForMaskedLMLoss with some shared implementation for both.

Also, as an aside I would love to see the Bart derived "decode_input_ids from labels" logic adapted to Whisper - but not sure I have the experience to know how.

Grad acc loss bug still applies to Whisper

As it is implemented now - you can (thankfully) customize the loss calc using "compute_loss_func" which was introduced in #34198 - and this is mandatory for anyone who want to avoid the grad acc loss described here and fixed in many PR's around the above mentioned efforts.
This is actually an open bug for Whisper which did not enjoy the common fixed_cross_entropy injection onto other models.

Thanks guys for all the great documentation on this, so much easier to try and contribute back!

@eustlb

cc @eustlb

Super welcome to anyone that wishes to fix it! 🤗

@ArthurZucker

@ArthurZucker I would like to work on this!

I do need some guidance though - Is my analysis of the "label"->"input decode ids" problem correct?
Should I create a new "loss type" or try and force Whisper to play nice with ForCausalLMLoss ? That would require many developers to adapt their Whisper DataCollators though - and I know we do not wish to break legacy support.

That is the main issue here basically, ForCausalLMLoss expects labels to be shifted, Whisper expects labels after the shift.

Thanks!

@yoadsn

@yoadsn @ArthurZucker Is this issue still open? I would like to work on it

Reviewed the PR sorry for the delay

@tanuj-rai

@tanuj-rai Our training scripts contain workarounds to this problem.
Although I wish it would be solved in the code base to help new-comers as well as anyone trying to base a model on top of the Whisper code.
LMK if I can help.

@yoadsn

@Tanuj-rai Our training scripts contain workarounds to this problem. Although I wish it would be solved in the code base to help new-comers as well as anyone trying to base a model on top of the Whisper code. LMK if I can help.

Sure @yoadsn, I would greatly appreciate your help.

I noticed #46068 was closed. I'd like to attempt a clean fix myself with a proper understanding of the approach. Could you give guidance on what the correct loss type should be and what went wrong with the previous PR?

@manak-sharma20

@manak-sharma20 I have since move on from this - but would love to see it taken care of.
Whisper is still alive and kicking Mid 2026, and many still fine tune it today.

At this point the training script we used internally has a solid workaround (custom loss function) which is aligned with how our collators structure the labels.

This is a multi prong problem - requires some decision making on breaking changes to the Whisper training scripts - which I could not make as an outsider. I no longer hold the deep understanding I had when opening this issue, and seems like the last person trying to fix this was AI-coding it without proper understanding judging by the "ai slop" label on the PR when it was rejected.

Still, my original post holds - representing the problems, and my suggestions on how to fix them:

Whisper uses the "new" (now quite old) loss function unification infra
Whisper not shifting labels just to shift them again - the expectation of who shifts the labels should be clear so that collators work with the loss function smoothly. Collators are in User Land, loss function is in Transformers land - and breaking changes are expected. This is more a project management decisions I feel than actual technical decision.

Good luck to All 👍

@ArthurZucker

Hey @ArthurZucker, came across this issue and wanted to take a shot at it...

I dug into the Whisper code and compared it against Bart, and noticed two things:

The **kwargs in forward() is untyped, Bart and most other models use Unpack[TransformersKwargs], but Whisper is just using bare **kwargs.
When labels are passed, Whisper doesn't set use_cache = False. Bart does this with a warning, and without it you can get incorrect loss if someone is training with gradient accumulation.
Wrote a fix for both and ran the test suite, 361 passed, no failures.

Wanted to check in before opening a PR though, since I saw there's a broader refactor discussion going on. Didn't want to step on anything or go in the wrong direction. Let me know if you'd like me to open the PR..!

@yoadsn

Thanks @yoadsn for the context! I'd like to tackle the loss_function unification fix specifically. @ArthurZucker could you confirm
ForMaskedLMLoss is the right approach before I start writing code?

Navigation Menu

Repository navigation

Adapting Whisper to the new loss_function attribute #36119

It does not accept kwargs

I does not use the global "loss_function" attr (introduced with #34191)

Some background on the above

Grad acc loss bug still applies to Whisper

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	if labels is not None:
	loss_fct = CrossEntropyLoss()
	# move labels to correct device to enable PP
	labels = labels.to(lm_logits.device)
	loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1))

Adapting Whisper to the new loss_function attribute #36119

Description

It does not accept kwargs

I does not use the global "loss_function" attr (introduced with #34191)

Some background on the above

Grad acc loss bug still applies to Whisper

Activity

Rocketknight1 commented on Feb 11, 2025

ArthurZucker commented on Feb 19, 2025

yoadsn commented on Feb 19, 2025

SuccessMoses commented on Apr 21, 2025

ArthurZucker commented on Jun 24, 2025

yoadsn commented on Jun 25, 2025

tanuj-rai commented on Jun 26, 2025

manak-sharma20 commented on Jun 17, 2026

yoadsn commented on Jun 17, 2026

Sahith59 commented on Jun 19, 2026

manak-sharma20 commented on Jun 19, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions