From 290a540ad730e0de7c1450e34fe4bd39e096dddd Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Tue, 10 Feb 2026 08:57:47 +0000 Subject: [PATCH] qwen3 updates - no sharding like in gemma3 - replaced custom rmsnorm, einsum classes with nnx - sharding test --- bonsai/models/qwen3/modeling.py | 268 +++++++++--------- bonsai/models/qwen3/params.py | 12 +- bonsai/models/qwen3/tests/run_model.py | 4 +- .../models/qwen3/tests/test_outputs_qwen3.py | 18 +- .../models/qwen3/tests/test_sharding_qwen3.py | 46 +++ pyproject.toml | 4 +- scripts/run_selective_tests.py | 10 + 7 files changed, 219 insertions(+), 143 deletions(-) create mode 100644 bonsai/models/qwen3/tests/test_sharding_qwen3.py diff --git a/bonsai/models/qwen3/modeling.py b/bonsai/models/qwen3/modeling.py index 7a038a84..0d2b1d65 100644 --- a/bonsai/models/qwen3/modeling.py +++ b/bonsai/models/qwen3/modeling.py @@ -14,7 +14,7 @@ import dataclasses import math -from functools import partial +from enum import Enum from typing import TypeAlias import jax @@ -24,55 +24,48 @@ from jax.sharding import PartitionSpec, get_abstract_mesh, reshard from jaxtyping import Array, ArrayLike -_K_MASK = jnp.finfo(jnp.bfloat16).min -ShardingSpec = PartitionSpec +LARGE_NEGATIVE = jnp.finfo(jnp.bfloat16).min + + +class ShardMode(Enum): + FSDP = "fsdp" + TP = "tp" @dataclasses.dataclass(slots=True, frozen=True) class ShardConfig: - emb_vd: ShardingSpec - emb_dv: ShardingSpec - q_weight_ndh: ShardingSpec - kv_weight_ndh: ShardingSpec - o_weight_nhd: ShardingSpec - ffw_weight_df: ShardingSpec - ffw_weight_fd: ShardingSpec - rms_norm: ShardingSpec - act_btd: ShardingSpec - act_btf: ShardingSpec - act_btnh: ShardingSpec + emb_vd: PartitionSpec | None = None + emb_dv: PartitionSpec | None = None + q_weight_ndh: PartitionSpec | None = None + kv_weight_ndh: PartitionSpec | None = None + o_weight_nhd: PartitionSpec | None = None + ffw_weight_df: PartitionSpec | None = None + ffw_weight_fd: PartitionSpec | None = None + decoder_norm: PartitionSpec | None = None + act_btd: PartitionSpec | None = None + act_btf: PartitionSpec | None = None + act_btnh: PartitionSpec | None = None @staticmethod def no_sharding(): - """Configuration with no sharding (all None).""" - return ShardConfig( - emb_vd=P(None, None), - emb_dv=P(None, None), - q_weight_ndh=P(None, None, None), - kv_weight_ndh=P(None, None, None), - o_weight_nhd=P(None, None, None), - ffw_weight_df=P(None, None), - ffw_weight_fd=P(None, None), - rms_norm=P(None), - act_btd=P(None, None, None), - act_btf=P(None, None, None), - act_btnh=P(None, None, None, None), - ) + return ShardConfig() @staticmethod - def default(): + 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 ShardConfig( - emb_vd=P("tp", "fsdp"), - emb_dv=P("fsdp", "tp"), - q_weight_ndh=P("tp", "fsdp", None), - kv_weight_ndh=P("tp", "fsdp", None), - o_weight_nhd=P("tp", None, "fsdp"), - ffw_weight_df=P("fsdp", "tp"), - ffw_weight_fd=P("tp", "fsdp"), - rms_norm=P("tp"), - act_btd=P("fsdp", None, "tp"), - act_btf=P("fsdp", None, "tp"), - act_btnh=P("fsdp", None, "tp", None), + emb_vd=P(tp, fsdp), + emb_dv=P(fsdp, tp), + q_weight_ndh=P(tp, fsdp, None), + kv_weight_ndh=P(tp, fsdp, None), + o_weight_nhd=P(tp, None, fsdp), + ffw_weight_df=P(fsdp, tp), + ffw_weight_fd=P(tp, fsdp), + decoder_norm=P(tp), + act_btd=P(fsdp, None, tp), + act_btf=P(fsdp, None, tp), + act_btnh=P(fsdp, None, tp, None), ) @@ -93,15 +86,14 @@ class ModelConfig: shd_cfg: ShardConfig = ShardConfig.no_sharding() @classmethod - def _from_param(cls, use_sharding: bool, **kwargs): - if use_sharding: - kwargs["shd_cfg"] = ShardConfig.default() + def _from_param(cls, use_fsdp: bool = False, use_tp: bool = False, **kwargs): + if use_fsdp or use_tp: + kwargs["shd_cfg"] = ShardConfig.default(use_fsdp=use_fsdp, use_tp=use_tp) return cls(**kwargs) @classmethod - def qwen3_0_6b(cls, use_sharding: bool = False): # qwen3-0.6B + def qwen3_0_6b(cls, use_fsdp: bool = False, use_tp: bool = False): # qwen3-0.6B return cls._from_param( - use_sharding, num_layers=28, vocab_size=151936, emb_dim=1024, @@ -114,12 +106,13 @@ def qwen3_0_6b(cls, use_sharding: bool = False): # qwen3-0.6B rope_scaling_factor=8.0, local_rope_theta=1e4, tie_word_embeddings=True, + use_fsdp=use_fsdp, + use_tp=use_tp, ) @classmethod - def qwen3_1_7b(cls, use_sharding: bool = False): # qwen3-1.7B + def qwen3_1_7b(cls, use_fsdp: bool = False, use_tp: bool = False): # qwen3-1.7B return cls._from_param( - use_sharding, num_layers=28, vocab_size=151936, emb_dim=2048, @@ -132,12 +125,13 @@ def qwen3_1_7b(cls, use_sharding: bool = False): # qwen3-1.7B rope_scaling_factor=8.0, local_rope_theta=1e4, tie_word_embeddings=True, + use_fsdp=use_fsdp, + use_tp=use_tp, ) @classmethod - def qwen3_4b(cls, use_sharding: bool = False): # qwen3-4B + def qwen3_4b(cls, use_fsdp: bool = False, use_tp: bool = False): # qwen3-4B return cls._from_param( - use_sharding, num_layers=36, vocab_size=151936, emb_dim=2560, @@ -150,12 +144,13 @@ def qwen3_4b(cls, use_sharding: bool = False): # qwen3-4B rope_scaling_factor=8.0, local_rope_theta=1e4, tie_word_embeddings=True, + use_fsdp=use_fsdp, + use_tp=use_tp, ) @classmethod - def qwen3_8b(cls, use_sharding: bool = False): # qwen3-8B + def qwen3_8b(cls, use_fsdp: bool = False, use_tp: bool = False): # qwen3-8B return cls._from_param( - use_sharding, num_layers=36, vocab_size=151936, emb_dim=4096, @@ -168,12 +163,13 @@ def qwen3_8b(cls, use_sharding: bool = False): # qwen3-8B rope_scaling_factor=8.0, local_rope_theta=1e4, tie_word_embeddings=False, + use_fsdp=use_fsdp, + use_tp=use_tp, ) @classmethod - def qwen3_14b(cls, use_sharding: bool = False): # qwen3-14B + def qwen3_14b(cls, use_fsdp: bool = False, use_tp: bool = False): # qwen3-14B return cls._from_param( - use_sharding, num_layers=40, vocab_size=151936, emb_dim=5120, @@ -186,12 +182,14 @@ def qwen3_14b(cls, use_sharding: bool = False): # qwen3-14B rope_scaling_factor=8.0, local_rope_theta=1e4, tie_word_embeddings=False, + use_fsdp=use_fsdp, + use_tp=use_tp, ) -def shard(x: jnp.ndarray, s: ShardingSpec): +def shard(x: jnp.ndarray, s: PartitionSpec | None): mesh = get_abstract_mesh() - if not mesh.empty and len(mesh.axis_names) > 0: + if not mesh.empty and len(mesh.axis_names) > 0 and s is not None: return reshard(x, s) return x @@ -199,28 +197,18 @@ def shard(x: jnp.ndarray, s: ShardingSpec): class LayerCache(nnx.Module): def __init__(self, cfg: ModelConfig, batch_size: int, cache_size: int, dtype: jnp.dtype): cache_shape = (batch_size, cache_size, cfg.num_kv_heads, cfg.head_dim) - self.k_cache = shard(nnx.Cache(jnp.zeros(cache_shape, dtype=dtype)), cfg.shd_cfg.act_btnh) - self.v_cache = shard(nnx.Cache(jnp.zeros(cache_shape, dtype=dtype)), cfg.shd_cfg.act_btnh) + kv_shd = cfg.shd_cfg.act_btnh + self.k_cache = nnx.Cache(jnp.zeros(cache_shape, dtype=dtype, out_sharding=kv_shd)) + self.v_cache = nnx.Cache(jnp.zeros(cache_shape, dtype=dtype, out_sharding=kv_shd)) self.size = self.k_cache.shape[1] - batch_sharding = P(cfg.shd_cfg.act_btnh[0]) if cfg.shd_cfg.act_btnh else P(None) - self.start_ind = shard(nnx.Variable(-1 * jnp.ones((batch_size,), dtype=jnp.int32)), batch_sharding) - self.cur_ind = nnx.Variable(jnp.zeros((), dtype=jnp.int32)) # scalar for compute efficiency. + start_ind_shd = None if kv_shd is None else P(kv_shd[0]) + self.start_ind = nnx.Variable(-1 * jnp.ones((batch_size,), dtype=jnp.int32, out_sharding=start_ind_shd)) + self.cur_ind = nnx.Variable(jnp.zeros((), dtype=jnp.int32)) Cache: TypeAlias = list[LayerCache] -class Einsum(nnx.Module): - def __init__(self, einsum_str: str, shape: tuple[int, ...], *, shd: ShardingSpec, rngs: nnx.Rngs): - self.einsum_str = einsum_str - self.shape = shape - self.w = shard(nnx.Param(nnx.initializers.normal()(rngs.params(), shape)), shd) - - @jax.named_scope("einsum") - def __call__(self, x: ArrayLike) -> Array: - return jnp.einsum(self.einsum_str, x, self.w[...]) - - def _generate_pos_embeddings( positions: jax.Array, head_dim: int, rope_theta: int = 1_000_000 ) -> tuple[jax.Array, jax.Array]: @@ -241,18 +229,6 @@ def apply_rope(x: jax.Array, sin: jax.Array, cos: jax.Array) -> jax.Array: return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1).astype(x.dtype) -class RMSNorm(nnx.Module): - def __init__(self, dim: int, cfg: ModelConfig, *, rngs: nnx.Rngs): - self.scale = shard(nnx.Param(nnx.initializers.ones_init()(rngs.params(), dim)), cfg.shd_cfg.rms_norm) - self.norm_eps = cfg.norm_eps - - @jax.named_scope("rms_norm") - def __call__(self, x: Array) -> Array: - dtype = x.dtype - rms = jnp.sqrt(jnp.mean(jnp.astype(x, jnp.float32) ** 2, axis=-1, keepdims=True) + self.norm_eps) - return jnp.astype(self.scale[...] * x / rms, dtype) - - def count_left_pads(x: jax.Array) -> int: """Count left padding tokens.""" return jnp.sum(jnp.cumsum(x != 0, axis=-1) == 0, -1) @@ -272,34 +248,48 @@ def compute_positions_from_segment_ids(seg_ids): class Attention(nnx.Module): def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): self.shd_cfg = cfg.shd_cfg - einsum_fn = partial(Einsum, rngs=rngs) - self.q_proj = einsum_fn( - "BTD,DNH->BTNH", (cfg.emb_dim, cfg.num_heads, cfg.head_dim), shd=self.shd_cfg.q_weight_ndh + self.q_proj = nnx.Einsum( + "BTD,DNH->BTNH", + (cfg.emb_dim, cfg.num_heads, cfg.head_dim), + kernel_metadata={"out_sharding": self.shd_cfg.q_weight_ndh}, + rngs=rngs, ) - self.k_proj = einsum_fn( - "BSD,DKH->BSKH", (cfg.emb_dim, cfg.num_kv_heads, cfg.head_dim), shd=self.shd_cfg.kv_weight_ndh + self.k_proj = nnx.Einsum( + "BSD,DKH->BSKH", + (cfg.emb_dim, cfg.num_kv_heads, cfg.head_dim), + kernel_metadata={"out_sharding": self.shd_cfg.kv_weight_ndh}, + rngs=rngs, ) - self.v_proj = einsum_fn( - "BSD,DKH->BSKH", (cfg.emb_dim, cfg.num_kv_heads, cfg.head_dim), shd=self.shd_cfg.kv_weight_ndh + self.v_proj = nnx.Einsum( + "BSD,DKH->BSKH", + (cfg.emb_dim, cfg.num_kv_heads, cfg.head_dim), + kernel_metadata={"out_sharding": self.shd_cfg.kv_weight_ndh}, + rngs=rngs, ) - self.o_proj = einsum_fn( - "BTNH,NHD->BTD", (cfg.num_heads, cfg.head_dim, cfg.emb_dim), shd=self.shd_cfg.o_weight_nhd + self.o_proj = nnx.Einsum( + "BTNH,NHD->BTD", + (cfg.num_heads, cfg.head_dim, cfg.emb_dim), + kernel_metadata={"out_sharding": self.shd_cfg.o_weight_nhd}, + rngs=rngs, ) - self.q_norm = RMSNorm(cfg.head_dim, cfg, rngs=rngs) - self.k_norm = RMSNorm(cfg.head_dim, cfg, rngs=rngs) + scale_metadata = {"out_sharding": P() if self.shd_cfg.q_weight_ndh is not None else None} + self.q_norm = nnx.RMSNorm(cfg.head_dim, epsilon=cfg.norm_eps, scale_metadata=scale_metadata, rngs=rngs) + self.k_norm = nnx.RMSNorm(cfg.head_dim, epsilon=cfg.norm_eps, scale_metadata=scale_metadata, rngs=rngs) self.n_rep = cfg.num_heads // cfg.num_kv_heads self.scale = cfg.head_dim**-0.5 @jax.named_scope("attention") def __call__(self, x: Array, cache: LayerCache | None, segment_ids: Array) -> Array: - query_proj = shard(self.q_norm(self.q_proj(x)), self.shd_cfg.act_btnh) # [B, T, N, H] - key_proj = shard(self.k_norm(self.k_proj(x)), self.shd_cfg.act_btnh) # [B, T, K, H] - value_proj = shard(self.v_proj(x), self.shd_cfg.act_btnh) # [B, T, K, H] + shd = self.shd_cfg.act_btnh + query_proj = self.q_norm(self.q_proj(x, out_sharding=shd)) # [B, T, N, H] + key_proj = self.k_norm(self.k_proj(x, out_sharding=shd)) # [B, T, K, H] + value_proj = self.v_proj(x, out_sharding=shd) # [B, T, K, H] # RoPE and Cache Logic left_pads = count_left_pads(segment_ids) - left_pads = shard(left_pads, P(self.shd_cfg.act_btnh[0])) + if self.shd_cfg.act_btnh is not None: + left_pads = shard(left_pads, P(self.shd_cfg.act_btnh[0])) cache.start_ind.set_value(jnp.where(cache.start_ind[...] < 0, left_pads, cache.start_ind[...])) position_ids = compute_positions_from_segment_ids(segment_ids) + cache.cur_ind[...] sin, cos = _generate_pos_embeddings(position_ids, self.head_dim) @@ -308,8 +298,8 @@ def __call__(self, x: Array, cache: LayerCache | None, segment_ids: Array) -> Ar # Update K/V cache [B, S, K, H] slice_indices = (0, cache.cur_ind[...], 0, 0) - cache.v_cache[...] = jax.lax.dynamic_update_slice(cache.v_cache[...], value_proj, slice_indices) - cache.k_cache[...] = jax.lax.dynamic_update_slice(cache.k_cache[...], key_proj, slice_indices) + cache.v_cache.set_value(jax.lax.dynamic_update_slice(cache.v_cache[...], value_proj, slice_indices)) + cache.k_cache.set_value(jax.lax.dynamic_update_slice(cache.k_cache[...], key_proj, slice_indices)) b, t, n, h = query_proj.shape @@ -326,50 +316,57 @@ def __call__(self, x: Array, cache: LayerCache | None, segment_ids: Array) -> Ar segment_mask = kv_segment_ids[:, None, :] == segment_ids[:, :, None] final_mask = causal_mask & segment_mask # (B, T, S) attn_mask = final_mask[:, :, :, None, None] - attn_logits = jnp.where(attn_mask, attn_logits, _K_MASK) + attn_logits = jnp.where(attn_mask, attn_logits, LARGE_NEGATIVE) # Softmax attn_weights = jax.nn.softmax(attn_logits.astype(jnp.float32), axis=2).astype(attn_logits.dtype) qkv = jnp.einsum("BTSKG,BSKH->BTKGH", attn_weights, cache.v_cache[...]) qkv = qkv.reshape((b, t, n, h)) - cache.cur_ind[...] = cache.cur_ind[...] + t - return shard(self.o_proj(qkv), self.shd_cfg.act_btd) + cache.cur_ind.set_value(cache.cur_ind[...] + t) + return self.o_proj(qkv, out_sharding=self.shd_cfg.act_btd) @property def head_dim(self): - return self.o_proj.shape[1] + return self.o_proj.kernel.shape[1] @property def num_heads(self): - return self.q_proj.shape[1] + return self.q_proj.kernel.shape[1] @property def num_kv_heads(self): - return self.k_proj.shape[1] + return self.k_proj.kernel.shape[1] class MLP(nnx.Module): def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): self.shd_cfg = cfg.shd_cfg - linear = partial(nnx.Linear, use_bias=False, rngs=rngs) - self.gate_proj = shard(linear(cfg.emb_dim, cfg.mlp_dim), self.shd_cfg.ffw_weight_df) - self.up_proj = shard(linear(cfg.emb_dim, cfg.mlp_dim), self.shd_cfg.ffw_weight_df) - self.down_proj = shard(linear(cfg.mlp_dim, cfg.emb_dim), self.shd_cfg.ffw_weight_fd) + kernel_metadata = {"out_sharding": self.shd_cfg.ffw_weight_df} + self.gate_proj = nnx.Linear( + cfg.emb_dim, cfg.mlp_dim, kernel_metadata=kernel_metadata, use_bias=False, rngs=rngs + ) + self.up_proj = nnx.Linear(cfg.emb_dim, cfg.mlp_dim, kernel_metadata=kernel_metadata, use_bias=False, rngs=rngs) + kernel_metadata = {"out_sharding": self.shd_cfg.ffw_weight_fd} + self.down_proj = nnx.Linear( + cfg.mlp_dim, cfg.emb_dim, kernel_metadata=kernel_metadata, use_bias=False, rngs=rngs + ) @jax.named_scope("feed_forward") def __call__(self, x: ArrayLike) -> Array: - activations = nnx.silu(self.gate_proj(x)) * self.up_proj(x) - activations = shard(activations, self.shd_cfg.act_btf) - outputs = self.down_proj(activations) + ux = self.up_proj(x, out_sharding=self.shd_cfg.act_btf) + gx = nnx.silu(self.gate_proj(x, out_sharding=self.shd_cfg.act_btf)) + outputs = self.down_proj(gx * ux, out_sharding=self.shd_cfg.act_btf) return outputs class DecoderLayer(nnx.Module): def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): - self.input_layernorm = RMSNorm(cfg.emb_dim, cfg, rngs=rngs) + scale_metadata = {"out_sharding": cfg.shd_cfg.decoder_norm} + norm_kwargs = dict(num_features=cfg.emb_dim, epsilon=cfg.norm_eps, scale_metadata=scale_metadata, rngs=rngs) + self.input_layernorm = nnx.RMSNorm(**norm_kwargs) self.attn = Attention(cfg=cfg, rngs=rngs) - self.post_attention_layernorm = RMSNorm(cfg.emb_dim, cfg, rngs=rngs) + self.post_attention_layernorm = nnx.RMSNorm(**norm_kwargs) self.mlp = MLP(cfg=cfg, rngs=rngs) def __call__(self, x: Array, cache: LayerCache | None, segment_ids: Array) -> Array: @@ -381,35 +378,44 @@ def __call__(self, x: Array, cache: LayerCache | None, segment_ids: Array) -> Ar class Qwen3(nnx.Module): def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): - self.embedder = shard( - nnx.Embed(num_embeddings=cfg.vocab_size, features=cfg.emb_dim, dtype=jnp.bfloat16, rngs=rngs), - cfg.shd_cfg.emb_vd, + embedding_metadata = {"out_sharding": cfg.shd_cfg.emb_vd} + self.embedder = nnx.Embed( + num_embeddings=cfg.vocab_size, + features=cfg.emb_dim, + dtype=jnp.bfloat16, + embedding_metadata=embedding_metadata, + rngs=rngs, ) - self.out_emb_shd = None if get_abstract_mesh().empty else cfg.shd_cfg.act_btd + self.out_emb_shd = cfg.shd_cfg.act_btd self.layers = nnx.List([DecoderLayer(cfg=cfg, rngs=rngs) for _ in range(cfg.num_layers)]) - self.final_norm = RMSNorm(cfg.emb_dim, cfg, rngs=rngs) - self.lm_head = Einsum( - einsum_str="BTD,DV->BTV", shape=(cfg.emb_dim, cfg.vocab_size), shd=cfg.shd_cfg.emb_dv, rngs=rngs + scale_metadata = {"out_sharding": cfg.shd_cfg.decoder_norm} + self.final_norm = nnx.RMSNorm(cfg.emb_dim, epsilon=cfg.norm_eps, scale_metadata=scale_metadata, rngs=rngs) + self.lm_head = nnx.Einsum( + "BTD,DV->BTV", + (cfg.emb_dim, cfg.vocab_size), + kernel_metadata={"out_sharding": cfg.shd_cfg.emb_dv}, + rngs=rngs, ) - def init_cache( - self, cfg: ModelConfig, batch_size: int, token_len: int, generate_steps: int, dtype: jnp.dtype = jnp.bfloat16 - ) -> Cache: - cache_size = 2 ** math.ceil(math.log2(max(token_len + generate_steps, 1))) # Pad for a sharding-friendly size. - return [LayerCache(cfg, batch_size, cache_size, dtype) for _ in range(cfg.num_layers)] - - def __call__(self, tokens, segment_ids, cache, num_right_pads): + def __call__(self, tokens, *, cache, segment_ids): x = self.embedder.embedding[...].at[(tokens,)].get(out_sharding=self.out_emb_shd) for i, layer in enumerate(self.layers): x = layer(x, cache[i], segment_ids) - logits = self.lm_head(self.final_norm(x)) + logits = self.lm_head(self.final_norm(x), out_sharding=self.out_emb_shd) return logits +def init_cache( + cfg: ModelConfig, batch_size: int, token_len: int, generate_steps: int, dtype: jnp.dtype = jnp.bfloat16 +) -> Cache: + cache_size = 2 ** math.ceil(math.log2(max(token_len + generate_steps, 1))) # Pad for a sharding-friendly size. + return [LayerCache(cfg, batch_size, cache_size, dtype) for _ in range(cfg.num_layers)] + + @jax.jit def forward(model: nnx.Module, cache: Cache, tokens: Array, pad_id: int) -> tuple[Array, nnx.Cache]: segment_ids = 1 * (tokens != pad_id) num_right_pads = count_right_pads(tokens, pad_id) - logits = model(tokens, segment_ids, cache, num_right_pads) + logits = model(tokens, cache=cache, segment_ids=segment_ids) target_ind = tokens.shape[-1] - num_right_pads - 1 return logits[:, target_ind], cache diff --git a/bonsai/models/qwen3/params.py b/bonsai/models/qwen3/params.py index 8f032d82..93f5c54f 100644 --- a/bonsai/models/qwen3/params.py +++ b/bonsai/models/qwen3/params.py @@ -36,10 +36,10 @@ class Transform(Enum): # Mapping of torch_keys -> (nnx_keys, (permute_rule, reshape_rule)). return { r"model\.embed_tokens\.weight": ("embedder.embedding", Transform.EMBED), - r"model\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": (r"layers.\1.attn.q_proj.w", Transform.ATTN_Q), - r"model\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": (r"layers.\1.attn.k_proj.w", Transform.ATTN_KV), - r"model\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": (r"layers.\1.attn.v_proj.w", Transform.ATTN_KV), - r"model\.layers\.([0-9]+)\.self_attn\.o_proj\.weight": (r"layers.\1.attn.o_proj.w", Transform.ATTN_OUT), + r"model\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": (r"layers.\1.attn.q_proj.kernel", Transform.ATTN_Q), + r"model\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": (r"layers.\1.attn.k_proj.kernel", Transform.ATTN_KV), + r"model\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": (r"layers.\1.attn.v_proj.kernel", Transform.ATTN_KV), + r"model\.layers\.([0-9]+)\.self_attn\.o_proj\.weight": (r"layers.\1.attn.o_proj.kernel", Transform.ATTN_OUT), # mlp r"model\.layers\.([0-9]+)\.mlp\.gate_proj\.weight": (r"layers.\1.mlp.gate_proj.kernel", Transform.LINEAR), r"model\.layers\.([0-9]+)\.mlp\.up_proj\.weight": (r"layers.\1.mlp.up_proj.kernel", Transform.LINEAR), @@ -54,7 +54,7 @@ class Transform(Enum): r"layers.\1.post_attention_layernorm.scale", Transform.SCALE, ), - r"lm_head\.weight": ("lm_head.w", Transform.LINEAR), + r"lm_head\.weight": ("lm_head.kernel", Transform.LINEAR), } @@ -93,6 +93,6 @@ def create_model_from_safe_tensors(file_dir: str, cfg: model_lib.ModelConfig) -> raise RuntimeError(f"Encountered {len(conversion_errors)} weight conversion errors. Log:\n{full_error_log}") if cfg.tie_word_embeddings: - state_dict["lm_head"]["w"] = state_dict["embedder"]["embedding"].T + state_dict["lm_head"]["kernel"] = state_dict["embedder"]["embedding"].T gc.collect() return nnx.merge(graph_def, state_dict) diff --git a/bonsai/models/qwen3/tests/run_model.py b/bonsai/models/qwen3/tests/run_model.py index 1603ec99..9679064e 100644 --- a/bonsai/models/qwen3/tests/run_model.py +++ b/bonsai/models/qwen3/tests/run_model.py @@ -43,7 +43,7 @@ def run_model(): config = modeling.ModelConfig.qwen3_0_6b(use_sharding=False) mesh, batch_shd = None, None - # Enable sharding below if you have mtuliple devices. + # Enable sharding below if you have multiple devices. # model_ckpt_path = snapshot_download("Qwen/Qwen3-4B") # config = modeling.ModelConfig.qwen3_4b(use_sharding=True) # mesh = jax.make_mesh((2, 2), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) @@ -61,7 +61,7 @@ def run_model(): generate_steps = 32 model = params.create_model_from_safe_tensors(model_ckpt_path, config) - cache = model.init_cache(config, batch_size, token_len, generate_steps, dtype=jnp.float32) + cache = modeling.init_cache(config, batch_size, token_len, generate_steps, dtype=jnp.float32) key = jax.random.key(0) sampler = Sampler(temperature=1.0, top_p=0.8, top_k=10) diff --git a/bonsai/models/qwen3/tests/test_outputs_qwen3.py b/bonsai/models/qwen3/tests/test_outputs_qwen3.py index d9999104..2d73c68f 100644 --- a/bonsai/models/qwen3/tests/test_outputs_qwen3.py +++ b/bonsai/models/qwen3/tests/test_outputs_qwen3.py @@ -2,7 +2,7 @@ import jax.numpy as jnp import numpy as np import torch -from absl.testing import absltest +from absl.testing import absltest, parameterized from flax import nnx from huggingface_hub import snapshot_download from jax.sharding import AxisType @@ -27,7 +27,7 @@ def setUpClass(cls): ## models cls.torch_model = Qwen3ForCausalLM.from_pretrained(model_name, dtype="auto").eval() - cls.bonsai_config = modeling.ModelConfig.qwen3_0_6b(use_sharding=False) + cls.bonsai_config = modeling.ModelConfig.qwen3_0_6b() model_ckpt_path = snapshot_download("Qwen/Qwen3-0.6B") cls.mesh = jax.make_mesh(((1, 1)), ("fsdp", "tp"), axis_types=(AxisType.Explicit, AxisType.Explicit)) jax.set_mesh(cls.mesh) @@ -122,7 +122,7 @@ def _process_hf_tokens(self, query: list[str]): return model_inputs def _init_nnx_cache(self, batch_size: int): - return self.nnx_model.init_cache( + return modeling.init_cache( cfg=self.bonsai_config, batch_size=batch_size, token_len=10, generate_steps=32, dtype=jnp.float32 ) @@ -391,5 +391,17 @@ def test_full_batched(self): self._check_batched_logits(torch_inputs["left_pads"], torch_logits, nnx_logits) +class TestModuleOtherConfigs(parameterized.TestCase): + @parameterized.named_parameters( + ("qwen3_0_6b", "qwen3_0_6b"), + ("qwen3_1_7b", "qwen3_1_7b"), + ("qwen3_4b", "qwen3_4b"), + ("qwen3_8b", "qwen3_8b"), + ("qwen3_14b", "qwen3_14b"), + ) + def test_config(self, config_name): + _ = getattr(modeling.ModelConfig, config_name)() + + if __name__ == "__main__": absltest.main() diff --git a/bonsai/models/qwen3/tests/test_sharding_qwen3.py b/bonsai/models/qwen3/tests/test_sharding_qwen3.py new file mode 100644 index 00000000..510a462d --- /dev/null +++ b/bonsai/models/qwen3/tests/test_sharding_qwen3.py @@ -0,0 +1,46 @@ +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest +from flax import nnx +from jax import P +from jax.sharding import AxisType, NamedSharding + +from bonsai.models.qwen3 import modeling + + +@unittest.skipIf(jax.device_count() < 4, "At least 4 devices required") +class TestSharding(absltest.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + + cls.mesh = jax.make_mesh(((2, 2)), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(cls.mesh) + + cls.bonsai_config = modeling.ModelConfig.qwen3_0_6b(True, True) + cls.bonsai_model = modeling.Qwen3(cls.bonsai_config, rngs=nnx.Rngs(0)) + + def test_full(self): + nm = self.bonsai_model + fsdp = modeling.ShardMode.FSDP.value + + batch_size = 2 # should be evenly divisible to num devices for fsdp axis + num_tokens = 128 + n_text = jax.device_put( + np.arange(batch_size * num_tokens).reshape(batch_size, -1), + device=P(fsdp), + ) + segment_ids = jnp.ones((batch_size, num_tokens), out_sharding=P(fsdp)) + cache = modeling.init_cache(self.bonsai_config, batch_size, num_tokens, 1, jnp.float32) + + out = nm(n_text, cache=cache, segment_ids=segment_ids) + assert isinstance(out.sharding, NamedSharding) + assert out.sharding.spec == self.bonsai_config.shd_cfg.act_btd + + +if __name__ == "__main__": + absltest.main() diff --git a/pyproject.toml b/pyproject.toml index 92cf2d82..cedfaeaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,9 @@ classifiers = [ dependencies = [ "jax >= 0.8.0, < 0.9.0", "jaxlib >= 0.8.0, < 0.9.0", - "flax >= 0.12.0, < 0.13.0", + # Set flax from source as long as bonsai has no stable release + # "flax >= 0.12.0, < 0.13.0", + "flax @ git+https://github.com/google/flax.git", "jaxtyping>=0.2.20", "jinja2>=3.0.0", "huggingface-hub>=0.20.0", diff --git a/scripts/run_selective_tests.py b/scripts/run_selective_tests.py index 339230d6..ba1e85c9 100644 --- a/scripts/run_selective_tests.py +++ b/scripts/run_selective_tests.py @@ -122,6 +122,16 @@ def main(): test_files = find_tests_in_directory(Path(target)) if test_files: subprocess.run(["uv", "run", "pytest", target, "-v"], check=True) + if any("sharding" in path.stem for path in test_files): + subprocess.run( + ["uv", "run", "pytest", target, "-v", "-k", "TestSharding"], + check=True, + env=os.environ + | { + "JAX_NUM_CPU_DEVICES": "4", + "JAX_PLATFORMS": "cpu", + }, + ) else: print(f"No tests to run for {target}")