Adding jit support to Qwen3VL, readme fixes - #170
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 introduces Just-In-Time (JIT) compilation support for the Qwen3VL model's vision component, enhancing its performance and enabling more efficient batched processing of visual inputs. The core vision embedding logic has been refactored to dynamically handle variable image grid dimensions within a JIT-compatible framework. Additionally, the changes include minor documentation updates to reflect the improved model capabilities and status. 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 successfully adds JIT compilation support to the Qwen3VL model's vision processing path, which is a great performance enhancement. The core change involves making the image grid dimensions a static argument by changing its type from a JAX array to a tuple of integers. This required some thoughtful refactoring to support batched inputs, which has been done well. The tests and documentation have been updated accordingly. I have a couple of minor suggestions to improve code conciseness.
| results = [] | ||
| for grid_t, grid_h, grid_w in grid_thw: | ||
| results.append(self._fast_pos_embed_interpolate_single(grid_t, grid_h, grid_w)) |
There was a problem hiding this comment.
For conciseness and better readability, this for-loop can be replaced with a list comprehension. This is a common Pythonic pattern.
| results = [] | |
| for grid_t, grid_h, grid_w in grid_thw: | |
| results.append(self._fast_pos_embed_interpolate_single(grid_t, grid_h, grid_w)) | |
| results = [self._fast_pos_embed_interpolate_single(grid_t, grid_h, grid_w) for grid_t, grid_h, grid_w in grid_thw] |
| embs = [] | ||
| for grid_t, grid_h, grid_w in grid_thw: | ||
| embs.append(self._rot_pos_emb_single(grid_t, grid_h, grid_w)) |
There was a problem hiding this comment.
Similar to the other comment, this loop can be simplified into a more concise list comprehension, which is generally preferred in Python for creating lists.
| embs = [] | |
| for grid_t, grid_h, grid_w in grid_thw: | |
| embs.append(self._rot_pos_emb_single(grid_t, grid_h, grid_w)) | |
| embs = [self._rot_pos_emb_single(grid_t, grid_h, grid_w) for grid_t, grid_h, grid_w in grid_thw] |
| grid_thw_np = np.array([[1, 16, 16]], dtype=np.int64) | ||
| grid_thw_pt = torch.tensor(grid_thw_np) | ||
| grid_thw_jax = jnp.array(grid_thw_np) | ||
| grid_thw_tuple = ((1, 16, 16),) |
There was a problem hiding this comment.
I think we can rewrite these 3 lines as following:
grid_thw_tuple = ((1, 16, 16),)
grid_thw_pt = torch.tensor(grid_thw_tuple, dtype=torch.long)and remove grid_thw_np as unused
| grid_thw_np = np.array([[grid_t, grid_h, grid_w]], dtype=np.int64) | ||
| grid_thw_pt = torch.tensor(grid_thw_np) | ||
| grid_thw_jax = jnp.array(grid_thw_np) | ||
| grid_thw_tuple = ((grid_t, grid_h, grid_w),) |
|
@coder0143 can you try to run the following sharding test to make sure the batched input issue is fixed. Thanks! @unittest.skipIf(jax.device_count() < 8, "Atleast 8 devices required")
class TestSharding(absltest.TestCase):
"""Test sharding with simulated 8-device mesh."""
@classmethod
def setUpClass(cls):
cls.mesh = jax.make_mesh((4, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit))
jax.set_mesh(cls.mesh)
cls.cfg_unsharded = get_test_config(use_fsdp=False, use_tp=False)
cls.cfg_sharded = get_test_config(use_fsdp=True, use_tp=True)
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
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 =((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_btd |
|
@vfdev-5 The fsdp test added and some small rewrites for unused vars done! |
| self.assertIsInstance(out.sharding, NamedSharding) | ||
| expected_logit_shd = P(cfg.text_config.shd_cfg.act_btd[0], None, None) | ||
| self.assertEqual(out.sharding.spec, expected_logit_shd) | ||
| print(f"Full vision+text forward: shape={out.shape}, sharding={out.sharding.spec} ✓") |
|
|
||
| def test_full_vision_text_forward(self): | ||
| """Test full vision+text forward pass with sharding (FSDP=4, TP=2).""" | ||
| mesh = jax.make_mesh((4, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) |
There was a problem hiding this comment.
we can set the mesh once in the class setUp
|
On it! |
|
@vfdev-5 Done, sharding file has been updated, removed prints, setup mesh of (4,2) across the file. |
vfdev-5
left a comment
vfdev-5
left a comment
There was a problem hiding this comment.
Thanks for the PR @coder0143
I'll run locally the sharding test_full to make sure there is nothing missing.
I left a minor comment
| inputs = self._create_dummy_image_input() | ||
| grid_thw_pt = inputs["image_grid_thw"] | ||
| grid_thw_jax = jnp.array(grid_thw_pt.numpy()) | ||
| grid_thw_tuple = tuple((int(x[0]), int(x[1]), int(x[2])) for x in grid_thw_pt) |
There was a problem hiding this comment.
Can't we do grid_thw_tuple = tuple(grid_thw_pt.long().tolist()) instead?
|
@vfdev-5 Fixed everywhere for |
| pixel_values = jnp.array(inputs_vision["pixel_values"].numpy()) | ||
| image_grid_thw = jnp.array(inputs_vision["image_grid_thw"].numpy()) | ||
| image_grid_thw_raw = inputs_vision["image_grid_thw"].numpy() | ||
| image_grid_thw = tuple((int(row[0]), int(row[1]), int(row[2])) for row in image_grid_thw_raw) |
There was a problem hiding this comment.
And here we can't use the previous trick (tuple(grid_thw_pt.long().tolist()))?
You do not need to define image_grid_thw_raw neither
|
I sincerely apologize for being impatient and rushing the maintainers. From now on I will double check and ping only when absolutely necessary, thankyou. |
|
@coder0143 please squash all commits into 1. I can confirm sharding tests are passing locally. |
|
I have combined all commits into a single commit, I am not touching the main Readme file (earlier I added a few newly added models to it), do check if any more changes have to be made @vfdev-5 , thanks! |
| if "pixel_values" in inputs_vision: | ||
| pixel_values = jnp.array(inputs_vision["pixel_values"].numpy()) | ||
| image_grid_thw = jnp.array(inputs_vision["image_grid_thw"].numpy()) | ||
| image_grid_thw = tuple(tuple(row) for row in inputs_vision["image_grid_thw"].long().tolist()) |
There was a problem hiding this comment.
I see that finally we could not make the original code simple as we still need to iterate over its structure. I wonder if we can accept the type of image_grid_thw as list[list[int]] instead of tuple[tuple[int, int, int]] ? Can this work with jit and the type checker etc?
There was a problem hiding this comment.
No, list[list[int]] won't work directly with static_argnums because lists are not hashable — JAX needs to hash static arguments for its compilation cache, and list raises TypeError: unhashable type: 'list'.
To make it simple we can do something like:
@functools.partial(jax.jit, static_argnums=(4,))
def _forward_vision_jit(model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids):
logits = model(input_ids, cache, pixel_values, image_grid_thw, token_type_ids)
return logits[:, -1, :], cache
def forward_vision(model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids):
"""Accepts list[list[int]] or tuple — converts to hashable tuple for JIT."""
grid_thw = tuple(tuple(row) for row in image_grid_thw)
return _forward_vision_jit(model, cache, input_ids, pixel_values, grid_thw, token_type_ids)Which makes the callers simple:
image_grid_thw = inputs["image_grid_thw"].long().tolist() # [[1, 64, 64]]
logits, cache = forward_vision(model, cache, ids, pixels, image_grid_thw, tti)@vfdev-5 Do tell if you want me to make the changes.
There was a problem hiding this comment.
Finally, I think it's fine to keep it everywhere as tuple of tuples. We have to do the conversion in tests only so it's ok
|
Just notifying, I have started working on Qwen3.5 and will make a PR soon |
|
Just a reminder.. |
|
@vfdev-5 Please do merge if possible |
|
@coder0143 PR needs rebasing as this branch has conflicts that must be resolved |
|
@vfdev-5 Yup, the earlier code was mine and was merged by chapman20j (is working on another team) earlier, the next code has been pushed here by me, I have resolved the conflicts by accepting the current(newer) changes which were verified via the tests. Thankyou for approving the changes @cgarciae ! |
|
@coder0143 seems that pre-commit is failing, see test. |
|
@cgarciae Thanks for reviewing, fixed the cache bug, pre-commit is running now and the rest of the tests are passing as well |
3 participants
Adds jit support to Qwen3VL model and fixes some readme files
Based on @vfdev-5 's suggestions, jit support and batched input support added for Qwen3VL model, the tests can be seen here: https://www.kaggle.com/code/prathamshahmldlds/test-bonsai-qwen3-vl (on T4x2 with mesh (1,2)).
@chapman20j and @jenriver do check and merge.
Fixes #161
Checklist
run_model.py,test_outputs.py, and/ormodel_validation_colab.ipynb).