Enable device map

[darshana1406]

What does this PR do?

Fixes #30858

• Added _no_split_modules = ["VideoLlavaVisionAttention"] to src/transformers/models/video_llava/modeling_video_llava.py
• Works on a 4 GPU NVIDIA GeForce RTX 2080 Ti setup.

Who can review?

@zucchini-nlp @amyeroberts

[amyeroberts]

Thanks for adding this feature @darshana1406!

Could you run the following tests in a multi-gpu environment and share the terminal output?

pytest tests/models/video_llava/test_modeling_ video_llava.py -vv -k "offload or parallelism"

[zucchini-nlp]

Awesome, thanks a lot for adding device map support!

[darshana1406]

@amyeroberts Here is the terminal output.

(videocon) darshana.s@gnode084:~/transformers$ python -m pytest tests/models/video_llava/test_modeling_video_llava.py -vv -k "offload or parallelism"
============================================================= test session starts =============================================================
platform linux -- Python 3.10.14, pytest-8.2.0, pluggy-1.5.0 -- /home2/darshana.s/miniconda3/envs/videocon/bin/python
cachedir: .pytest_cache
rootdir: /home2/darshana.s/transformers
configfile: pyproject.toml
plugins: anyio-4.3.0
collected 125 items / 121 deselected / 4 selected                                                                                             

tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_cpu_offload <- tests/test_modeling_common.py PASSED [ 25%]
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_bin <- tests/test_modeling_common.py PASSED [ 50%]
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_safetensors <- tests/test_modeling_common.py PASSED [ 75%]
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_model_parallelism <- tests/test_modeling_common.py PASSED [100%]

============================================================== warnings summary ===============================================================
../miniconda3/envs/videocon/lib/python3.10/site-packages/_pytest/config/__init__.py:1448
  /home2/darshana.s/miniconda3/envs/videocon/lib/python3.10/site-packages/_pytest/config/__init__.py:1448: PytestConfigWarning: Unknown config option: doctest_glob
  
    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_cpu_offload
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_cpu_offload
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_bin
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_bin
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_safetensors
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_model_parallelism
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_model_parallelism
  /home2/darshana.s/miniconda3/envs/videocon/lib/python3.10/site-packages/accelerate/utils/modeling.py:1142: DeprecationWarning: The 'warn' method is deprecated, use 'warning' instead
    logger.warn(

tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_cpu_offload
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_bin
tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_disk_offload_safetensors
  /home2/darshana.s/miniconda3/envs/videocon/lib/python3.10/site-packages/accelerate/utils/modeling.py:1363: UserWarning: Current model requires 262176 bytes of buffer for offloaded layers, which seems does not fit any GPU's remaining memory. If you are experiencing a OOM later, please consider using offload_buffers=True.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=============================================== 4 passed, 121 deselected, 11 warnings in 25.69s ===============================================

[amyeroberts]

Thanks for adding this really useful feature @darshana1406! ❤️

[darshana1406]

Thank you for guiding me! @zucchini-nlp @amyeroberts