Fix: bamba error handling kwargs with forward pass

[Ssukriti]

This PR adds a small change to handle additional kwargs passed to the forward function of BambaModel architecture.

Fixes a behavior when tuning Bamba models - when HF Trainer would pass additional LossArgs to BambaForCausalLM.forward() , which are passed to BambaModel.forward , and need to be ignored.
This would fix the error - num_items_in_batch is an unexpected arg , while tuning the model.

Additional Context:
LlamaForCausalLM class passes all kwargs to self.model.forward , hence they need to be handled.
In the future, we would add support for flashAttentionKwargs in BambaModel.forward() .

cc: @fabianlim

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.

[fabianlim]

@ArthurZucker @molbap this is just a temporary minor fix to make the model work with trainer.Trainer. We have not properly handled the FlashAttentionKwargs yet because we plan to get to that when we support padding free. But that work is involved, and for now we want to make sure the model is able to tune properly

cc: @ani300 @raghukiran1224