Feature/qwen3 vl - #133
Conversation
…sine similarity in tests
…ue to gated repo.
| x = self.linear_fc1(x, out_sharding=P(None, None)) | ||
| x = nnx.gelu(x, approximate=True) | ||
| return self.linear_fc2(x, out_sharding=P(None, None)) |
There was a problem hiding this comment.
This will always remove the sharding on the outputs of the layers. Could you update this to have appropriate shardings?
| ) | ||
| self.num_grid_per_side = int(config.num_position_embeddings**0.5) | ||
| self.blocks = nnx.List([Qwen3VLVisionBlock(config, rngs=rngs) for _ in range(config.depth)]) | ||
| self.merger = Qwen3VLPatchMerger(config, use_postshuffle_norm=False, rngs=rngs) |
There was a problem hiding this comment.
can you make the use_postshuffle_norm part of the config? This would allow the user to change this through the config.
There was a problem hiding this comment.
We could also remove it as an input to the Qwen3VLPatchMerger layer because it would be accessible through the config.
There was a problem hiding this comment.
I have inspected this, the issue is that the merger has use_postshuffle_norm=False but the deepstack_merger_list has it true, so its better to keep it this way
| EOS_TOKEN_ID = 151643 | ||
|
|
||
|
|
||
| def generate(model, cache, input_ids, max_new_tokens: int = 50): |
There was a problem hiding this comment.
Generally Bonsai's convention is to define the actual model functions in the model's __call__ function, and wrap it in jitted forward (example [gemma3]), and run the forward function in run_model example.
Could you update to follow this pattern?
There was a problem hiding this comment.
Also, let's omit the performance measuring parts to make the code more simple. gemma3 is a good example.
| k = k.transpose(0, 2, 1, 3) | ||
| v = v.transpose(0, 2, 1, 3) | ||
|
|
||
| attn_weights = jnp.matmul(q, k.transpose(0, 1, 3, 2)) * self.scale |
There was a problem hiding this comment.
Make sure to specify an outsharding here.
| image_grid_thw: Array, | ||
| token_type_ids: Array, | ||
| ) -> Tuple[Array, Cache]: | ||
| """Forward pass with vision inputs (not JIT - vision has data-dependent shapes).""" |
There was a problem hiding this comment.
I think jit should still be used for this function. We can pre-process images to have the same shapes.
jenriver
left a comment
jenriver
left a comment
There was a problem hiding this comment.
Thanks for this contribution, this is a really great addition!
One thing we could do later is to add jit to forward_vision after setting a default shape and padding, but we can leave it as a TODO for later.
|
Cool, if some small change has to be made, do inform before merging, thanks for reviewing @jenriver . |
|
Hey! We just updated our contribution standards to keep the repo as lean as possible. Since this is already in great shape, could you do one quick final pass to squash your changes into a single commit and prune any extra comments or boilerplate? We’re moving to the convention detailed here: #159. Thanks for the help! |
|
On it! |
|
I ran the command |
|
Hello @jenriver can you test and merge? |
|
Hi @coder0143. If your main branch is up to date, you could try |
|
Cool, will stage changes then commit and open a new PR. |
| from jax._src.mesh import AxisType | ||
| from jax.sharding import PartitionSpec as P | ||
|
|
||
| jax.config.update("jax_platform_name", "cpu") |
There was a problem hiding this comment.
Adding this line and os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8" only changes the runtime before any array allocation and this configuration remains until we exit the python interpreter. If this code is run first, then we will run other tests on cpu as well (which may be unwanted) and if another code was run first on an accelerator then this line wont have any effect:
>>> import jax
>>> print(jax.numpy.array([1]).device)
cuda:0
>>> jax.config.update("jax_platform_name", "cpu")
>>> print(jax.numpy.array([1]).device)
cuda:0I would say a better way to run sharding tests can be like here:
|
Closing since we merged in the other qwen3vl with fewer commits. |
6 participants
Resolves #93
Reference
Ported Qwen3-VL model. Some things to be discussed @chapman20j @jenriver :
(vision features are variable, don't know if we have to manually setup jit compilation (using jax.lax api), also, jit compilation is quite expensive with vision inputs).
Checklist
run_model.pyfor model usage,test_outputs.pyand/ormodel_validation_colab.ipynbfor quality).