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

[molbap]

Observed issue

Found out when running multi-gpu slow tests in #33560 .

Line 479 exactly of the mamba2 modeling file

transformers/src/transformers/models/mamba2/modeling_mamba2.py

Lines 472 to 480 in 8efc06e

	# Discretize x into dB
	# [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
	hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
	dBx = dB * hidden_states[..., None]
	# State calculation
	cache_params.ssm_states[self.layer_idx].copy_(
	cache_params.ssm_states[self.layer_idx] * dA + dBx
	)

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.