DeepSpeed

4e668fce - [AutoSP] (Sequence Parallelism) support for Multimodal Models (ViT + LLM) (#7984)

[AutoSP] (Sequence Parallelism) support for Multimodal Models (ViT + LLM) (#7984) ## Description Hello DeepSpeed Team! 👋 This PR directly addresses the **"Multimodal model support"** goal outlined in the **DeepSpeed Roadmap Q2 2026 (#7861)**. It introduces **AutoSP (Sequence Parallelism) support for Multimodal Models (ViT + LLM)** out of the box. As noted in the roadmap, multimodal models handle significantly longer sequence lengths, making SP critical. This PR automates the injection of DeepSpeed Ulysses-based sequence parallelism into multimodal architectures, removing the need for manual and error-prone engineering efforts. This is a consolidated PR of several incremental features developed and thoroughly tested in my fork. ### 🎯 Related Issue - Addresses the Multimodal model support item in **#7861 (DeepSpeed Roadmap Q2 2026)**. - Builds upon the AutoSP foundation introduced in #7860. ### 🌟 Key Features & Contributions 1. **AutoSP Scaffolding & Detector (`auto_wrap_model_for_sp`)**: - Introduced a scanning utility to automatically detect ViT encoders and LLM decoders within a multimodal model. - Automatically wraps LLM decoder attention layers with DeepSpeed's existing `DistributedAttention`. 2. **ViT Sequence Parallelism (`UlyssesSPViTAttention`)**: - Implemented a Ulysses-style `Gather-Compute-Scatter` sequence parallel wrapper tailored for non-causal ViT attention layers. - Significantly reduces the memory footprint of ViT Feed-Forward Networks (FFN) and LayerNorms across the sequence dimension. 3. **Cross-Modal Fusion Adapters (Phase 2)**: - Handled the complex sequence scatter/gather at the vision-language boundary to ensure the LLM decoder receives uniformly sharded fused sequences. - Supported architectures include: - **LLaVA** (`LlavaFusionAdapter`): Visual token splice replacing image placeholders. - **InternVL** (`InternVLFusionAdapter`): `IMG_CONTEXT` token splice. - **Qwen2-VL** (`Qwen2VLFusionAdapter`): Vision_start/end bounded splice. ### 🧪 Testing & Validation To ensure this PR does not break any existing functionality and is numerically sound, comprehensive tests have been added: - **Numerical Equivalence Tests**: Added multi-GPU tests (`tests/unit/sequence_parallelism/test_autosp_equivalence.py`) verifying that the SP-wrapped path across N ranks produces the **exact same numerical results** as the equivalent single-device (non-SP) computation. - **Integration Tests**: End-to-end mock integration tests validating the full pipeline from ViT to fusion adapter. - **Benchmarks Provided**: Included a multimodal SP benchmark script (`benchmarks/autosp/bench_multimodal_sp.py`) to easily verify throughput scaling and peak GPU memory reduction. *(All tests pass cleanly on 2 GPUs with `NCCL_P2P_DISABLE=1`)* ### 🚧 Known Limitations & Future Work To be fully transparent, there are a few limitations in the current design that I plan to improve in follow-up iterations (or would love guidance on from the team): 1. **Manual Wrapping for Fusion Layers**: While ViT and LLM attentions are wrapped automatically, the vision projection layer currently requires manual wrapping with `ModalityFusionSPAdapter` due to varying HF model implementations. Fully automating Phase 2 is a logical next step. 2. **ViT SP Trade-off**: The current `UlyssesSPViTAttention` uses a Gather-Compute-Scatter approach. While it successfully reduces FFN memory by $1/N$, it still computes the full attention matrix on every rank. A true All-to-All sequence-to-head transposition for Opaque ViT layers is something I am actively exploring. 3. **Padding Attention Mask**: When `fused_len % world_size != 0`, zero-padding is applied. Currently, the global `attention_mask` is not automatically intercepted and patched, which might require user attention during inference. --- I would deeply appreciate any feedback or suggestions from the maintainers! I am more than happy to make any required adjustments, refactorings, or add further test cases to get this perfectly aligned with the Q2 roadmap and DeepSpeed's standards. Thank you for your time reviewing this! 🚀 --------- Signed-off-by: nathon-lee <leejianwoo@gmail.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: nathon-lee <248585198+nathon-lee@users.noreply.github.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>