Feature/Qwen3vl (single commit) - #161
Conversation
Summary of ChangesHello @coder0143, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates the Qwen3-VL Vision Language Model into the JAX ecosystem, providing a complete and functional implementation. It focuses on enabling efficient and scalable inference through detailed sharding configurations and a robust mechanism for loading pretrained weights. The addition is thoroughly validated with tests ensuring numerical consistency and proper distributed behavior. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive JAX implementation of the Qwen3-VL model, including the model architecture, weight conversion utilities from PyTorch, and an extensive test suite. The code is well-structured and the addition is a valuable contribution to the repository. My review focuses on improving maintainability and fixing a minor documentation issue. I've identified a dependency on an internal JAX API that should be replaced, an unused variable that can be removed for code clarity, and a broken link in the README.
| if jax._src.config.abstract_mesh_context_manager.value is not None: | ||
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2), out_sharding=self.shd_cfg.act_btd) * self.scale | ||
| else: | ||
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2)) * self.scale |
There was a problem hiding this comment.
The code uses an internal JAX API (jax._src.config.abstract_mesh_context_manager.value) to check if it's running within a mesh context. Relying on internal APIs is fragile and can break with future JAX updates. It's better to use the public jax.sharding.get_abstract_mesh() API for this check, similar to the pattern used in the ShardedLinear module.
| if jax._src.config.abstract_mesh_context_manager.value is not None: | |
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2), out_sharding=self.shd_cfg.act_btd) * self.scale | |
| else: | |
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2)) * self.scale | |
| mesh = get_abstract_mesh() | |
| if not mesh.empty and len(mesh.axis_names) > 0: | |
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2), out_sharding=self.shd_cfg.act_btd) * self.scale | |
| else: | |
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2)) * self.scale |
|
|
||
| ### Running this model | ||
|
|
||
| Run Qwen3 in action, implemented in [900 lines of code](bonsai/models/qwen3_vl/modeling.py) in JAX. |
There was a problem hiding this comment.
The relative link to modeling.py is incorrect. Since this README.md is in the same directory as modeling.py, the path should be relative to the current directory.
| Run Qwen3 in action, implemented in [900 lines of code](bonsai/models/qwen3_vl/modeling.py) in JAX. | |
| Run Qwen3 in action, implemented in [900 lines of code](modeling.py) in JAX. |
| key_mapping = _get_key_and_transform_mapping(config.text_config.tie_word_embeddings) | ||
|
|
||
| conversion_errors = [] | ||
| loaded_keys = set() |
There was a problem hiding this comment.
The loaded_keys variable is initialized here and updated on line 369, but its value is never read or used. This unused variable should be removed to improve code clarity. The update on line 369 should also be removed.
|
Shall I make these changes? |
|
Seems like we have a large overlap with qwen3 model, maybe we can either put vl version to qwen3 or import qwen3 code here? |
|
Hello @vfdev-5 , there is an overlap, but I think we should keep them separate. Only 3 things are common for the text decoder parts: Some things are fundamentally different and changing them might break multimodal reasoning:
|
|
I will make the final changes along with testing and pack it all into a single commit |
|
|
||
| # --- Sharding Configuration --- # | ||
| @dataclass(slots=True, frozen=True) | ||
| class VisionShardingCfg: |
There was a problem hiding this comment.
Let's unify the sharding config names, recent update was to adopt this naming:
bonsai/bonsai/models/gemma3/modeling.py
Line 50 in 73a80c2
Also let's implement no sharding option as no partition, same as
bonsai/bonsai/models/gemma3/modeling.py
Lines 50 to 66 in 73a80c2
Let's set ShardMode names instead of hardcoded "fsdp", "tp":
bonsai/bonsai/models/gemma3/modeling.py
Lines 39 to 41 in 73a80c2
Ideally, these names should be even configurable, but this is for another PR.
|
Cool, on it! |
|
Final changes based on all the reviews along with testing done, the code can be merged @jenriver @chapman20j , @vfdev-5 , thanks for reviewing my PR |
| cache: Cache, | ||
| input_ids: Array, | ||
| pixel_values: Array, | ||
| image_grid_thw: Array, |
There was a problem hiding this comment.
@coder0143 What is the structure of image_grid_thw? I saw in run_model it was reported as [[1, 64, 64]]. In the code 3 values only used :
grid_h, grid_w = int(grid_thw[0, 1]), int(grid_thw[0, 2])
grid_t = int(grid_thw[0, 0])
Should not it be static tuple of int instead of being a jax array? And to confirm, we are not missing the batch dim?
There was a problem hiding this comment.
@vfdev-5 yup you are correct, I'm making changes on my fork, will you review the changes, or you can open a PR yourself to support jit for forward_vision.
There was a problem hiding this comment.
I think there is a bug in Qwen3VLVisionModel, the line:
hidden_states = hidden_states + pos_embeds[:seq_len]The error message:
E TypeError: add got incompatible shapes for broadcasting: (784, 64), (196, 64).
if we preprocess a message with 4 chat messages (dicts), we will get the following shapes for pixel values and thw grid:
>>> inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt")
>>> inputs["pixel_values"].shape
torch.Size([1024, 1536])
>>> inputs["image_grid_thw"]
tensor([[ 1, 16, 16],
[ 1, 16, 16],
[ 1, 16, 16],
[ 1, 16, 16]])
>>> inputs["image_grid_thw"].shape
torch.Size([4, 3])
There was a problem hiding this comment.
Here is the test I would like to run in test_sharding_qwen3vl.py
def test_full(self):
"""Test full model can be created with sharding enabled."""
rngs = nnx.Rngs(0)
fsdp = modeling.ShardMode.FSDP.value
model = modeling.Qwen3VLForConditionalGeneration(self.cfg_sharded, rngs=rngs)
config = model.config
# TODO: enable the test once fixed the bug in Qwen3VLVisionModel when batch size > 1
batch_size = 4 # should be evenly divisible to num devices for fsdp axis
num_tokens = 128
key = jax.random.key(0)
patch_size = config.vision_config.patch_size
img_size = patch_size * patch_size
n_img = jax.random.uniform(
key,
(batch_size * img_size, img_size * 3 * 2),
dtype=jnp.float32,
minval=-1,
maxval=1,
out_sharding=P(fsdp),
)
n_text = jax.device_put(
np.arange(batch_size * num_tokens).reshape(batch_size, -1),
device=P(fsdp),
)
token_type_ids = np.zeros((batch_size, num_tokens), dtype=int)
token_type_ids[:, 12:98] = 1
n_tti = jax.device_put(
token_type_ids,
device=P(fsdp),
)
image_grid_thw = jnp.asarray([[1, patch_size, patch_size]] * batch_size)
cache = modeling.init_cache(config, batch_size, num_tokens, 1, jnp.float32)
out = model(n_text, n_img, image_grid_thw, cache=cache, token_type_ids=n_tti)
assert isinstance(out.sharding, NamedSharding)
assert out.sharding.spec == config.text_config.shd_cfg.act_btd3 participants
Resolves #133
Reference
Adding qwen3vl model in single commit, btw do update the main readme with all the newer models post merging. To @chapman20j @jenriver
Checklist
run_model.pyfor model usage,test_outputs.pyand/ormodel_validation_colab.ipynbfor quality).