Closed
Description
Feature request
Feature Request
transformers
models can be easily loaded across multiple devices using device_map="auto"
. This will automatically allocate weights across available devices e.g. GPUs and offload any weights onto CPU, then disk as necessary. This is useful when doing inference with large models.
To enable this, _no_split_modules
has to be defined in the model's pretrained model class e.g. like here for LLaMa. This defines layers which should not be split across devices, and should contain as few layers as possible.
Steps to add
- Pick a model to work on and open a PR - comment on this issue to say which model you're working on
- Define
_no_split_modules
in the PreTrainedModel subclass. Try with_no_split_modules = []
first - Enable testing
- Ensure the following tests are not skipped for the model:
test_disk_offload_bin
,test_disk_offload_safetensors
,test_cpu_offload
,test_model_parallelism
,test_model_parallel_beam_search
- Run the tests in a multi-gpu environment
pytest tests/models/{MODEL_NAME}/test_modeling_{MODEL_NAME}.py -vv -k "offload or parallelism"
- Ensure the following tests are not skipped for the model:
Models
- Align
- Altclip
- Beit Enable multi-device for more models #30379
- Bit
- Blip
- Chinese_clip
- Convnext Enable multi-device for some models #30207
- Convnextv2 Enable multi-device for some models #30207
- Cvt Enable multi-device for some models #30207
- Data2vec
- Depth_anything
- Dinat
- Dinov2
- Donut
- Dpt
- Efficientformer
- Efficientnet Enable multi-device for efficientnet #29989
- Flava
- Focalnet Enable multi-device for some models #30207
- Git
- Glpn Enable multi-device for some models #30207
- Groupvit
- Imagegpt Enable multi-device for some models #30207
- Layoutlmv3
- Levit Enable multi-device for some models #30207
- Mask2former
- Maskformer
- Maskformer
- Mgp_str Enable multi-device for some models #30207
- Mobilenet_v1 Enable multi-device for some models #30207
- Mobilenet_v2 Enable multi-device for some models #30207
- Mobilevit Enable multi-device for some models #30207
- Mobilevitv2
- Nat
- Oneformer
- Perceiver
- Poolformer Enable multi-device for some models #30207
- Pvt
- Regnet Enable multi-device for some models #30207
- Resnet Enable multi-device for some models #30207
- Sam Enable multi-device for some models #30207
- Segformer
- Swiftformer Enable multi-device for some models #30207
- Swin Enable multi-device for some models #30207
- Swin2sr
- Swinv2 Enable multi-device for some models #30207
- Timesformer
- Timm_backbone
- Trocr Enable multi-device for some models #30207
- Tvlt
- Tvp
- Upernet Enable multi-device for some models #30207
- Videomae
- Vit_mae
- Vit_msn
- Vitmatte Enable multi-device for more models #30379
- Vivit Enable multi-device for more models #30379
- X_clip
- Yolos Enable multi-device for some models #30207
Motivation
Enable a powerful HF feature for all of our vision models
Your contribution
Ping me for review 🤗