Fix YOLOS device mismatch with device_map="auto" (#46886)
* Fix YOLOS device mismatch with device_map="auto"
* Set _no_split_modules for YOLOS and revert .to() workaround
YOLOS was missing the proper _no_split_modules entries, which caused
infer_auto_device_map to split inside YolosLayer when using
device_map='auto'. This led to device mismatch errors on multi-GPU
setups because residual connections inside YolosLayer were split
across devices.
This change:
- Sets _no_split_modules = ["YolosEmbeddings", "YolosLayer"] to match
the pattern used by ViT, DeiT, DINOv2, BEiT, and other vision models.
- Reverts the .to(hidden_states.device) workaround, which was a
symptom-level fix rather than the structural fix.
Fixes #46869
---------
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
Co-authored-by: Pablo Montalvo <39954772+molbap@users.noreply.github.com>
