diff --git a/bonsai/models/qwen3_vl/README.md b/bonsai/models/qwen3_vl/README.md new file mode 100644 index 00000000..e4f3dcb5 --- /dev/null +++ b/bonsai/models/qwen3_vl/README.md @@ -0,0 +1,40 @@ +# Qwen3-VL in JAX + +This directory contains a pure JAX implementation of the [Qwen3-VL SOTA Vision Language Model](https://github.com/QwenLM/Qwen3-VL), using the [Flax NNX](https://flax.readthedocs.io/en/stable/index.html) API. + + +## Model Configuration Support Status + +| Model Name | Config Support Status | +| :--- | :--- | +| **Dense Models** | | +| [Qwen3-VL-2B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-2B-Instruct) | **✅ Supported** | +| [Qwen3-VL-2B-Thinking](https://huggingface.co/Qwen/Qwen3-VL-2B-Thinking) | **✅ Supported** | +| [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct) | **✅ Supported** | +| [Qwen3-VL-4B-Thinking](https://huggingface.co/Qwen/Qwen3-VL-4B-Thinking) | **✅ Supported** | +| [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) | **✅ Supported** | +| [Qwen3-VL-8B-Thinking](https://huggingface.co/Qwen/Qwen3-VL-8B-Thinking) | **✅ Supported** | +| [Qwen3-VL-32B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-32B-Instruct) | **✅ Supported** | +| [Qwen3-VL-32B-Thinking](https://huggingface.co/Qwen/Qwen3-VL-32B-Thinking) | **✅ Supported** | +| [NVIDIA-Cosmos-Reason2-2B](https://huggingface.co/nvidia/Cosmos-Reason2-2B) | **✅ Supported** | +| [NVIDIA-Cosmos-Reason2-8B](https://huggingface.co/nvidia/Cosmos-Reason2-8B) | **✅ Supported** | + +| **MoE Models** | | +| [Qwen3-VL-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | **🟡 Not started** | +| [Qwen3-VL-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | **🟡 Not started** | + + +### Running this model + +Run Qwen3 in action, implemented in [900 lines of code](modeling.py) in JAX. + +```sh +python3 -m bonsai.models.qwen3_vl.tests.run_model +``` + + +## How to contribute to this model + +We welcome contributions! You can contribute to this model via the following: +* Add a model config variant from the above `🟡 Not started` to `class ModelConfig` in [modeling.py](modeling.py). Make sure your code is runnable on at least one hardware before creating a PR. +* Got some hardware? Run [run_model.py](tests/run_model.py) the existing configs above on hardwares marked `❔ Needs check`. Mark as `✅ Runs` or `⛔️ Not supported`. diff --git a/bonsai/models/qwen3_vl/__init__.py b/bonsai/models/qwen3_vl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/qwen3_vl/modeling.py b/bonsai/models/qwen3_vl/modeling.py new file mode 100644 index 00000000..994267d8 --- /dev/null +++ b/bonsai/models/qwen3_vl/modeling.py @@ -0,0 +1,1075 @@ +import math +from dataclasses import dataclass +from typing import Optional, Tuple, TypeAlias +from enum import Enum + +import jax +import jax.numpy as jnp +from flax import nnx +from jax import Array, P +from jax.sharding import PartitionSpec, get_abstract_mesh, reshard + +_K_MASK = jnp.finfo(jnp.bfloat16).min + + +class ShardMode(Enum): + FSDP = "fsdp" + TP = "tp" + + +# --- Sharding Configuration --- # +@dataclass(slots=True, frozen=True) +class VisionShardingConfig: + """Sharding configuration for vision encoder components.""" + + attn_qkv_kernel: PartitionSpec + attn_proj_kernel: PartitionSpec + mlp_fc1_kernel: PartitionSpec + mlp_fc2_kernel: PartitionSpec + layer_norm: PartitionSpec + activation: PartitionSpec + patch_embed_kernel: PartitionSpec + pos_embed: PartitionSpec + + @staticmethod + def no_sharding(): + return VisionShardingConfig( + attn_qkv_kernel=P(None, None), + attn_proj_kernel=P(None, None), + mlp_fc1_kernel=P(None, None), + mlp_fc2_kernel=P(None, None), + layer_norm=P(None), + activation=P(None, None), + patch_embed_kernel=P(None, None, None, None, None), + pos_embed=P(None, None), + ) + + @staticmethod + def default(use_fsdp: bool, use_tp: bool): + fsdp = ShardMode.FSDP.value if use_fsdp else None + tp = ShardMode.TP.value if use_tp else None + return VisionShardingConfig( + attn_qkv_kernel=P(fsdp, tp), + attn_proj_kernel=P(tp, fsdp), + mlp_fc1_kernel=P(fsdp, tp), + mlp_fc2_kernel=P(tp, fsdp), + layer_norm=P(tp), + activation=P(fsdp, tp), + patch_embed_kernel=P(None, None, None, None, tp), + pos_embed=P(None, tp), + ) + + +@dataclass(slots=True, frozen=True) +class TextShardingConfig: + """Sharding configuration for text decoder components.""" + + q_weight: PartitionSpec + kv_weight: PartitionSpec + o_weight: PartitionSpec + mlp_gate_up_kernel: PartitionSpec + mlp_down_kernel: PartitionSpec + rms_norm: PartitionSpec + embed_kernel: PartitionSpec + cache: PartitionSpec + act_btd: PartitionSpec + act_btf: PartitionSpec + act_btnh: PartitionSpec + act_bhsd: PartitionSpec # (batch, heads, seq, dim) - post-transpose attention + attn_logit_shd: PartitionSpec # (batch, heads, q_seq, k_seq) - attention logits + + @staticmethod + def no_sharding(): + return TextShardingConfig( + q_weight=P(None, None), + kv_weight=P(None, None), + o_weight=P(None, None), + mlp_gate_up_kernel=P(None, None), + mlp_down_kernel=P(None, None), + rms_norm=P(None), + embed_kernel=P(None, None), + cache=P(None, None, None, None), + act_btd=P(None, None, None), + act_btf=P(None, None, None), + act_btnh=P(None, None, None, None), + act_bhsd=P(None, None, None, None), + attn_logit_shd=P(None, None, None, None), + ) + + @staticmethod + def default(use_fsdp: bool, use_tp: bool): + fsdp = ShardMode.FSDP.value if use_fsdp else None + tp = ShardMode.TP.value if use_tp else None + return TextShardingConfig( + q_weight=P(fsdp, tp), + kv_weight=P(fsdp, tp), + o_weight=P(tp, fsdp), + mlp_gate_up_kernel=P(fsdp, tp), + mlp_down_kernel=P(tp, fsdp), + rms_norm=P(tp), + embed_kernel=P(tp, fsdp), + cache=P(fsdp, None, tp, None), + act_btd=P(fsdp, None, tp), + act_btf=P(fsdp, None, tp), + act_btnh=P(fsdp, None, tp, None), + act_bhsd=P(fsdp, tp, None, None), # transpose of act_btnh + attn_logit_shd=P(fsdp, tp, None, None), # TP on heads for attn logits + ) + + +def shard(x, spec: PartitionSpec): + """Reshard tensor according to partition spec if mesh is available.""" + mesh = get_abstract_mesh() + if not mesh.empty and len(mesh.axis_names) > 0: + return reshard(x, spec) + return x + + +# --- Sharded Layer Components (Gemma3 Pattern) --- # +class ShardedLinear(nnx.Module): + """Linear layer with explicit sharding on matmul output. + + Unlike standard Linear, sharding is applied only at call time via out_sharding, + not during initialization. This allows model creation without an active mesh. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + *, + use_bias: bool = True, + kernel_sharding: PartitionSpec = None, # Stored but not used in init + dtype=None, + rngs: nnx.Rngs, + ): + # Initialize without sharding - sharding happens at call time + kernel_initializer = jax.nn.initializers.lecun_normal() + self.kernel = nnx.Param(kernel_initializer(rngs.params(), (in_dim, out_dim), dtype=dtype)) + self.use_bias = use_bias + if use_bias: + self.bias = nnx.Param(jnp.zeros((out_dim,), dtype=dtype)) + else: + self.bias = None + + def __call__(self, x, *, out_sharding: PartitionSpec): + # Apply sharding on output - handles both mesh and no-mesh contexts + mesh = get_abstract_mesh() + if mesh.empty or len(mesh.axis_names) == 0: + # No mesh - skip sharding + result = jnp.matmul(x, self.kernel[...]) + else: + result = jnp.matmul(x, self.kernel[...], out_sharding=out_sharding) + if self.use_bias and self.bias is not None: + result = result + self.bias[...] + return result + + +class ShardedEmbedding(nnx.Embed): + """Embedding layer with explicit sharding on gather output.""" + + def __call__(self, inputs: Array, *, out_sharding: PartitionSpec) -> Array: + if not jnp.issubdtype(inputs.dtype, jnp.integer): + raise ValueError("Input type must be an integer or unsigned integer.") + (embedding,) = self.promote_dtype((self.embedding[...],), dtype=self.dtype, inexact=False) + if self.num_embeddings == 1: + return jnp.broadcast_to(embedding, (*inputs.shape, self.features)) + # Apply sharding on output - handles both mesh and no-mesh contexts + mesh = get_abstract_mesh() + if mesh.empty or len(mesh.axis_names) == 0: + return embedding[inputs] + return embedding.at[inputs].get(out_sharding=out_sharding) + + def attend(self, query: Array, *, out_sharding: PartitionSpec) -> Array: + """Compute logits by matrix multiply with embedding weights.""" + query, embedding = self.promote_dtype((query, self.embedding[...]), dtype=self.dtype) + mesh = get_abstract_mesh() + if mesh.empty or len(mesh.axis_names) == 0: + return jnp.matmul(query, embedding.T) + return jnp.matmul(query, embedding.T, out_sharding=out_sharding) + + +@dataclass(frozen=True) +class Qwen3VLVisionConfig: + """Vision encoder configuration for Qwen3-VL.""" + + depth: int = 24 + hidden_size: int = 1024 + intermediate_size: int = 4096 + num_heads: int = 16 + in_channels: int = 3 + patch_size: int = 16 + temporal_patch_size: int = 2 + spatial_merge_size: int = 2 + out_hidden_size: int = 2048 + num_position_embeddings: int = 2304 + deepstack_visual_indexes: tuple = (5, 11, 17) + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-6 + rope_theta: float = 10000.0 + shd_cfg: VisionShardingConfig = VisionShardingConfig.no_sharding() + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_heads + + @classmethod + def qwen3vl_2b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = ( + VisionShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else VisionShardingConfig.no_sharding() + ) + return cls( + depth=24, + hidden_size=1024, + intermediate_size=4096, + num_heads=16, + out_hidden_size=2048, + deepstack_visual_indexes=(5, 11, 17), + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_4b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = ( + VisionShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else VisionShardingConfig.no_sharding() + ) + return cls( + depth=24, + hidden_size=1024, + intermediate_size=4096, + num_heads=16, + out_hidden_size=2560, + deepstack_visual_indexes=(5, 11, 17), + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_8b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = ( + VisionShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else VisionShardingConfig.no_sharding() + ) + return cls( + depth=27, + hidden_size=1152, + intermediate_size=4304, + num_heads=16, + out_hidden_size=4096, + deepstack_visual_indexes=(8, 16, 24), + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_32b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = ( + VisionShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else VisionShardingConfig.no_sharding() + ) + return cls( + depth=27, + hidden_size=1152, + intermediate_size=4304, + num_heads=16, + out_hidden_size=5120, + deepstack_visual_indexes=(8, 16, 24), + shd_cfg=shd, + ) + + +@dataclass(frozen=True) +class Qwen3VLTextConfig: + """Text decoder configuration for Qwen3-VL.""" + + vocab_size: int = 151936 + hidden_size: int = 2048 + intermediate_size: int = 6144 + num_hidden_layers: int = 28 + num_attention_heads: int = 16 + num_key_value_heads: int = 8 + head_dim: int = 128 + hidden_act: str = "silu" + rms_norm_eps: float = 1e-6 + rope_theta: float = 5_000_000 + mrope_section: tuple = (24, 20, 20) # T, H, W partitions of head_dim + attention_bias: bool = False + tie_word_embeddings: bool = True + shd_cfg: TextShardingConfig = TextShardingConfig.no_sharding() + + @classmethod + def qwen3vl_2b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = TextShardingConfig.default(use_fsdp, use_tp) if (use_fsdp or use_tp) else TextShardingConfig.no_sharding() + return cls( + hidden_size=2048, + intermediate_size=6144, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=8, + tie_word_embeddings=True, + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_4b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = TextShardingConfig.default(use_fsdp, use_tp) if (use_fsdp or use_tp) else TextShardingConfig.no_sharding() + return cls( + hidden_size=2560, + intermediate_size=9728, + num_hidden_layers=36, + num_attention_heads=32, + num_key_value_heads=8, + tie_word_embeddings=True, + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_8b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = TextShardingConfig.default(use_fsdp, use_tp) if (use_fsdp or use_tp) else TextShardingConfig.no_sharding() + return cls( + hidden_size=4096, + intermediate_size=12288, + num_hidden_layers=36, + num_attention_heads=32, + num_key_value_heads=8, + tie_word_embeddings=False, + shd_cfg=shd, + ) + + @classmethod + def qwen3vl_32b(cls, use_fsdp: bool = False, use_tp: bool = False): + shd = TextShardingConfig.default(use_fsdp, use_tp) if (use_fsdp or use_tp) else TextShardingConfig.no_sharding() + return cls( + hidden_size=5120, + intermediate_size=25600, + num_hidden_layers=64, + num_attention_heads=64, + num_key_value_heads=8, + tie_word_embeddings=False, + shd_cfg=shd, + ) + + +@dataclass(frozen=True) +class Qwen3VLConfig: + """Combined configuration for Qwen3-VL model.""" + + vision_config: Qwen3VLVisionConfig + text_config: Qwen3VLTextConfig + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + + @classmethod + def qwen3vl_2b(cls, use_fsdp: bool = False, use_tp: bool = False): + """Qwen3-VL 2B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_2b(use_fsdp, use_tp), + text_config=Qwen3VLTextConfig.qwen3vl_2b(use_fsdp, use_tp), + ) + + @classmethod + def qwen3vl_4b(cls, use_fsdp: bool = False, use_tp: bool = False): + """Qwen3-VL 4B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_4b(use_fsdp, use_tp), + text_config=Qwen3VLTextConfig.qwen3vl_4b(use_fsdp, use_tp), + ) + + @classmethod + def qwen3vl_8b(cls, use_fsdp: bool = False, use_tp: bool = False): + """Qwen3-VL 8B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_8b(use_fsdp, use_tp), + text_config=Qwen3VLTextConfig.qwen3vl_8b(use_fsdp, use_tp), + ) + + @classmethod + def qwen3vl_32b(cls, use_fsdp: bool = False, use_tp: bool = False): + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_32b(use_fsdp, use_tp), + text_config=Qwen3VLTextConfig.qwen3vl_32b(use_fsdp, use_tp), + ) + + +class RMSNorm(nnx.Module): + """Root Mean Square Layer Normalization.""" + + def __init__(self, dim: int, eps: float = 1e-6, *, rngs: nnx.Rngs): + self.eps = eps + self.weight = nnx.Param(jnp.ones((dim,), dtype=jnp.float32)) + + def __call__(self, x: Array) -> Array: + x_f32 = x.astype(jnp.float32) + rms = jax.lax.rsqrt(jnp.mean(x_f32**2, axis=-1, keepdims=True) + self.eps) + out = (x_f32 * rms) * self.weight[...] + return out.astype(x.dtype) + + +class Qwen3VLPatchEmbed(nnx.Module): + """3D Convolutional patch embedding for vision input using nnx.Conv.""" + + def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs): + self.config = config + kernel = (config.temporal_patch_size, config.patch_size, config.patch_size) + self.proj = nnx.Conv( + in_features=config.in_channels, + out_features=config.hidden_size, + kernel_size=kernel, + strides=kernel, + padding="VALID", # No padding - matches PyTorch Conv3d default + use_bias=True, + rngs=rngs, + ) + + def __call__(self, hidden_states: Array) -> Array: + # Input: (num_patches, in_channels * temporal_patch_size * patch_size * patch_size) + cfg = self.config + seq_len = hidden_states.shape[0] + + hidden_states = hidden_states.reshape( + seq_len, cfg.in_channels, cfg.temporal_patch_size, cfg.patch_size, cfg.patch_size + ) + # (seq, C, D, H, W) -> (seq, D, H, W, C) + hidden_states = hidden_states.transpose(0, 2, 3, 4, 1) + + # Apply conv: input (seq, D, H, W, C) -> output (seq , hidden_size) + return self.proj(hidden_states).reshape(seq_len, cfg.hidden_size) + + +class Qwen3VLVisionMLP(nnx.Module): + """Vision encoder MLP with GELU activation.""" + + def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs): + self.shd_cfg = config.shd_cfg + self.linear_fc1 = ShardedLinear( + config.hidden_size, + config.intermediate_size, + use_bias=True, + kernel_sharding=self.shd_cfg.mlp_fc1_kernel, + rngs=rngs, + ) + self.linear_fc2 = ShardedLinear( + config.intermediate_size, + config.hidden_size, + use_bias=True, + kernel_sharding=self.shd_cfg.mlp_fc2_kernel, + rngs=rngs, + ) + + def __call__(self, x: Array) -> Array: + # Vision operates on (seq, hidden) without batch - no sharding benefit + x = self.linear_fc1(x, out_sharding=self.shd_cfg.mlp_fc1_kernel) + x = nnx.gelu(x, approximate=True) + return self.linear_fc2(x, out_sharding=self.shd_cfg.mlp_fc2_kernel) + + +class Qwen3VLVisionAttention(nnx.Module): + """Vision encoder multi-head attention with RoPE.""" + + def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs): + self.shd_cfg = config.shd_cfg + self.num_heads = config.num_heads + self.head_dim = config.head_dim + hidden_size = config.hidden_size + self.qkv = ShardedLinear( + hidden_size, 3 * hidden_size, use_bias=True, kernel_sharding=self.shd_cfg.attn_qkv_kernel, rngs=rngs + ) + self.proj = ShardedLinear( + hidden_size, hidden_size, use_bias=True, kernel_sharding=self.shd_cfg.attn_proj_kernel, rngs=rngs + ) + self.scale = self.head_dim**-0.5 + + def apply_rope(self, cos: Array, sin: Array, q: Array, k: Array): + """Apply RoPE using shared function.""" + # cos/sin: (seq, head_dim) -> (seq, 1, head_dim) for broadcasting with (seq, heads, head_dim) + cos = cos[:, None, :] + sin = sin[:, None, :] + q = apply_rotary_pos_emb(q, cos, sin) + k = apply_rotary_pos_emb(k, cos, sin) + return q, k + + def __call__(self, hidden_states: Array, position_embeddings: Tuple[Array, Array]) -> Array: + seq_len = hidden_states.shape[0] + cos, sin = position_embeddings # (seq_len, head_dim) + # Use no sharding for QKV since we reshape immediately after + qkv_out = self.qkv(hidden_states, out_sharding=P(None, None)) + qkv = qkv_out.reshape(seq_len, 3, self.num_heads, self.head_dim) + q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] # (seq, heads, head_dim) + + q, k = self.apply_rope(cos, sin, q, k) + + q, k, v = q.transpose(1, 0, 2), k.transpose(1, 0, 2), v.transpose(1, 0, 2) + attn_weights = jnp.matmul(q, k.transpose(0, 2, 1)) * self.scale + attn_weights = jax.nn.softmax(attn_weights.astype(jnp.float32), axis=-1).astype(q.dtype) + out = jnp.matmul(attn_weights, v).transpose(1, 0, 2).reshape(seq_len, -1) + return self.proj(out, out_sharding=P(None, None)) + + +class Qwen3VLVisionBlock(nnx.Module): + """Single transformer block for vision encoder.""" + + def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs): + self.norm1 = nnx.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps, rngs=rngs) + self.norm2 = nnx.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps, rngs=rngs) + self.attn = Qwen3VLVisionAttention(config, rngs=rngs) + self.mlp = Qwen3VLVisionMLP(config, rngs=rngs) + + def __call__(self, hidden_states: Array, position_embeddings: Tuple[Array, Array]) -> Array: + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class Qwen3VLPatchMerger(nnx.Module): + """Merge spatial patches after vision encoding.""" + + def __init__(self, config: Qwen3VLVisionConfig, use_postshuffle_norm: bool = False, *, rngs: nnx.Rngs): + self.config = config + self.shd_cfg = config.shd_cfg + merge_factor = config.spatial_merge_size**2 + hidden_merged = config.hidden_size * merge_factor + norm_dim = hidden_merged if use_postshuffle_norm else config.hidden_size + self.norm = nnx.LayerNorm(norm_dim, epsilon=config.layer_norm_eps, rngs=rngs) + self.use_postshuffle_norm = use_postshuffle_norm + self.linear_fc1 = ShardedLinear( + hidden_merged, hidden_merged, use_bias=True, kernel_sharding=self.shd_cfg.mlp_fc1_kernel, rngs=rngs + ) + self.linear_fc2 = ShardedLinear( + hidden_merged, config.out_hidden_size, use_bias=True, kernel_sharding=self.shd_cfg.mlp_fc2_kernel, rngs=rngs + ) + + def __call__(self, x: Array) -> Array: + if not self.use_postshuffle_norm: + x = self.norm(x) + merge_factor = self.config.spatial_merge_size**2 + n_patches = x.shape[0] // merge_factor + # Remove sharding before reshape (vision ops don't benefit from sharding on seq dim) + x = shard(x, P(None, None)) + x = x.reshape(n_patches, -1) + if self.use_postshuffle_norm: + x = self.norm(x) + x = self.linear_fc1(x, out_sharding=P(None, None)) + x = nnx.gelu(x) + return self.linear_fc2(x, out_sharding=P(None, None)) + + +class Qwen3VLVisionModel(nnx.Module): + """Complete vision encoder with deepstack feature extraction.""" + + def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs): + self.config = config + self.shd_cfg = config.shd_cfg + self.patch_embed = Qwen3VLPatchEmbed(config, rngs=rngs) + self.pos_embed = nnx.Embed( + num_embeddings=config.num_position_embeddings, features=config.hidden_size, rngs=rngs + ) + 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) + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nnx.List( + [ + Qwen3VLPatchMerger(config, use_postshuffle_norm=True, rngs=rngs) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + def _fast_pos_embed_interpolate(self, grid_thw: Array) -> Array: + """Bilinear interpolation for position embeddings, matching PyTorch.""" + grid_h, grid_w = int(grid_thw[0, 1]), int(grid_thw[0, 2]) + + # Create interpolation indices + h_idxs = jnp.linspace(0, self.num_grid_per_side - 1, grid_h) + w_idxs = jnp.linspace(0, self.num_grid_per_side - 1, grid_w) + + h_floor = jnp.floor(h_idxs).astype(jnp.int32) + w_floor = jnp.floor(w_idxs).astype(jnp.int32) + h_ceil = jnp.clip(h_floor + 1, 0, self.num_grid_per_side - 1) + w_ceil = jnp.clip(w_floor + 1, 0, self.num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # 2D grid indices for 4 corners + base_h = h_floor * self.num_grid_per_side + base_h_ceil = h_ceil * self.num_grid_per_side + + idx00 = (base_h[:, None] + w_floor[None, :]).flatten() + idx01 = (base_h[:, None] + w_ceil[None, :]).flatten() + idx10 = (base_h_ceil[:, None] + w_floor[None, :]).flatten() + idx11 = (base_h_ceil[:, None] + w_ceil[None, :]).flatten() + + # Weights for bilinear interpolation + w00 = ((1 - dh)[:, None] * (1 - dw)[None, :]).flatten() + w01 = ((1 - dh)[:, None] * dw[None, :]).flatten() + w10 = (dh[:, None] * (1 - dw)[None, :]).flatten() + w11 = (dh[:, None] * dw[None, :]).flatten() + + # Lookup and interpolate + pos_embeds = ( + self.pos_embed(idx00) * w00[:, None] + + self.pos_embed(idx01) * w01[:, None] + + self.pos_embed(idx10) * w10[:, None] + + self.pos_embed(idx11) * w11[:, None] + ) + + # Apply spatial merge permutation + merge_size = self.config.spatial_merge_size + grid_t = int(grid_thw[0, 0]) + + # Reshape: (H*W, D) -> (H, W, D) + pos_embeds = pos_embeds.reshape(grid_h, grid_w, -1) + + # Repeat for temporal dimension + if grid_t > 1: + pos_embeds = jnp.tile(pos_embeds[None], (grid_t, 1, 1, 1)) # (T, H, W, D) + else: + pos_embeds = pos_embeds[None] # (1, H, W, D) + + # Permute for spatial merge: (T, H, W, D) -> (T, H//m, m, W//m, m, D) -> (T, H//m, W//m, m, m, D) + merged_h, merged_w = grid_h // merge_size, grid_w // merge_size + pos_embeds = pos_embeds.reshape(grid_t, merged_h, merge_size, merged_w, merge_size, -1) + pos_embeds = pos_embeds.transpose(0, 1, 3, 2, 4, 5) # (T, merged_h, merged_w, m, m, D) + pos_embeds = pos_embeds.reshape(-1, pos_embeds.shape[-1]) # Flatten to (seq, D) + + return pos_embeds + + def _rot_pos_emb(self, grid_thw: Array) -> Tuple[Array, Array]: + """Compute rotary position embeddings matching PyTorch rot_pos_emb.""" + merge_size = self.config.spatial_merge_size + grid_h, grid_w = int(grid_thw[0, 1]), int(grid_thw[0, 2]) + grid_t = int(grid_thw[0, 0]) + + # Compute merged dimensions + merged_h, merged_w = grid_h // merge_size, grid_w // merge_size + + # Compute position indices for each patch + block_rows = jnp.arange(merged_h) + block_cols = jnp.arange(merged_w) + intra_row = jnp.arange(merge_size) + intra_col = jnp.arange(merge_size) + + # Full resolution positions + row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] + col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] + + # Expand and reshape to match spatial merge order: (merged_h, merged_w, m, m) -> (merged_h, merged_w, m, m) + row_idx = jnp.broadcast_to(row_idx, (merged_h, merged_w, merge_size, merge_size)).reshape(-1) + col_idx = jnp.broadcast_to(col_idx, (merged_h, merged_w, merge_size, merge_size)).reshape(-1) + + # Repeat for temporal dimension + if grid_t > 1: + row_idx = jnp.tile(row_idx, grid_t) + col_idx = jnp.tile(col_idx, grid_t) + + # Create frequency table - PyTorch uses rotary_dim = head_dim // 2 + # And inv_freq has length rotary_dim // 2 = head_dim // 4 + max_hw = max(grid_h, grid_w) + head_dim = self.config.head_dim + rotary_dim = head_dim // 2 # = 32 for head_dim=64 + inv_freq_dim = rotary_dim // 2 # = 16 + inv_freq = 1.0 / (self.config.rope_theta ** (jnp.arange(0, rotary_dim, 2, dtype=jnp.float32) / rotary_dim)) + seq_positions = jnp.arange(max_hw, dtype=jnp.float32) + freq_table = jnp.outer(seq_positions, inv_freq) # (max_hw, rotary_dim//2) = (max_hw, 16) + + # Lookup embeddings for row and col: (seq, rotary_dim//2) each + row_emb = freq_table[row_idx] # (seq, 16) + col_emb = freq_table[col_idx] # (seq, 16) + + # Concatenate row and col: (seq, 32) + emb = jnp.concatenate([row_emb, col_emb], axis=-1) + + # Double the embedding (matching PyTorch cat): (seq, 64) + emb = jnp.concatenate([emb, emb], axis=-1) + + # Apply cos/sin + cos = jnp.cos(emb) + sin = jnp.sin(emb) + return cos, sin + + def __call__(self, hidden_states: Array, grid_thw: Array) -> Tuple[Array, list[Array]]: + hidden_states = self.patch_embed(hidden_states) + seq_len = hidden_states.shape[0] + + # Position embeddings with bilinear interpolation + pos_embeds = self._fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds[:seq_len] + + # RoPE embeddings + cos, sin = self._rot_pos_emb(grid_thw) + position_embeddings = (cos[:seq_len], sin[:seq_len]) + + deepstack_features = [] + for layer_idx, block in enumerate(self.blocks): + hidden_states = block(hidden_states, position_embeddings) + if layer_idx in self.deepstack_visual_indexes: + ds_idx = list(self.deepstack_visual_indexes).index(layer_idx) + deepstack_features.append(self.deepstack_merger_list[ds_idx](hidden_states)) + + hidden_states = self.merger(hidden_states) + return hidden_states, deepstack_features + + +class LayerCache(nnx.Module): + """KV-cache for a single decoder layer.""" + + def __init__(self, config: Qwen3VLTextConfig, batch_size: int, cache_size: int, dtype: jnp.dtype = jnp.bfloat16): + cache_shape = (batch_size, cache_size, config.num_key_value_heads, config.head_dim) + shd = config.shd_cfg + self.k_cache = nnx.Cache(shard(jnp.zeros(cache_shape, dtype=dtype), shd.cache)) + self.v_cache = nnx.Cache(shard(jnp.zeros(cache_shape, dtype=dtype), shd.cache)) + self.size = cache_size + self.cur_ind = nnx.Variable(jnp.zeros((), dtype=jnp.int32)) + + +Cache: TypeAlias = list[LayerCache] + + +def init_cache( + config: Qwen3VLConfig, batch_size: int, token_len: int, generate_steps: int, dtype: jnp.dtype = jnp.bfloat16 +) -> Cache: + """Initialize KV-cache for all layers.""" + cache_size = 2 ** math.ceil(math.log2(max(token_len + generate_steps, 1))) + return [ + LayerCache(config.text_config, batch_size, cache_size, dtype) + for _ in range(config.text_config.num_hidden_layers) + ] + + +class Qwen3VLMLP(nnx.Module): + """SiLU-gated MLP for text decoder.""" + + def __init__(self, config: Qwen3VLTextConfig, *, rngs: nnx.Rngs): + self.shd_cfg = config.shd_cfg + self.gate_proj = ShardedLinear( + config.hidden_size, + config.intermediate_size, + use_bias=False, + kernel_sharding=self.shd_cfg.mlp_gate_up_kernel, + rngs=rngs, + ) + self.up_proj = ShardedLinear( + config.hidden_size, + config.intermediate_size, + use_bias=False, + kernel_sharding=self.shd_cfg.mlp_gate_up_kernel, + rngs=rngs, + ) + self.down_proj = ShardedLinear( + config.intermediate_size, + config.hidden_size, + use_bias=False, + kernel_sharding=self.shd_cfg.mlp_down_kernel, + rngs=rngs, + ) + + def __call__(self, x: Array) -> Array: + gate = self.gate_proj(x, out_sharding=self.shd_cfg.act_btf) + up = self.up_proj(x, out_sharding=self.shd_cfg.act_btf) + activations = nnx.silu(gate) * up + return self.down_proj(activations, out_sharding=self.shd_cfg.act_btd) + + +def _generate_rope(positions: Array, head_dim: int, rope_theta: float) -> Tuple[Array, Array]: + """Generate RoPE cos/sin embeddings.""" + fraction = jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim + timescale = rope_theta**fraction + sinusoid_inp = jnp.einsum( + "bt,k->btk", positions.astype(jnp.float32), 1.0 / timescale, precision=jax.lax.Precision.HIGHEST + ) + return jnp.sin(sinusoid_inp), jnp.cos(sinusoid_inp) + + +def rotate_half(x: Array) -> Array: + """Rotate half the hidden dims of the input: (-x2, x1).""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return jnp.concatenate([-x2, x1], axis=-1) + + +def apply_rotary_pos_emb(x: Array, cos: Array, sin: Array) -> Array: + """Apply rotary position embeddings using rotate_half pattern. + + This is the standard RoPE formula: (x * cos) + (rotate_half(x) * sin) + Works for both vision (seq, heads, dim) and text (batch, seq, heads, dim). + """ + return (x * cos) + (rotate_half(x) * sin) + + +def _apply_rope(x: Array, sin: Array, cos: Array) -> Array: + """Apply rotary position embeddings for text model. + + Note: sin/cos have shape (batch, seq, head_dim//2) and need broadcasting. + """ + # Expand sin/cos for heads dimension: (batch, seq, 1, head_dim//2) + sin, cos = sin[:, :, None, :], cos[:, :, None, :] + # Duplicate to full head_dim for rotate_half pattern + cos_full = jnp.concatenate([cos, cos], axis=-1) + sin_full = jnp.concatenate([sin, sin], axis=-1) + return apply_rotary_pos_emb(x, cos_full, sin_full) + + +def repeat_kv(hidden_states: Array, n_rep: int) -> Array: + """Repeat KV heads for GQA using reshape+tile for sharding compatibility.""" + if n_rep == 1: + return hidden_states + b, t, kv_heads, head_dim = hidden_states.shape + # Reshape to add a dimension, tile, then reshape back + # This avoids jnp.repeat which requires out_sharding under mesh context + hidden_states = hidden_states[:, :, :, None, :] # (b, t, kv_heads, 1, head_dim) + hidden_states = jnp.tile(hidden_states, (1, 1, 1, n_rep, 1)) # (b, t, kv_heads, n_rep, head_dim) + return hidden_states.reshape(b, t, kv_heads * n_rep, head_dim) + + +class Qwen3VLAttention(nnx.Module): + """Text decoder attention with GQA and Q/K normalization.""" + + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int, *, rngs: nnx.Rngs): + self.config = config + self.shd_cfg = config.shd_cfg + self.layer_idx = layer_idx + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.n_rep = self.num_heads // self.num_kv_heads + self.scale = self.head_dim**-0.5 + + self.q_proj = ShardedLinear( + config.hidden_size, + self.num_heads * self.head_dim, + use_bias=config.attention_bias, + kernel_sharding=self.shd_cfg.q_weight, + rngs=rngs, + ) + self.k_proj = ShardedLinear( + config.hidden_size, + self.num_kv_heads * self.head_dim, + use_bias=config.attention_bias, + kernel_sharding=self.shd_cfg.kv_weight, + rngs=rngs, + ) + self.v_proj = ShardedLinear( + config.hidden_size, + self.num_kv_heads * self.head_dim, + use_bias=config.attention_bias, + kernel_sharding=self.shd_cfg.kv_weight, + rngs=rngs, + ) + self.o_proj = ShardedLinear( + self.num_heads * self.head_dim, + config.hidden_size, + use_bias=config.attention_bias, + kernel_sharding=self.shd_cfg.o_weight, + rngs=rngs, + ) + self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps, rngs=rngs) + self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps, rngs=rngs) + + def __call__(self, x: Array, cache: LayerCache, sin: Array, cos: Array, mask: Array | None) -> Array: + batch, seq_len, _ = x.shape + shd = self.shd_cfg + + q_out = self.q_proj(x, out_sharding=shd.act_btd) + k_out = self.k_proj(x, out_sharding=shd.act_btd) + v_out = self.v_proj(x, out_sharding=shd.act_btd) + + q = shard(self.q_norm(q_out.reshape(batch, seq_len, self.num_heads, self.head_dim)), shd.act_btnh) + k = shard(self.k_norm(k_out.reshape(batch, seq_len, self.num_kv_heads, self.head_dim)), shd.act_btnh) + v = shard(v_out.reshape(batch, seq_len, self.num_kv_heads, self.head_dim), shd.act_btnh) + + q = _apply_rope(q, sin, cos) + k = _apply_rope(k, sin, cos) + + # Update cache - cast to cache dtype + cache_dtype = cache.k_cache[...].dtype + k_cached = k.astype(cache_dtype) + v_cached = v.astype(cache_dtype) + slice_indices = (0, cache.cur_ind[...], 0, 0) + cache.k_cache[...] = jax.lax.dynamic_update_slice(cache.k_cache[...], k_cached, slice_indices) + cache.v_cache[...] = jax.lax.dynamic_update_slice(cache.v_cache[...], v_cached, slice_indices) + + k = shard(repeat_kv(cache.k_cache[...], self.n_rep), shd.act_btnh) + v = shard(repeat_kv(cache.v_cache[...], self.n_rep), shd.act_btnh) + + # Transpose to (B, heads, T, dim) and re-shard for attention + q = shard(q.transpose(0, 2, 1, 3), shd.act_bhsd) + k = shard(k.transpose(0, 2, 1, 3), shd.act_bhsd) + v = shard(v.transpose(0, 2, 1, 3), shd.act_bhsd) + + # Attention: TP shards heads (axis 1), safe for any seq_len + attn_weights = shard( + jnp.matmul(q, k.transpose(0, 1, 3, 2)) * self.scale, + shd.attn_logit_shd, + ) + + if mask is not None: + attn_weights = jnp.where(mask, attn_weights, _K_MASK) + attn_weights = jax.nn.softmax(attn_weights.astype(jnp.float32), axis=-1).astype(q.dtype) + attn_out = shard(jnp.matmul(attn_weights, v), shd.act_bhsd) + attn_out = attn_out.transpose(0, 2, 1, 3).reshape(batch, seq_len, -1) + + cache.cur_ind[...] = cache.cur_ind[...] + seq_len + return self.o_proj(attn_out, out_sharding=shd.act_btd) + + +class Qwen3VLDecoderLayer(nnx.Module): + """Single decoder layer for text model.""" + + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int, *, rngs: nnx.Rngs): + self.self_attn = Qwen3VLAttention(config, layer_idx, rngs=rngs) + self.mlp = Qwen3VLMLP(config, rngs=rngs) + self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps, rngs=rngs) + self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps, rngs=rngs) + + def __call__(self, x: Array, cache: LayerCache, sin: Array, cos: Array, mask: Array | None) -> Array: + x = x + self.self_attn(self.input_layernorm(x), cache, sin, cos, mask) + x = x + self.mlp(self.post_attention_layernorm(x)) + return x + + +class Qwen3VLTextModel(nnx.Module): + """Text decoder model.""" + + def __init__(self, config: Qwen3VLTextConfig, *, rngs: nnx.Rngs): + self.config = config + self.shd_cfg = config.shd_cfg + self.embed_tokens = ShardedEmbedding(num_embeddings=config.vocab_size, features=config.hidden_size, rngs=rngs) + self.layers = nnx.List([Qwen3VLDecoderLayer(config, i, rngs=rngs) for i in range(config.num_hidden_layers)]) + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, rngs=rngs) + + def __call__(self, inputs_embeds: Array, cache: Cache, sin: Array, cos: Array, mask: Array | None) -> Array: + hidden_states = inputs_embeds # Already sharded when passed in + for i, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, cache[i], sin, cos, mask) + return self.norm(hidden_states) + + +def merge_modalities(img_emb: Array, text_emb: Array, token_mask: Array) -> Array: + """Merge image embeddings into text sequence at masked positions.""" + img_indices = jnp.cumsum(token_mask) - 1 + safe_indices = jnp.clip(img_indices, 0, img_emb.shape[0] - 1) + aligned_images = img_emb[safe_indices] + return jnp.where(token_mask[:, None], aligned_images, text_emb) + + +def batched_merge_modalities(img_emb: Array, text_emb: Array, token_mask: Array) -> Array: + """Batched version of merge_modalities.""" + return jax.vmap(merge_modalities)(img_emb, text_emb, token_mask) + + +def make_causal_mask(cache: LayerCache, seq_len: int) -> Array: + """Create causal attention mask.""" + cache_size = cache.size + cur_pos = cache.cur_ind[...] + seq_arange = jnp.arange(seq_len) + cache_arange = jnp.arange(cache_size) + mask = (seq_arange[:, None] + cur_pos) >= cache_arange[None, :] + return mask[None, None, :, :] + + +class Qwen3VLForConditionalGeneration(nnx.Module): + """Qwen3-VL model with language modeling head.""" + + def __init__(self, config: Qwen3VLConfig, *, rngs: nnx.Rngs): + self.config = config + self.model = Qwen3VLModel(config, rngs=rngs) + if config.text_config.tie_word_embeddings: + self.lm_head = None + else: + self.lm_head = ShardedLinear( + config.text_config.hidden_size, + config.text_config.vocab_size, + use_bias=False, + kernel_sharding=config.text_config.shd_cfg.embed_kernel, + rngs=rngs, + ) + + def __call__( + self, + input_ids: Array, + cache: Cache, + pixel_values: Optional[Array] = None, + image_grid_thw: Optional[Array] = None, + token_type_ids: Optional[Array] = None, + ) -> Array: + """Forward pass with KV-cache.""" + batch, seq_len = input_ids.shape + shd = self.config.text_config.shd_cfg + + # Generate position embeddings + positions = jnp.arange(seq_len)[None, :] + cache[0].cur_ind[...] + positions = jnp.broadcast_to(positions, (batch, seq_len)) + sin, cos = _generate_rope(positions, self.config.text_config.head_dim, self.config.text_config.rope_theta) + + # Create causal mask + mask = make_causal_mask(cache[0], seq_len) + + # Get text embeddings using ShardedEmbedding with explicit out_sharding + inputs_embeds = self.model.language_model.embed_tokens(input_ids, out_sharding=shd.act_btd) + + # Merge vision if provided + if pixel_values is not None and image_grid_thw is not None and token_type_ids is not None: + vision_embeds, _ = self.model.visual(pixel_values, image_grid_thw) + # Broadcast vision_embeds for batch + vision_embeds_batched = jnp.broadcast_to( + vision_embeds[None], (batch, vision_embeds.shape[0], vision_embeds.shape[1]) + ) + # Unshard inputs before vmap (vmap requires uniform sharding on mapped axis) + inputs_embeds_unsharded = shard(inputs_embeds, P(None, None, None)) + token_type_ids_unsharded = shard(token_type_ids, P(None, None)) + merged = batched_merge_modalities(vision_embeds_batched, inputs_embeds_unsharded, token_type_ids_unsharded) + # Reshard output + inputs_embeds = shard(merged, shd.act_btd) + + hidden_states = self.model.language_model(inputs_embeds, cache, sin, cos, mask) + + # Compute logits + logits_sharding = P(shd.act_btd[0], None, None) # (batch, seq, vocab) + if self.lm_head is not None: + logits = self.lm_head(hidden_states, out_sharding=logits_sharding) + else: + # Use attend() method for tied embeddings + logits = self.model.language_model.embed_tokens.attend(hidden_states, out_sharding=logits_sharding) + + return logits + + +class Qwen3VLModel(nnx.Module): + """Qwen3-VL backbone model.""" + + def __init__(self, config: Qwen3VLConfig, *, rngs: nnx.Rngs): + self.config = config + self.visual = Qwen3VLVisionModel(config.vision_config, rngs=rngs) + self.language_model = Qwen3VLTextModel(config.text_config, rngs=rngs) + + +@jax.jit +def forward(model: Qwen3VLForConditionalGeneration, cache: Cache, input_ids: Array) -> Tuple[Array, Cache]: + """JIT-compiled forward pass for text-only generation.""" + logits = model(input_ids, cache) + return logits[:, -1, :], cache + + +# TODO: Add jit compilation support by fixing image size +def forward_vision( + model: Qwen3VLForConditionalGeneration, + cache: Cache, + input_ids: Array, + pixel_values: Array, + image_grid_thw: Array, + token_type_ids: Array, +) -> Tuple[Array, Cache]: + """Forward pass with vision inputs (not JIT - vision has data-dependent shapes).""" + logits = model(input_ids, cache, pixel_values, image_grid_thw, token_type_ids) + return logits[:, -1, :], cache diff --git a/bonsai/models/qwen3_vl/params.py b/bonsai/models/qwen3_vl/params.py new file mode 100644 index 00000000..f0ff1790 --- /dev/null +++ b/bonsai/models/qwen3_vl/params.py @@ -0,0 +1,387 @@ +import gc +import re +from enum import Enum +from pathlib import Path + +import jax +import safetensors +from flax import nnx + +from bonsai.models.qwen3_vl import modeling as model_lib + + +class Transform(Enum): + # (Permute_axes, reshape, reshape_first) + DEFAULT = (None, None, False) + BIAS = (None, None, False) + LINEAR = ((1, 0), None, False) # PyTorch (out, in) -> JAX (in, out) + CONV3D = ((2, 3, 4, 1, 0), None, False) # (out, in, T, H, W) -> (T, H, W, in, out) + EMBED = (None, None, False) + ATTN_Q = None # Special handling needed + ATTN_KV = None # Special handling needed + + +def _get_vision_key_mapping(): + """Mapping for vision encoder weights.""" + return { + # Patch embedding - nnx.Conv uses (D, H, W, in_channels, out_channels) kernel + # PyTorch Conv3d: (out_channels, in_channels, D, H, W) + r"^model\.visual\.patch_embed\.proj\.weight$": ( + "model.visual.patch_embed.proj.kernel", + Transform.CONV3D, + ), + r"^model\.visual\.patch_embed\.proj\.bias$": ( + "model.visual.patch_embed.proj.bias", + Transform.BIAS, + ), + # Position embedding + r"^model\.visual\.pos_embed\.weight$": ( + "model.visual.pos_embed.embedding", + Transform.EMBED, + ), + # Vision blocks + r"^model\.visual\.blocks\.(\d+)\.norm1\.weight$": ( + r"model.visual.blocks.\1.norm1.scale", + Transform.DEFAULT, + ), + r"^model\.visual\.blocks\.(\d+)\.norm1\.bias$": ( + r"model.visual.blocks.\1.norm1.bias", + Transform.BIAS, + ), + r"^model\.visual\.blocks\.(\d+)\.norm2\.weight$": ( + r"model.visual.blocks.\1.norm2.scale", + Transform.DEFAULT, + ), + r"^model\.visual\.blocks\.(\d+)\.norm2\.bias$": ( + r"model.visual.blocks.\1.norm2.bias", + Transform.BIAS, + ), + # Vision attention - fused QKV + r"^model\.visual\.blocks\.(\d+)\.attn\.qkv\.weight$": ( + r"model.visual.blocks.\1.attn.qkv.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.blocks\.(\d+)\.attn\.qkv\.bias$": ( + r"model.visual.blocks.\1.attn.qkv.bias", + Transform.BIAS, + ), + r"^model\.visual\.blocks\.(\d+)\.attn\.proj\.weight$": ( + r"model.visual.blocks.\1.attn.proj.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.blocks\.(\d+)\.attn\.proj\.bias$": ( + r"model.visual.blocks.\1.attn.proj.bias", + Transform.BIAS, + ), + # Vision MLP + r"^model\.visual\.blocks\.(\d+)\.mlp\.linear_fc1\.weight$": ( + r"model.visual.blocks.\1.mlp.linear_fc1.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.blocks\.(\d+)\.mlp\.linear_fc1\.bias$": ( + r"model.visual.blocks.\1.mlp.linear_fc1.bias", + Transform.BIAS, + ), + r"^model\.visual\.blocks\.(\d+)\.mlp\.linear_fc2\.weight$": ( + r"model.visual.blocks.\1.mlp.linear_fc2.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.blocks\.(\d+)\.mlp\.linear_fc2\.bias$": ( + r"model.visual.blocks.\1.mlp.linear_fc2.bias", + Transform.BIAS, + ), + # Merger + r"^model\.visual\.merger\.norm\.weight$": ( + "model.visual.merger.norm.scale", + Transform.DEFAULT, + ), + r"^model\.visual\.merger\.norm\.bias$": ( + "model.visual.merger.norm.bias", + Transform.BIAS, + ), + r"^model\.visual\.merger\.linear_fc1\.weight$": ( + "model.visual.merger.linear_fc1.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.merger\.linear_fc1\.bias$": ( + "model.visual.merger.linear_fc1.bias", + Transform.BIAS, + ), + r"^model\.visual\.merger\.linear_fc2\.weight$": ( + "model.visual.merger.linear_fc2.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.merger\.linear_fc2\.bias$": ( + "model.visual.merger.linear_fc2.bias", + Transform.BIAS, + ), + # Deepstack mergers + r"^model\.visual\.deepstack_merger_list\.(\d+)\.norm\.weight$": ( + r"model.visual.deepstack_merger_list.\1.norm.scale", + Transform.DEFAULT, + ), + r"^model\.visual\.deepstack_merger_list\.(\d+)\.norm\.bias$": ( + r"model.visual.deepstack_merger_list.\1.norm.bias", + Transform.BIAS, + ), + r"^model\.visual\.deepstack_merger_list\.(\d+)\.linear_fc1\.weight$": ( + r"model.visual.deepstack_merger_list.\1.linear_fc1.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.deepstack_merger_list\.(\d+)\.linear_fc1\.bias$": ( + r"model.visual.deepstack_merger_list.\1.linear_fc1.bias", + Transform.BIAS, + ), + r"^model\.visual\.deepstack_merger_list\.(\d+)\.linear_fc2\.weight$": ( + r"model.visual.deepstack_merger_list.\1.linear_fc2.kernel", + Transform.LINEAR, + ), + r"^model\.visual\.deepstack_merger_list\.(\d+)\.linear_fc2\.bias$": ( + r"model.visual.deepstack_merger_list.\1.linear_fc2.bias", + Transform.BIAS, + ), + } + + +def _get_text_key_mapping(tie_word_embeddings: bool = True): + """Mapping for text decoder weights. + + Args: + tie_word_embeddings: If True, lm_head.weight maps to embed_tokens.embedding. + If False, lm_head.weight maps to lm_head.kernel. + """ + mapping = { + # Token embedding (note: PyTorch key is model.language_model.embed_tokens) + r"^model\.language_model\.embed_tokens\.weight$": ( + "model.language_model.embed_tokens.embedding", + Transform.EMBED, + ), + # Decoder layers - attention + r"^model\.language_model\.layers\.(\d+)\.self_attn\.q_proj\.weight$": ( + r"model.language_model.layers.\1.self_attn.q_proj.kernel", + Transform.LINEAR, + ), + r"^model\.language_model\.layers\.(\d+)\.self_attn\.k_proj\.weight$": ( + r"model.language_model.layers.\1.self_attn.k_proj.kernel", + Transform.LINEAR, + ), + r"^model\.language_model\.layers\.(\d+)\.self_attn\.v_proj\.weight$": ( + r"model.language_model.layers.\1.self_attn.v_proj.kernel", + Transform.LINEAR, + ), + r"^model\.language_model\.layers\.(\d+)\.self_attn\.o_proj\.weight$": ( + r"model.language_model.layers.\1.self_attn.o_proj.kernel", + Transform.LINEAR, + ), + # Q/K norms + r"^model\.language_model\.layers\.(\d+)\.self_attn\.q_norm\.weight$": ( + r"model.language_model.layers.\1.self_attn.q_norm.weight", + Transform.DEFAULT, + ), + r"^model\.language_model\.layers\.(\d+)\.self_attn\.k_norm\.weight$": ( + r"model.language_model.layers.\1.self_attn.k_norm.weight", + Transform.DEFAULT, + ), + # Decoder layers - MLP + r"^model\.language_model\.layers\.(\d+)\.mlp\.gate_proj\.weight$": ( + r"model.language_model.layers.\1.mlp.gate_proj.kernel", + Transform.LINEAR, + ), + r"^model\.language_model\.layers\.(\d+)\.mlp\.up_proj\.weight$": ( + r"model.language_model.layers.\1.mlp.up_proj.kernel", + Transform.LINEAR, + ), + r"^model\.language_model\.layers\.(\d+)\.mlp\.down_proj\.weight$": ( + r"model.language_model.layers.\1.mlp.down_proj.kernel", + Transform.LINEAR, + ), + # Decoder layers - norms + r"^model\.language_model\.layers\.(\d+)\.input_layernorm\.weight$": ( + r"model.language_model.layers.\1.input_layernorm.weight", + Transform.DEFAULT, + ), + r"^model\.language_model\.layers\.(\d+)\.post_attention_layernorm\.weight$": ( + r"model.language_model.layers.\1.post_attention_layernorm.weight", + Transform.DEFAULT, + ), + # Final norm + r"^model\.language_model\.norm\.weight$": ( + "model.language_model.norm.weight", + Transform.DEFAULT, + ), + } + + # lm_head mapping depends on whether embeddings are tied + if tie_word_embeddings: + # For tied embeddings: lm_head.weight maps to embed_tokens.embedding + mapping[r"^lm_head\.weight$"] = ( + "model.language_model.embed_tokens.embedding", + Transform.EMBED, + ) + else: + # For untied embeddings: lm_head.weight maps to lm_head.kernel (ShardedLinear) + mapping[r"^lm_head\.weight$"] = ( + "lm_head.kernel", + Transform.LINEAR, + ) + + return mapping + + +def _get_key_and_transform_mapping(tie_word_embeddings: bool = True): + """Combined key mapping for all model components. + + Args: + tie_word_embeddings: If True, lm_head.weight maps to embed_tokens.embedding. + If False, lm_head.weight maps to lm_head.kernel. + """ + mapping = {} + mapping.update(_get_vision_key_mapping()) + mapping.update(_get_text_key_mapping(tie_word_embeddings)) + return mapping + + +def _torch_key_to_jax_key(mapping: dict, source_key: str) -> tuple[str | None, Transform | None]: + """Map a PyTorch key to JAX key with transform specification.""" + matches = [ + (re.sub(pat, repl, source_key), transform) + for pat, (repl, transform) in mapping.items() + if re.match(pat, source_key) + ] + if not matches: + return None, None + if len(matches) > 1: + raise ValueError(f"Multiple mappings found for {source_key}: {[m[0] for m in matches]}") + return matches[0] + + +def _stoi(s: str) -> int | str: + """Convert string to int if possible.""" + try: + return int(s) + except ValueError: + return s + + +def _apply_transform(tensor, transform: Transform | None): + """Apply weight transformation to tensor.""" + if transform is None or transform.value is None: + return tensor + + permute, reshape, reshape_first = transform.value + + if reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if permute is not None: + tensor = tensor.transpose(permute) + if not reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + + return tensor + + +def _assign_weights( + keys: list, + tensor, + state_dict: dict, + torch_key: str, + transform: Transform | None, + sharding_dict: dict | None, +): + """Recursively assign transformed weight to state dict.""" + key, *rest = keys + + if not rest: + # Apply transformation + tensor = _apply_transform(tensor, transform) + + # Validate shape + if tensor.shape != state_dict[key].shape: + raise ValueError(f"Shape mismatch for {torch_key}: got {tensor.shape}, expected {state_dict[key].shape}") + + # Place on device with optional sharding + if sharding_dict is not None: + state_dict[key] = jax.device_put(tensor, sharding_dict[key]) + else: + state_dict[key] = jax.device_put(tensor) + else: + next_sharding = sharding_dict[key] if sharding_dict is not None else None + _assign_weights(rest, tensor, state_dict[key], torch_key, transform, next_sharding) + + +def create_model_from_safe_tensors( + file_dir: str, + config: model_lib.Qwen3VLConfig, + mesh: jax.sharding.Mesh | None = None, + model_filename: str | None = None, +) -> model_lib.Qwen3VLForConditionalGeneration: + """Load pretrained weights from safetensors and create model. + + Args: + file_dir: Path to directory containing .safetensors files. + config: Model configuration. + mesh: Optional JAX mesh for sharding. + model_filename: Optional specific filename to load (e.g., "model.safetensors"). + + Returns: + Qwen3VLForConditionalGeneration model with loaded weights. + + Raises: + ValueError: If no safetensors files found or weight conversion fails. + """ + path = Path(file_dir).expanduser() + if model_filename: + files = [path / model_filename] + if not files[0].exists(): + raise ValueError(f"Specified file not found: {files[0]}") + else: + files = list(path.glob("*.safetensors")) + if not files: + raise ValueError(f"No safetensors files found in {file_dir}") + + # Create model with abstract state (no actual arrays) + model = nnx.eval_shape(lambda: model_lib.Qwen3VLForConditionalGeneration(config, rngs=nnx.Rngs(params=0))) + graph_def, abs_state = nnx.split(model) + state_dict = abs_state.to_pure_dict() + + # Get sharding if mesh provided + sharding = nnx.get_named_sharding(abs_state, mesh).to_pure_dict() if mesh is not None else None + + # Key mapping - depends on whether embeddings are tied + key_mapping = _get_key_and_transform_mapping(config.text_config.tie_word_embeddings) + + conversion_errors = [] + + for f in files: + with safetensors.safe_open(f, framework="numpy") as sf: + for torch_key in sf.keys(): + tensor = sf.get_tensor(torch_key) + + jax_key, transform = _torch_key_to_jax_key(key_mapping, torch_key) + if jax_key is None: + # Skip unmapped keys (e.g., rotary buffers) + continue + + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights(keys, tensor, state_dict, torch_key, transform, sharding) + except Exception as e: + full_jax_key = ".".join(str(k) for k in keys) + conversion_errors.append( + f"Failed to assign '{torch_key}' to '{full_jax_key}': {type(e).__name__}: {e}" + ) + + gc.collect() + + if conversion_errors: + error_log = "\n".join(conversion_errors) + raise RuntimeError(f"Encountered {len(conversion_errors)} weight conversion errors:\n{error_log}") + + # Handle tied embeddings + if config.text_config.tie_word_embeddings: + # lm_head shares weights with embed_tokens + # The model implementation handles this at inference time + pass + + gc.collect() + return nnx.merge(graph_def, state_dict) diff --git a/bonsai/models/qwen3_vl/tests/run_model.py b/bonsai/models/qwen3_vl/tests/run_model.py new file mode 100644 index 00000000..b43d6181 --- /dev/null +++ b/bonsai/models/qwen3_vl/tests/run_model.py @@ -0,0 +1,173 @@ +import time + +import jax +import jax.numpy as jnp +import numpy as np +from huggingface_hub import snapshot_download +from jax._src.mesh import AxisType +from transformers import AutoProcessor + +from bonsai.models.qwen3_vl import modeling +from bonsai.models.qwen3_vl import params + +# ============================================================================ +# MESH CONFIGURATION - Modify these to enable sharding + +# Set USE_SHARDING=True and MESH_SHAPE to enable distributed inference +# Examples: +# MESH_SHAPE = (2, 4) -> 2 FSDP devices x 4 TP devices (8 total) +# MESH_SHAPE = (1, 4) -> Pure tensor parallelism on 4 devices +# MESH_SHAPE = (4, 1) -> Pure FSDP on 4 devices +USE_SHARDING = False +MESH_SHAPE = (1, 1) # (fsdp, tp) - only used when USE_SHARDING=True +# ============================================================================ + +MODEL_ID = "Qwen/Qwen3-VL-2B-Instruct" +EOS_TOKEN_ID = 151643 + + +def generate_with_vision( + model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids, max_new_tokens: int = 50 +): + """Generation with vision inputs.""" + batch_size, seq_len = input_ids.shape + generated_tokens = [] + + # Prefill with vision + logits, cache = modeling.forward_vision(model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids) + + next_token = jnp.argmax(logits, axis=-1, keepdims=True) + generated_tokens.append(next_token) + + # Decode loop (text-only after prefill) + for step in range(max_new_tokens - 1): + logits, cache = modeling.forward(model, cache, next_token) + next_token = jnp.argmax(logits, axis=-1, keepdims=True) + generated_tokens.append(next_token) + + # Stop on EOS - use numpy to avoid sharding issues with indexing + if int(np.array(next_token)[0, 0]) == EOS_TOKEN_ID: + break + + # Use numpy for concatenation to avoid sharding mismatches + all_tokens = [np.array(input_ids)] + [np.array(t) for t in generated_tokens] + return jnp.array(np.concatenate(all_tokens, axis=1)) + + +def main(): + print("=" * 60) + print("Qwen3-VL JAX Fast Generation Demo") + print("=" * 60) + + # Setup mesh context if sharding enabled + use_fsdp = USE_SHARDING and MESH_SHAPE[0] > 1 + use_tp = USE_SHARDING and MESH_SHAPE[1] > 1 + + if USE_SHARDING: + print(f"\nSharding enabled: MESH_SHAPE={MESH_SHAPE}, use_fsdp={use_fsdp}, use_tp={use_tp}") + mesh = jax.make_mesh(MESH_SHAPE, ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + else: + print("\nSharding disabled (running on single device)") + + # Download model + print("\n1. Downloading model...") + model_path = snapshot_download(MODEL_ID) + + # Load Flax model + print("\n2. Loading Flax model...") + start = time.time() + flax_config = modeling.Qwen3VLConfig.qwen3vl_2b(use_fsdp=use_fsdp, use_tp=use_tp) + flax_model = params.create_model_from_safe_tensors(model_path, flax_config) + print(f" Model loaded in {time.time() - start:.2f}s") + + # Load processor + processor = AutoProcessor.from_pretrained(model_path) + + print("\n" + "=" * 60) + print("3. Example : Image + Text Generation") + print("=" * 60) + + messages_vision = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.jpeg", + }, + {"type": "text", "text": "What is in this image? Answer in detail"}, + ], + } + ] + + print("\n3. Processing image input...") + inputs_vision = processor.apply_chat_template( + messages_vision, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt", + ) + + input_ids_vision = jnp.array(inputs_vision["input_ids"].numpy()) + seq_len_vision = input_ids_vision.shape[1] + print(f" Input IDs: {seq_len_vision} tokens") + + # Check for vision inputs + 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()) + + # Create token_type_ids (1 for image tokens, 0 for text) + # Image token ID is 151655 + token_type_ids = (input_ids_vision == flax_config.image_token_id).astype(jnp.int32) + + print(f" Pixel values: {pixel_values.shape}") + print(f" Image grid THW: {image_grid_thw}") + print(f" Image tokens in sequence: {int(token_type_ids.sum())}") + + # Initialize cache + batch_vision = MESH_SHAPE[0] if USE_SHARDING else 1 + if USE_SHARDING and input_ids_vision.shape[0] < batch_vision: + input_ids_vision = jnp.concatenate( + [input_ids_vision, jnp.zeros((batch_vision - 1, input_ids_vision.shape[1]), dtype=jnp.int32)], axis=0 + ) + token_type_ids = jnp.concatenate( + [token_type_ids, jnp.zeros((batch_vision - 1, token_type_ids.shape[1]), dtype=jnp.int32)], axis=0 + ) + cache_vision = modeling.init_cache(flax_config, batch_vision, seq_len_vision, generate_steps=200) + + print("\n4. Generating with vision...") + generated_ids_vision = generate_with_vision( + flax_model, + cache_vision, + input_ids_vision, + pixel_values, + image_grid_thw, + token_type_ids, + max_new_tokens=200, + ) + + # Decode + generated_ids_trimmed_vision = generated_ids_vision[:, seq_len_vision:] + output_text_vision = processor.batch_decode( + np.asarray(generated_ids_trimmed_vision), + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) + + print("\n" + "-" * 40) + print("Q: What is in this image?") + print(f"A: {output_text_vision[0]}") + print("-" * 40) + else: + print(" No pixel_values in inputs - vision processing may have failed.") + + print("\n" + "=" * 60) + print("Demo complete!") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/bonsai/models/qwen3_vl/tests/test_outputs_qwen3vl.py b/bonsai/models/qwen3_vl/tests/test_outputs_qwen3vl.py new file mode 100644 index 00000000..1c544bb8 --- /dev/null +++ b/bonsai/models/qwen3_vl/tests/test_outputs_qwen3vl.py @@ -0,0 +1,820 @@ +import os +import shutil +import tempfile +import gc # Added for memory management + +# 1. CRITICAL: Prevent JAX from pre-allocating all memory, allowing PyTorch to coexist +os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" + +import jax +import jax.numpy as jnp +import numpy as np +import torch +from absl.testing import absltest +from huggingface_hub import snapshot_download +from safetensors.torch import save_model +from transformers import AutoProcessor, Qwen3VLConfig, Qwen3VLForConditionalGeneration +from PIL import Image + +from bonsai.models.qwen3_vl import modeling as model_lib +from bonsai.models.qwen3_vl import params + +# Set highest precision for testing +jax.config.update("jax_default_matmul_precision", "highest") + + +class PretrainedModelMixin: + """ + Singleton Mixin to load the heavy pretrained model exactly once for the entire + test suite, rather than reloading it for every test class. + """ + + _models_loaded = False + model_path = None + pt_model = None + flax_config = None + flax_model = None + processor = None + + @classmethod + def load_models_once(cls): + """Loads models if they haven't been loaded yet.""" + if cls._models_loaded: + return + + print("\n[Mixin] Loading Qwen3-VL-2B models (Shared instance)...") + MODEL_ID = "Qwen/Qwen3-VL-2B-Instruct" + + # Download and Load + cls.model_path = snapshot_download(MODEL_ID) + cls.pt_model = Qwen3VLForConditionalGeneration.from_pretrained(cls.model_path, dtype=torch.float32).eval() + + cls.flax_config = model_lib.Qwen3VLConfig.qwen3vl_2b() + cls.flax_model = params.create_model_from_safe_tensors(cls.model_path, cls.flax_config) + cls.processor = AutoProcessor.from_pretrained(cls.model_path) + + cls._models_loaded = True + print("[Mixin] Models loaded successfully.\n") + + def tearDown(self): + """ + Aggressive cleanup after every individual test method to + prevent intermediate tensors from piling up. + """ + super().tearDown() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + jax.clear_caches() + gc.collect() + + +def get_test_config_torch() -> Qwen3VLConfig: + return Qwen3VLConfig( + vision_config={ + "depth": 2, + "hidden_size": 64, + "intermediate_size": 128, + "num_heads": 4, + "in_channels": 3, + "patch_size": 8, + "temporal_patch_size": 2, + "spatial_merge_size": 2, + "out_hidden_size": 128, + "num_position_embeddings": 256, + "deepstack_visual_indexes": [0, 1], + "hidden_act": "gelu_pytorch_tanh", + }, + text_config={ + "vocab_size": 1000, + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 32, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "rope_theta": 5_000_000, + "rope_scaling": { + "mrope_interleaved": True, + "mrope_section": [12, 10, 10], + "rope_type": "default", + }, + "attention_bias": False, + "tie_word_embeddings": True, + }, + image_token_id=151655, + video_token_id=151656, + ) + + +def get_test_config_flax() -> model_lib.Qwen3VLConfig: + return model_lib.Qwen3VLConfig( + vision_config=model_lib.Qwen3VLVisionConfig( + depth=2, + hidden_size=64, + intermediate_size=128, + num_heads=4, + in_channels=3, + patch_size=8, + temporal_patch_size=2, + spatial_merge_size=2, + out_hidden_size=128, + num_position_embeddings=256, + deepstack_visual_indexes=(0, 1), + ), + text_config=model_lib.Qwen3VLTextConfig( + vocab_size=1000, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + mrope_section=(12, 10, 10), + tie_word_embeddings=True, + ), + ) + + +class TestForwardPass(absltest.TestCase): + """Test forward pass using small manually-initialized models.""" + + def setUp(self): + super().setUp() + self.test_dir = tempfile.mkdtemp() + + self.pt_config = get_test_config_torch() + self.pt_model = Qwen3VLForConditionalGeneration(config=self.pt_config) + self.model_filename = "qwen3vl_test.safetensors" + self.model_ckpt_path = os.path.join(self.test_dir, self.model_filename) + save_model(self.pt_model, self.model_ckpt_path) + + self.flax_config = get_test_config_flax() + self.flax_model = params.create_model_from_safe_tensors( + self.test_dir, self.flax_config, model_filename=self.model_filename + ) + + self.pt_model.eval() + self.batch_size = 1 + self.seq_len = 10 + + def tearDown(self): + # Recursively remove the temporary directory and all its contents + shutil.rmtree(self.test_dir) + gc.collect() # Explicit GC + super().tearDown() + + def test_text_embeddings(self): + """Compare embedding layer output.""" + pt_embed = self.pt_model.model.language_model.embed_tokens + flax_embed = self.flax_model.model.language_model.embed_tokens + + key = jax.random.PRNGKey(0) + input_ids = jax.random.randint(key, (self.batch_size, self.seq_len), 0, 100) + pt_ids = torch.tensor(np.asarray(input_ids), dtype=torch.long) + + with torch.inference_mode(): + pt_out = pt_embed(pt_ids) + flax_out = flax_embed(input_ids, out_sharding=model_lib.P(None, None, None)) + + np.testing.assert_allclose(np.asarray(flax_out), pt_out.numpy(), rtol=1e-7, atol=1e-7) + + def test_rms_norm(self): + """Compare RMSNorm layer.""" + pt_norm = self.pt_model.model.language_model.norm + flax_norm = self.flax_model.model.language_model.norm + + key = jax.random.PRNGKey(0) + hidden_size = self.flax_config.text_config.hidden_size + jx = jax.random.normal(key, (self.batch_size, self.seq_len, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_out = pt_norm(tx) + flax_out = flax_norm(jx) + + np.testing.assert_allclose(np.asarray(flax_out), pt_out.numpy(), rtol=1e-6, atol=1e-6) + + def test_text_mlp(self): + """Compare gated MLP layer.""" + pt_mlp = self.pt_model.model.language_model.layers[0].mlp + flax_mlp = self.flax_model.model.language_model.layers[0].mlp + + key = jax.random.PRNGKey(0) + hidden_size = self.flax_config.text_config.hidden_size + jx = jax.random.normal(key, (self.batch_size, self.seq_len, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_out = pt_mlp(tx) + # Reshape to (batch * seq, hidden) for ShardedLinear, then reshape back + jx_flat = jx.reshape(-1, hidden_size) + flax_out_flat = flax_mlp(jx_flat) + flax_out = flax_out_flat.reshape(self.batch_size, self.seq_len, -1) + + np.testing.assert_allclose(np.asarray(flax_out), pt_out.numpy(), rtol=1e-7, atol=1e-7) + + +class TestVisionComponentsEquivalence(absltest.TestCase): + """Test vision encoder components with PyTorch comparison.""" + + def setUp(self): + super().setUp() + self.test_dir = tempfile.mkdtemp() + + self.pt_config = get_test_config_torch() + self.pt_model = Qwen3VLForConditionalGeneration(config=self.pt_config) + self.model_filename = "qwen3vl_vision_test.safetensors" + self.model_ckpt_path = os.path.join(self.test_dir, self.model_filename) + save_model(self.pt_model, self.model_ckpt_path) + + self.flax_config = get_test_config_flax() + self.flax_model = params.create_model_from_safe_tensors( + self.test_dir, self.flax_config, model_filename=self.model_filename + ) + + self.pt_model.eval() + + def tearDown(self): + shutil.rmtree(self.test_dir) + gc.collect() + super().tearDown() + + def test_vision_mlp(self): + """Compare vision MLP output.""" + pt_mlp = self.pt_model.model.visual.blocks[0].mlp + flax_mlp = self.flax_model.model.visual.blocks[0].mlp + + hidden_size = 64 + key = jax.random.PRNGKey(0) + jx = jax.random.normal(key, (16, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_out = pt_mlp(tx) + flax_out = flax_mlp(jx) + + np.testing.assert_allclose(np.asarray(flax_out), pt_out.numpy(), rtol=1e-5, atol=1e-5) + + def test_vision_layernorm(self): + """Compare vision LayerNorm output.""" + pt_norm = self.pt_model.model.visual.blocks[0].norm1 + flax_norm = self.flax_model.model.visual.blocks[0].norm1 + + hidden_size = 64 + key = jax.random.PRNGKey(0) + jx = jax.random.normal(key, (16, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_out = pt_norm(tx) + flax_out = flax_norm(jx) + + np.testing.assert_allclose(np.asarray(flax_out), pt_out.numpy(), rtol=1e-6, atol=1e-6) + + def test_vision_attention_qkv(self): + """Compare vision attention QKV projections.""" + pt_attn = self.pt_model.model.visual.blocks[0].attn + flax_attn = self.flax_model.model.visual.blocks[0].attn + + hidden_size = 64 + seq_len = 16 + key = jax.random.PRNGKey(0) + jx = jax.random.normal(key, (seq_len, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + # Compare QKV projection output + with torch.inference_mode(): + pt_qkv = pt_attn.qkv(tx).numpy() + flax_qkv = np.array(flax_attn.qkv(jx, out_sharding=model_lib.P(None, None))) + + np.testing.assert_allclose(flax_qkv, pt_qkv, rtol=1e-6, atol=1e-6) + + def test_vision_block_components(self): + """Compare vision block layer norm and MLP separately.""" + pt_block = self.pt_model.model.visual.blocks[0] + flax_block = self.flax_model.model.visual.blocks[0] + + hidden_size = 64 + seq_len = 16 + key = jax.random.PRNGKey(42) + jx = jax.random.normal(key, (seq_len, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + # Test norm1 + with torch.inference_mode(): + pt_norm1 = pt_block.norm1(tx).numpy() + flax_norm1 = np.array(flax_block.norm1(jx)) + np.testing.assert_allclose(flax_norm1, pt_norm1, rtol=1e-6, atol=1e-6) + + # Test norm2 + with torch.inference_mode(): + pt_norm2 = pt_block.norm2(tx).numpy() + flax_norm2 = np.array(flax_block.norm2(jx)) + np.testing.assert_allclose(flax_norm2, pt_norm2, rtol=1e-6, atol=1e-6) + + # Test MLP + with torch.inference_mode(): + pt_mlp = pt_block.mlp(tx).numpy() + flax_mlp = np.array(flax_block.mlp(jx)) + np.testing.assert_allclose(flax_mlp, pt_mlp, rtol=1e-5, atol=1e-5) + + def test_vision_attention_with_rope(self): + """Compare vision attention output with RoPE applied.""" + pt_attn = self.pt_model.model.visual.blocks[0].attn + flax_attn = self.flax_model.model.visual.blocks[0].attn + + hidden_size = 64 + seq_len = 16 + head_dim = hidden_size // 4 # 4 heads + + key = jax.random.PRNGKey(42) + jx = jax.random.normal(key, (seq_len, hidden_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + # Create position embeddings (cos, sin) for RoPE + # Using same formula as vision encoder + rotary_dim = head_dim // 2 + inv_freq = 1.0 / (10000.0 ** (np.arange(0, rotary_dim, 2, dtype=np.float32) / rotary_dim)) + positions = np.arange(seq_len, dtype=np.float32) + freqs = np.outer(positions, inv_freq) # (seq, rotary_dim//2) + emb = np.concatenate([freqs, freqs], axis=-1) # (seq, rotary_dim) + emb = np.concatenate([emb, emb], axis=-1) # (seq, head_dim) + cos_np, sin_np = np.cos(emb), np.sin(emb) + + cos_jax = jnp.array(cos_np) + sin_jax = jnp.array(sin_np) + cos_pt = torch.tensor(cos_np) + sin_pt = torch.tensor(sin_np) + + with torch.inference_mode(): + cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32) + pt_out = pt_attn(tx, cu_seqlens=cu_seqlens, position_embeddings=(cos_pt, sin_pt)).numpy() + flax_out = np.array(flax_attn(jx, position_embeddings=(cos_jax, sin_jax))) + + np.testing.assert_allclose(flax_out, pt_out, rtol=1e-4, atol=1e-4) + + def test_patch_embed(self): + """Compare patch embedding output with dummy config.""" + pt_patch_embed = self.pt_model.model.visual.patch_embed + flax_patch_embed = self.flax_model.model.visual.patch_embed + + # Create input matching the expected format + # patch_size=8, temporal_patch_size=2, so input per patch is: 3 * 2 * 8 * 8 = 384 + cfg = self.flax_config.vision_config + per_patch_size = cfg.in_channels * cfg.temporal_patch_size * cfg.patch_size * cfg.patch_size + num_patches = 16 # Arbitrary number of patches + + key = jax.random.PRNGKey(0) + jx = jax.random.normal(key, (num_patches, per_patch_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_out = pt_patch_embed(tx).numpy() + flax_out = np.array(flax_patch_embed(jx)) + + np.testing.assert_allclose(flax_out, pt_out, rtol=1e-5, atol=1e-5) + + def test_rope_embedding(self): + """Compare RoPE embedding computation with dummy config.""" + pt_visual = self.pt_model.model.visual + flax_visual = self.flax_model.model.visual + + # Create grid_thw for a small image: 1 frame, 16x16 grid + 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) + + with torch.inference_mode(): + pt_rope = pt_visual.rot_pos_emb(grid_thw_pt) + pt_emb = torch.cat([pt_rope, pt_rope], dim=-1) + pt_cos = pt_emb.cos().numpy() + pt_sin = pt_emb.sin().numpy() + + flax_cos, flax_sin = flax_visual._rot_pos_emb(grid_thw_jax) + + np.testing.assert_allclose(np.array(flax_cos), pt_cos, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(np.array(flax_sin), pt_sin, rtol=1e-5, atol=1e-5) + + def test_full_vision_encoder(self): + """Compare full vision encoder output with dummy config.""" + pt_visual = self.pt_model.model.visual + flax_visual = self.flax_model.model.visual + + # Create dummy vision input + # grid_thw: 1 frame, 16x16 grid (spatial_merge_size=2, so 8x8 merged patches = 64) + cfg = self.flax_config.vision_config + per_patch_size = cfg.in_channels * cfg.temporal_patch_size * cfg.patch_size * cfg.patch_size + grid_t, grid_h, grid_w = 1, 16, 16 + num_patches = grid_t * grid_h * grid_w # 256 patches before merge + + 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) + + key = jax.random.PRNGKey(42) + jx = jax.random.normal(key, (num_patches, per_patch_size), dtype=jnp.float32) + tx = torch.tensor(np.asarray(jx), dtype=torch.float32) + + with torch.inference_mode(): + pt_result = pt_visual(tx, grid_thw_pt) + # Handle both tuple and BaseModelOutputWithDeepstackFeatures + # HF now returns pre-merger output in last_hidden_state, so we apply merger manually + if hasattr(pt_result, "last_hidden_state"): + pt_hidden = pt_result.last_hidden_state + # Apply merger to get merged output matching Flax + pt_out = pt_visual.merger(pt_hidden).numpy() + else: + pt_out = pt_result[0].numpy() + + flax_out, flax_deepstack = flax_visual(jx, grid_thw_jax) + flax_out = np.array(flax_out) + + # After spatial merge (2x2), 256 patches become 64 + expected_merged_patches = num_patches // (cfg.spatial_merge_size**2) + self.assertEqual(flax_out.shape[0], expected_merged_patches) + self.assertEqual(pt_out.shape[0], expected_merged_patches) + self.assertEqual(flax_out.shape[1], cfg.out_hidden_size) + self.assertEqual(pt_out.shape[1], cfg.out_hidden_size) + + np.testing.assert_allclose(flax_out, pt_out, rtol=1e-3, atol=5e-3) + + +class TestKVCache(absltest.TestCase): + """Test KV-cache functionality.""" + + def test_cache_initialization(self): + """Test cache is initialized with correct shapes.""" + config = get_test_config_flax() + cache = model_lib.init_cache(config, batch_size=2, token_len=10, generate_steps=5) + + self.assertEqual(len(cache), config.text_config.num_hidden_layers) + + for layer_cache in cache: + k_shape = layer_cache.k_cache[:].shape + self.assertEqual(k_shape[0], 2) # batch + self.assertEqual(k_shape[2], config.text_config.num_key_value_heads) + self.assertEqual(k_shape[3], config.text_config.head_dim) + + +# --- UPDATED CLASSES USING MIXIN --- + + +class TestPretrained2B(PretrainedModelMixin, absltest.TestCase): + """Test pretrained Qwen3-VL-2B model. Inherits from Mixin to share model.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.load_models_once() + + def test_embedding_output(self): + """Check embedding outputs match.""" + text = "Hello" + inputs = self.processor(text=text, return_tensors="pt") + input_ids_pt = inputs["input_ids"] + input_ids_jax = jnp.array(inputs["input_ids"].numpy()) + + with torch.inference_mode(): + pt_out = self.pt_model.model.language_model.embed_tokens(input_ids_pt).numpy() + flax_out = np.array( + self.flax_model.model.language_model.embed_tokens(input_ids_jax, out_sharding=model_lib.P(None, None, None)) + ) + + np.testing.assert_allclose(flax_out, pt_out, rtol=1e-6, atol=1e-6) + + def test_text_forward_pass(self): + """Full text forward pass comparison.""" + text = "Hello" + inputs = self.processor(text=text, return_tensors="pt") + input_ids_jax = jnp.array(inputs["input_ids"].numpy()) + + with torch.inference_mode(): + pt_logits = self.pt_model(input_ids=inputs["input_ids"]).logits.numpy() + + batch, seq_len = input_ids_jax.shape + cache = model_lib.init_cache(self.flax_config, batch, seq_len, generate_steps=10) + flax_logits = np.array(self.flax_model(input_ids_jax, cache)) + + np.testing.assert_allclose(flax_logits, pt_logits, rtol=1e-5, atol=1e-2) + + def test_text_decoder_layer_output(self): + """Compare text decoder layer 0 output.""" + text = "Hello world" + inputs = self.processor(text=text, return_tensors="pt") + input_ids_pt = inputs["input_ids"] + input_ids_jax = jnp.array(input_ids_pt.numpy()) + + # Get hidden states after layer 0 using hooks + pt_hidden_after_layer0 = [] + + def hook(module, input, output): + pt_hidden_after_layer0.append(output[0].detach()) + + handle = self.pt_model.model.language_model.layers[0].register_forward_hook(hook) + + with torch.inference_mode(): + _ = self.pt_model(input_ids=input_ids_pt) + + handle.remove() + pt_layer0_out = pt_hidden_after_layer0[0].numpy() + + # Flax: run through embed + layer 0 manually + batch, seq_len = input_ids_jax.shape + cache = model_lib.init_cache(self.flax_config, batch, seq_len, 10) + + flax_embeds = self.flax_model.model.language_model.embed_tokens( + input_ids_jax, out_sharding=model_lib.P(None, None, None) + ) + + # Generate position embeddings + positions = jnp.arange(seq_len)[None, :] + sin, cos = model_lib._generate_rope( + positions, self.flax_config.text_config.head_dim, self.flax_config.text_config.rope_theta + ) + + # Run through layer 0 + mask = model_lib.make_causal_mask(cache[0], seq_len) + flax_layer0_out = self.flax_model.model.language_model.layers[0](flax_embeds, cache[0], sin, cos, mask) + + np.testing.assert_allclose(np.array(flax_layer0_out.squeeze(0)), pt_layer0_out, rtol=1e-4, atol=5e-3) + + def test_jit_forward(self): + """Test JIT-compiled forward.""" + text = "Hello" + inputs = self.processor(text=text, return_tensors="pt") + input_ids = jnp.array(inputs["input_ids"].numpy()) + + cache = model_lib.init_cache(self.flax_config, 1, input_ids.shape[1], 20) + logits, _ = model_lib.forward(self.flax_model, cache, input_ids) + + self.assertEqual(logits.shape, (1, self.flax_config.text_config.vocab_size)) + + def test_generation_step_with_numeric_check(self): + """Test generation step with numeric output verification.""" + # Use processor to get proper token ids + text = "Hello" + inputs = self.processor(text=text, return_tensors="pt") + input_ids_pt = inputs["input_ids"] + input_ids_jax = jnp.array(input_ids_pt.numpy()) + + seq_len = input_ids_jax.shape[1] + cache = model_lib.init_cache(self.flax_config, 1, seq_len, 20) + + # Prefill - compare with PyTorch + with torch.inference_mode(): + pt_logits = self.pt_model(input_ids=input_ids_pt).logits[:, -1, :].numpy() + + flax_logits, cache = model_lib.forward(self.flax_model, cache, input_ids_jax) + flax_logits_np = np.array(flax_logits) + + # Verify numeric equivalence + np.testing.assert_allclose(flax_logits_np, pt_logits, rtol=1e-4, atol=1e-2) + + # Verify cache position + self.assertEqual(int(cache[0].cur_ind.get_value()), seq_len) + + # Decode step + next_token_flax = jnp.argmax(flax_logits, axis=-1, keepdims=True) + # next_token_pt = torch.tensor(np.array(next_token_flax), dtype=torch.long) + + flax_logits2, cache = model_lib.forward(self.flax_model, cache, next_token_flax) + + # Verify cache position updated + self.assertEqual(int(cache[0].cur_ind.get_value()), seq_len + 1) + + # Verify logits are valid (not NaN/Inf) + self.assertFalse(np.any(np.isnan(np.array(flax_logits2)))) + self.assertFalse(np.any(np.isinf(np.array(flax_logits2)))) + + +class TestVisionEncoderPretrained(PretrainedModelMixin, absltest.TestCase): + """Test vision encoder with pretrained weights. Inherits from Mixin.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.load_models_once() + + def _create_dummy_image_input(self): + """Create dummy image input for testing.""" + image = Image.new("RGB", (256, 256), color=(128, 128, 128)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "What?"}, + ], + } + ] + return self.processor.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" + ) + + def test_patch_embed_output(self): + """Check patch embedding output matches.""" + inputs = self._create_dummy_image_input() + pixel_values_pt = inputs["pixel_values"] + pixel_values_jax = jnp.array(pixel_values_pt.numpy()) + + with torch.inference_mode(): + pt_out = self.pt_model.model.visual.patch_embed(pixel_values_pt).numpy() + flax_out = np.array(self.flax_model.model.visual.patch_embed(pixel_values_jax)) + + np.testing.assert_allclose(flax_out, pt_out, rtol=1e-5, atol=1e-5) + + def test_position_embedding_output(self): + """Check position embedding interpolation output matches.""" + inputs = self._create_dummy_image_input() + grid_thw_pt = inputs["image_grid_thw"] + grid_thw_jax = jnp.array(grid_thw_pt.numpy()) + + with torch.inference_mode(): + pt_pos = self.pt_model.model.visual.fast_pos_embed_interpolate(grid_thw_pt).numpy() + flax_pos = np.array(self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_jax)) + + np.testing.assert_allclose(flax_pos, pt_pos, rtol=1e-5, atol=3e-5) + + def test_rope_embedding(self): + """Check RoPE embedding output matches.""" + inputs = self._create_dummy_image_input() + grid_thw_pt = inputs["image_grid_thw"] + grid_thw_jax = jnp.array(grid_thw_pt.numpy()) + + with torch.inference_mode(): + pt_rope = self.pt_model.model.visual.rot_pos_emb(grid_thw_pt) + pt_emb = torch.cat([pt_rope, pt_rope], dim=-1) + pt_cos = pt_emb.cos().numpy() + pt_sin = pt_emb.sin().numpy() + + flax_cos, flax_sin = self.flax_model.model.visual._rot_pos_emb(grid_thw_jax) + + np.testing.assert_allclose(np.array(flax_cos), pt_cos, rtol=1e-6, atol=2e-5) + np.testing.assert_allclose(np.array(flax_sin), pt_sin, rtol=1e-6, atol=2e-5) + + def test_vision_patch_plus_pos_output(self): + """Check patch + position embedding output matches.""" + inputs = self._create_dummy_image_input() + pixel_values_pt = inputs["pixel_values"] + grid_thw_pt = inputs["image_grid_thw"] + pixel_values_jax = jnp.array(pixel_values_pt.numpy()) + grid_thw_jax = jnp.array(grid_thw_pt.numpy()) + + # Get patch + pos embeddings + with torch.inference_mode(): + pt_patches = self.pt_model.model.visual.patch_embed(pixel_values_pt) + pt_pos = self.pt_model.model.visual.fast_pos_embed_interpolate(grid_thw_pt) + pt_hidden = (pt_patches + pt_pos).numpy() + + flax_patches = self.flax_model.model.visual.patch_embed(pixel_values_jax) + flax_pos = self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_jax) + flax_hidden = np.array(flax_patches + flax_pos) + + np.testing.assert_allclose(flax_hidden, pt_hidden, rtol=1e-5, atol=1e-5) + + def test_full_vision_output(self): + """Check full vision encoder output matches.""" + inputs = self._create_dummy_image_input() + pixel_values_pt = inputs["pixel_values"] + grid_thw_pt = inputs["image_grid_thw"] + pixel_values_jax = jnp.array(pixel_values_pt.numpy()) + grid_thw_jax = jnp.array(grid_thw_pt.numpy()) + + with torch.inference_mode(): + pt_result = self.pt_model.model.visual(pixel_values_pt, grid_thw_pt) + # Handle both tuple and BaseModelOutputWithDeepstackFeatures + # HF now returns pre-merger output in last_hidden_state, so we apply merger manually + if hasattr(pt_result, "last_hidden_state"): + pt_hidden = pt_result.last_hidden_state + pt_out = self.pt_model.model.visual.merger(pt_hidden).numpy() + else: + pt_out = pt_result[0].numpy() + flax_out, _ = self.flax_model.model.visual(pixel_values_jax, grid_thw_jax) + + self.assertEqual( + np.array(flax_out).shape, + pt_out.shape, + f"Shape mismatch: Flax {np.array(flax_out).shape} vs PT {pt_out.shape}.", + ) + np.testing.assert_allclose(np.array(flax_out), pt_out, rtol=1e-5, atol=8e-3) + + +class TestVisionTextGeneration(PretrainedModelMixin, absltest.TestCase): + """Test vision + text generation with pretrained model. Inherits from Mixin.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.load_models_once() + + def test_vision_forward_with_numeric_check(self): + """Test vision forward pass with numeric verification.""" + image = Image.new("RGB", (256, 256), color=(100, 150, 200)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "Describe"}, + ], + } + ] + inputs = self.processor.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" + ) + + pixel_values_pt = inputs["pixel_values"] + grid_thw_pt = inputs["image_grid_thw"] + input_ids_pt = inputs["input_ids"] + + pixel_values_jax = jnp.array(pixel_values_pt.numpy()) + grid_thw_jax = jnp.array(grid_thw_pt.numpy()) + input_ids_jax = jnp.array(input_ids_pt.numpy()) + + # Create token_type_ids: 1 for image tokens, 0 for text + image_token_id = self.flax_config.image_token_id + token_type_ids_jax = (input_ids_jax == image_token_id).astype(jnp.int32) + + # PyTorch forward + with torch.inference_mode(): + pt_out = self.pt_model( + input_ids=input_ids_pt, + pixel_values=pixel_values_pt, + image_grid_thw=grid_thw_pt, + ) + pt_logits = pt_out.logits[:, -1, :].numpy() + + # Flax forward + batch, seq_len = input_ids_jax.shape + cache = model_lib.init_cache(self.flax_config, batch, seq_len, 20) + + flax_logits, cache = model_lib.forward_vision( + self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_jax, token_type_ids_jax + ) + + # Vision+text forward has larger tolerance due to accumulated numerical diffs + # Top logits should still match for generation correctness + flax_logits_np = np.array(flax_logits) + + # Check top-k token predictions match + pt_top5 = np.argsort(pt_logits[0])[-5:] + flax_top5 = np.argsort(flax_logits_np[0])[-5:] + overlap = len(set(pt_top5) & set(flax_top5)) + self.assertTrue(overlap >= 4, f"Out of Top-5 tokens, top 4 must overlap, overlap: {overlap}") + + # Check overall correlation + corr = np.corrcoef(pt_logits.flatten(), flax_logits_np.flatten())[0, 1] + self.assertGreater(corr, 0.95, f"Logits should be highly correlated, got {corr}") + + def test_generation_with_vision_input(self): + """Test generation step with vision input.""" + image = Image.new("RGB", (256, 256), color=(50, 100, 150)) + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "What is this?"}, + ], + } + ] + inputs = self.processor.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" + ) + + pixel_values_jax = jnp.array(inputs["pixel_values"].numpy()) + grid_thw_jax = jnp.array(inputs["image_grid_thw"].numpy()) + input_ids_jax = jnp.array(inputs["input_ids"].numpy()) + + # Create token_type_ids: 1 for image tokens, 0 for text + image_token_id = self.flax_config.image_token_id + token_type_ids_jax = (input_ids_jax == image_token_id).astype(jnp.int32) + + batch, seq_len = input_ids_jax.shape + cache = model_lib.init_cache(self.flax_config, batch, seq_len, 30) + + # Prefill with vision + logits, cache = model_lib.forward_vision( + self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_jax, token_type_ids_jax + ) + + # Verify cache position + self.assertEqual(int(cache[0].cur_ind.get_value()), seq_len) + + # Generate a few tokens + for i in range(3): + next_token = jnp.argmax(logits, axis=-1, keepdims=True) + logits, cache = model_lib.forward(self.flax_model, cache, next_token) + + # Verify logits are valid + self.assertFalse(np.any(np.isnan(np.array(logits)))) + self.assertFalse(np.any(np.isinf(np.array(logits)))) + + # Verify cache position updates + self.assertEqual(int(cache[0].cur_ind.get_value()), seq_len + i + 1) + + +if __name__ == "__main__": + absltest.main() diff --git a/bonsai/models/qwen3_vl/tests/test_sharding_qwen3vl.py b/bonsai/models/qwen3_vl/tests/test_sharding_qwen3vl.py new file mode 100644 index 00000000..64891c18 --- /dev/null +++ b/bonsai/models/qwen3_vl/tests/test_sharding_qwen3vl.py @@ -0,0 +1,271 @@ +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest +import unittest + +from flax import nnx +from jax._src.mesh import AxisType +from jax.sharding import PartitionSpec as P + +from bonsai.models.qwen3_vl import modeling + + +def get_test_config(use_fsdp=False, use_tp=False): + """Get a small test config with optional sharding.""" + return modeling.Qwen3VLConfig( + vision_config=modeling.Qwen3VLVisionConfig( + depth=2, + hidden_size=64, + intermediate_size=128, + num_heads=4, + in_channels=3, + patch_size=14, + temporal_patch_size=2, + spatial_merge_size=2, + out_hidden_size=128, + num_position_embeddings=256, + deepstack_visual_indexes=(0,), + shd_cfg=modeling.VisionShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else modeling.VisionShardingConfig.no_sharding(), + ), + text_config=modeling.Qwen3VLTextConfig( + vocab_size=1000, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=32, + shd_cfg=modeling.TextShardingConfig.default(use_fsdp, use_tp) + if (use_fsdp or use_tp) + else modeling.TextShardingConfig.no_sharding(), + ), + ) + + +@unittest.skipIf(jax.device_count() < 4, "Atleast 4 devices required") +class TestSharding(absltest.TestCase): + """Test sharding with simulated 8-device mesh.""" + + @classmethod + def setUpClass(cls): + print(f"JAX devices: {jax.devices()}") + print(f"Device count: {len(jax.devices())}") + assert len(jax.devices()) == 8, f"Expected 8 simulated devices, got {len(jax.devices())}" + + def test_sharding_config_creation(self): + """Test that sharding configs are created correctly.""" + # No sharding + cfg_unsharded = modeling.Qwen3VLConfig.qwen3vl_2b() + self.assertEqual(cfg_unsharded.text_config.shd_cfg.q_weight, P(None, None)) + + # With sharding + cfg_sharded = modeling.Qwen3VLConfig.qwen3vl_2b(use_fsdp=True, use_tp=True) + self.assertEqual(cfg_sharded.text_config.shd_cfg.q_weight, P("fsdp", "tp")) + self.assertEqual(cfg_sharded.vision_config.shd_cfg.attn_qkv_kernel, P("fsdp", "tp")) + + def test_text_mlp_sharded_vs_unsharded(self): + """Test text MLP output is numerically equivalent with/without sharding.""" + # Setup mesh for sharded test + mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + + # Create unsharded model and input + cfg_unsharded = get_test_config(use_fsdp=False, use_tp=False) + rngs = nnx.Rngs(42) + mlp_unsharded = modeling.Qwen3VLMLP(cfg_unsharded.text_config, rngs=rngs) + + # Capture weights from unsharded model using [...] indexing + gate_kernel = np.array(mlp_unsharded.gate_proj.kernel[...]) + up_kernel = np.array(mlp_unsharded.up_proj.kernel[...]) + down_kernel = np.array(mlp_unsharded.down_proj.kernel[...]) + + # Use batch=2 so it's divisible by fsdp=2 + x = jnp.ones((2, 8, 128), dtype=jnp.float32) # (batch, seq, hidden) + out_unsharded = mlp_unsharded(x) + + # Create sharded model with same weights + jax.set_mesh(mesh) + cfg_sharded = get_test_config(use_fsdp=True, use_tp=True) + rngs = nnx.Rngs(42) + mlp_sharded = modeling.Qwen3VLMLP(cfg_sharded.text_config, rngs=rngs) + + # Copy weights to sharded model using [...] indexing + mlp_sharded.gate_proj.kernel[...] = jnp.array(gate_kernel) + mlp_sharded.up_proj.kernel[...] = jnp.array(up_kernel) + mlp_sharded.down_proj.kernel[...] = jnp.array(down_kernel) + + # Recreate input inside mesh context + x_sharded = jnp.ones((2, 8, 128), dtype=jnp.float32) + out_sharded = mlp_sharded(x_sharded) + + # Numerical comparison + np.testing.assert_allclose( + np.array(out_unsharded), + np.array(out_sharded), + rtol=1e-5, + atol=1e-5, + err_msg="Sharded MLP output differs from unsharded", + ) + print("Text MLP: sharded vs unsharded numerical match ✓") + + # Reset mesh + jax.set_mesh(jax.make_mesh((1,), ("dummy",), axis_types=(AxisType.Explicit,))) + + def test_vision_mlp_sharded_vs_unsharded(self): + """Test vision MLP output is numerically equivalent with/without sharding.""" + mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + + # Create unsharded model + cfg_unsharded = get_test_config(use_fsdp=False, use_tp=False) + rngs = nnx.Rngs(42) + mlp_unsharded = modeling.Qwen3VLVisionMLP(cfg_unsharded.vision_config, rngs=rngs) + + fc1_kernel = np.array(mlp_unsharded.linear_fc1.kernel[...]) + fc2_kernel = np.array(mlp_unsharded.linear_fc2.kernel[...]) + fc1_bias = np.array(mlp_unsharded.linear_fc1.bias[...]) + fc2_bias = np.array(mlp_unsharded.linear_fc2.bias[...]) + + x = jnp.ones((16, 64), dtype=jnp.float32) # (seq, hidden) + out_unsharded = mlp_unsharded(x) + + # Create sharded model + jax.set_mesh(mesh) + cfg_sharded = get_test_config(use_fsdp=True, use_tp=True) + rngs = nnx.Rngs(42) + mlp_sharded = modeling.Qwen3VLVisionMLP(cfg_sharded.vision_config, rngs=rngs) + + mlp_sharded.linear_fc1.kernel[...] = jnp.array(fc1_kernel) + mlp_sharded.linear_fc2.kernel[...] = jnp.array(fc2_kernel) + mlp_sharded.linear_fc1.bias[...] = jnp.array(fc1_bias) + mlp_sharded.linear_fc2.bias[...] = jnp.array(fc2_bias) + + # Recreate input inside mesh context + x_sharded = jnp.ones((16, 64), dtype=jnp.float32) + out_sharded = mlp_sharded(x_sharded) + + np.testing.assert_allclose( + np.array(out_unsharded), + np.array(out_sharded), + rtol=1e-5, + atol=1e-5, + err_msg="Sharded Vision MLP output differs from unsharded", + ) + print("Vision MLP: sharded vs unsharded numerical match ✓") + + jax.set_mesh(jax.make_mesh((1,), ("dummy",), axis_types=(AxisType.Explicit,))) + + def test_full_model_creation_with_sharding(self): + """Test full model can be created with sharding enabled.""" + mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + + cfg = get_test_config(use_fsdp=True, use_tp=True) + rngs = nnx.Rngs(0) + + model = modeling.Qwen3VLForConditionalGeneration(cfg, rngs=rngs) + + # Just verify model was created successfully + self.assertIsNotNone(model) + print("Full model created with sharding ✓") + + jax.set_mesh(jax.make_mesh((1,), ("dummy",), axis_types=(AxisType.Explicit,))) + + def test_text_model_forward_with_sharding(self): + """Test text model can run forward pass with sharding enabled.""" + mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + + # Create sharded model + cfg = get_test_config(use_fsdp=True, use_tp=True) + rngs = nnx.Rngs(42) + text_model = modeling.Qwen3VLTextModel(cfg.text_config, rngs=rngs) + + # Create input - batch=2 divisible by fsdp=2 + batch, seq_len = 2, 8 + inputs_embeds = jnp.ones((batch, seq_len, 128), dtype=jnp.float32) + cache = modeling.init_cache(cfg, batch_size=batch, token_len=seq_len, generate_steps=4) + positions = jnp.arange(seq_len)[None, :].repeat(batch, axis=0) + sin, cos = modeling._generate_rope(positions, cfg.text_config.head_dim, cfg.text_config.rope_theta) + + # Forward pass should complete without error + output = text_model(inputs_embeds, cache, sin, cos, mask=None) + + self.assertEqual(output.shape, (batch, seq_len, 128)) + print("Text Model: forward pass with sharding ✓") + + jax.set_mesh(jax.make_mesh((1,), ("dummy",), axis_types=(AxisType.Explicit,))) + + def test_vision_attention_sharded_vs_unsharded(self): + """Test vision attention output is numerically equivalent with/without sharding.""" + mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) + + # Create unsharded model + cfg_unsharded = get_test_config(use_fsdp=False, use_tp=False) + rngs = nnx.Rngs(42) + attn_unsharded = modeling.Qwen3VLVisionAttention(cfg_unsharded.vision_config, rngs=rngs) + + # Get weights + graphdef, state = nnx.split(attn_unsharded) + flat_state = nnx.to_flat_state(state) + state_arrays = {k: np.array(v[...]) for k, v in zip(flat_state.paths, flat_state.leaves)} + + # Create input - (seq, hidden) + seq_len = 16 + hidden_size = cfg_unsharded.vision_config.hidden_size + x = jnp.ones((seq_len, hidden_size), dtype=jnp.float32) + + # Create RoPE cos/sin for vision + # VisionAttention expects (seq_len, head_dim) for cos/sin + head_dim = cfg_unsharded.vision_config.head_dim + # Simple RoPE: just use arange positions with full head_dim + positions = jnp.arange(seq_len, dtype=jnp.float32) + theta = 10000.0 + freqs = 1.0 / (theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) + angles = jnp.outer(positions, freqs) # (seq_len, head_dim/2) + # Repeat to match full head_dim + angles = jnp.concatenate([angles, angles], axis=-1) # (seq_len, head_dim) + cos_vals = jnp.cos(angles) + sin_vals = jnp.sin(angles) + + out_unsharded = attn_unsharded(x, (cos_vals, sin_vals)) + + # Create sharded model + jax.set_mesh(mesh) + cfg_sharded = get_test_config(use_fsdp=True, use_tp=True) + rngs = nnx.Rngs(42) + attn_sharded = modeling.Qwen3VLVisionAttention(cfg_sharded.vision_config, rngs=rngs) + + # Copy weights + sharded_graphdef, sharded_state = nnx.split(attn_sharded) + sharded_flat_state = nnx.to_flat_state(sharded_state) + for k, v in zip(sharded_flat_state.paths, sharded_flat_state.leaves): + if k in state_arrays: + v[...] = jnp.array(state_arrays[k]) + + # Recreate ALL inputs inside mesh context + x_sharded = jnp.ones((seq_len, hidden_size), dtype=jnp.float32) + positions_sharded = jnp.arange(seq_len, dtype=jnp.float32) + freqs_sharded = 1.0 / (theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) + angles_sharded = jnp.outer(positions_sharded, freqs_sharded) + angles_sharded = jnp.concatenate([angles_sharded, angles_sharded], axis=-1) + cos_sharded = jnp.cos(angles_sharded) + sin_sharded = jnp.sin(angles_sharded) + + out_sharded = attn_sharded(x_sharded, (cos_sharded, sin_sharded)) + + np.testing.assert_allclose( + np.array(out_unsharded), + np.array(out_sharded), + rtol=1e-4, + atol=1e-4, + err_msg="Sharded VisionAttention output differs from unsharded", + ) + print("Vision Attention: sharded vs unsharded numerical match ✓") + + jax.set_mesh(jax.make_mesh((1,), ("dummy",), axis_types=(AxisType.Explicit,))) + + +if __name__ == "__main__": + absltest.main()