fix: unpackaging error in Custom Mixture of Experts model when aux_loss_enabled is set to True.

[Jonathanjordan21]

What does this PR do?

This PR fixes #2038.

Fix unpackaging error due to additional aux_loss returned by concatenated_forward function when aux_loss_enabled=True.

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[kashif]

thanks @Jonathanjordan21

perhaps its more ellegant to do:

reference_chosen_logps, reference_rejected_logps = self.concatenated_forward(self.ref_model, padded_batch)[:2]

what do you think?

[Jonathanjordan21]

@kashif seems good. I actually just followed the earlier code which calculate the policy losses in get_batch_loss_metrics function.

1440        forward_output = self.concatenated_forward(model, batch)
1441        (
1442            policy_chosen_logps,
1443            policy_rejected_logps,
1444            policy_chosen_logits,
1445            policy_rejected_logits,
1446            policy_nll_loss,
1447       ) = forward_output[:5]
1448       if self.aux_loss_enabled:
1449            aux_loss = forward_output[5]

[kashif]

yeah... i should have just done the above but happy if you do it!

[kashif]

you might need to run pre-commit run --all-files in the root of the TRL folder fix any formatting issues

[HuggingFaceDocBuilderDev]

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.