fix models that fail in test_model_parallelism

[faaany]

What does this PR do?

This PR fixes the following failed tests:

pytest -rA tests/models/clip/test_modeling_clip.py::CLIPModelTest::test_model_parallelism
pytest -rA tests/models/m2m_100/test_modeling_m2m_100.py::M2M100ModelTest::test_model_parallelism
pytest -rA tests/models/mamba/test_modeling_mamba.py::MambaModelTest::test_model_parallelism
pytest -rA tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_model_parallelism
pytest -rA tests/models/qdqbert/test_modeling_qdqbert.py::QDQBertModelTest::test_model_parallelism
pytest -rA tests/models/udop/test_modeling_udop.py::UdopModelTest::test_model_parallelism

@ydshieh

[faaany]

There are 2 further tests that fail:

pytest -rA tests/models/cohere/test_modeling_cohere.py::CohereModelTest::test_model_parallelism
pytest -rA tests/models/udop/test_modeling_udop.py::UdopEncoderOnlyModelTest::test_model_parallelism

Do you have an idea why they fail? I tried to fix them, but I am not very confident about my fix. So I would like to hear about your thoughts on this first.

[faaany]

Hi @amyeroberts , could you help review this PR? Thank a lot!

[ydshieh]

Hi @faaany . Thanks for the fix.

There is a work #30836 which has some overlaps with this PR? Might be better to wait that merged and see if there is any remaining to fix?

[faaany]

Hi @faaany . Thanks for the fix.

There is a work #30836 which has some overlaps with this PR? Might be better to wait that merged and see if there is any remaining to fix?

Yes, indeed! Sure, let's wait. Thx!

[ydshieh]

Hi @faaany That PR is merged. You can try to rebase (and resolve the conflicts) to see if there is anything remaining to commit 🙏

[faaany]

Hi @ydshieh, I tested the current master branch and all failing tests got fixed in PR #30806! Thx and close this one.