Skip to content

Remove the redundant shift during the loss computation in the Moshi m…#36928

Open
glynpu wants to merge 1 commit into
huggingface:mainfrom
glynpu:moshi_fix
Open

Remove the redundant shift during the loss computation in the Moshi m…#36928
glynpu wants to merge 1 commit into
huggingface:mainfrom
glynpu:moshi_fix

Conversation

@glynpu

@glynpu glynpu commented Mar 24, 2025

Copy link
Copy Markdown

What does this PR do?

Correct the loss computation process in the Moshi model to apply the shift only once, as it is currently being applied twice.

Because the class name MoshiForCausalLM contains 'ForCausalLM', according to the mapping rules in LOSS_MAPPING, the self.loss_function used in the forward function of MoshiForCausalLM should be ForCausalLMLoss. As a result, logits and labels are shifted twice: once before calling self.loss_function and once inside self.loss_function. This leads to tokens < n - 1 predicting n instead of the expected behavior where tokens < n predict n.

This PR removes the shift before the self.loss_function call.

References:
LOSS_MAPPING:

LOSS_MAPPING = {
"ForCausalLM": ForCausalLMLoss,
"ForMaskedLM": ForMaskedLMLoss,
"ForQuestionAnswering": ForQuestionAnsweringLoss,
"ForSequenceClassification": ForSequenceClassificationLoss,
"ForTokenClassification": ForTokenClassification,

ForCausalLMLoss:

def ForCausalLMLoss(
logits,
labels,
vocab_size: int,
num_items_in_batch: int = None,
ignore_index: int = -100,
shift_labels=None,
**kwargs,
):
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
if shift_labels is None:
labels = labels.to(logits.device)
# Shift so that tokens < n predict n
labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
logits = logits.view(-1, vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(logits.device)
loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)
return loss

@github-actions
github-actions Bot marked this pull request as draft March 24, 2025 12:49
@github-actions

Copy link
Copy Markdown
Contributor

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

@glynpu
glynpu marked this pull request as ready for review March 24, 2025 12:56
@github-actions
github-actions Bot requested a review from eustlb March 24, 2025 12:56
@Rocketknight1

Copy link
Copy Markdown
Member

cc @eustlb!

@eustlb

eustlb commented Jul 7, 2025

Copy link
Copy Markdown
Collaborator

Hey @glynpu, thank you so much for your patience.
I am allocating time this week to a bug fixing sprint on Moshi, I'll review this PR at the same time 🤗

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants