Unhandled 'num_items_in_batch' in Mistral model

[gheinrich]

System Info

• Transformer version: 4.46.0
• Model: nvidia/Mistral-NeMo-Minitron-8B-Base

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

When calling the forward method on the NeMo Mistral model, the following exception occurs:

[rank2]:   File "/lustre/fsw/portfolios/llmservice/users/gheinrich/anaconda3/envs/vila/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1582, in _call_impl
[rank2]:     result = forward_call(*args, **kwargs)
[rank2]: TypeError: MistralForCausalLM.forward() got an unexpected keyword argument 'num_items_in_batch'

Expected behavior

The forward() method should use num_items_in_batch for the loss calculation.

[arivero]

and careful here, not to break multi gpu... With Llama 3.1 I get usually an error when dividing loss/num_items_in_batch, they tend to go to different gpus.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.