From fe6b7009c2ab8ac101a58afc682b1b1d3d71481b Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Thu, 22 Jan 2026 13:42:29 +0900 Subject: [PATCH 01/18] feat(llama32): implement llama3.2 model architecture feat(llama32): enhance LlamaShardCfg with activation and logits partition specs --- bonsai/models/llama32/modeling.py | 640 ++++++++++++++++++ bonsai/models/llama32/params.py | 169 +++++ bonsai/models/llama32/tests/run_model.py | 120 ++++ bonsai/models/llama32/tests/run_model_base.py | 113 ++++ 4 files changed, 1042 insertions(+) create mode 100644 bonsai/models/llama32/modeling.py create mode 100644 bonsai/models/llama32/params.py create mode 100644 bonsai/models/llama32/tests/run_model.py create mode 100644 bonsai/models/llama32/tests/run_model_base.py diff --git a/bonsai/models/llama32/modeling.py b/bonsai/models/llama32/modeling.py new file mode 100644 index 00000000..ddb1ea4b --- /dev/null +++ b/bonsai/models/llama32/modeling.py @@ -0,0 +1,640 @@ +# Copyright 2026 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import math +from functools import partial +from typing import cast, TypeAlias +from enum import Enum + +import jax +from jax import P +from jax.sharding import PartitionSpec +import jax.numpy as jnp +from jaxtyping import Array, ArrayLike +from flax import nnx +from flax.nnx.nn.linear import default_embed_init + + +class ShardMode(Enum): + """ Sharding Modes for Model Parameters """ + FSDP = "fsdp" + TP = "tp" + + +@dataclasses.dataclass(slots=True, frozen=True) +class LlamaShardCfg: + # Embedding + emb_weight: PartitionSpec + activation: PartitionSpec + logits: PartitionSpec + + # Attention + q_proj: PartitionSpec + k_proj: PartitionSpec + v_proj: PartitionSpec + o_proj: PartitionSpec + + attn_logits: PartitionSpec + attn_out: PartitionSpec + + cache: PartitionSpec + + # MLP + gate_proj: PartitionSpec + up_proj: PartitionSpec + down_proj: PartitionSpec + + # Head + lm_head: PartitionSpec + + @classmethod + def no_sharding(cls) -> "LlamaShardCfg": + return cls.default(use_fsdp=False, use_tp=False) + + @classmethod + def default(cls, use_fsdp: bool, use_tp: bool) -> "LlamaShardCfg": + fsdp = ShardMode.FSDP.value if use_fsdp else None + tp = ShardMode.TP.value if use_tp else None + + return cls( + emb_weight=P(None, tp), + activation=P(fsdp, None, tp), + logits=P(fsdp, None, tp), + q_proj=P(fsdp, tp), + k_proj=P(fsdp, tp), + v_proj=P(fsdp, tp), + o_proj=P(tp, fsdp), + attn_logits=P(fsdp, None, None, tp, None), + attn_out=P(fsdp, None, tp, None, None), + cache=P(fsdp, None, tp, None), + gate_proj=P(fsdp, tp), + up_proj=P(fsdp, tp), + down_proj=P(tp, fsdp), + lm_head=P(None, tp), + ) + + +@dataclasses.dataclass(frozen=True) +class RopeScalingConfig: + factor: float + low_freq_factor: float = 1.0 + high_freq_factor: float = 4.0 + original_max_position_embeddings: int = 8192 + + + +@dataclasses.dataclass(frozen=True) +class ModelConfig: + vocab_size: int + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + head_dim: int + num_key_value_heads: int + max_position_embeddings: int + rms_norm_eps: float + rope_theta: float + rope_scaling: RopeScalingConfig | None + tie_word_embeddings: bool + shd_cfg: LlamaShardCfg + dtype: jnp.dtype = jnp.bfloat16 + + @classmethod + def llama3_2_1b(cls, use_fsdp: bool, use_tp: bool) -> "ModelConfig": + return cls( + vocab_size=128256, + hidden_size=2048, + intermediate_size=8192, + num_hidden_layers=16, + num_attention_heads=32, + head_dim=64, + num_key_value_heads=8, + max_position_embeddings=131072, + rms_norm_eps=1e-5, + rope_theta=500000.0, + rope_scaling=RopeScalingConfig(factor=32.0), + tie_word_embeddings=True, + shd_cfg=LlamaShardCfg.default(use_fsdp, use_tp), + ) + + @classmethod + def llama3_2_3b(cls, use_fsdp: bool, use_tp: bool) -> "ModelConfig": + return cls( + vocab_size=128256, + hidden_size=3072, + intermediate_size=8192, + num_hidden_layers=28, + num_attention_heads=24, + head_dim=128, + num_key_value_heads=8, + max_position_embeddings=131072, + rms_norm_eps=1e-5, + rope_theta=500000.0, + rope_scaling=RopeScalingConfig(factor=32.0), + tie_word_embeddings=True, + shd_cfg=LlamaShardCfg.default(use_fsdp, use_tp), + ) + + +class LayerCache(nnx.Module): + """Key-Value Cache for Attention Layers""" + + def __init__(self, cfg: ModelConfig, batch_size: int, cache_size: int, dtype: jnp.dtype): + cache_shape = (batch_size, cache_size, cfg.num_key_value_heads, cfg.head_dim) + kv_shd = cfg.shd_cfg.cache + 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] + self.start_ind = nnx.Variable(-1 * jnp.ones((batch_size,), dtype=jnp.int32, out_sharding=P(kv_shd[0]))) + self.cur_ind = nnx.Variable(jnp.zeros((), dtype=jnp.int32)) + + +Cache: TypeAlias = list[LayerCache] + + +class LlamaRMSNorm(nnx.Module): + """Root Mean Square Layer Normalization""" + + def __init__(self, hidden_size: int, eps: float, rngs: nnx.Rngs): + self.scale = nnx.Param(nnx.initializers.ones_init()(rngs.params(), (hidden_size,))) + self.norm_eps = eps + + @jax.named_scope("rms_norm") + def __call__(self, x: Array) -> Array: + dtype = x.dtype + x_fp32 = x.astype(jnp.float32) + variance = jnp.mean(jnp.square(x_fp32), axis=-1, keepdims=True) + inv_rms = jax.lax.rsqrt(variance + self.norm_eps) + + return (self.scale[...] * x_fp32 * inv_rms).astype(dtype) + +# TODO: Replace with nnx.Linear once explicit sharding is supported. +class ShardedLinear(nnx.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + sharding: PartitionSpec, + *, + use_bias: bool = True, + dtype=jnp.bfloat16, + rngs: nnx.Rngs, + ): + kernel_initializer = jax.nn.initializers.lecun_normal() + self.kernel = nnx.Param( + kernel_initializer(rngs.params(), (in_dim, out_dim), dtype=dtype, out_sharding=sharding) + ) + if use_bias: + self.bias = nnx.Param(jnp.zeros((out_dim,), dtype=dtype)) + else: + self.bias = None + + def __call__(self, x: ArrayLike, *, out_sharding: PartitionSpec | None = None) -> Array: + out = jnp.matmul(x, self.kernel[...], out_sharding=out_sharding) + if self.bias is None: + return out + return out + self.bias[...] + + +class ShardedEmbedding(nnx.Embed): + """Sharded Embedding Layer""" + + def __call__(self, inputs: Array, *, out_sharding: PartitionSpec | None = None) -> 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)) + return embedding.at[inputs].get(out_sharding=out_sharding) + + def decode(self, query: Array, *, out_sharding: PartitionSpec | None = None) -> Array: + query, embedding = self.promote_dtype((query, self.embedding[...]), dtype=self.dtype) + return jnp.dot(query, embedding.T, out_sharding=out_sharding) + + + +class LlamaMLP(nnx.Module): + """Feed Forward Network""" + + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): + super().__init__() + self.config = cfg + self.hidden_size = cfg.hidden_size + self.intermediate_size = cfg.intermediate_size + self.gate_proj = ShardedLinear( + self.hidden_size, + self.intermediate_size, + sharding=cfg.shd_cfg.gate_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + self.up_proj = ShardedLinear(self.hidden_size, self.intermediate_size, sharding=cfg.shd_cfg.up_proj, use_bias=False, dtype=cfg.dtype, rngs=rngs) + self.down_proj = ShardedLinear(self.intermediate_size, self.hidden_size, sharding=cfg.shd_cfg.down_proj, use_bias=False, dtype=cfg.dtype, rngs=rngs) + + + @jax.named_scope('feed_forward') + def __call__(self, x: ArrayLike) -> Array: + act_shd = self.config.shd_cfg.activation + gated = nnx.silu(self.gate_proj(x, out_sharding=act_shd)) * self.up_proj(x, out_sharding=act_shd) + return self.down_proj(gated, out_sharding=act_shd) + + +def _generate_pos_embeddings( + positions: Array, + head_dim: int, + rope_theta: float = 500000.0, + rope_scaling: RopeScalingConfig | None = None, +) -> tuple[Array, Array]: + """Generate sin and cos for rotary position embeddings.""" + if rope_scaling is None: + rope_scaling = RopeScalingConfig(factor=1.0) + + inv_freq = 1.0 / ( + rope_theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim) + ) + + factor = rope_scaling.factor + low_freq_factor = rope_scaling.low_freq_factor + high_freq_factor = rope_scaling.high_freq_factor + old_context_len = rope_scaling.original_max_position_embeddings + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + + wavelen = (2.0 * math.pi) / inv_freq + inv_freq_llama = jnp.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq) + smooth_factor = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + smoothed_inv_freq = (1.0 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama + is_medium_freq = jnp.logical_and( + wavelen >= high_freq_wavelen, + wavelen <= low_freq_wavelen, + ) + inv_freq_llama = jnp.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + + positions_f32 = positions.astype(jnp.float32) + sinusoid_inp = jnp.einsum( + "BT,k->BTk", positions_f32, inv_freq_llama, precision=jax.lax.Precision.HIGHEST + ) + return jnp.sin(sinusoid_inp), jnp.cos(sinusoid_inp) + + +def apply_rope(x: Array, sin: Array, cos: Array) -> Array: + """Apply rotary position embeddings to input tensor.""" + if x.ndim != 4 or sin.ndim != 3 or cos.ndim != 3: + raise ValueError( + "apply_rope expects x.ndim == 4 and sin.ndim == cos.ndim == 3; " + f"got x.ndim={x.ndim}, sin.ndim={sin.ndim}, cos.ndim={cos.ndim}" + ) + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + # [B, T, head_dim] -> [B, h, T, head_dim] + sin, cos = sin[:, :, None, :], cos[:, :, None, :] + return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1).astype(x.dtype) + + +def count_left_pads(x: Array) -> Array: + """Count left padding tokens per batch element.""" + return jnp.sum(jnp.cumsum(x != 0, axis=-1) == 0, -1) + + +def count_right_pads(x: Array, pad_id: int) -> Array: + """Count right padding tokens per batch element.""" + all_pad = jnp.all(x == pad_id, axis=1) + right_pad = jnp.argmin(jnp.flip(x == pad_id, axis=1).astype(jnp.int32), axis=1) + max_len = jnp.full_like(right_pad, x.shape[1]) + return jnp.where(all_pad, max_len, right_pad) + + +def count_right_pads_from_mask(attn_mask: Array) -> Array: + """Count right padding tokens from a 0/1 attention mask.""" + mask = attn_mask.astype(jnp.int32) + all_pad = jnp.all(mask == 0, axis=1) + right_pad = jnp.argmax(jnp.flip(mask, axis=1), axis=1) + max_len = jnp.full_like(right_pad, mask.shape[1]) + return jnp.where(all_pad, max_len, right_pad) + +def compute_positions_from_segment_ids(seg_ids: Array) -> Array: + """Compute position ids from segment ids.""" + seg_ids = seg_ids.astype(jnp.int32) + pad_sentinel = 2**30 + + def step(carry: tuple[Array, Array], seg_id: Array) -> tuple[tuple[Array, Array], Array]: + prev_seg, prev_pos = carry + is_pad = seg_id == 0 + is_new = seg_id != prev_seg + zero = jnp.zeros_like(seg_id) + pos = jnp.where(is_pad, zero, jnp.where(is_new, zero, prev_pos + 1)) + pad_val = jnp.full_like(seg_id, pad_sentinel) + out = jnp.where(is_pad, pad_val, pos) + new_prev_seg = jnp.where(is_pad, zero, seg_id) + new_prev_pos = jnp.where(is_pad, zero, pos) + return (new_prev_seg, new_prev_pos), out + + base = jnp.zeros_like(seg_ids[:, 0]) + init = (base, base) + _, out = jax.lax.scan(step, init, seg_ids.T) + return cast(Array, out.T) + + +def sharded_attention( + q: Array, + k: Array, + v: Array, + attn_mask: Array | None, + scale: float, + *, + attn_logit_sharding: PartitionSpec, + out_sharding: PartitionSpec, +) -> Array: + """Compute scaled dot-product attention with optional masking.""" + attn_logits = jnp.einsum("BTKGH,BSKH->BTSKG", q, k, out_sharding=attn_logit_sharding) * scale + + if attn_mask is not None: + if attn_mask.ndim == 2: + attn_mask = attn_mask[:, None, :, None, None] # [B, 1, S, 1, 1] + elif attn_mask.ndim == 3: + attn_mask = attn_mask[:, :, :, None, None] # [B, T, S, 1, 1] + else: + raise ValueError(f"attn_mask must be rank-2 or rank-3, got {attn_mask.ndim}") + attn_logits = jnp.where(attn_mask, attn_logits, jnp.finfo(attn_logits.dtype).min) + + attn_weights = jax.nn.softmax(attn_logits.astype(jnp.float32), axis=2).astype(attn_logits.dtype) + attn_output = jnp.einsum("BTSKG,BSKH->BTKGH", attn_weights, v, out_sharding=out_sharding) + + return attn_output.astype(q.dtype) + + + + +class LlamaAttention(nnx.Module): + """Multi-Head Self Attention""" + + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): + self.config = cfg + self.hidden_size = cfg.hidden_size + self.head_dim = cfg.head_dim + self.num_heads = cfg.num_attention_heads + self.num_kv_heads = cfg.num_key_value_heads + self.n_rep = self.num_heads // self.num_kv_heads + self.scale = self.head_dim**-0.5 + + self.q_proj = ShardedLinear( + self.hidden_size, + self.num_heads * self.head_dim, + sharding=cfg.shd_cfg.q_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + self.k_proj = ShardedLinear( + self.hidden_size, + self.num_kv_heads * self.head_dim, + sharding=cfg.shd_cfg.k_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + self.v_proj = ShardedLinear( + self.hidden_size, + self.num_kv_heads * self.head_dim, + sharding=cfg.shd_cfg.v_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + self.o_proj = ShardedLinear( + self.num_heads * self.head_dim, + self.hidden_size, + sharding=cfg.shd_cfg.o_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + + def _make_cache_mask(self, cache: LayerCache, t: int) -> Array: + """Create attention mask for cached key-value states.""" + + q_pos = cache.cur_ind[...] + jnp.arange(t, dtype=jnp.int32)[None, :] - cache.start_ind[...][:, None] + ts = jnp.arange(cache.size, dtype=jnp.int32) + kv_valid = (ts[None, :] >= cache.start_ind[...][:, None]) & (ts[None, :] < cache.cur_ind[...] + t) + k_pos = ts[None, :] - cache.start_ind[...][:, None] + causal_mask = k_pos[:, None, :] <= q_pos[:, :, None] + return causal_mask & kv_valid[:, None, :] + + def _make_stateless_mask(self, segment_ids: Array, t: int) -> Array: + """Create attention mask without cache.""" + q_pos = jnp.arange(t, dtype=jnp.int32)[None, :] + k_pos = jnp.arange(t, dtype=jnp.int32)[None, :] + causal_mask = k_pos[:, None, :] <= q_pos[:, :, None] + segment_mask = segment_ids[:, :, None] == segment_ids[:, None, :] + return causal_mask & segment_mask + + @jax.named_scope("attention") + def __call__( + self, + x: Array, + segment_ids: Array, + attn_mask: Array | None, + cache: LayerCache | None = None, + ) -> Array: + b, t, _ = x.shape + + # Project to Q, K, V and reshape to [B, T, N/K, H] + act_shd = self.config.shd_cfg.activation + q = self.q_proj(x, out_sharding=act_shd).reshape((b, t, self.num_heads, self.head_dim)) + k = self.k_proj(x, out_sharding=act_shd).reshape((b, t, self.num_kv_heads, self.head_dim)) + v = self.v_proj(x, out_sharding=act_shd).reshape((b, t, self.num_kv_heads, self.head_dim)) + + # Apply RoPE + position_ids = compute_positions_from_segment_ids(segment_ids) + if cache is not None: + left_pads = count_left_pads(segment_ids) + cache.start_ind[...] = jnp.where(cache.start_ind[...] < 0, left_pads, cache.start_ind[...]) + position_ids = position_ids + cache.cur_ind[...] + + sin, cos = _generate_pos_embeddings( + position_ids, self.head_dim, self.config.rope_theta, self.config.rope_scaling + ) + q = apply_rope(q, sin, cos) + k = apply_rope(k, sin, cos) + + if cache is not None: + # Update cache with new keys and values + slice_indices = (0, cache.cur_ind[...], 0, 0) + cache_dtype = cache.k_cache[...].dtype + k = k.astype(cache_dtype) + v = v.astype(cache_dtype) + cache.k_cache[...] = jax.lax.dynamic_update_slice(cache.k_cache[...], k, slice_indices) + cache.v_cache[...] = jax.lax.dynamic_update_slice(cache.v_cache[...], v, slice_indices) + k = cache.k_cache[...] + v = cache.v_cache[...] + + # Reshape query for GQA: [B, T, K, n_rep, H] + query_proj_gqa = q.reshape((b, t, self.num_kv_heads, self.n_rep, self.head_dim)) + + if attn_mask is None: + if cache is not None: + attn_mask = self._make_cache_mask(cache, t) + else: + attn_mask = self._make_stateless_mask(segment_ids, t) + + # Attention (with optional sharding) -> [B, T, K, n_rep, H] + attn_output = sharded_attention( + query_proj_gqa, + k, + v, + attn_mask, + self.scale, + attn_logit_sharding=self.config.shd_cfg.attn_logits, + out_sharding=self.config.shd_cfg.attn_out, + ) + + if cache is not None: + cache.cur_ind[...] = cache.cur_ind[...] + t + + # Reshape back to [B, T, N * H] for the output projection. + attn_output = attn_output.reshape((b, t, self.num_heads * self.head_dim)) + + # Output projection: [B, T, D] + return self.o_proj(attn_output, out_sharding=act_shd) + + +class LlamaDecoderLayer(nnx.Module): + """Llama Decoder Layer""" + + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): + self.self_attn = LlamaAttention(cfg, rngs=rngs) + self.mlp = LlamaMLP(cfg, rngs=rngs) + self.input_layernorm = LlamaRMSNorm(cfg.hidden_size, cfg.rms_norm_eps, rngs=rngs) + self.post_attention_layernorm = LlamaRMSNorm(cfg.hidden_size, cfg.rms_norm_eps, rngs=rngs) + + @jax.named_scope("decoder_layer") + def __call__( + self, + x: Array, + segment_ids: Array, + attn_mask: Array | None = None, + cache: LayerCache | None = None, + ) -> Array: + # Self-Attention Block + normed_x = self.input_layernorm(x) + attn_output = self.self_attn(normed_x, segment_ids, attn_mask, cache=cache) + x = x + attn_output + + # Feed-Forward Block + normed_x = self.post_attention_layernorm(x) + mlp_output = self.mlp(normed_x) + x = x + mlp_output + + return x + + +class Llama(nnx.Module): + """Llama Model""" + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): + self.config = cfg + embed_init = partial(default_embed_init, out_sharding=cfg.shd_cfg.emb_weight) + self.embedder = ShardedEmbedding( + num_embeddings=cfg.vocab_size, + features=cfg.hidden_size, + dtype=cfg.dtype, + embedding_init=embed_init, + rngs=rngs, + ) + self.layers = nnx.List([LlamaDecoderLayer(cfg, rngs=rngs) for _ in range(cfg.num_hidden_layers)]) + self.final_norm = LlamaRMSNorm(cfg.hidden_size, cfg.rms_norm_eps, rngs=rngs) + if cfg.tie_word_embeddings: + self.lm_head = None + else: + self.lm_head = ShardedLinear( + cfg.hidden_size, + cfg.vocab_size, + sharding=cfg.shd_cfg.lm_head, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + + def init_cache( + self, + cfg: ModelConfig, + batch_size: int, + token_len: int, + generate_steps: int, + max_cache_len: int = 4096, + ) -> Cache: + target_len = min(max_cache_len, token_len + generate_steps) + cache_size = 2 ** math.ceil(math.log2(max(target_len, 1))) + cache_size = min(cache_size, max_cache_len) + cache_dtype = self.layers[0].self_attn.k_proj.kernel[...].dtype + return [LayerCache(cfg, batch_size, cache_size, cache_dtype) for _ in range(cfg.num_hidden_layers)] + + def __call__( + self, + tokens: Array, + segment_ids: Array, + cache: Cache | None, + attn_mask: Array | None = None, + ) -> Array: + x = self.embedder(tokens, out_sharding=self.config.shd_cfg.activation) + for i, layer in enumerate(self.layers): + layer_cache = cache[i] if cache is not None else None + x = layer(x, segment_ids, attn_mask=attn_mask, cache=layer_cache) + hidden = self.final_norm(x) + if self.config.tie_word_embeddings: + logits = self.embedder.decode(hidden, out_sharding=self.config.shd_cfg.logits) + else: + assert self.lm_head is not None + logits = self.lm_head(hidden, out_sharding=self.config.shd_cfg.logits) + return logits + + +@jax.jit +def forward( + model: nnx.Module, + cache: Cache, + tokens: Array, + pad_id: int, + attention_mask: Array | None = None, + segment_ids: Array | None = None, +) -> tuple[Array, Cache]: + + # Use attention_mask when available because pad_id can equal eos_id, which would + # misclassify real tokens as padding if we rely on tokens != pad_id. + if attention_mask is not None and segment_ids is not None: + raise ValueError("Provide only one of attention_mask or segment_ids.") + + + if segment_ids is None: + if attention_mask is None: + segment_ids = 1 * (tokens != pad_id) + pad_mask = segment_ids + num_right_pads = count_right_pads(tokens, pad_id) + else: + segment_ids = attention_mask.astype(jnp.int32) + pad_mask = segment_ids + num_right_pads = count_right_pads_from_mask(pad_mask) + else: + segment_ids = segment_ids.astype(jnp.int32) + pad_mask = (segment_ids != 0).astype(jnp.int32) + num_right_pads = count_right_pads_from_mask(pad_mask) + + logits = model(tokens, segment_ids, cache, attn_mask=None) + target_ind = tokens.shape[-1] - num_right_pads - 1 + batch_idx = jnp.arange(tokens.shape[0]) + return logits[batch_idx, target_ind], cache diff --git a/bonsai/models/llama32/params.py b/bonsai/models/llama32/params.py new file mode 100644 index 00000000..c4f16c2a --- /dev/null +++ b/bonsai/models/llama32/params.py @@ -0,0 +1,169 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import re +from enum import Enum + +import jax +import safetensors +from etils import epath +from flax import nnx + +from bonsai.models.llama32 import modeling as model_lib + + +def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig): + class Transform(Enum): + """Transformations for model parameters""" + + BIAS = None + LINEAR = ((1, 0), None, False) + EMBED = None + SCALE = None + + mapping = { + r"model\.embed_tokens\.weight": ("embedder.embedding", Transform.EMBED), + r"model\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": ( + r"layers.\1.self_attn.q_proj.kernel", + Transform.LINEAR, + ), + r"model\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": ( + r"layers.\1.self_attn.k_proj.kernel", + Transform.LINEAR, + ), + r"model\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": ( + r"layers.\1.self_attn.v_proj.kernel", + Transform.LINEAR, + ), + r"model\.layers\.([0-9]+)\.self_attn\.o_proj\.weight": ( + r"layers.\1.self_attn.o_proj.kernel", + Transform.LINEAR, + ), + 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), + r"model\.layers\.([0-9]+)\.mlp\.down_proj\.weight": (r"layers.\1.mlp.down_proj.kernel", Transform.LINEAR), + r"model\.layers\.([0-9]+)\.input_layernorm\.weight": (r"layers.\1.input_layernorm.scale", Transform.SCALE), + r"model\.layers\.([0-9]+)\.post_attention_layernorm\.weight": ( + r"layers.\1.post_attention_layernorm.scale", + Transform.SCALE, + ), + r"model\.norm\.weight": ("final_norm.scale", Transform.SCALE), + } + if not cfg.tie_word_embeddings: + mapping[r"lm_head\.weight"] = ("lm_head.kernel", Transform.LINEAR) + return mapping + + +def _torch_key_to_jax_key(mapping, source_key): + subs = [ + (re.sub(pat, repl, source_key), reshape) + for pat, (repl, reshape) in mapping.items() + if re.match(pat, source_key) + ] + if len(subs) == 0: + return None, None + if len(subs) != 1: + raise ValueError(f"Expected at most one key match for {source_key}, found {len(subs)}: {subs}") + return subs[0] + + +def _assign_weights(keys, tensor, state_dict, st_key, transform, sharding_dict): + """Recursively descend into state_dict and assign the (possibly permuted/reshaped) tensor.""" + key, *rest = keys + if not rest: + if transform is not None: + permute, reshape, reshape_first = transform + if reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if permute: + tensor = tensor.transpose(permute) + if not reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if tensor.shape != state_dict[key].shape: + raise ValueError(f"Shape mismatch for {st_key}: {tensor.shape} vs {state_dict[key].shape}") + 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], st_key, transform, next_sharding) + + +def _stoi(s): + try: + return int(s) + except ValueError: + return s + + +def create_model_from_safe_tensors( + file_dir: str, cfg: model_lib.ModelConfig, mesh: jax.sharding.Mesh | None = None +) -> model_lib.Llama: + """Load tensors from the safetensors file and create a Llama model (memory-optimized).""" + files = list(epath.Path(file_dir).expanduser().glob("*.safetensors")) + if not files: + raise ValueError(f"No safetensors found in {file_dir}") + + llama = nnx.eval_shape(lambda: model_lib.Llama(cfg, rngs=nnx.Rngs(params=0))) + graph_def, abs_state = nnx.split(llama) + state_dict = abs_state.to_pure_dict() + sharding = nnx.get_named_sharding(abs_state, mesh).to_pure_dict() if mesh is not None else None + + key_mapping = _get_key_and_transform_mapping(cfg) + conversion_errors = [] + unexpected_biases = [] + + for f in files: + with safetensors.safe_open(f, framework="numpy") as sf: + for torch_key in sf.keys(): + if torch_key.endswith(".bias"): + unexpected_biases.append(torch_key) + continue + + # When embeddings are tied, lm_head weights are derived from embedder. + if cfg.tie_word_embeddings and torch_key == "lm_head.weight": + continue + + tensor = sf.get_tensor(torch_key) + jax_key, transform = _torch_key_to_jax_key(key_mapping, torch_key) + if jax_key is None: + continue + if transform is None: + raise ValueError(f"Missing transform for {torch_key}") + + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights(keys, tensor, state_dict, torch_key, transform.value, 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 unexpected_biases: + bias_list = "\n".join(unexpected_biases) + raise RuntimeError(f"Unexpected bias parameters found for Llama (biases are disabled):\n{bias_list}") + + if conversion_errors: + full_error_log = "\n".join(conversion_errors) + raise RuntimeError(f"Encountered {len(conversion_errors)} weight conversion errors. Log:\n{full_error_log}") + + if cfg.tie_word_embeddings and "lm_head" in state_dict: + state_dict["lm_head"]["kernel"] = state_dict["embedder"]["embedding"].T + + gc.collect() + return nnx.merge(graph_def, state_dict) diff --git a/bonsai/models/llama32/tests/run_model.py b/bonsai/models/llama32/tests/run_model.py new file mode 100644 index 00000000..b75eb31b --- /dev/null +++ b/bonsai/models/llama32/tests/run_model.py @@ -0,0 +1,120 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import time + +import jax +import jax.numpy as jnp +import numpy as np +from huggingface_hub import snapshot_download +from jax.sharding import AxisType + +from transformers import AutoTokenizer + +from bonsai.models.llama32 import modeling, params +from bonsai.utils import Sampler + + +def tokenize(tokenizer, prompts: list[str], shd=None): + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "left" + lines = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True + ) + for prompt in prompts + ] + batch = tokenizer(lines, padding=True, return_tensors="np", add_special_tokens=False) + input_ids = jnp.array(batch["input_ids"], out_sharding=shd) + attention_mask = jnp.array(batch["attention_mask"], out_sharding=shd) + return input_ids, attention_mask + + +def run_model(): + + # Choose a checkpoint and config; defaults to the 1B Instruct variant. + model_id = "meta-llama/Llama-3.2-1B-Instruct" + try: + access_token = os.environ["HF_TOKEN"] + except KeyError: + print("\nError: HF_TOKEN is not set.", file=sys.stderr) + print("Please set the HF_TOKEN environment variable and retry.", file=sys.stderr) + sys.exit(1) + + model_ckpt_path = snapshot_download(model_id, token=access_token) + + # Default: no sharding (single-device friendly). + config = modeling.ModelConfig.llama3_2_1b(use_fsdp=False, use_tp=False) + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + batch_shd = None + + prompts = [ + "Summarize what a tokenizer does in one paragraph.", + "Write a short, friendly explanation of gradient descent for beginners.", + ] + + tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd) + batch_size, token_len = tokens.shape + + generate_steps = 64 + model = params.create_model_from_safe_tensors(model_ckpt_path, config, mesh) + cache = model.init_cache(config, batch_size, token_len, generate_steps) + + key = jax.random.key(0) + sampler = Sampler(temperature=1.0, top_p=0.9, top_k=50) + jit_sampler = jax.jit(sampler) + + # prefill + logits, cache = modeling.forward(model, cache, tokens, pad_id, attention_mask=attention_mask) + key, subkey = jax.random.split(key) + next_tokens = jit_sampler(logits, key=subkey) + + # decode + tokens_list = [next_tokens] + finished = jnp.zeros((batch_size,), dtype=jnp.bool_) + start = time.time() + for _ in range(generate_steps): + logits, cache = modeling.forward(model, cache, next_tokens, pad_id) + key, subkey = jax.random.split(key) + next_tokens = jit_sampler(logits, key=subkey) + finished = finished | (next_tokens.squeeze(-1) == tokenizer.eos_token_id) + tokens_list.append(next_tokens) + if finished.all(): + break + + elapsed = time.time() - start + all_output_tokens = jax.device_get(jnp.concatenate(tokens_list, axis=-1)) + print(f"Generated {all_output_tokens.shape[1]} tokens in {elapsed:.3f}s") + for i, prompt in enumerate(prompts): + print(f"User:\n {prompt}") + seq_tokens = all_output_tokens[i] + eos_idx = np.where(seq_tokens == tokenizer.eos_token_id)[0] + if eos_idx.size > 0: + seq_tokens = seq_tokens[: eos_idx[0]] + decoded = tokenizer.decode(seq_tokens, skip_special_tokens=True) + print(f"Answer:\n {decoded}\n\n") + + +if __name__ == "__main__": + run_model() + + +__all__ = ["run_model"] diff --git a/bonsai/models/llama32/tests/run_model_base.py b/bonsai/models/llama32/tests/run_model_base.py new file mode 100644 index 00000000..35e03b5b --- /dev/null +++ b/bonsai/models/llama32/tests/run_model_base.py @@ -0,0 +1,113 @@ +# Copyright 2026 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import time + +import jax +import jax.numpy as jnp +import numpy as np +from huggingface_hub import snapshot_download +from jax.sharding import AxisType +from transformers import AutoTokenizer + +from bonsai.models.llama32 import modeling, params +from bonsai.utils import Sampler + + +def tokenize(tokenizer, prompts: list[str], shd=None): + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "left" + batch = tokenizer(prompts, padding=True, return_tensors="np") + input_ids = jnp.array(batch["input_ids"], out_sharding=shd) + attention_mask = jnp.array(batch["attention_mask"], out_sharding=shd) + return input_ids, attention_mask + + +def run_model(): + + # Choose a checkpoint and config; defaults to the 1B base variant. + model_id = "meta-llama/Llama-3.2-1B" + try: + access_token = os.environ["HF_TOKEN"] + except KeyError: + print("\nError: HF_TOKEN is not set.", file=sys.stderr) + print("Please set the HF_TOKEN environment variable and retry.", file=sys.stderr) + sys.exit(1) + + model_ckpt_path = snapshot_download(model_id, token=access_token) + + # Default: no sharding (single-device friendly). + config = modeling.ModelConfig.llama3_2_1b(use_fsdp=False, use_tp=False) + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + batch_shd = None + + prompts = [ + "The capital of France is", + "The definition of a tokenizer in NLP is strictly defined as follows: A tokenizer is an algorithm that", + ] + + tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd) + batch_size, token_len = tokens.shape + + generate_steps = 64 + model = params.create_model_from_safe_tensors(model_ckpt_path, config, mesh) + cache = model.init_cache(config, batch_size, token_len, generate_steps) + + key = jax.random.key(0) + sampler = Sampler(temperature=1.0, top_p=0.9, top_k=50) + jit_sampler = jax.jit(sampler) + + # prefill + logits, cache = modeling.forward(model, cache, tokens, pad_id, attention_mask=attention_mask) + key, subkey = jax.random.split(key) + next_tokens = jit_sampler(logits, key=subkey) + + # decode + tokens_list = [next_tokens] + finished = jnp.zeros((batch_size,), dtype=jnp.bool_) + start = time.time() + for _ in range(generate_steps): + logits, cache = modeling.forward(model, cache, next_tokens, pad_id) + key, subkey = jax.random.split(key) + next_tokens = jit_sampler(logits, key=subkey) + finished = finished | (next_tokens.squeeze(-1) == tokenizer.eos_token_id) + tokens_list.append(next_tokens) + if finished.all(): + break + + elapsed = time.time() - start + all_output_tokens = jax.device_get(jnp.concatenate(tokens_list, axis=-1)) + print(f"Generated {all_output_tokens.shape[1]} tokens in {elapsed:.3f}s") + for i, prompt in enumerate(prompts): + print(f"Prompt:\n {prompt}") + seq_tokens = all_output_tokens[i] + eos_idx = np.where(seq_tokens == tokenizer.eos_token_id)[0] + if eos_idx.size > 0: + seq_tokens = seq_tokens[: eos_idx[0]] + decoded = tokenizer.decode(seq_tokens, skip_special_tokens=True) + print(f"Completion:\n {decoded}\n\n") + + +if __name__ == "__main__": + run_model() + + +__all__ = ["run_model"] From 4c7a5c2f32dac659577891b77880ff1aa9403ace Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Thu, 22 Jan 2026 13:51:46 +0900 Subject: [PATCH 02/18] docs(llama32): add README --- bonsai/models/llama32/README.md | 54 +++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 bonsai/models/llama32/README.md diff --git a/bonsai/models/llama32/README.md b/bonsai/models/llama32/README.md new file mode 100644 index 00000000..70eecce5 --- /dev/null +++ b/bonsai/models/llama32/README.md @@ -0,0 +1,54 @@ +# Llama 3.2 in JAX + +This directory contains a pure JAX implementation of the +[Llama 3.2 language model](https://huggingface.co/meta-llama), +using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API. + +Note: You need a Hugging Face access token to download model weights. +Set an environment variable `HF_TOKEN` before running any scripts that fetch checkpoints. + +```sh +export HF_TOKEN="your_hf_access_token" +``` + +Some Llama models are gated. Make sure you have accepted the license in the +Hugging Face UI for the specific model you want to use. + +## Model Configuration Support Status + +| Model Name | Config Support Status | +| :--- | :--- | +| [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) | **✅ Supported** | +| [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) | **✅ Supported** | +| [Llama-3.2-3B](https://huggingface.co/meta-llama/Llama-3.2-3B) | **✅ Supported** | +| [Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | **✅ Supported** | + +## Running this model + +```sh +# Base model +python3 -m bonsai.models.llama32.tests.run_model_base + +# Instruct model +python3 -m bonsai.models.llama32.tests.run_model +``` + +## Output parity tests + +These tests compare JAX outputs against Hugging Face PyTorch outputs and require `HF_TOKEN`. + +```sh +python3 -m bonsai.models.llama32.tests.test_outputs_llama32 +``` + +## References + +* Paper: [The Llama 3 Herd of Models](https://arxiv.org/abs/2407.21783) +* Model code: [Hugging Face Transformers (LlamaModel)](https://github.com/huggingface/transformers/tree/main/src/transformers/models/llama) + +## How to contribute to this model + +We welcome contributions! You can contribute via the following: + +* Add a model config variant to `ModelConfig` in [modeling.py](modeling.py). +* Run [run_model.py](tests/run_model.py) and report whether the variant runs on your hardware. From 76e5106e033dd26a62b10261d659c3d9a6a6d532 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Thu, 22 Jan 2026 13:43:12 +0900 Subject: [PATCH 03/18] test(llama32): add tests for Llama-3.2 model output and padding behavior --- .../llama32/tests/test_outputs_llama32.py | 258 ++++++++++++++++++ .../llama32/tests/test_padding_llama32.py | 114 ++++++++ 2 files changed, 372 insertions(+) create mode 100644 bonsai/models/llama32/tests/test_outputs_llama32.py create mode 100644 bonsai/models/llama32/tests/test_padding_llama32.py diff --git a/bonsai/models/llama32/tests/test_outputs_llama32.py b/bonsai/models/llama32/tests/test_outputs_llama32.py new file mode 100644 index 00000000..53e7d9f4 --- /dev/null +++ b/bonsai/models/llama32/tests/test_outputs_llama32.py @@ -0,0 +1,258 @@ +import dataclasses +import os +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import torch +from absl.testing import absltest +from flax import nnx +from huggingface_hub import snapshot_download +from jax.sharding import AxisType +from transformers import AutoModelForCausalLM, AutoTokenizer + +from bonsai.models.llama32 import modeling, params + +# used to set highest precision on matrix multiplication for testing +jax.config.update("jax_default_matmul_precision", "highest") + + +def check_hf_token(): + try: + access_token = os.environ["HF_TOKEN"] + AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", token=access_token) + except Exception as e: + print("Failed to access HF_TOKEN or download tokenizer:") + print(e) + return True + return False + + +@unittest.skipIf(check_hf_token(), "Skipping Llama32 output tests due to HF_TOKEN failure.") +class TestOutputsLlama32(absltest.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.model_name = "meta-llama/Llama-3.2-1B-Instruct" + cls.torch_device = "cuda" if torch.cuda.is_available() else "cpu" + access_token = os.environ["HF_TOKEN"] + + cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name, token=access_token) + if cls.tokenizer.pad_token_id is None: + cls.tokenizer.pad_token = cls.tokenizer.eos_token + cls.tokenizer.padding_side = "left" + cls.pad_id = cls.tokenizer.pad_token_id + + cls.torch_model = ( + AutoModelForCausalLM.from_pretrained(cls.model_name, token=access_token, dtype=torch.float32) + .to(device=cls.torch_device, dtype=torch.float32) + .eval() + ) + + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + cls.mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(cls.mesh) + + cls.llama_config = dataclasses.replace( + modeling.ModelConfig.llama3_2_1b(use_fsdp=False, use_tp=False), + dtype=jnp.float32, + ) + model_ckpt_path = snapshot_download(cls.model_name, token=access_token) + graph_def, state = nnx.split( + params.create_model_from_safe_tensors(model_ckpt_path, cls.llama_config, mesh=cls.mesh) + ) + state = jax.tree.map(lambda x: x.astype(jnp.float32) if isinstance(x, jax.Array) else x, state) + cls.llama_model = nnx.merge(graph_def, state) + + cls.batch_size = 4 + cls.num_input_tokens = 6 + cls.relaxed_tol = 1e-3 + + def _make_torch_input(self): + messages = [{"role": "user", "content": "Summarize what a tokenizer does in one paragraph."}] + prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = self.tokenizer([prompt], padding=False, return_tensors="pt", add_special_tokens=False) + return {k: v.to(device=self.torch_device) for k, v in batch.items()} + + def _make_hidden_states(self, shape): + jx = jax.random.normal(jax.random.key(0), shape=shape) + tx = torch.tensor(np.array(jx, dtype=np.float32)) + return jx, tx + + def test_embedder(self): + nm = self.llama_model.embedder + tm = self.torch_model.model.embed_tokens + + tx = torch.randint(0, self.torch_model.config.vocab_size, size=(self.batch_size, self.num_input_tokens)) + jx = jnp.array(tx.cpu().detach().numpy()) + + jy = nm(jx) + ty = tm(tx) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_rms_norm(self): + nm = self.llama_model.layers[0].input_layernorm + tm = self.torch_model.model.layers[0].input_layernorm + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx) + ty = tm(tx) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_q_proj(self): + nm = self.llama_model.layers[0].self_attn.q_proj + tm = self.torch_model.model.layers[0].self_attn.q_proj.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_attention_heads, self.llama_config.head_dim + ) + ty = tm(tx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_attention_heads, self.llama_config.head_dim + ) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_k_proj(self): + nm = self.llama_model.layers[0].self_attn.k_proj + tm = self.torch_model.model.layers[0].self_attn.k_proj.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_key_value_heads, self.llama_config.head_dim + ) + ty = tm(tx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_key_value_heads, self.llama_config.head_dim + ) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_v_proj(self): + nm = self.llama_model.layers[0].self_attn.v_proj + tm = self.torch_model.model.layers[0].self_attn.v_proj.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_key_value_heads, self.llama_config.head_dim + ) + ty = tm(tx).reshape( + self.batch_size, self.num_input_tokens, self.llama_config.num_key_value_heads, self.llama_config.head_dim + ) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_o_proj(self): + nm = self.llama_model.layers[0].self_attn.o_proj + tm = self.torch_model.model.layers[0].self_attn.o_proj.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx) + ty = tm(tx) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_mlp(self): + nm = self.llama_model.layers[0].mlp + tm = self.torch_model.model.layers[0].mlp.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm(jx) + ty = tm(tx) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_lm_head(self): + nm = self.llama_model.embedder + tm = self.torch_model.lm_head.to(torch.float32) + + shape = (self.batch_size, self.num_input_tokens, self.llama_config.hidden_size) + jx, tx = self._make_hidden_states(shape) + + jy = nm.decode(jx) + ty = tm(tx) + torch.testing.assert_close( + torch.tensor(np.array(jy, dtype=np.float32)), + ty, + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_full_logits(self): + t_inputs = self._make_torch_input() + n_tokens = jnp.array(t_inputs["input_ids"].detach().cpu().numpy()) + attention_mask = jnp.array(t_inputs["attention_mask"].detach().cpu().numpy()) + segment_ids = attention_mask.astype(jnp.int32) + + with torch.no_grad(): + t_logits = self.torch_model(**t_inputs).logits + + n_logits = self.llama_model(n_tokens, segment_ids, cache=None, attn_mask=None) + np.testing.assert_allclose(n_logits, t_logits.detach().cpu().numpy(), rtol=5e-2, atol=5e-2) + + def test_forward_logits(self): + t_inputs = self._make_torch_input() + n_tokens = jnp.array(t_inputs["input_ids"].detach().cpu().numpy()) + attention_mask = jnp.array(t_inputs["attention_mask"].detach().cpu().numpy()) + batch_size, token_len = n_tokens.shape + cache = self.llama_model.init_cache(self.llama_config, batch_size, token_len, generate_steps=1) + + with torch.no_grad(): + t_logits = self.torch_model(**t_inputs).logits + + n_logits, _ = modeling.forward(self.llama_model, cache, n_tokens, self.pad_id, attention_mask=attention_mask) + np.testing.assert_allclose(n_logits, t_logits[:, -1].detach().cpu().numpy(), rtol=1e-4, atol=1e-4) + + +if __name__ == "__main__": + absltest.main() diff --git a/bonsai/models/llama32/tests/test_padding_llama32.py b/bonsai/models/llama32/tests/test_padding_llama32.py new file mode 100644 index 00000000..a61e473f --- /dev/null +++ b/bonsai/models/llama32/tests/test_padding_llama32.py @@ -0,0 +1,114 @@ +import numpy as np +import jax +import jax.numpy as jnp +from flax import nnx +from jax.sharding import AxisType + +from bonsai.models.llama32 import modeling + + +def _tiny_config() -> modeling.ModelConfig: + return modeling.ModelConfig( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + head_dim=4, + num_key_value_heads=1, + max_position_embeddings=32, + rms_norm_eps=1e-5, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=True, + shd_cfg=modeling.LlamaShardCfg.no_sharding(), + dtype=jnp.float32, + ) + + +def test_forward_uses_per_sample_right_padding(): + cfg = _tiny_config() + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) + + tokens = jnp.array( + [ + [1, 2, 3, 0, 0], + [4, 5, 6, 7, 0], + ], + dtype=jnp.int32, + ) + attention_mask = jnp.array( + [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 0], + ], + dtype=jnp.int32, + ) + + cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + logits, _ = modeling.forward(model, cache, tokens, pad_id=0, attention_mask=attention_mask) + + ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + full_logits = model(tokens, attention_mask.astype(jnp.int32), ref_cache, attn_mask=None) + target_ind = jnp.sum(attention_mask, axis=1) - 1 + expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] + + np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) + + +def test_compute_positions_from_segment_ids_packed(): + seg_ids = jnp.array( + [ + [1, 1, 1, 2, 2, 0, 0], + [3, 3, 0, 4, 4, 4, 0], + [0, 0, 0, 0, 0, 0, 0], + ], + dtype=jnp.int32, + ) + positions = modeling.compute_positions_from_segment_ids(seg_ids) + pad_val = 2**30 + expected = jnp.array( + [ + [0, 1, 2, 0, 1, pad_val, pad_val], + [0, 1, pad_val, 0, 1, 2, pad_val], + [pad_val, pad_val, pad_val, pad_val, pad_val, pad_val, pad_val], + ], + dtype=jnp.int32, + ) + np.testing.assert_array_equal(np.array(positions), np.array(expected)) + + +def test_forward_accepts_segment_ids_packed(): + cfg = _tiny_config() + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) + + tokens = jnp.array( + [ + [1, 2, 3, 4, 0, 0], + [5, 6, 7, 8, 9, 0], + ], + dtype=jnp.int32, + ) + segment_ids = jnp.array( + [ + [1, 1, 2, 2, 0, 0], + [3, 3, 3, 4, 4, 0], + ], + dtype=jnp.int32, + ) + + cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + logits, _ = modeling.forward(model, cache, tokens, pad_id=0, segment_ids=segment_ids) + + ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + full_logits = model(tokens, segment_ids, ref_cache, attn_mask=None) + target_ind = jnp.sum(segment_ids != 0, axis=1) - 1 + expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] + + np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) From f52f7f547387593be1e9e4a655020957aa123cf4 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Thu, 22 Jan 2026 21:15:24 +0900 Subject: [PATCH 04/18] test(llama32): add unit tests for sharding behavior in Llama model --- .../llama32/tests/test_sharding_llama32.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 bonsai/models/llama32/tests/test_sharding_llama32.py diff --git a/bonsai/models/llama32/tests/test_sharding_llama32.py b/bonsai/models/llama32/tests/test_sharding_llama32.py new file mode 100644 index 00000000..b6077ea4 --- /dev/null +++ b/bonsai/models/llama32/tests/test_sharding_llama32.py @@ -0,0 +1,66 @@ +import jax +import jax.numpy as jnp +from absl.testing import absltest +from flax import nnx +from jax import P +from jax.sharding import AxisType + +from bonsai.models.llama32 import modeling + + +def _tiny_sharded_config() -> modeling.ModelConfig: + return modeling.ModelConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + head_dim=8, + num_key_value_heads=2, + max_position_embeddings=128, + rms_norm_eps=1e-5, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=True, + shd_cfg=modeling.LlamaShardCfg.default(use_fsdp=True, use_tp=True), + dtype=jnp.float32, + ) + + +class TestShardingLlama32(absltest.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + cls.mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(cls.mesh) + + cls.cfg = _tiny_sharded_config() + cls.model = modeling.Llama(cls.cfg, rngs=nnx.Rngs(params=0)) + + def test_forward_sharded_inputs(self): + fsdp = modeling.ShardMode.FSDP.value + tokens = jnp.array( + [ + [1, 2, 3, 0], + [4, 5, 0, 0], + ], + dtype=jnp.int32, + out_sharding=P(fsdp), + ) + attention_mask = jnp.array( + [ + [1, 1, 1, 0], + [1, 1, 0, 0], + ], + dtype=jnp.int32, + out_sharding=P(fsdp), + ) + segment_ids = attention_mask.astype(jnp.int32) + + cache = self.model.init_cache(self.cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + _ = self.model(tokens, segment_ids, cache, attn_mask=None) + + +if __name__ == "__main__": + absltest.main() From 126f82fded6fefda749e80740b9509201ecc664a Mon Sep 17 00:00:00 2001 From: Jen Ha Date: Thu, 22 Jan 2026 17:26:03 -0800 Subject: [PATCH 05/18] ruff format --- bonsai/models/llama32/modeling.py | 45 ++++++++++--------- bonsai/models/llama32/tests/run_model.py | 5 +-- bonsai/models/llama32/tests/run_model_base.py | 1 - 3 files changed, 26 insertions(+), 25 deletions(-) diff --git a/bonsai/models/llama32/modeling.py b/bonsai/models/llama32/modeling.py index ddb1ea4b..c376bb3c 100644 --- a/bonsai/models/llama32/modeling.py +++ b/bonsai/models/llama32/modeling.py @@ -28,7 +28,8 @@ class ShardMode(Enum): - """ Sharding Modes for Model Parameters """ + """Sharding Modes for Model Parameters""" + FSDP = "fsdp" TP = "tp" @@ -94,7 +95,6 @@ class RopeScalingConfig: original_max_position_embeddings: int = 8192 - @dataclasses.dataclass(frozen=True) class ModelConfig: vocab_size: int @@ -181,6 +181,7 @@ def __call__(self, x: Array) -> Array: return (self.scale[...] * x_fp32 * inv_rms).astype(dtype) + # TODO: Replace with nnx.Linear once explicit sharding is supported. class ShardedLinear(nnx.Module): def __init__( @@ -225,7 +226,6 @@ def decode(self, query: Array, *, out_sharding: PartitionSpec | None = None) -> return jnp.dot(query, embedding.T, out_sharding=out_sharding) - class LlamaMLP(nnx.Module): """Feed Forward Network""" @@ -242,11 +242,24 @@ def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): dtype=cfg.dtype, rngs=rngs, ) - self.up_proj = ShardedLinear(self.hidden_size, self.intermediate_size, sharding=cfg.shd_cfg.up_proj, use_bias=False, dtype=cfg.dtype, rngs=rngs) - self.down_proj = ShardedLinear(self.intermediate_size, self.hidden_size, sharding=cfg.shd_cfg.down_proj, use_bias=False, dtype=cfg.dtype, rngs=rngs) - + self.up_proj = ShardedLinear( + self.hidden_size, + self.intermediate_size, + sharding=cfg.shd_cfg.up_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) + self.down_proj = ShardedLinear( + self.intermediate_size, + self.hidden_size, + sharding=cfg.shd_cfg.down_proj, + use_bias=False, + dtype=cfg.dtype, + rngs=rngs, + ) - @jax.named_scope('feed_forward') + @jax.named_scope("feed_forward") def __call__(self, x: ArrayLike) -> Array: act_shd = self.config.shd_cfg.activation gated = nnx.silu(self.gate_proj(x, out_sharding=act_shd)) * self.up_proj(x, out_sharding=act_shd) @@ -263,9 +276,7 @@ def _generate_pos_embeddings( if rope_scaling is None: rope_scaling = RopeScalingConfig(factor=1.0) - inv_freq = 1.0 / ( - rope_theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim) - ) + inv_freq = 1.0 / (rope_theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) factor = rope_scaling.factor low_freq_factor = rope_scaling.low_freq_factor @@ -277,9 +288,7 @@ def _generate_pos_embeddings( wavelen = (2.0 * math.pi) / inv_freq inv_freq_llama = jnp.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq) - smooth_factor = (old_context_len / wavelen - low_freq_factor) / ( - high_freq_factor - low_freq_factor - ) + smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) smoothed_inv_freq = (1.0 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama is_medium_freq = jnp.logical_and( wavelen >= high_freq_wavelen, @@ -288,9 +297,7 @@ def _generate_pos_embeddings( inv_freq_llama = jnp.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) positions_f32 = positions.astype(jnp.float32) - sinusoid_inp = jnp.einsum( - "BT,k->BTk", positions_f32, inv_freq_llama, precision=jax.lax.Precision.HIGHEST - ) + sinusoid_inp = jnp.einsum("BT,k->BTk", positions_f32, inv_freq_llama, precision=jax.lax.Precision.HIGHEST) return jnp.sin(sinusoid_inp), jnp.cos(sinusoid_inp) @@ -328,6 +335,7 @@ def count_right_pads_from_mask(attn_mask: Array) -> Array: max_len = jnp.full_like(right_pad, mask.shape[1]) return jnp.where(all_pad, max_len, right_pad) + def compute_positions_from_segment_ids(seg_ids: Array) -> Array: """Compute position ids from segment ids.""" seg_ids = seg_ids.astype(jnp.int32) @@ -379,8 +387,6 @@ def sharded_attention( return attn_output.astype(q.dtype) - - class LlamaAttention(nnx.Module): """Multi-Head Self Attention""" @@ -546,6 +552,7 @@ def __call__( class Llama(nnx.Module): """Llama Model""" + def __init__(self, cfg: ModelConfig, *, rngs: nnx.Rngs): self.config = cfg embed_init = partial(default_embed_init, out_sharding=cfg.shd_cfg.emb_weight) @@ -613,13 +620,11 @@ def forward( attention_mask: Array | None = None, segment_ids: Array | None = None, ) -> tuple[Array, Cache]: - # Use attention_mask when available because pad_id can equal eos_id, which would # misclassify real tokens as padding if we rely on tokens != pad_id. if attention_mask is not None and segment_ids is not None: raise ValueError("Provide only one of attention_mask or segment_ids.") - if segment_ids is None: if attention_mask is None: segment_ids = 1 * (tokens != pad_id) diff --git a/bonsai/models/llama32/tests/run_model.py b/bonsai/models/llama32/tests/run_model.py index b75eb31b..03527d7c 100644 --- a/bonsai/models/llama32/tests/run_model.py +++ b/bonsai/models/llama32/tests/run_model.py @@ -33,9 +33,7 @@ def tokenize(tokenizer, prompts: list[str], shd=None): tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" lines = [ - tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True - ) + tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True) for prompt in prompts ] batch = tokenizer(lines, padding=True, return_tensors="np", add_special_tokens=False) @@ -45,7 +43,6 @@ def tokenize(tokenizer, prompts: list[str], shd=None): def run_model(): - # Choose a checkpoint and config; defaults to the 1B Instruct variant. model_id = "meta-llama/Llama-3.2-1B-Instruct" try: diff --git a/bonsai/models/llama32/tests/run_model_base.py b/bonsai/models/llama32/tests/run_model_base.py index 35e03b5b..65346ab8 100644 --- a/bonsai/models/llama32/tests/run_model_base.py +++ b/bonsai/models/llama32/tests/run_model_base.py @@ -38,7 +38,6 @@ def tokenize(tokenizer, prompts: list[str], shd=None): def run_model(): - # Choose a checkpoint and config; defaults to the 1B base variant. model_id = "meta-llama/Llama-3.2-1B" try: From bfbe0e05de6694a10506b5f956ea71c771c91a22 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Fri, 23 Jan 2026 13:11:00 +0900 Subject: [PATCH 06/18] chore(llama3.2): rename llama32 directory to llama3_2 --- bonsai/models/{llama32 => llama3_2}/README.md | 6 +++--- bonsai/models/{llama32 => llama3_2}/modeling.py | 0 bonsai/models/{llama32 => llama3_2}/params.py | 2 +- bonsai/models/{llama32 => llama3_2}/tests/run_model.py | 2 +- bonsai/models/{llama32 => llama3_2}/tests/run_model_base.py | 2 +- .../tests/test_outputs_llama3_2.py} | 2 +- .../tests/test_padding_llama3_2.py} | 2 +- .../tests/test_sharding_llama3_2.py} | 2 +- 8 files changed, 9 insertions(+), 9 deletions(-) rename bonsai/models/{llama32 => llama3_2}/README.md (91%) rename bonsai/models/{llama32 => llama3_2}/modeling.py (100%) rename bonsai/models/{llama32 => llama3_2}/params.py (99%) rename bonsai/models/{llama32 => llama3_2}/tests/run_model.py (98%) rename bonsai/models/{llama32 => llama3_2}/tests/run_model_base.py (98%) rename bonsai/models/{llama32/tests/test_outputs_llama32.py => llama3_2/tests/test_outputs_llama3_2.py} (99%) rename bonsai/models/{llama32/tests/test_padding_llama32.py => llama3_2/tests/test_padding_llama3_2.py} (98%) rename bonsai/models/{llama32/tests/test_sharding_llama32.py => llama3_2/tests/test_sharding_llama3_2.py} (97%) diff --git a/bonsai/models/llama32/README.md b/bonsai/models/llama3_2/README.md similarity index 91% rename from bonsai/models/llama32/README.md rename to bonsai/models/llama3_2/README.md index 70eecce5..1f65e3e4 100644 --- a/bonsai/models/llama32/README.md +++ b/bonsai/models/llama3_2/README.md @@ -27,10 +27,10 @@ Hugging Face UI for the specific model you want to use. ```sh # Base model -python3 -m bonsai.models.llama32.tests.run_model_base +python3 -m bonsai.models.llama3_2.tests.run_model_base # Instruct model -python3 -m bonsai.models.llama32.tests.run_model +python3 -m bonsai.models.llama3_2.tests.run_model ``` ## Output parity tests @@ -38,7 +38,7 @@ python3 -m bonsai.models.llama32.tests.run_model These tests compare JAX outputs against Hugging Face PyTorch outputs and require `HF_TOKEN`. ```sh -python3 -m bonsai.models.llama32.tests.test_outputs_llama32 +python3 -m bonsai.models.llama3_2.tests.test_outputs_llama3_2 ``` ## References diff --git a/bonsai/models/llama32/modeling.py b/bonsai/models/llama3_2/modeling.py similarity index 100% rename from bonsai/models/llama32/modeling.py rename to bonsai/models/llama3_2/modeling.py diff --git a/bonsai/models/llama32/params.py b/bonsai/models/llama3_2/params.py similarity index 99% rename from bonsai/models/llama32/params.py rename to bonsai/models/llama3_2/params.py index c4f16c2a..aa21fdc1 100644 --- a/bonsai/models/llama32/params.py +++ b/bonsai/models/llama3_2/params.py @@ -21,7 +21,7 @@ from etils import epath from flax import nnx -from bonsai.models.llama32 import modeling as model_lib +from bonsai.models.llama3_2 import modeling as model_lib def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig): diff --git a/bonsai/models/llama32/tests/run_model.py b/bonsai/models/llama3_2/tests/run_model.py similarity index 98% rename from bonsai/models/llama32/tests/run_model.py rename to bonsai/models/llama3_2/tests/run_model.py index 03527d7c..759e6204 100644 --- a/bonsai/models/llama32/tests/run_model.py +++ b/bonsai/models/llama3_2/tests/run_model.py @@ -24,7 +24,7 @@ from transformers import AutoTokenizer -from bonsai.models.llama32 import modeling, params +from bonsai.models.llama3_2 import modeling, params from bonsai.utils import Sampler diff --git a/bonsai/models/llama32/tests/run_model_base.py b/bonsai/models/llama3_2/tests/run_model_base.py similarity index 98% rename from bonsai/models/llama32/tests/run_model_base.py rename to bonsai/models/llama3_2/tests/run_model_base.py index 65346ab8..04b0e6ec 100644 --- a/bonsai/models/llama32/tests/run_model_base.py +++ b/bonsai/models/llama3_2/tests/run_model_base.py @@ -23,7 +23,7 @@ from jax.sharding import AxisType from transformers import AutoTokenizer -from bonsai.models.llama32 import modeling, params +from bonsai.models.llama3_2 import modeling, params from bonsai.utils import Sampler diff --git a/bonsai/models/llama32/tests/test_outputs_llama32.py b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py similarity index 99% rename from bonsai/models/llama32/tests/test_outputs_llama32.py rename to bonsai/models/llama3_2/tests/test_outputs_llama3_2.py index 53e7d9f4..cd73aca8 100644 --- a/bonsai/models/llama32/tests/test_outputs_llama32.py +++ b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py @@ -12,7 +12,7 @@ from jax.sharding import AxisType from transformers import AutoModelForCausalLM, AutoTokenizer -from bonsai.models.llama32 import modeling, params +from bonsai.models.llama3_2 import modeling, params # used to set highest precision on matrix multiplication for testing jax.config.update("jax_default_matmul_precision", "highest") diff --git a/bonsai/models/llama32/tests/test_padding_llama32.py b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py similarity index 98% rename from bonsai/models/llama32/tests/test_padding_llama32.py rename to bonsai/models/llama3_2/tests/test_padding_llama3_2.py index a61e473f..7b7b1024 100644 --- a/bonsai/models/llama32/tests/test_padding_llama32.py +++ b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py @@ -4,7 +4,7 @@ from flax import nnx from jax.sharding import AxisType -from bonsai.models.llama32 import modeling +from bonsai.models.llama3_2 import modeling def _tiny_config() -> modeling.ModelConfig: diff --git a/bonsai/models/llama32/tests/test_sharding_llama32.py b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py similarity index 97% rename from bonsai/models/llama32/tests/test_sharding_llama32.py rename to bonsai/models/llama3_2/tests/test_sharding_llama3_2.py index b6077ea4..c7117e97 100644 --- a/bonsai/models/llama32/tests/test_sharding_llama32.py +++ b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py @@ -5,7 +5,7 @@ from jax import P from jax.sharding import AxisType -from bonsai.models.llama32 import modeling +from bonsai.models.llama3_2 import modeling def _tiny_sharded_config() -> modeling.ModelConfig: From 313fda9c26c46bf3942b72fb2ef62699c2de1766 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Fri, 23 Jan 2026 13:19:09 +0900 Subject: [PATCH 07/18] refactor(llama3.2): replace deprecated flax.nnx.State with nnx.to_pure_dict --- bonsai/models/llama3_2/params.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bonsai/models/llama3_2/params.py b/bonsai/models/llama3_2/params.py index aa21fdc1..870cf091 100644 --- a/bonsai/models/llama3_2/params.py +++ b/bonsai/models/llama3_2/params.py @@ -119,8 +119,8 @@ def create_model_from_safe_tensors( llama = nnx.eval_shape(lambda: model_lib.Llama(cfg, rngs=nnx.Rngs(params=0))) graph_def, abs_state = nnx.split(llama) - state_dict = abs_state.to_pure_dict() - sharding = nnx.get_named_sharding(abs_state, mesh).to_pure_dict() if mesh is not None else None + state_dict = nnx.to_pure_dict(abs_state) + sharding = nnx.to_pure_dict(nnx.get_named_sharding(abs_state, mesh)) if mesh is not None else None key_mapping = _get_key_and_transform_mapping(cfg) conversion_errors = [] From 618b339a11d2299034015db84683a4c06be911a7 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Fri, 23 Jan 2026 21:26:08 +0900 Subject: [PATCH 08/18] test(llama32): update tolerance values in output tests for improved precision --- .../llama3_2/tests/test_outputs_llama3_2.py | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py index cd73aca8..cb1a0b08 100644 --- a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py @@ -67,7 +67,6 @@ def setUpClass(cls): cls.batch_size = 4 cls.num_input_tokens = 6 - cls.relaxed_tol = 1e-3 def _make_torch_input(self): messages = [{"role": "user", "content": "Summarize what a tokenizer does in one paragraph."}] @@ -92,9 +91,9 @@ def test_embedder(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_rms_norm(self): @@ -109,9 +108,9 @@ def test_rms_norm(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_q_proj(self): @@ -130,9 +129,9 @@ def test_q_proj(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_k_proj(self): @@ -151,9 +150,9 @@ def test_k_proj(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_v_proj(self): @@ -172,9 +171,9 @@ def test_v_proj(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_o_proj(self): @@ -189,9 +188,9 @@ def test_o_proj(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_mlp(self): @@ -206,9 +205,9 @@ def test_mlp(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_lm_head(self): @@ -223,9 +222,9 @@ def test_lm_head(self): torch.testing.assert_close( torch.tensor(np.array(jy, dtype=np.float32)), ty, - rtol=self.relaxed_tol, - atol=self.relaxed_tol, - check_dtype=False, + rtol=1e-5, + atol=1e-5, + check_dtype=True, ) def test_full_logits(self): @@ -238,7 +237,7 @@ def test_full_logits(self): t_logits = self.torch_model(**t_inputs).logits n_logits = self.llama_model(n_tokens, segment_ids, cache=None, attn_mask=None) - np.testing.assert_allclose(n_logits, t_logits.detach().cpu().numpy(), rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(n_logits, t_logits.detach().cpu().numpy(), rtol=5e-4, atol=1e-3) def test_forward_logits(self): t_inputs = self._make_torch_input() From c5603bd37a3823319f51bfeb496c7b4a5a0d54f2 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Fri, 30 Jan 2026 21:15:09 +0900 Subject: [PATCH 09/18] docs(llama3.2): update Flax NNX API link to stable version --- bonsai/models/llama3_2/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bonsai/models/llama3_2/README.md b/bonsai/models/llama3_2/README.md index 1f65e3e4..ba64c919 100644 --- a/bonsai/models/llama3_2/README.md +++ b/bonsai/models/llama3_2/README.md @@ -2,7 +2,7 @@ This directory contains a pure JAX implementation of the [Llama 3.2 language model](https://huggingface.co/meta-llama), -using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API. +using the [Flax NNX](https://flax.readthedocs.io/en/stable/index.html) API. Note: You need a Hugging Face access token to download model weights. Set an environment variable `HF_TOKEN` before running any scripts that fetch checkpoints. From b4cf559315c77e8a2f2e736537eb76b5081b0745 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:03:35 +0900 Subject: [PATCH 10/18] fix(llama3.2): correct sharding axes for GQA and fix explicit mesh errors --- bonsai/models/llama3_2/modeling.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/bonsai/models/llama3_2/modeling.py b/bonsai/models/llama3_2/modeling.py index c376bb3c..28b8beb7 100644 --- a/bonsai/models/llama3_2/modeling.py +++ b/bonsai/models/llama3_2/modeling.py @@ -20,7 +20,7 @@ import jax from jax import P -from jax.sharding import PartitionSpec +from jax.sharding import PartitionSpec, get_abstract_mesh, reshard import jax.numpy as jnp from jaxtyping import Array, ArrayLike from flax import nnx @@ -87,6 +87,12 @@ def default(cls, use_fsdp: bool, use_tp: bool) -> "LlamaShardCfg": ) +def shard(x: jnp.ndarray, s: PartitionSpec): + mesh = get_abstract_mesh() + if not mesh.empty and len(mesh.axis_names) > 0: + return reshard(x, s) + return x + @dataclasses.dataclass(frozen=True) class RopeScalingConfig: factor: float @@ -485,6 +491,9 @@ def __call__( cache_dtype = cache.k_cache[...].dtype k = k.astype(cache_dtype) v = v.astype(cache_dtype) + cache_shd = self.config.shd_cfg.cache + k = shard(k, cache_shd) + v = shard(v, cache_shd) cache.k_cache[...] = jax.lax.dynamic_update_slice(cache.k_cache[...], k, slice_indices) cache.v_cache[...] = jax.lax.dynamic_update_slice(cache.v_cache[...], v, slice_indices) k = cache.k_cache[...] From a678b0164b6da818e809d640dd0c143152d5f6d7 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:04:14 +0900 Subject: [PATCH 11/18] docs(llama3.2): enhance docstring for compute_positions_from_segment_ids to clarify packed sequences support --- bonsai/models/llama3_2/modeling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bonsai/models/llama3_2/modeling.py b/bonsai/models/llama3_2/modeling.py index 28b8beb7..05841283 100644 --- a/bonsai/models/llama3_2/modeling.py +++ b/bonsai/models/llama3_2/modeling.py @@ -343,7 +343,7 @@ def count_right_pads_from_mask(attn_mask: Array) -> Array: def compute_positions_from_segment_ids(seg_ids: Array) -> Array: - """Compute position ids from segment ids.""" + """Compute position ids from segment ids with support for packed sequences.""" seg_ids = seg_ids.astype(jnp.int32) pad_sentinel = 2**30 From ee9e595dab8f90b14442753d096b30b5976ba016 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:08:43 +0900 Subject: [PATCH 12/18] test(llama3.2): consolidate test configs --- .../llama3_2/tests/test_sharding_llama3_2.py | 22 +++------------- bonsai/models/llama3_2/tests/test_utils.py | 26 +++++++++++++++++++ 2 files changed, 29 insertions(+), 19 deletions(-) create mode 100644 bonsai/models/llama3_2/tests/test_utils.py diff --git a/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py index c7117e97..59afbd03 100644 --- a/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py @@ -6,25 +6,9 @@ from jax.sharding import AxisType from bonsai.models.llama3_2 import modeling +from bonsai.models.llama3_2.tests.test_utils import tiny_config - -def _tiny_sharded_config() -> modeling.ModelConfig: - return modeling.ModelConfig( - vocab_size=64, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - head_dim=8, - num_key_value_heads=2, - max_position_embeddings=128, - rms_norm_eps=1e-5, - rope_theta=10000.0, - rope_scaling=None, - tie_word_embeddings=True, - shd_cfg=modeling.LlamaShardCfg.default(use_fsdp=True, use_tp=True), - dtype=jnp.float32, - ) +jax.config.update("jax_num_cpu_devices", 8) class TestShardingLlama32(absltest.TestCase): @@ -35,7 +19,7 @@ def setUpClass(cls): cls.mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) jax.set_mesh(cls.mesh) - cls.cfg = _tiny_sharded_config() + cls.cfg = tiny_config(use_sharding=True) cls.model = modeling.Llama(cls.cfg, rngs=nnx.Rngs(params=0)) def test_forward_sharded_inputs(self): diff --git a/bonsai/models/llama3_2/tests/test_utils.py b/bonsai/models/llama3_2/tests/test_utils.py new file mode 100644 index 00000000..d6019c31 --- /dev/null +++ b/bonsai/models/llama3_2/tests/test_utils.py @@ -0,0 +1,26 @@ +import jax.numpy as jnp + +from bonsai.models.llama3_2 import modeling + +def tiny_config(*, use_sharding: bool = False) -> modeling.ModelConfig: + """Create a minimal Llama3.2 model configuration for testing purposes""" + return modeling.ModelConfig( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + head_dim=4, + num_key_value_heads=1, + max_position_embeddings=32, + rms_norm_eps=1e-5, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=True, + shd_cfg=( + modeling.LlamaShardCfg.default(use_fsdp=True, use_tp=True) + if use_sharding + else modeling.LlamaShardCfg.no_sharding() + ), + dtype=jnp.float32, + ) From 4c8ba67b74ad5cc7deaf6f069cbcf992519ecaca Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:11:30 +0900 Subject: [PATCH 13/18] test(llama3.2): refactor padding tests --- .../llama3_2/tests/test_padding_llama3_2.py | 200 ++++++++---------- 1 file changed, 93 insertions(+), 107 deletions(-) diff --git a/bonsai/models/llama3_2/tests/test_padding_llama3_2.py b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py index 7b7b1024..bea68ed5 100644 --- a/bonsai/models/llama3_2/tests/test_padding_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py @@ -1,114 +1,100 @@ import numpy as np import jax import jax.numpy as jnp +from absl.testing import absltest from flax import nnx from jax.sharding import AxisType from bonsai.models.llama3_2 import modeling - - -def _tiny_config() -> modeling.ModelConfig: - return modeling.ModelConfig( - vocab_size=16, - hidden_size=8, - intermediate_size=16, - num_hidden_layers=1, - num_attention_heads=2, - head_dim=4, - num_key_value_heads=1, - max_position_embeddings=32, - rms_norm_eps=1e-5, - rope_theta=10000.0, - rope_scaling=None, - tie_word_embeddings=True, - shd_cfg=modeling.LlamaShardCfg.no_sharding(), - dtype=jnp.float32, - ) - - -def test_forward_uses_per_sample_right_padding(): - cfg = _tiny_config() - fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value - mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) - jax.set_mesh(mesh) - model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) - - tokens = jnp.array( - [ - [1, 2, 3, 0, 0], - [4, 5, 6, 7, 0], - ], - dtype=jnp.int32, - ) - attention_mask = jnp.array( - [ - [1, 1, 1, 0, 0], - [1, 1, 1, 1, 0], - ], - dtype=jnp.int32, - ) - - cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) - logits, _ = modeling.forward(model, cache, tokens, pad_id=0, attention_mask=attention_mask) - - ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) - full_logits = model(tokens, attention_mask.astype(jnp.int32), ref_cache, attn_mask=None) - target_ind = jnp.sum(attention_mask, axis=1) - 1 - expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] - - np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) - - -def test_compute_positions_from_segment_ids_packed(): - seg_ids = jnp.array( - [ - [1, 1, 1, 2, 2, 0, 0], - [3, 3, 0, 4, 4, 4, 0], - [0, 0, 0, 0, 0, 0, 0], - ], - dtype=jnp.int32, - ) - positions = modeling.compute_positions_from_segment_ids(seg_ids) - pad_val = 2**30 - expected = jnp.array( - [ - [0, 1, 2, 0, 1, pad_val, pad_val], - [0, 1, pad_val, 0, 1, 2, pad_val], - [pad_val, pad_val, pad_val, pad_val, pad_val, pad_val, pad_val], - ], - dtype=jnp.int32, - ) - np.testing.assert_array_equal(np.array(positions), np.array(expected)) - - -def test_forward_accepts_segment_ids_packed(): - cfg = _tiny_config() - fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value - mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) - jax.set_mesh(mesh) - model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) - - tokens = jnp.array( - [ - [1, 2, 3, 4, 0, 0], - [5, 6, 7, 8, 9, 0], - ], - dtype=jnp.int32, - ) - segment_ids = jnp.array( - [ - [1, 1, 2, 2, 0, 0], - [3, 3, 3, 4, 4, 0], - ], - dtype=jnp.int32, - ) - - cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) - logits, _ = modeling.forward(model, cache, tokens, pad_id=0, segment_ids=segment_ids) - - ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) - full_logits = model(tokens, segment_ids, ref_cache, attn_mask=None) - target_ind = jnp.sum(segment_ids != 0, axis=1) - 1 - expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] - - np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) +from bonsai.models.llama3_2.tests.test_utils import tiny_config + + +class TestPaddingLlama32(absltest.TestCase): + def test_forward_uses_per_sample_right_padding(self): + cfg = tiny_config(use_sharding=False) + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) + + tokens = jnp.array( + [ + [1, 2, 3, 0, 0], + [4, 5, 6, 7, 0], + ], + dtype=jnp.int32, + ) + attention_mask = jnp.array( + [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 0], + ], + dtype=jnp.int32, + ) + + cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + logits, _ = modeling.forward(model, cache, tokens, pad_id=0, attention_mask=attention_mask) + + ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + full_logits = model(tokens, attention_mask.astype(jnp.int32), ref_cache, attn_mask=None) + target_ind = jnp.sum(attention_mask, axis=1) - 1 + expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] + + np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) + + def test_compute_positions_from_segment_ids_packed(self): + seg_ids = jnp.array( + [ + [1, 1, 1, 2, 2, 0, 0], + [3, 3, 0, 4, 4, 4, 0], + [0, 0, 0, 0, 0, 0, 0], + ], + dtype=jnp.int32, + ) + positions = modeling.compute_positions_from_segment_ids(seg_ids) + pad_val = 2**30 + expected = jnp.array( + [ + [0, 1, 2, 0, 1, pad_val, pad_val], + [0, 1, pad_val, 0, 1, 2, pad_val], + [pad_val, pad_val, pad_val, pad_val, pad_val, pad_val, pad_val], + ], + dtype=jnp.int32, + ) + np.testing.assert_array_equal(np.array(positions), np.array(expected)) + + def test_forward_accepts_segment_ids_packed(self): + cfg = tiny_config(use_sharding=False) + fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value + mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) + jax.set_mesh(mesh) + model = modeling.Llama(cfg, rngs=nnx.Rngs(params=0)) + + tokens = jnp.array( + [ + [1, 2, 3, 4, 0, 0], + [5, 6, 7, 8, 9, 0], + ], + dtype=jnp.int32, + ) + segment_ids = jnp.array( + [ + [1, 1, 2, 2, 0, 0], + [3, 3, 3, 4, 4, 0], + ], + dtype=jnp.int32, + ) + + cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + logits, _ = modeling.forward(model, cache, tokens, pad_id=0, segment_ids=segment_ids) + + ref_cache = model.init_cache(cfg, batch_size=tokens.shape[0], token_len=tokens.shape[1], generate_steps=1) + full_logits = model(tokens, segment_ids, ref_cache, attn_mask=None) + target_ind = jnp.sum(segment_ids != 0, axis=1) - 1 + expected = full_logits[jnp.arange(tokens.shape[0]), target_ind] + + np.testing.assert_allclose(np.array(logits), np.array(expected), rtol=1e-6, atol=1e-6) + + +if __name__ == "__main__": + absltest.main() From 74074391962c2cd907eff6cdfc9c15505ed7b101 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:11:38 +0900 Subject: [PATCH 14/18] test(llama3.2): add attention mask tests for padding and future tokens --- .../llama3_2/tests/test_outputs_llama3_2.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py index cb1a0b08..fd324b8d 100644 --- a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py @@ -193,6 +193,44 @@ def test_o_proj(self): check_dtype=True, ) + def test_attention_mask_blocks_padding(self): + nm = self.llama_model.layers[0].self_attn + batch_size = 2 + token_len = self.num_input_tokens + hidden_size = self.llama_config.hidden_size + + key = jax.random.key(0) + x = jax.random.normal(key, (batch_size, token_len, hidden_size), dtype=jnp.float32) + segment_ids = jnp.ones((batch_size, token_len), dtype=jnp.int32) + segment_ids = segment_ids.at[0, :2].set(0) + segment_ids = segment_ids.at[1, :1].set(0) + pad_mask = (segment_ids == 0)[..., None].astype(x.dtype) + x_alt = x + pad_mask * 5.0 + + y = nm(x, segment_ids, attn_mask=None, cache=None) + y_alt = nm(x_alt, segment_ids, attn_mask=None, cache=None) + + valid_mask = np.array(segment_ids == 1)[..., None].astype(np.float32) + y_masked = np.array(y) * valid_mask + y_alt_masked = np.array(y_alt) * valid_mask + np.testing.assert_allclose(y_masked, y_alt_masked, rtol=1e-5, atol=1e-5) + + def test_attention_mask_blocks_future_tokens(self): + nm = self.llama_model.layers[0].self_attn + batch_size = 2 + token_len = self.num_input_tokens + hidden_size = self.llama_config.hidden_size + + key = jax.random.key(1) + x = jax.random.normal(key, (batch_size, token_len, hidden_size), dtype=jnp.float32) + segment_ids = jnp.ones((batch_size, token_len), dtype=jnp.int32) + x_alt = x.at[:, -1, :].add(5.0) + + y = nm(x, segment_ids, attn_mask=None, cache=None) + y_alt = nm(x_alt, segment_ids, attn_mask=None, cache=None) + + np.testing.assert_allclose(np.array(y)[:, :-1, :], np.array(y_alt)[:, :-1, :], rtol=1e-5, atol=1e-5) + def test_mlp(self): nm = self.llama_model.layers[0].mlp tm = self.torch_model.model.layers[0].mlp.to(torch.float32) From fee4a82d540b4a525a5cb7b778cd3a0b112f9520 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Sun, 1 Feb 2026 19:48:41 +0900 Subject: [PATCH 15/18] refactor(llama3.2): enhance run_model script with argument parsing and remove deprecated base model test --- bonsai/models/llama3_2/README.md | 14 ++- bonsai/models/llama3_2/tests/run_model.py | 75 ++++++++++-- .../models/llama3_2/tests/run_model_base.py | 112 ------------------ 3 files changed, 75 insertions(+), 126 deletions(-) delete mode 100644 bonsai/models/llama3_2/tests/run_model_base.py diff --git a/bonsai/models/llama3_2/README.md b/bonsai/models/llama3_2/README.md index ba64c919..b6ddcd7c 100644 --- a/bonsai/models/llama3_2/README.md +++ b/bonsai/models/llama3_2/README.md @@ -26,11 +26,17 @@ Hugging Face UI for the specific model you want to use. ## Running this model ```sh -# Base model -python3 -m bonsai.models.llama3_2.tests.run_model_base - -# Instruct model +# Instruct model (default: 1B) python3 -m bonsai.models.llama3_2.tests.run_model + +# Base model (1B) +python3 -m bonsai.models.llama3_2.tests.run_model --base + +# Base model (3B) +python3 -m bonsai.models.llama3_2.tests.run_model --size 3B --base + +# Instruct model (3B) +python3 -m bonsai.models.llama3_2.tests.run_model --size 3B ``` ## Output parity tests diff --git a/bonsai/models/llama3_2/tests/run_model.py b/bonsai/models/llama3_2/tests/run_model.py index 759e6204..a35eb3fc 100644 --- a/bonsai/models/llama3_2/tests/run_model.py +++ b/bonsai/models/llama3_2/tests/run_model.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import argparse +import dataclasses import os import sys import time +from typing import Literal import jax import jax.numpy as jnp @@ -28,23 +31,69 @@ from bonsai.utils import Sampler -def tokenize(tokenizer, prompts: list[str], shd=None): +@dataclasses.dataclass(frozen=True) +class Args: + model_size: Literal["1B", "3B"] + use_base_model: bool = False + + +def _parse_args(argv: list[str]) -> Args: + def _size(value: str) -> str: + value = value.strip().upper() + if value not in ("1B", "3B"): + raise argparse.ArgumentTypeError("size must be 1B or 3B") + return value + + parser = argparse.ArgumentParser(description="Run a small Llama 3.2 inference example.") + parser.add_argument( + "--size", + type=_size, + default="1B", + help="Model size to load. Choices: 1B or 3B. Default: 1B.", + ) + parser.add_argument( + "--base", + action="store_true", + help="Use the base (non-instruct) checkpoint. Default uses Instruct.", + ) + parsed = parser.parse_args(argv) + return Args(model_size=parsed.size, use_base_model=parsed.base) + + +def tokenize(tokenizer, prompts: list[str], shd=None, use_chat_template: bool = True): if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" - lines = [ - tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True) - for prompt in prompts - ] - batch = tokenizer(lines, padding=True, return_tensors="np", add_special_tokens=False) + use_template = use_chat_template and getattr(tokenizer, "chat_template", None) is not None + + if use_template: + lines = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True + ) + for prompt in prompts + ] + else: + lines = prompts + + batch = tokenizer(lines, padding=True, return_tensors="np", add_special_tokens=not use_template) input_ids = jnp.array(batch["input_ids"], out_sharding=shd) attention_mask = jnp.array(batch["attention_mask"], out_sharding=shd) return input_ids, attention_mask def run_model(): + args = _parse_args(sys.argv[1:]) + # Choose a checkpoint and config; defaults to the 1B Instruct variant. - model_id = "meta-llama/Llama-3.2-1B-Instruct" + model_size = args.model_size + use_base_model = args.use_base_model + + model_id = ( + f"meta-llama/Llama-3.2-{model_size}" + if use_base_model + else f"meta-llama/Llama-3.2-{model_size}-Instruct" + ) try: access_token = os.environ["HF_TOKEN"] except KeyError: @@ -55,20 +104,26 @@ def run_model(): model_ckpt_path = snapshot_download(model_id, token=access_token) # Default: no sharding (single-device friendly). - config = modeling.ModelConfig.llama3_2_1b(use_fsdp=False, use_tp=False) + config_fn = modeling.ModelConfig.llama3_2_1b if model_size == "1B" else modeling.ModelConfig.llama3_2_3b + config = config_fn(use_fsdp=False, use_tp=False) fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) jax.set_mesh(mesh) batch_shd = None - prompts = [ + instruct_prompts = [ "Summarize what a tokenizer does in one paragraph.", "Write a short, friendly explanation of gradient descent for beginners.", ] + base_prompts = [ + "The capital of Japan is", + "The definition of a tokenizer in NLP is strictly defined as follows: A tokenizer is an algorithm that", + ] + prompts = base_prompts if use_base_model else instruct_prompts tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd) + tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd, use_chat_template=not use_base_model) batch_size, token_len = tokens.shape generate_steps = 64 diff --git a/bonsai/models/llama3_2/tests/run_model_base.py b/bonsai/models/llama3_2/tests/run_model_base.py deleted file mode 100644 index 04b0e6ec..00000000 --- a/bonsai/models/llama3_2/tests/run_model_base.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright 2026 The JAX Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -import sys -import time - -import jax -import jax.numpy as jnp -import numpy as np -from huggingface_hub import snapshot_download -from jax.sharding import AxisType -from transformers import AutoTokenizer - -from bonsai.models.llama3_2 import modeling, params -from bonsai.utils import Sampler - - -def tokenize(tokenizer, prompts: list[str], shd=None): - if tokenizer.pad_token_id is None: - tokenizer.pad_token = tokenizer.eos_token - tokenizer.padding_side = "left" - batch = tokenizer(prompts, padding=True, return_tensors="np") - input_ids = jnp.array(batch["input_ids"], out_sharding=shd) - attention_mask = jnp.array(batch["attention_mask"], out_sharding=shd) - return input_ids, attention_mask - - -def run_model(): - # Choose a checkpoint and config; defaults to the 1B base variant. - model_id = "meta-llama/Llama-3.2-1B" - try: - access_token = os.environ["HF_TOKEN"] - except KeyError: - print("\nError: HF_TOKEN is not set.", file=sys.stderr) - print("Please set the HF_TOKEN environment variable and retry.", file=sys.stderr) - sys.exit(1) - - model_ckpt_path = snapshot_download(model_id, token=access_token) - - # Default: no sharding (single-device friendly). - config = modeling.ModelConfig.llama3_2_1b(use_fsdp=False, use_tp=False) - fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value - mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) - jax.set_mesh(mesh) - batch_shd = None - - prompts = [ - "The capital of France is", - "The definition of a tokenizer in NLP is strictly defined as follows: A tokenizer is an algorithm that", - ] - - tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) - pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd) - batch_size, token_len = tokens.shape - - generate_steps = 64 - model = params.create_model_from_safe_tensors(model_ckpt_path, config, mesh) - cache = model.init_cache(config, batch_size, token_len, generate_steps) - - key = jax.random.key(0) - sampler = Sampler(temperature=1.0, top_p=0.9, top_k=50) - jit_sampler = jax.jit(sampler) - - # prefill - logits, cache = modeling.forward(model, cache, tokens, pad_id, attention_mask=attention_mask) - key, subkey = jax.random.split(key) - next_tokens = jit_sampler(logits, key=subkey) - - # decode - tokens_list = [next_tokens] - finished = jnp.zeros((batch_size,), dtype=jnp.bool_) - start = time.time() - for _ in range(generate_steps): - logits, cache = modeling.forward(model, cache, next_tokens, pad_id) - key, subkey = jax.random.split(key) - next_tokens = jit_sampler(logits, key=subkey) - finished = finished | (next_tokens.squeeze(-1) == tokenizer.eos_token_id) - tokens_list.append(next_tokens) - if finished.all(): - break - - elapsed = time.time() - start - all_output_tokens = jax.device_get(jnp.concatenate(tokens_list, axis=-1)) - print(f"Generated {all_output_tokens.shape[1]} tokens in {elapsed:.3f}s") - for i, prompt in enumerate(prompts): - print(f"Prompt:\n {prompt}") - seq_tokens = all_output_tokens[i] - eos_idx = np.where(seq_tokens == tokenizer.eos_token_id)[0] - if eos_idx.size > 0: - seq_tokens = seq_tokens[: eos_idx[0]] - decoded = tokenizer.decode(seq_tokens, skip_special_tokens=True) - print(f"Completion:\n {decoded}\n\n") - - -if __name__ == "__main__": - run_model() - - -__all__ = ["run_model"] From c047b0bffc837cad267c4785526038b8b4273a79 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Mon, 2 Feb 2026 10:44:37 +0900 Subject: [PATCH 16/18] refactor(llama3.2): add no mesh support Ref: #141 --- bonsai/models/llama3_2/modeling.py | 54 +++++++++++++---------- bonsai/models/llama3_2/tests/run_model.py | 15 +++---- 2 files changed, 35 insertions(+), 34 deletions(-) diff --git a/bonsai/models/llama3_2/modeling.py b/bonsai/models/llama3_2/modeling.py index 05841283..872ba0ab 100644 --- a/bonsai/models/llama3_2/modeling.py +++ b/bonsai/models/llama3_2/modeling.py @@ -37,35 +37,38 @@ class ShardMode(Enum): @dataclasses.dataclass(slots=True, frozen=True) class LlamaShardCfg: # Embedding - emb_weight: PartitionSpec - activation: PartitionSpec - logits: PartitionSpec + emb_weight: PartitionSpec | None = None + activation: PartitionSpec | None = None + logits: PartitionSpec | None = None # Attention - q_proj: PartitionSpec - k_proj: PartitionSpec - v_proj: PartitionSpec - o_proj: PartitionSpec + q_proj: PartitionSpec | None = None + k_proj: PartitionSpec | None = None + v_proj: PartitionSpec | None = None + o_proj: PartitionSpec | None = None - attn_logits: PartitionSpec - attn_out: PartitionSpec + attn_logits: PartitionSpec | None = None + attn_out: PartitionSpec | None = None - cache: PartitionSpec + cache: PartitionSpec | None = None # MLP - gate_proj: PartitionSpec - up_proj: PartitionSpec - down_proj: PartitionSpec + gate_proj: PartitionSpec | None = None + up_proj: PartitionSpec | None = None + down_proj: PartitionSpec | None = None # Head - lm_head: PartitionSpec + lm_head: PartitionSpec | None = None @classmethod def no_sharding(cls) -> "LlamaShardCfg": - return cls.default(use_fsdp=False, use_tp=False) + return LlamaShardCfg() @classmethod def default(cls, use_fsdp: bool, use_tp: bool) -> "LlamaShardCfg": + if not (use_fsdp and use_tp): + return cls.no_sharding() + fsdp = ShardMode.FSDP.value if use_fsdp else None tp = ShardMode.TP.value if use_tp else None @@ -87,7 +90,9 @@ def default(cls, use_fsdp: bool, use_tp: bool) -> "LlamaShardCfg": ) -def shard(x: jnp.ndarray, s: PartitionSpec): +def shard(x: jnp.ndarray, s: PartitionSpec | None): + if s is None: + return x mesh = get_abstract_mesh() if not mesh.empty and len(mesh.axis_names) > 0: return reshard(x, s) @@ -164,7 +169,8 @@ def __init__(self, cfg: ModelConfig, batch_size: int, cache_size: int, dtype: jn 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] - self.start_ind = nnx.Variable(-1 * jnp.ones((batch_size,), dtype=jnp.int32, out_sharding=P(kv_shd[0]))) + 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)) @@ -194,7 +200,7 @@ def __init__( self, in_dim: int, out_dim: int, - sharding: PartitionSpec, + sharding: PartitionSpec | None, *, use_bias: bool = True, dtype=jnp.bfloat16, @@ -372,8 +378,8 @@ def sharded_attention( attn_mask: Array | None, scale: float, *, - attn_logit_sharding: PartitionSpec, - out_sharding: PartitionSpec, + attn_logit_sharding: PartitionSpec | None, + out_sharding: PartitionSpec | None, ) -> Array: """Compute scaled dot-product attention with optional masking.""" attn_logits = jnp.einsum("BTKGH,BSKH->BTSKG", q, k, out_sharding=attn_logit_sharding) * scale @@ -476,7 +482,7 @@ def __call__( position_ids = compute_positions_from_segment_ids(segment_ids) if cache is not None: left_pads = count_left_pads(segment_ids) - cache.start_ind[...] = jnp.where(cache.start_ind[...] < 0, left_pads, cache.start_ind[...]) + cache.start_ind.set_value(jnp.where(cache.start_ind[...] < 0, left_pads, cache.start_ind[...])) position_ids = position_ids + cache.cur_ind[...] sin, cos = _generate_pos_embeddings( @@ -494,8 +500,8 @@ def __call__( cache_shd = self.config.shd_cfg.cache k = shard(k, cache_shd) v = shard(v, cache_shd) - cache.k_cache[...] = jax.lax.dynamic_update_slice(cache.k_cache[...], k, slice_indices) - cache.v_cache[...] = jax.lax.dynamic_update_slice(cache.v_cache[...], v, slice_indices) + cache.k_cache.set_value(jax.lax.dynamic_update_slice(cache.k_cache[...], k, slice_indices)) + cache.v_cache.set_value(jax.lax.dynamic_update_slice(cache.v_cache[...], v, slice_indices)) k = cache.k_cache[...] v = cache.v_cache[...] @@ -520,7 +526,7 @@ def __call__( ) if cache is not None: - cache.cur_ind[...] = cache.cur_ind[...] + t + cache.cur_ind.set_value(cache.cur_ind[...] + t) # Reshape back to [B, T, N * H] for the output projection. attn_output = attn_output.reshape((b, t, self.num_heads * self.head_dim)) diff --git a/bonsai/models/llama3_2/tests/run_model.py b/bonsai/models/llama3_2/tests/run_model.py index a35eb3fc..4d13fa01 100644 --- a/bonsai/models/llama3_2/tests/run_model.py +++ b/bonsai/models/llama3_2/tests/run_model.py @@ -23,7 +23,6 @@ import jax.numpy as jnp import numpy as np from huggingface_hub import snapshot_download -from jax.sharding import AxisType from transformers import AutoTokenizer @@ -60,7 +59,7 @@ def _size(value: str) -> str: return Args(model_size=parsed.size, use_base_model=parsed.base) -def tokenize(tokenizer, prompts: list[str], shd=None, use_chat_template: bool = True): +def tokenize(tokenizer, prompts: list[str], use_chat_template: bool = True): if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" @@ -77,8 +76,8 @@ def tokenize(tokenizer, prompts: list[str], shd=None, use_chat_template: bool = lines = prompts batch = tokenizer(lines, padding=True, return_tensors="np", add_special_tokens=not use_template) - input_ids = jnp.array(batch["input_ids"], out_sharding=shd) - attention_mask = jnp.array(batch["attention_mask"], out_sharding=shd) + input_ids = jnp.array(batch["input_ids"]) + attention_mask = jnp.array(batch["attention_mask"]) return input_ids, attention_mask @@ -106,10 +105,6 @@ def run_model(): # Default: no sharding (single-device friendly). config_fn = modeling.ModelConfig.llama3_2_1b if model_size == "1B" else modeling.ModelConfig.llama3_2_3b config = config_fn(use_fsdp=False, use_tp=False) - fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value - mesh = jax.make_mesh((1, 1), (fsdp, tp), axis_types=(AxisType.Explicit, AxisType.Explicit)) - jax.set_mesh(mesh) - batch_shd = None instruct_prompts = [ "Summarize what a tokenizer does in one paragraph.", @@ -123,11 +118,11 @@ def run_model(): tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - tokens, attention_mask = tokenize(tokenizer, prompts, batch_shd, use_chat_template=not use_base_model) + tokens, attention_mask = tokenize(tokenizer, prompts, use_chat_template=not use_base_model) batch_size, token_len = tokens.shape generate_steps = 64 - model = params.create_model_from_safe_tensors(model_ckpt_path, config, mesh) + model = params.create_model_from_safe_tensors(model_ckpt_path, config) cache = model.init_cache(config, batch_size, token_len, generate_steps) key = jax.random.key(0) From a32f929072b2b1b6c95bc6b24f11040fd1379b20 Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Mon, 2 Feb 2026 10:47:04 +0900 Subject: [PATCH 17/18] style(llama3.2): clean up code formatting --- bonsai/models/llama3_2/modeling.py | 1 + bonsai/models/llama3_2/tests/run_model.py | 6 +----- bonsai/models/llama3_2/tests/test_utils.py | 1 + 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/bonsai/models/llama3_2/modeling.py b/bonsai/models/llama3_2/modeling.py index 872ba0ab..a975c5e6 100644 --- a/bonsai/models/llama3_2/modeling.py +++ b/bonsai/models/llama3_2/modeling.py @@ -98,6 +98,7 @@ def shard(x: jnp.ndarray, s: PartitionSpec | None): return reshard(x, s) return x + @dataclasses.dataclass(frozen=True) class RopeScalingConfig: factor: float diff --git a/bonsai/models/llama3_2/tests/run_model.py b/bonsai/models/llama3_2/tests/run_model.py index 4d13fa01..3e3b4038 100644 --- a/bonsai/models/llama3_2/tests/run_model.py +++ b/bonsai/models/llama3_2/tests/run_model.py @@ -88,11 +88,7 @@ def run_model(): model_size = args.model_size use_base_model = args.use_base_model - model_id = ( - f"meta-llama/Llama-3.2-{model_size}" - if use_base_model - else f"meta-llama/Llama-3.2-{model_size}-Instruct" - ) + model_id = f"meta-llama/Llama-3.2-{model_size}" if use_base_model else f"meta-llama/Llama-3.2-{model_size}-Instruct" try: access_token = os.environ["HF_TOKEN"] except KeyError: diff --git a/bonsai/models/llama3_2/tests/test_utils.py b/bonsai/models/llama3_2/tests/test_utils.py index d6019c31..ff56e046 100644 --- a/bonsai/models/llama3_2/tests/test_utils.py +++ b/bonsai/models/llama3_2/tests/test_utils.py @@ -2,6 +2,7 @@ from bonsai.models.llama3_2 import modeling + def tiny_config(*, use_sharding: bool = False) -> modeling.ModelConfig: """Create a minimal Llama3.2 model configuration for testing purposes""" return modeling.ModelConfig( From c8404a06be6073a56d652c82ad35584a01b655da Mon Sep 17 00:00:00 2001 From: Moriyuki SUZUKI Date: Mon, 2 Feb 2026 11:18:24 +0900 Subject: [PATCH 18/18] style(llama3.2): rename test classes --- bonsai/models/llama3_2/tests/test_outputs_llama3_2.py | 4 ++-- bonsai/models/llama3_2/tests/test_padding_llama3_2.py | 2 +- bonsai/models/llama3_2/tests/test_sharding_llama3_2.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py index fd324b8d..e49b025f 100644 --- a/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_outputs_llama3_2.py @@ -29,8 +29,8 @@ def check_hf_token(): return False -@unittest.skipIf(check_hf_token(), "Skipping Llama32 output tests due to HF_TOKEN failure.") -class TestOutputsLlama32(absltest.TestCase): +@unittest.skipIf(check_hf_token(), "Skipping Llama3.2 output tests due to HF_TOKEN failure.") +class TestOutputsLlama3_2(absltest.TestCase): @classmethod def setUpClass(cls): super().setUpClass() diff --git a/bonsai/models/llama3_2/tests/test_padding_llama3_2.py b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py index bea68ed5..e3bd7e1b 100644 --- a/bonsai/models/llama3_2/tests/test_padding_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_padding_llama3_2.py @@ -9,7 +9,7 @@ from bonsai.models.llama3_2.tests.test_utils import tiny_config -class TestPaddingLlama32(absltest.TestCase): +class TestPaddingLlama3_2(absltest.TestCase): def test_forward_uses_per_sample_right_padding(self): cfg = tiny_config(use_sharding=False) fsdp, tp = modeling.ShardMode.FSDP.value, modeling.ShardMode.TP.value diff --git a/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py index 59afbd03..bdd6a6d2 100644 --- a/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py +++ b/bonsai/models/llama3_2/tests/test_sharding_llama3_2.py @@ -11,7 +11,7 @@ jax.config.update("jax_num_cpu_devices", 8) -class TestShardingLlama32(absltest.TestCase): +class TestShardingLlama3_2(absltest.TestCase): @classmethod def setUpClass(cls): super().setUpClass()