Mamba 2 Multi-GPU errors out on generation with parallel beam search

[molbap]

Observed issue

Found out when running multi-gpu slow tests in https://github.com/huggingface/transformers/pull/33560 .

Line 479 exactly of the mamba2 modeling file

https://github.com/huggingface/transformers/blob/8efc06ee1863bd6e34e8adb7b10901da87c66818/src/transformers/models/mamba2/modeling_mamba2.py#L472-L480

Will raise the following for the test tests/models/mamba2/test_modeling_mamba2.py::Mamba2ModelTest::test_model_parallel_beam_search

E           RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

Mitigation

Move the smallest tensors of that operation to the device of the largest. Will do that soon, but if anyone wants to jump on the issue happy to review as well.

[github-actions[bot]]

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.