deepspeedai / deepspeedai/DeepSpeed

[REQUEST] Ulysses support for VLMs

Open
#7,834 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Feature description
UlyssesSPAttentionHF.register_with_transformers assumes that a text-only model will be used. I am hoping to use it with a VLM (for example, Qwen/Qwen3-VL-8B-Instruct.)

UlyssesSPAttentionHF.register_with_transformers assumes certain keys in the config, for example: attn_head_count=hf_model_config.num_attention_heads,. However for Qwen3VL, the correct key would be hf_model_config.text_config.num_attention_heads. This would be easy enough to change to make it specific to the language model (or the vision encoder, as necessary.)

However the broader problem is that this new attention implementation is architecture dependent (ie. depends on the number of heads) and is used to overwrite the global attention implementation:

# So instead we hack `ALL_ATTENTION_FUNCTIONS` to override all existing keys with our implementation, since it only gets used at the point of calling the attention and that's what we want, all other code branches relying on the original core `attn_implementation` will still be executed. This is what we called "Being John Malkovich"
for key in ALL_ATTENTION_FUNCTIONS.keys():
    ALL_ATTENTION_FUNCTIONS[key] = uattn_wrapper

This enables loading the model with no changes - ie. AutoModel.from_pretrained. However the vision encoder has a different architecture with a different number of heads. When it is loaded, it will have the same sequence parallel attention as the language model, which would be incorrect. Further, it is likely unnecessary to implement sequence parallel attention for the vision encoder because it is much smaller than the language model. But because the global attention implementation is overwritten with the Ulysses attention, changing the attention implementation for the language model will change it for the vision encoder as well.

Describe the solution you'd like
The ability to use UlyssesSPAttentionHF.register_with_transformers to enable sequence parallel for the language model, but not the vision encoder, of a VLM.

Describe alternatives you've considered
For selecting the correct config keys, an additional parameter could be added to UlyssesSPAttentionHF.register_with_transformers that allows for specifying the config key that contains the architectural details for the new attention function.

For the global overwriting of the attention implementations, it's less clear to me. Maybe the vision encoder should be loaded before the global overwriting, and then the language model is loaded after? And then the two are connected back together? However this would mean messing with the model loading.

I am happy to make the necessary changes and submit a PR, but direction on the best approach to solving these issues would be appreciated.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in deepspeed/runtime/sequence_parallel/ulysses_sp.py at UlyssesSPAttentionHF.register_with_transformers and trace how ALL_ATTENTION_FUNCTIONS is overwritten. Compare the language-model and vision-encoder configuration paths in Qwen/Qwen3-VL-8B-Instruct, then determine how registration can target only the language model. Done means the VLM uses Ulysses attention for its language model without applying it to the vision encoder; the issue requests maintainer direction on the exact loading approach.

Written by the indexing model from the issue text.

Assessment

Tech stack
huggingface, python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.