diff --git a/transformers-5.5.0.dev0-py3-none-any.whl b/transformers-5.5.0.dev0-py3-none-any.whl new file mode 100644 index 0000000000..65e8b8cd23 Binary files /dev/null and b/transformers-5.5.0.dev0-py3-none-any.whl differ diff --git a/unsloth/models/gemma4.py b/unsloth/models/gemma4.py new file mode 100644 index 0000000000..727ed9f761 --- /dev/null +++ b/unsloth/models/gemma4.py @@ -0,0 +1,87 @@ +# Copyright © 2026 Apple Inc. + +from dataclasses import dataclass +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten, tree_unflatten + +from . import gemma4_text +from .base import BaseModelArgs + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + text_config: dict + + @classmethod + def from_dict(cls, params): + if "text_config" not in params: + return cls(model_type = params["model_type"], text_config = params) + return super().from_dict(params) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.language_model = gemma4_text.Model( + gemma4_text.ModelArgs.from_dict(args.text_config) + ) + + def __call__( + self, + inputs: mx.array, + cache = None, + input_embeddings: Optional[mx.array] = None, + ): + return self.language_model( + inputs, + cache = cache, + input_embeddings = input_embeddings, + ) + + def sanitize(self, weights): + weights = tree_unflatten(list(weights.items())) + + if "model" in weights: + model_weights = weights["model"] + else: + model_weights = weights + + for key in [ + "vision_tower", + "embed_vision", + "audio_tower", + "embed_audio", + ]: + model_weights.pop(key, None) + + if "language_model" in model_weights: + source_lm_weights = dict(tree_flatten(model_weights["language_model"])) + else: + source_lm_weights = dict(tree_flatten(model_weights)) + + lm_weights = {} + for key, value in source_lm_weights.items(): + if key.startswith("model.") or key.startswith("lm_head."): + lm_weights[key] = value + else: + lm_weights[f"model.{key}"] = value + + lm_head = model_weights.get("lm_head", weights.get("lm_head")) + if isinstance(lm_head, dict) and "weight" in lm_head: + lm_weights["lm_head.weight"] = lm_head["weight"] + + lm_weights = self.language_model.sanitize(lm_weights) + return {f"language_model.{key}": value for key, value in lm_weights.items()} + + @property + def layers(self): + return self.language_model.layers + + def make_cache(self): + return self.language_model.make_cache() diff --git a/unsloth/models/gemma4_text.py b/unsloth/models/gemma4_text.py new file mode 100644 index 0000000000..fd6af0307b --- /dev/null +++ b/unsloth/models/gemma4_text.py @@ -0,0 +1,754 @@ +# Copyright © 2026 Apple Inc. + +from dataclasses import dataclass +from functools import partial +from typing import Any, Dict, Optional, Union + +import mlx.core as mx +import mlx.nn as nn + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import KVCache, RotatingKVCache +from .rope_utils import initialize_rope + + +def _gelu_pytorch_tanh(x: mx.array) -> mx.array: + return nn.gelu_approx(x) + + +ACT2FN = { + "gelu": nn.gelu, + "gelu_pytorch_tanh": _gelu_pytorch_tanh, + "silu": nn.silu, +} + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + rms_norm_eps: float + vocab_size: int + max_position_embeddings: int = 131072 + layer_types: Optional[list[str]] = None + sliding_window: int = 512 + rope_parameters: Optional[Dict[str, Dict[str, Any]]] = None + hidden_size_per_layer_input: int = 0 + vocab_size_per_layer_input: int = 262144 + num_global_key_value_heads: Optional[int] = None + global_head_dim: Optional[int] = None + attention_k_eq_v: bool = False + num_kv_shared_layers: int = 0 + use_double_wide_mlp: bool = False + enable_moe_block: bool = False + num_experts: Optional[int] = None + top_k_experts: Optional[int] = None + moe_intermediate_size: Optional[int] = None + hidden_activation: str = "gelu_pytorch_tanh" + tie_word_embeddings: bool = True + final_logit_softcapping: Optional[float] = None + attention_bias: bool = False + attention_dropout: float = 0.0 + use_bidirectional_attention: Optional[str] = None + + def __post_init__(self): + if self.layer_types is None: + self.layer_types = [ + "sliding_attention" if (i + 1) % 6 else "full_attention" + for i in range(self.num_hidden_layers) + ] + if self.layer_types[-1] != "full_attention": + self.layer_types[-1] = "full_attention" + if self.num_global_key_value_heads is None: + self.num_global_key_value_heads = self.num_key_value_heads + if self.global_head_dim is None: + self.global_head_dim = self.head_dim + if self.rope_parameters is None: + self.rope_parameters = { + "sliding_attention": { + "rope_type": "default", + "rope_theta": 10_000.0, + }, + "full_attention": { + "rope_type": "proportional", + "partial_rotary_factor": 0.25, + "rope_theta": 1_000_000.0, + }, + } + + +class Gemma4RMSNorm(nn.Module): + def __init__(self, dims: int, eps: float = 1e-6, with_scale: bool = True): + super().__init__() + self.eps = eps + self.with_scale = with_scale + if self.with_scale: + self.weight = mx.ones((dims,)) + + def __call__(self, x: mx.array) -> mx.array: + y = x.astype(mx.float32) + mean_squared = mx.mean(y * y, axis = -1, keepdims = True) + self.eps + y = y * mx.rsqrt(mean_squared) + if self.with_scale: + y = y * self.weight.astype(mx.float32) + return y.astype(x.dtype) + + +class Float32RoPE(nn.Module): + def __init__(self, rope: nn.Module): + super().__init__() + self.rope = rope + + def __call__(self, x: mx.array, offset: Union[int, mx.array] = 0) -> mx.array: + y = self.rope(x.astype(mx.float32), offset = offset) + return y.astype(x.dtype) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + first_kv_shared_layer_idx = args.num_hidden_layers - args.num_kv_shared_layers + is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 + use_double_wide_mlp = args.use_double_wide_mlp and is_kv_shared_layer + hidden_dim = args.intermediate_size * (2 if use_double_wide_mlp else 1) + + self.gate_proj = nn.Linear(args.hidden_size, hidden_dim, bias = False) + self.up_proj = nn.Linear(args.hidden_size, hidden_dim, bias = False) + self.down_proj = nn.Linear(hidden_dim, args.hidden_size, bias = False) + self.act = ACT2FN[args.hidden_activation] + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x)) + + +class Router(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.norm = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps, with_scale = False + ) + self.proj = nn.Linear(args.hidden_size, args.num_experts, bias = False) + self.scale = mx.ones((args.hidden_size,)) + self.per_expert_scale = mx.ones((args.num_experts,)) + self._root_size = args.hidden_size**-0.5 + + def __call__(self, x: mx.array): + x = self.norm(x) + x = x * self._root_size + x = x * self.scale + + expert_scores = self.proj(x) + router_probs = mx.softmax(expert_scores, axis = -1) + + top_k_indices = mx.argpartition( + -expert_scores, kth = self.args.top_k_experts - 1, axis = -1 + )[..., : self.args.top_k_experts] + + top_k_weights = mx.take_along_axis(router_probs, top_k_indices, axis = -1) + top_k_weights = top_k_weights / mx.sum(top_k_weights, axis = -1, keepdims = True) + top_k_weights = top_k_weights * self.per_expert_scale[top_k_indices] + return top_k_indices, top_k_weights + + +class GeGLU(nn.Module): + def __call__(self, x, gate): + return nn.gelu_approx(gate) * x + + +class Experts(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + try: + from .switch_layers import SwitchGLU + except ImportError: + raise ImportError( + "Gemma4 MoE requires mlx-lm >= 0.31. Please upgrade: pip install -U mlx-lm" + ) + + self.switch_glu = SwitchGLU( + input_dims = args.hidden_size, + hidden_dims = args.moe_intermediate_size, + num_experts = args.num_experts, + activation = GeGLU(), + bias = False, + ) + + def __call__( + self, + x: mx.array, + top_k_indices: mx.array, + top_k_weights: mx.array, + ) -> mx.array: + B, S, H = x.shape + K = top_k_indices.shape[-1] + + x_flat = x.reshape(B * S, H) + indices_flat = top_k_indices.reshape(B * S, K) + + expert_out = self.switch_glu(x_flat, indices_flat) + + weights = top_k_weights.reshape(B * S, K)[..., None] + return (expert_out * weights).sum(axis = -2).reshape(B, S, H) + + +def build_rope(args: ModelArgs, layer_type: str, head_dim: int): + rope_config = args.rope_parameters[layer_type] + rope_type = rope_config.get("rope_type", "default") + rope_theta = rope_config.get("rope_theta", 10_000.0) + + if rope_type == "proportional": + partial_rotary_factor = rope_config.get("partial_rotary_factor", 1.0) + rope_angles = int(partial_rotary_factor * head_dim // 2) + # Use full head_dim RoPE but with zero inv_freq for NoPE dimensions, + # matching HF's rotate_half pairing: (0, head_dim//2), (1, head_dim//2+1), ... + return Float32RoPE(ProportionalRoPE(head_dim, rope_angles, base = rope_theta)) + + return Float32RoPE( + initialize_rope( + dims = head_dim, + base = rope_theta, + traditional = False, + scaling_config = rope_config, + max_position_embeddings = args.max_position_embeddings, + ) + ) + + +class ProportionalRoPE(nn.Module): + """RoPE with partial rotation matching HF's rotate_half pairing. + + Rotates `rope_angles` pairs out of `head_dim // 2` total pairs. + Non-rotated pairs get cos=1, sin=0 (identity). + Pairing follows HF convention: (i, i + head_dim//2). + """ + + def __init__(self, head_dim: int, rope_angles: int, base: float = 10000.0): + super().__init__() + self.head_dim = head_dim + self.rope_angles = rope_angles + + inv_freq_rotated = 1.0 / ( + base ** (mx.arange(0, 2 * rope_angles, 2, dtype = mx.float32) / head_dim) + ) + nope_angles = head_dim // 2 - rope_angles + if nope_angles > 0: + self._inv_freq = mx.concatenate( + [inv_freq_rotated, mx.zeros(nope_angles, dtype = mx.float32)] + ) + else: + self._inv_freq = inv_freq_rotated + + def __call__(self, x: mx.array, offset: int = 0) -> mx.array: + # x shape: (B, n_heads, L, head_dim) + seq_len = x.shape[-2] + positions = mx.arange(seq_len, dtype = mx.float32) + offset + + # (L, head_dim//2) + freqs = mx.outer(positions, self._inv_freq) + # (L, head_dim) — interleaved cos/sin + cos = mx.cos(freqs) + sin = mx.sin(freqs) + + # HF-style rotate_half: split at head_dim//2 + half = self.head_dim // 2 + x1 = x[..., :half] + x2 = x[..., half:] + out = mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis = -1) + return out + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.layer_type = args.layer_types[layer_idx] + self.is_sliding = self.layer_type == "sliding_attention" + self.is_kv_shared_layer = ( + layer_idx >= (args.num_hidden_layers - args.num_kv_shared_layers) > 0 + ) + + self.n_heads = args.num_attention_heads + self.n_kv_heads = ( + args.num_key_value_heads + if self.is_sliding or not args.attention_k_eq_v + else args.num_global_key_value_heads + ) + self.head_dim = ( + args.head_dim + if self.is_sliding or not args.global_head_dim + else args.global_head_dim + ) + self.scale = 1.0 + self.use_alternative_attention = args.attention_k_eq_v and not self.is_sliding + + self.q_proj = nn.Linear( + args.hidden_size, + self.n_heads * self.head_dim, + bias = args.attention_bias, + ) + self.k_proj = nn.Linear( + args.hidden_size, + self.n_kv_heads * self.head_dim, + bias = args.attention_bias, + ) + self.v_proj = ( + None + if self.use_alternative_attention + else nn.Linear( + args.hidden_size, + self.n_kv_heads * self.head_dim, + bias = args.attention_bias, + ) + ) + self.o_proj = nn.Linear( + self.n_heads * self.head_dim, + args.hidden_size, + bias = args.attention_bias, + ) + + self.q_norm = Gemma4RMSNorm(self.head_dim, eps = args.rms_norm_eps) + self.k_norm = Gemma4RMSNorm(self.head_dim, eps = args.rms_norm_eps) + self.v_norm = Gemma4RMSNorm( + self.head_dim, eps = args.rms_norm_eps, with_scale = False + ) + self.rope = build_rope(args, self.layer_type, self.head_dim) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + position_offset: int = 0, + ) -> mx.array: + batch_size, seq_len, _ = x.shape + offset = position_offset + + queries = self.q_proj(x).reshape( + batch_size, seq_len, self.n_heads, self.head_dim + ) + queries = self.q_norm(queries).transpose(0, 2, 1, 3) + queries = self.rope(queries, offset = offset) + + _cache_not_empty = cache is not None and not ( + cache.empty() if hasattr(cache, "empty") else len(cache) == 0 + ) + if self.is_kv_shared_layer and _cache_not_empty: + keys, values = cache.state + else: + raw_keys = self.k_proj(x).reshape( + batch_size, seq_len, self.n_kv_heads, self.head_dim + ) + raw_values = ( + raw_keys + if self.v_proj is None + else self.v_proj(x).reshape( + batch_size, seq_len, self.n_kv_heads, self.head_dim + ) + ) + + keys = self.k_norm(raw_keys).transpose(0, 2, 1, 3) + keys = self.rope(keys, offset = offset) + + values = self.v_norm(raw_values).transpose(0, 2, 1, 3) + + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, + keys, + values, + cache = cache, + scale = self.scale, + mask = mask, + ) + output = output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, -1) + return self.o_proj(output) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.hidden_size_per_layer_input = args.hidden_size_per_layer_input + self.self_attn = Attention(args, layer_idx) + self.mlp = MLP(args, layer_idx) + + self.input_layernorm = Gemma4RMSNorm(args.hidden_size, eps = args.rms_norm_eps) + self.post_attention_layernorm = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + self.pre_feedforward_layernorm = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + self.post_feedforward_layernorm = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + self.layer_scalar = mx.ones((1,)) + + # MoE + self.enable_moe = args.enable_moe_block + if self.enable_moe: + self.router = Router(args) + self.experts = Experts(args) + self.post_feedforward_layernorm_1 = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + self.post_feedforward_layernorm_2 = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + self.pre_feedforward_layernorm_2 = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + + if self.hidden_size_per_layer_input: + self.act = ACT2FN[args.hidden_activation] + self.per_layer_input_gate = nn.Linear( + args.hidden_size, args.hidden_size_per_layer_input, bias = False + ) + self.per_layer_projection = nn.Linear( + args.hidden_size_per_layer_input, args.hidden_size, bias = False + ) + self.post_per_layer_input_norm = Gemma4RMSNorm( + args.hidden_size, eps = args.rms_norm_eps + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + per_layer_input: Optional[mx.array] = None, + position_offset: int = 0, + ) -> mx.array: + residual = x + h = self.input_layernorm(x) + h = self.self_attn(h, mask, cache, position_offset = position_offset) + h = self.post_attention_layernorm(h) + h = residual + h + + residual = h + + if self.enable_moe: + h1 = self.pre_feedforward_layernorm(h) + h1 = self.mlp(h1) + h1 = self.post_feedforward_layernorm_1(h1) + + top_k_indices, top_k_weights = self.router(h) + h2 = self.pre_feedforward_layernorm_2(h) + h2 = self.experts(h2, top_k_indices, top_k_weights) + h2 = self.post_feedforward_layernorm_2(h2) + + h = h1 + h2 + else: + h = self.pre_feedforward_layernorm(h) + h = self.mlp(h) + + h = self.post_feedforward_layernorm(h) + h = residual + h + + if self.hidden_size_per_layer_input: + residual = h + h = self.per_layer_input_gate(h) + h = self.act(h) + h = h * per_layer_input + h = self.per_layer_projection(h) + h = self.post_per_layer_input_norm(h) + h = residual + h + + return h * self.layer_scalar + + +@partial(mx.compile, shapeless = True) +def logit_softcap(softcap, x): + out = mx.tanh(x / softcap) + out = out * softcap + return out + + +class Gemma4Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + self.hidden_size_per_layer_input = args.hidden_size_per_layer_input + self.first_kv_shared_layer_idx = ( + args.num_hidden_layers - args.num_kv_shared_layers + ) + self.embed_scale = args.hidden_size**0.5 + self.per_layer_embed_scale = args.hidden_size_per_layer_input**0.5 + self.per_layer_projection_scale = args.hidden_size**-0.5 + self.per_layer_input_scale = 2.0**-0.5 + + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args = args, layer_idx = i) + for i in range(args.num_hidden_layers) + ] + self.norm = Gemma4RMSNorm(args.hidden_size, eps = args.rms_norm_eps) + + if self.hidden_size_per_layer_input: + self.embed_tokens_per_layer = nn.Embedding( + args.vocab_size_per_layer_input, + args.num_hidden_layers * args.hidden_size_per_layer_input, + ) + self.per_layer_model_projection = nn.Linear( + args.hidden_size, + args.num_hidden_layers * args.hidden_size_per_layer_input, + bias = False, + ) + self.per_layer_projection_norm = Gemma4RMSNorm( + args.hidden_size_per_layer_input, + eps = args.rms_norm_eps, + ) + + concrete_layers = args.layer_types[: self.first_kv_shared_layer_idx] + concrete_layer_types = set(concrete_layers) + for layer_type in args.layer_types[self.first_kv_shared_layer_idx :]: + if layer_type not in concrete_layer_types: + raise ValueError( + "num_kv_shared_layers requires at least one earlier " + f"{layer_type!r} layer before the shared suffix." + ) + self.layer_idx_to_cache_idx = [] + for i, layer_type in enumerate(args.layer_types): + if i < self.first_kv_shared_layer_idx: + self.layer_idx_to_cache_idx.append(i) + continue + + shared_idx = ( + len(concrete_layers) - 1 - concrete_layers[::-1].index(layer_type) + ) + self.layer_idx_to_cache_idx.append(shared_idx) + + self.first_full_idx = next( + ( + self.layer_idx_to_cache_idx[i] + for i, layer_type in enumerate(args.layer_types) + if layer_type == "full_attention" + ), + None, + ) + self.first_sliding_idx = next( + ( + self.layer_idx_to_cache_idx[i] + for i, layer_type in enumerate(args.layer_types) + if layer_type == "sliding_attention" + ), + None, + ) + + def get_input_embeddings(self, input_ids: mx.array) -> mx.array: + return self.embed_tokens(input_ids) * self.embed_scale + + def get_per_layer_inputs( + self, + input_ids: Optional[mx.array], + input_embeddings: Optional[mx.array], + ) -> mx.array: + if input_ids is None: + if input_embeddings is None: + raise ValueError( + "Either input ids or input embeddings are required for Gemma4 per-layer inputs." + ) + exact_matches = mx.all( + input_embeddings[:, :, None, :] + == self.embed_tokens.weight[None, None, :, :] * self.embed_scale, + axis = -1, + ) + if not mx.all(mx.sum(exact_matches, axis = -1) == 1): + raise ValueError( + "Gemma4 input embeddings must exactly match embed_tokens when " + "input ids are omitted." + ) + input_ids = mx.argmax(exact_matches, axis = -1).astype(mx.int32) + + tokens = mx.where( + input_ids < self.args.vocab_size_per_layer_input, + input_ids, + mx.zeros_like(input_ids), + ) + result = self.embed_tokens_per_layer(tokens) * self.per_layer_embed_scale + return result.reshape( + *input_ids.shape, + self.args.num_hidden_layers, + self.args.hidden_size_per_layer_input, + ) + + def project_per_layer_inputs( + self, + inputs_embeds: mx.array, + per_layer_inputs: mx.array, + ) -> mx.array: + per_layer_projection = ( + self.per_layer_model_projection(inputs_embeds) + * self.per_layer_projection_scale + ) + per_layer_projection = per_layer_projection.reshape( + *inputs_embeds.shape[:-1], + self.args.num_hidden_layers, + self.args.hidden_size_per_layer_input, + ) + per_layer_projection = self.per_layer_projection_norm(per_layer_projection) + return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale + + def __call__( + self, + inputs: Optional[mx.array], + cache = None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + if input_embeddings is None: + h = self.get_input_embeddings(inputs) + else: + h = input_embeddings + + per_layer_inputs = None + if self.hidden_size_per_layer_input: + per_layer_inputs = self.get_per_layer_inputs(inputs, h) + per_layer_inputs = self.project_per_layer_inputs(h, per_layer_inputs) + + if cache is None: + if self.first_kv_shared_layer_idx < self.num_hidden_layers: + # Must create real caches so template layers store KV + # for shared layers to reuse — even without external cache + cache = [] + for layer_type in self.args.layer_types[ + : self.first_kv_shared_layer_idx + ]: + if layer_type == "full_attention": + cache.append(KVCache()) + else: + cache.append( + RotatingKVCache(max_size = self.args.sliding_window, keep = 0) + ) + else: + cache = [None] * self.num_hidden_layers + + global_mask = ( + None + if self.first_full_idx is None + else create_attention_mask(h, cache[self.first_full_idx]) + ) + sliding_mask = ( + None + if self.first_sliding_idx is None + else create_attention_mask( + h, + cache[self.first_sliding_idx], + window_size = self.args.sliding_window, + ) + ) + global_offset = ( + 0 + if self.first_full_idx is None or cache[self.first_full_idx] is None + else cache[self.first_full_idx].offset + ) + sliding_offset = ( + 0 + if self.first_sliding_idx is None or cache[self.first_sliding_idx] is None + else cache[self.first_sliding_idx].offset + ) + + for i, layer in enumerate(self.layers): + layer_type = self.args.layer_types[i] + mask = global_mask if layer_type == "full_attention" else sliding_mask + position_offset = ( + global_offset if layer_type == "full_attention" else sliding_offset + ) + per_layer_input = ( + None if per_layer_inputs is None else per_layer_inputs[:, :, i, :] + ) + cache_entry = ( + None if cache is None else cache[self.layer_idx_to_cache_idx[i]] + ) + h = layer( + h, + mask = mask, + cache = cache_entry, + per_layer_input = per_layer_input, + position_offset = position_offset, + ) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Gemma4Model(args) + self.tie_word_embeddings = False + self.final_logit_softcapping = args.final_logit_softcapping + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias = False) + + def __call__( + self, + inputs: Optional[mx.array], + cache = None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + out = self.model(inputs, cache = cache, input_embeddings = input_embeddings) + if self.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + if self.final_logit_softcapping is not None: + out = logit_softcap(self.final_logit_softcapping, out) + return out + + def sanitize(self, weights): + if "lm_head.weight" not in weights: + self.tie_word_embeddings = True + self.pop("lm_head") + + sanitized = {} + for k, v in weights.items(): + if "rotary_emb" in k: + continue + + if k.endswith(".experts.down_proj"): + k = k.replace( + ".experts.down_proj", ".experts.switch_glu.down_proj.weight" + ) + sanitized[k] = v + continue + + if k.endswith(".experts.gate_up_proj"): + gate_key = k.replace( + ".experts.gate_up_proj", ".experts.switch_glu.gate_proj.weight" + ) + up_key = k.replace( + ".experts.gate_up_proj", ".experts.switch_glu.up_proj.weight" + ) + v = v.swapaxes(-1, -2) + mid_dim = v.shape[-1] // 2 + sanitized[gate_key] = v[..., :mid_dim].swapaxes(-1, -2) + sanitized[up_key] = v[..., mid_dim:].swapaxes(-1, -2) + continue + + sanitized[k] = v + return sanitized + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + caches = [] + for layer_type in self.args.layer_types[: self.model.first_kv_shared_layer_idx]: + if layer_type == "full_attention": + caches.append(KVCache()) + else: + caches.append( + RotatingKVCache(max_size = self.args.sliding_window, keep = 0) + ) + return caches diff --git a/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py b/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py new file mode 100644 index 0000000000..01c64357c3 --- /dev/null +++ b/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py @@ -0,0 +1,1653 @@ +import argparse +import codecs +import contextlib +import functools +import json +import time +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_reduce +from mlx_lm.generate import maybe_quantize_kv_cache as mlx_maybe_quantize_kv_cache +from mlx_lm.sample_utils import make_logits_processors, make_sampler +from tqdm import tqdm +from transformers import PreTrainedTokenizer + +from .models import cache +from .prompt_utils import apply_chat_template +from .turboquant import TurboQuantKVCache, turboquant_enabled +from .utils import ( + StoppingCriteria, + ThinkingBudgetCriteria, + group_images_by_shape, + load, + prepare_inputs, +) + +DEFAULT_MODEL_PATH = "mlx-community/nanoLLaVA-1.5-8bit" +DEFAULT_IMAGE = None +DEFAULT_AUDIO = None +DEFAULT_PROMPT = "What are these?" +DEFAULT_MAX_TOKENS = 256 +DEFAULT_TEMPERATURE = 0.0 +DEFAULT_TOP_P = 1.0 +DEFAULT_SEED = 0 +DEFAULT_TOP_K = 0 +DEFAULT_MIN_P = 0.0 +DEFAULT_REPETITION_CONTEXT_SIZE = 20 +DEFAULT_KV_GROUP_SIZE = 64 +DEFAULT_KV_QUANT_SCHEME = "uniform" +DEFAULT_COMPLETION_BATCH_SIZE = 32 +DEFAULT_PREFILL_BATCH_SIZE = 8 +DEFAULT_THINKING_START_TOKEN = "" +DEFAULT_THINKING_END_TOKEN = "" +DEFAULT_QUANTIZED_KV_START = 5000 +DEFAULT_PREFILL_STEP_SIZE = 2048 + + +def parse_arguments(): + parser = argparse.ArgumentParser( + description = "Generate text from an image using a model." + ) + parser.add_argument( + "--model", + type = str, + default = DEFAULT_MODEL_PATH, + help = "The path to the local model directory or Hugging Face repo.", + ) + parser.add_argument( + "--adapter-path", + type = str, + default = None, + help = "The path to the adapter weights.", + ) + parser.add_argument( + "--image", + type = str, + nargs = "+", + default = DEFAULT_IMAGE, + help = "URL or path of the image to process.", + ) + parser.add_argument( + "--audio", + type = str, + nargs = "+", + default = DEFAULT_AUDIO, + help = "URL or path of the audio to process.", + ) + parser.add_argument( + "--resize-shape", + type = int, + nargs = "+", + default = None, + help = "Resize shape for the image.", + ) + parser.add_argument( + "--prompt", + type = str, + nargs = "+", + default = DEFAULT_PROMPT, + help = "Message to be processed by the model.", + ) + parser.add_argument( + "--system", + type = str, + default = None, + help = "System message for the model.", + ) + parser.add_argument( + "--max-tokens", + type = int, + default = DEFAULT_MAX_TOKENS, + help = "Maximum number of tokens to generate.", + ) + parser.add_argument( + "--temperature", + type = float, + default = DEFAULT_TEMPERATURE, + help = "Temperature for sampling.", + ) + parser.add_argument("--chat", action = "store_true", help = "Chat in multi-turn style.") + parser.add_argument("--verbose", action = "store_false", help = "Detailed output.") + parser.add_argument( + "--eos-tokens", + type = str, + nargs = "+", + default = None, + help = "EOS tokens to add to the tokenizer.", + ) + parser.add_argument( + "--max-kv-size", + type = int, + default = None, + help = "Maximum KV size for the prompt cache.", + ) + parser.add_argument( + "--kv-bits", + type = float, + default = None, + help = "Number of bits to quantize the KV cache to.", + ) + parser.add_argument( + "--kv-quant-scheme", + type = str, + choices = ("uniform", "turboquant"), + default = DEFAULT_KV_QUANT_SCHEME, + help = "KV cache quantization backend. Fractional --kv-bits values use " + "TurboQuant automatically.", + ) + parser.add_argument( + "--kv-group-size", + type = int, + default = DEFAULT_KV_GROUP_SIZE, + help = "Group size for uniform KV cache quantization.", + ) + parser.add_argument( + "--quantized-kv-start", + type = int, + default = DEFAULT_QUANTIZED_KV_START, + help = "Start index for the quantized KV cache.", + ) + parser.add_argument( + "--skip-special-tokens", + action = "store_true", + help = "Skip special tokens in the detokenizer.", + ) + parser.add_argument( + "--force-download", + action = "store_true", + help = "Force download the model from Hugging Face.", + ) + parser.add_argument( + "--revision", + type = str, + default = "main", + help = "The specific model version to use (branch, tag, commit).", + ) + parser.add_argument( + "--trust-remote-code", + action = "store_true", + help = "Trust remote code when loading the model.", + ) + parser.add_argument( + "--quantize-activations", + "-qa", + action = "store_true", + help = "Enable activation quantization for QQLinear layers. " + "Only supported for models quantized with 'nvfp4' or 'mxfp8' modes.", + ) + parser.add_argument( + "--processor-kwargs", + type = json.loads, + default = {}, + help = "Extra processor kwargs as JSON. " + 'Example: --processor-kwargs \'{"cropping": false, "max_patches": 3}\'', + ) + parser.add_argument( + "--prefill-step-size", + type = int, + default = DEFAULT_PREFILL_STEP_SIZE, + help = "Number of tokens to process per prefill step. " + "Lower values reduce peak memory usage but may be slower. " + "Try 512 or 256 if you hit GPU memory errors during prefill.", + ) + parser.add_argument( + "--enable-thinking", + action = "store_true", + help = "Enable thinking mode in the chat template (e.g. for Qwen3.5).", + ) + parser.add_argument( + "--thinking-budget", + type = int, + default = None, + help = "Maximum number of thinking tokens before forcing the end-of-thinking token.", + ) + parser.add_argument( + "--thinking-start-token", + type = str, + default = DEFAULT_THINKING_START_TOKEN, + help = "Token that marks the start of a thinking block (default: %(default)s).", + ) + parser.add_argument( + "--thinking-end-token", + type = str, + default = DEFAULT_THINKING_END_TOKEN, + help = "Token that marks the end of a thinking block (default: %(default)s).", + ) + + return parser.parse_args() + + +def normalize_resize_shape( + values: Optional[Sequence[int]], +) -> Optional[Tuple[int, int]]: + if values is None: + return None + if not ( + isinstance(values, Sequence) + and not isinstance(values, (str, bytes)) + and len(values) in (1, 2) + and all(type(value) is int for value in values) + ): + raise ValueError("resize_shape must contain 1 or 2 integers") + return (values[0], values[0]) if len(values) == 1 else tuple(values) + + +# A stream on the default device just for generation +generation_stream = mx.new_stream(mx.default_device()) + + +def maybe_quantize_kv_cache( + prompt_cache, + quantized_kv_start, + kv_group_size, + kv_bits, + kv_quant_scheme: str = DEFAULT_KV_QUANT_SCHEME, +): + if kv_bits is None: + return + + if turboquant_enabled(kv_bits, kv_quant_scheme): + + def quantize_entry(entry): + if isinstance(entry, TurboQuantKVCache): + return entry + if isinstance(entry, cache.RotatingKVCache): + return entry + if isinstance(entry, cache.KVCache): + if entry.offset == 0: + # Empty: replace so update_and_fetch quantizes on the fly + return TurboQuantKVCache(bits = kv_bits) + if entry.offset < quantized_kv_start: + return entry + return TurboQuantKVCache.from_cache(entry, bits = kv_bits) + if isinstance(entry, cache.CacheList): + entry.caches = [quantize_entry(sub_entry) for sub_entry in entry.caches] + return entry + if isinstance(entry, list): + for i, sub_entry in enumerate(entry): + entry[i] = quantize_entry(sub_entry) + return entry + if isinstance(entry, tuple): + return tuple(quantize_entry(sub_entry) for sub_entry in entry) + return entry + + # Skip the last layer (before final norm/LM head) — it's highly + # sensitive to quantization in deep models (e.g. gemma-4-31b). + last_idx = len(prompt_cache) - 1 if len(prompt_cache) > 2 else -1 + for index, layer_cache in enumerate(prompt_cache): + if index == last_idx: + continue + prompt_cache[index] = quantize_entry(layer_cache) + return + + mlx_maybe_quantize_kv_cache( + prompt_cache, + quantized_kv_start = quantized_kv_start, + kv_group_size = kv_group_size, + kv_bits = int(kv_bits), + ) + + +@contextlib.contextmanager +def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): + """ + A context manager to temporarily change the wired limit. + + Note, the wired limit should not be changed during an async eval. If an + async eval could be running pass in the streams to synchronize with prior + to exiting the context manager. + """ + if not mx.metal.is_available(): + yield + return + + model_bytes = tree_reduce( + lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0 + ) + max_rec_size = mx.device_info()["max_recommended_working_set_size"] + if model_bytes > 0.9 * max_rec_size: + model_mb = model_bytes // 2**20 + max_rec_mb = max_rec_size // 2**20 + print( + f"[WARNING] Generating with a model that requires {model_mb} MB " + f"which is close to the maximum recommended size of {max_rec_mb} " + "MB. This can be slow. See the documentation for possible work-arounds: " + "https://github.com/ml-explore/mlx-lm/tree/main#large-models" + ) + old_limit = mx.set_wired_limit(max_rec_size) + try: + yield + finally: + if streams is not None: + for s in streams: + mx.synchronize(s) + else: + mx.synchronize() + mx.set_wired_limit(old_limit) + + +@dataclass +class GenerationResult: + text: str = "" + token: Optional[int] = None + logprobs: Optional[List[float]] = None + prompt_tokens: int = 0 + generation_tokens: int = 0 + total_tokens: int = 0 + prompt_tps: float = 0.0 + generation_tps: float = 0.0 + peak_memory: float = 0.0 + + +class PromptCacheState: + """Holds KV cache and token history across conversation turns. + + Pass this to stream_generate via the ``prompt_cache_state`` kwarg to + reuse the KV cache from previous turns. Only the new tokens (after + the common prefix) are processed, avoiding redundant prefill. + """ + + def __init__(self): + self.cache: Optional[List[Any]] = None + self.token_ids: Optional[List[int]] = None + + def find_prefix_length(self, new_ids: list) -> int: + """Return the number of leading tokens that match the cached ids.""" + if self.token_ids is None: + return 0 + max_len = min(len(self.token_ids), len(new_ids)) + for i in range(max_len): + if self.token_ids[i] != new_ids[i]: + return i + return max_len + + def update(self, token_ids: list, kv_cache: list): + """Store the full token sequence and corresponding KV cache.""" + self.token_ids = list(token_ids) + self.cache = kv_cache + + +def generate_step( + input_ids: mx.array, + model: nn.Module, + pixel_values, + mask, + *, + max_tokens: int = DEFAULT_MAX_TOKENS, + temperature: float = DEFAULT_TEMPERATURE, + repetition_penalty: Optional[float] = None, + repetition_context_size: Optional[int] = DEFAULT_REPETITION_CONTEXT_SIZE, + top_p: float = DEFAULT_TOP_P, + min_p: float = DEFAULT_MIN_P, + top_k: int = DEFAULT_TOP_K, + logit_bias: Optional[Dict[int, float]] = None, + prompt_cache: Optional[List[Any]] = None, + max_kv_size: Optional[int] = None, + kv_bits: Optional[float] = None, + kv_group_size: int = DEFAULT_KV_GROUP_SIZE, + kv_quant_scheme: str = DEFAULT_KV_QUANT_SCHEME, + quantized_kv_start: int = DEFAULT_QUANTIZED_KV_START, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None, + prefill_step_size: Optional[int] = DEFAULT_PREFILL_STEP_SIZE, + **kwargs, +) -> Generator[Tuple[mx.array, mx.array], None, None]: + """ + A generator producing token ids based on the given prompt from the model. + + Args: + input_ids (mx.array): The input prompt token ids. + model (nn.Module): The model to use for generation. + pixel_values: The pixel values for vision models (optional). + mask: The attention mask (optional). + max_tokens (int): Maximum number of tokens to generate. + temperature (float): The temperature for sampling, if 0 the argmax is used. + repetition_penalty (float, optional): The penalty factor for repeating + tokens. + repetition_context_size (int, optional): The number of tokens to + consider for repetition penalty. + top_p (float, optional): Nucleus sampling, higher means model considers + more less likely words. + min_p (float, optional): Minimum probability threshold relative to the + highest-probability token. + top_k (int, optional): Restrict sampling to the top-k tokens. + logit_bias (dictionary, optional): Additive logit bias. + prompt_cache (list, optional): Pre-existing KV cache for the prompt. + max_kv_size (int, optional): Maximum KV cache size. + kv_bits (float, optional): Number of bits for KV cache quantization. + kv_group_size (int): Group size for uniform KV cache quantization. + kv_quant_scheme (str): KV cache quantization backend. + quantized_kv_start (int): Start index for quantized KV cache. + sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a + token from a vector of log probabilities. + logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional): + A list of functions that take tokens and logits and return the processed + logits. + prefill_step_size (int): Number of tokens to process per prefill step. + Chunked prefill processes prompts in smaller chunks to reduce peak + memory usage. + + Yields: + Generator[Tuple[mx.array, mx.array], None, None]: A generator producing + one token and a vector of log probabilities. + """ + + quantize_cache_fn = functools.partial( + maybe_quantize_kv_cache, + quantized_kv_start = quantized_kv_start, + kv_group_size = kv_group_size, + kv_bits = kv_bits, + kv_quant_scheme = kv_quant_scheme, + ) + + if sampler is None: + sampler = make_sampler( + temp = temperature, + top_p = top_p, + min_p = min_p, + top_k = top_k, + ) + + processors = make_logits_processors( + logit_bias, repetition_penalty, repetition_context_size + ) + if logits_processors is not None: + processors.extend(logits_processors) + + y = input_ids + tokens = mx.array([], dtype = input_ids.dtype) + + thinking_budget_criteria = kwargs.pop("thinking_budget_criteria", None) + + # Create the KV cache for generation + if prompt_cache is None: + prompt_cache = cache.make_prompt_cache( + model.language_model, + max_kv_size = max_kv_size, + ) + + def _step(y, inputs_embeds = None): + nonlocal tokens, kwargs + + with mx.stream(generation_stream): + if "decoder_input_ids" in kwargs: + outputs = model.language_model( + cache = prompt_cache, + **kwargs, + ) + else: + outputs = model.language_model( + y, + inputs_embeds = inputs_embeds, + cache = prompt_cache, + **kwargs, + ) + + logits = outputs.logits[:, -1, :] + + if len(processors) > 0 and len(y) > 0: + tokens = mx.concat([tokens, y.flatten()]) + + for processor in processors: + logits = processor(tokens, logits) + + quantize_cache_fn(prompt_cache) + + logprobs = logits - mx.logsumexp(logits) + y = sampler(logprobs) + + if outputs.cross_attention_states is not None: + kwargs = {"cross_attention_states": outputs.cross_attention_states} + elif outputs.encoder_outputs is not None: + kwargs = {"encoder_outputs": outputs.encoder_outputs} + else: + kwargs = {} + + return y, logprobs.squeeze(0) + + with mx.stream(generation_stream): + # Get input embeddings (handles both multimodal and text-only) + embedding_output = model.get_input_embeddings( + input_ids, pixel_values, mask = mask, **kwargs + ) + + inputs_embeds = embedding_output.inputs_embeds + + kwargs.update( + { + k: v + for k, v in embedding_output.to_dict().items() + if k != "inputs_embeds" and v is not None + } + ) + if getattr(model, "no_chunked_prefill", False): + prefill_step_size = None + if prefill_step_size is not None and inputs_embeds.shape[1] > prefill_step_size: + # Chunked prefill with embeddings + total_tokens = inputs_embeds.shape[1] + with tqdm(total = total_tokens, desc = "Prefill", unit = "tok") as pbar: + while inputs_embeds.shape[1] > 1: + n_to_process = min(prefill_step_size, inputs_embeds.shape[1] - 1) + model.language_model( + inputs = input_ids[:, :n_to_process], + inputs_embeds = inputs_embeds[:, :n_to_process], + cache = prompt_cache, + n_to_process = n_to_process, + **kwargs, + ) + quantize_cache_fn(prompt_cache) + mx.eval([c.state for c in prompt_cache]) + inputs_embeds = inputs_embeds[:, n_to_process:] + input_ids = input_ids[:, n_to_process:] + mx.clear_cache() + pbar.update(n_to_process) + + input_ids = input_ids[:, -1:] + + y, logprobs = _step(input_ids, inputs_embeds = inputs_embeds) + + mx.async_eval(y) + + n = 0 + while True: + if n != max_tokens: + next_y, next_logprobs = _step(y[None]) + mx.async_eval(next_y) + if n == 0: + mx.eval(y) + if n == max_tokens: + break + + yield y.item(), logprobs + if n % 256 == 0: + mx.clear_cache() + + if thinking_budget_criteria is not None: + next_y = thinking_budget_criteria.apply_forced_token(next_y) + y, logprobs = next_y, next_logprobs + n += 1 + + +def stream_generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + **kwargs, +) -> Union[str, Generator[str, None, None]]: + """ + A generator producing text based on the given prompt from the model. + + Args: + model (nn.Module): The model to use for generation. + processor (PreTrainedTokenizer): The tokenizer/processor. + prompt (str): The input prompt text. + image (Union[str, List[str]], optional): Image path(s) or URL(s). + audio (Union[str, List[str]], optional): Audio file path(s). + prefill_step_size (int, optional): Number of tokens to process per prefill + step. When set, enables chunked prefill which processes long prompts in + smaller chunks to reduce peak memory usage. + kwargs: Additional options passed to :func:`generate_step`. + See :func:`generate_step` for more details. + + Yields: + Generator[GenerationResult]: A generator producing GenerationResult objects + containing the generated text, tokens, and statistics. + """ + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Set up thinking budget criteria if requested + thinking_budget = kwargs.pop("thinking_budget", None) + thinking_end_token = kwargs.pop("thinking_end_token", DEFAULT_THINKING_END_TOKEN) + thinking_start_token = kwargs.pop( + "thinking_start_token", DEFAULT_THINKING_START_TOKEN + ) + enable_thinking = kwargs.pop("enable_thinking", False) + + # Skip special tokens + skip_special_tokens = kwargs.pop("skip_special_tokens", False) + skip_special_token_ids = ( + set(tokenizer.all_special_ids) + if skip_special_tokens and hasattr(tokenizer, "all_special_ids") + else [] + ) + + add_special_tokens = ( + getattr(processor, "chat_template", None) is None + if model.config.model_type in ["gemma3", "gemma3n", "gemma4"] + else True + ) + + resize_shape = normalize_resize_shape(kwargs.pop("resize_shape", None)) + image_token_index = getattr(model.config, "image_token_index", None) + vision_cache = kwargs.pop("vision_cache", None) + + if kwargs.get("input_ids", None) is not None: + input_ids = kwargs.pop("input_ids") + pixel_values = kwargs.pop("pixel_values", None) + mask = kwargs.pop("mask", None) + else: + inputs = prepare_inputs( + processor, + images = image, + audio = audio, + prompts = prompt, + image_token_index = image_token_index, + resize_shape = resize_shape, + add_special_tokens = add_special_tokens, + **kwargs, + ) + input_ids = inputs.get("input_ids", None) + pixel_values = inputs.get("pixel_values", None) + mask = inputs.get("attention_mask", None) + data_kwargs = { + k: v + for k, v in inputs.items() + if k not in ["input_ids", "pixel_values", "attention_mask"] + } + kwargs.update(data_kwargs) + + # Vision feature caching: reuse cached image features across turns + if vision_cache is not None and image is not None and pixel_values is not None: + cached = vision_cache.get(image) + if cached is not None: + kwargs["cached_image_features"] = cached + elif hasattr(model, "encode_image"): + features = model.encode_image(pixel_values) + mx.eval(features) + vision_cache.put(image, features) + kwargs["cached_image_features"] = features + + # Prompt cache reuse: skip common prefix from previous turn + prompt_cache_state = kwargs.pop("prompt_cache_state", None) + reused_prefix_len = 0 + full_input_ids_list = input_ids.flatten().tolist() + + if prompt_cache_state is not None and prompt_cache_state.cache is not None: + prefix_len = prompt_cache_state.find_prefix_length(full_input_ids_list) + if prefix_len > 0 and prefix_len < input_ids.shape[1]: + reused_prefix_len = prefix_len + # Trim to only new tokens + input_ids = input_ids[:, prefix_len:] + if mask is not None: + mask = mask[:, prefix_len:] + # Only skip vision if no image tokens in the new (trimmed) tokens + image_token_id = getattr(model.config, "image_token_id", None) or getattr( + model.config, "image_token_index", None + ) + new_ids = input_ids.flatten().tolist() + has_image_in_new = image_token_id is not None and image_token_id in new_ids + if not has_image_in_new: + pixel_values = None + kwargs.pop("cached_image_features", None) + # Reuse the saved KV cache (trimmed to prefix length) + kv_cache = prompt_cache_state.cache + # Trim cache to prefix_len in case it includes generated tokens + for c in kv_cache: + if hasattr(c, "keys") and c.keys is not None: + cached_len = c.keys.shape[2] + if cached_len > prefix_len: + c.keys = c.keys[:, :, :prefix_len, :] + c.values = c.values[:, :, :prefix_len, :] + if hasattr(c, "offset"): + c.offset = prefix_len + kwargs["prompt_cache"] = kv_cache + + if thinking_budget is not None: + thinking_start_token_id = tokenizer.encode( + thinking_start_token, add_special_tokens = False + )[-1] + enable_thinking = enable_thinking and ( + thinking_start_token_id in input_ids.flatten().tolist() + ) + tokenizer.thinking_budget_criteria = ThinkingBudgetCriteria( + tokenizer = tokenizer, + thinking_budget = thinking_budget, + thinking_end_token = thinking_end_token, + thinking_start_token = thinking_start_token, + enable_thinking = enable_thinking, + ) + kwargs["thinking_budget_criteria"] = tokenizer.thinking_budget_criteria + else: + tokenizer.thinking_budget_criteria = None + + # Ensure we have a prompt_cache we can track for reuse + if "prompt_cache" not in kwargs: + kwargs["prompt_cache"] = cache.make_prompt_cache( + model.language_model, + max_kv_size = kwargs.get("max_kv_size", None), + ) + tracked_cache = kwargs["prompt_cache"] + + total_prompt_tokens = reused_prefix_len + input_ids.size + + with wired_limit(model, [generation_stream]): + detokenizer = processor.detokenizer + detokenizer.reset() + thinking_criteria = getattr(tokenizer, "thinking_budget_criteria", None) + gen = generate_step(input_ids, model, pixel_values, mask, **kwargs) + tic = time.perf_counter() + + generated_tokens = [] + for n, (token, logprobs) in enumerate(gen): + if n == 0: + prompt_time = time.perf_counter() - tic + prompt_tps = total_prompt_tokens / prompt_time + tic = time.perf_counter() + + generated_tokens.append(token) + + # Check thinking budget and force token if needed + if thinking_criteria is not None: + thinking_criteria(token) + + # Stop generation if the token is in the eos_token_ids + if tokenizer.stopping_criteria(token): + break + + detokenizer.add_token(token, skip_special_token_ids = skip_special_token_ids) + + # Yield the last segment if streaming + yield GenerationResult( + text = detokenizer.last_segment, + token = token, + logprobs = logprobs, + prompt_tokens = total_prompt_tokens, + generation_tokens = n + 1, + total_tokens = total_prompt_tokens + n + 1, + prompt_tps = prompt_tps, + generation_tps = (n + 1) / (time.perf_counter() - tic), + peak_memory = mx.get_peak_memory() / 1e9, + ) + + detokenizer.finalize() + yield GenerationResult( + text = detokenizer.last_segment, + token = token, + logprobs = logprobs, + prompt_tokens = total_prompt_tokens, + generation_tokens = n + 1, + total_tokens = total_prompt_tokens + n + 1, + prompt_tps = prompt_tps, + generation_tps = (n + 1) / (time.perf_counter() - tic), + peak_memory = mx.get_peak_memory() / 1e9, + ) + + # Save cache state for potential reuse on next turn + if prompt_cache_state is not None: + all_ids = full_input_ids_list + [ + t.item() if hasattr(t, "item") else t for t in generated_tokens + ] + prompt_cache_state.update(all_ids, tracked_cache) + + # Cleanup after generation + mx.clear_cache() + + +def generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + verbose: bool = False, + **kwargs, +) -> GenerationResult: + """ + Generate text from the model. + + Args: + model (nn.Module): The language model. + tokenizer (PreTrainedTokenizer): The tokenizer. + prompt (str): The string prompt. + temperature (float): The temperature for sampling (default 0). + max_tokens (int): The maximum number of tokens (default 100). + verbose (bool): If ``True``, print tokens and timing information + (default ``False``). + formatter (Optional[Callable]): A function which takes a token and a + probability and displays it. + repetition_penalty (float, optional): The penalty factor for repeating tokens. + repetition_context_size (int, optional): The number of tokens to consider for repetition penalty. + """ + + if verbose: + print("=" * 10) + files = [] + if image is not None: + files.extend(image) + if audio is not None: + files.extend(audio) + if kwargs.get("video") is not None: + files.extend(kwargs.get("video")) + + print(f"Files: {files}", "\n") + + print("Prompt:", prompt) + + text = "" + last_response = None + + eos_tokens = kwargs.get("eos_tokens", None) + stopping_criteria = kwargs.get("stopping_criteria", None) + + # Get the tokenizer + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Add custom EOS tokens to the stopping criteria + if eos_tokens is not None: + tokenizer.stopping_criteria.add_eos_token_ids(eos_tokens) + + # Use custom stopping criteria + elif stopping_criteria is not None: + if isinstance(stopping_criteria, StoppingCriteria) or callable( + stopping_criteria + ): + tokenizer.stopping_criteria = stopping_criteria + else: + raise ValueError( + "stopping_criteria must be an instance of StoppingCriteria or a callable" + ) + else: + tokenizer.stopping_criteria.reset(model.config.eos_token_id) + + for response in stream_generate(model, processor, prompt, image, audio, **kwargs): + if verbose: + print(response.text, end = "", flush = True) + text += response.text + last_response = response + + if verbose: + print("\n" + "=" * 10) + if len(text) == 0: + print("No text generated for this prompt") + return GenerationResult( + text = text, + token = None, + logprobs = None, + prompt_tokens = 0, + generation_tokens = 0, + total_tokens = 0, + prompt_tps = 0.0, + generation_tps = 0.0, + peak_memory = mx.get_peak_memory() / 1e9, + ) + print( + f"Prompt: {last_response.prompt_tokens} tokens, " + f"{last_response.prompt_tps:.3f} tokens-per-sec" + ) + print( + f"Generation: {last_response.generation_tokens} tokens, " + f"{last_response.generation_tps:.3f} tokens-per-sec" + ) + print(f"Peak memory: {last_response.peak_memory:.3f} GB") + + return GenerationResult( + text = text, + token = last_response.token, + logprobs = last_response.logprobs, + prompt_tokens = last_response.prompt_tokens, + generation_tokens = last_response.generation_tokens, + total_tokens = last_response.total_tokens, + prompt_tps = last_response.prompt_tps, + generation_tps = last_response.generation_tps, + peak_memory = last_response.peak_memory, + ) + + +@dataclass +class BatchGenerationResult: + """ + Result of batch generation with optional image size tracking. + + Attributes: + texts: Generated text for each sample + tokens: Last generated token for each sample + logprobs: Log probabilities for each sample + prompt_tokens: Number of prompt tokens per sample + generation_tokens: Number of generated tokens per sample + total_tokens: Total tokens (prompt + generation) per sample + prompt_tps: Prompt tokens per second per sample + generation_tps: Generation tokens per second per sample + peak_memory: Peak memory usage in GB + image_sizes: Original (height, width) for each image (for tracking) + """ + + texts: List[str] + tokens: List[Optional[int]] + logprobs: List[Optional[List[float]]] + prompt_tokens: List[int] + generation_tokens: List[int] + total_tokens: List[int] + prompt_tps: List[float] + generation_tps: List[float] + peak_memory: float = 0.0 + image_sizes: Optional[List[Tuple[int, int]]] = None + + +def _left_pad_prompts(prompts, max_length = None): + if max_length is None: + max_length = max(len(p) for p in prompts) + + return mx.array([[0] * (max_length - len(p)) + p for p in prompts]) + + +def _make_cache(model, left_padding): + """ + Convert a list of regular caches into their corresponding + batch-aware caches. + """ + + def to_batch_cache(c): + if isinstance(c, cache.KVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.ChunkedKVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.SimpleKVCache): + return cache.BatchKVCache(left_padding) + elif isinstance(c, cache.ArraysCache): + c.left_padding = mx.array(left_padding) + return c + elif isinstance(c, cache.RotatingKVCache): + if c.keep > 0: + raise ValueError("RotatingKVCache with keep tokens is not supported.") + return cache.BatchRotatingKVCache(c.max_size, left_padding) + elif isinstance(c, cache.CacheList): + return cache.CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches)) + elif isinstance(c, tuple): + return cache.CacheList(*(to_batch_cache(sub_c) for sub_c in c)) + else: + raise ValueError(f"{type(c)} does not yet support batching") + + if hasattr(model, "make_cache"): + model_cache = model.make_cache() + return [to_batch_cache(c) for c in model_cache] + else: + return [cache.BatchKVCache(left_padding) for _ in model.layers] + + +@dataclass +class BatchStats: + """ + An data object to hold generation stats. + + Args: + prompt_tokens (int): The number of prompt tokens processed. + prompt_tps (float): The prompt processing tokens-per-second. + prompt_time (float): The time in seconds spent in prompt processing. + generation_tokens (int): The number of generated tokens. + generation_tps (float): The tokens-per-second for generation. + generation_time (float): The time in seconds spent in generation . + peak_memory (float): The peak memory used so far in GB. + """ + + prompt_tokens: int = 0 + prompt_tps: float = 0 + prompt_time: float = 0 + generation_tokens: int = 0 + generation_tps: float = 0 + generation_time: float = 0 + peak_memory: float = 0 + + +@dataclass +class BatchResponse: + """ + An data object to hold a batch generation response. + + Args: + texts: (List[str]): The generated text for each prompt. + stats (BatchStats): Statistics about the generation. + image_sizes: (Optional[List[Tuple[int, int]]]): Original (height, width) + for each image. Useful for tracking which images produced which responses + and for debugging padding/batching behavior. + """ + + texts: List[str] + stats: BatchStats + image_sizes: Optional[List[Tuple[int, int]]] = None + + +@dataclass +class Batch: + uids: List[int] + y: mx.array + logprobs: mx.array + max_tokens: List[int] + num_tokens: List[int] + cache: List[Any] + + def __len__(self): + return len(self.uids) + + def filter(self, keep_idx: List[int]): + self.uids = [self.uids[k] for k in keep_idx] + self.max_tokens = [self.max_tokens[k] for k in keep_idx] + self.num_tokens = [self.num_tokens[k] for k in keep_idx] + keep_idx = mx.array(keep_idx, mx.int32) + self.y = self.y[keep_idx] + self.logprobs = self.logprobs[keep_idx] + for c in self.cache: + c.filter(keep_idx) + + def extend(self, other): + self.uids.extend(other.uids) + self.y = mx.concatenate([self.y, other.y]) + self.logprobs = mx.concatenate([self.logprobs, other.logprobs]) + self.num_tokens.extend(other.num_tokens) + self.max_tokens.extend(other.max_tokens) + for c, o in zip(self.cache, other.cache): + c.extend(o) + + +class BatchGenerator: + @dataclass + class Response: + uid: int + token: int + logprobs: mx.array + finish_reason: Optional[str] + + def __init__( + self, + model, + processor, + max_tokens: int = DEFAULT_MAX_TOKENS, + stop_tokens: Optional[set] = None, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + completion_batch_size: int = DEFAULT_COMPLETION_BATCH_SIZE, + prefill_batch_size: int = DEFAULT_PREFILL_BATCH_SIZE, + prefill_step_size: Optional[int] = DEFAULT_PREFILL_STEP_SIZE, + prompt_cache = None, + ): + self.model = model + self.unprocessed_prompts = [] + self.max_tokens = max_tokens + self.processor = processor + self.tokenizer = ( + processor.tokenizer if hasattr(processor, "tokenizer") else processor + ) + self.sampler = sampler or (lambda x: mx.argmax(x, axis = -1)) + self.uid_count = 0 + self.prefill_step_size = prefill_step_size + self.prefill_batch_size = prefill_batch_size + self.completion_batch_size = completion_batch_size + self.prompt_cache = prompt_cache + self._stats = BatchStats() + + self.tokenizer.stopping_criteria.add_eos_token_ids(stop_tokens) + + self.active_batch = None + + def insert(self, prompts, max_tokens: Union[List[int], int, None] = None): + uids = [] + + if max_tokens is None or isinstance(max_tokens, int): + max_tokens = [max_tokens or self.max_tokens] * len(prompts) + + for p, m in zip(prompts, max_tokens): + self.unprocessed_prompts.append((self.uid_count, p, m)) + uids.append(self.uid_count) + self.uid_count += 1 + # Sort in ascending order of length + self.unprocessed_prompts = sorted( + self.unprocessed_prompts, key = lambda x: len(x[1]) + ) + return uids + + def _process_prompts(self, prompts, **kwargs) -> Batch: + uids, inputs, max_tokens = zip(*prompts) + lengths = [len(p) for p in inputs] + max_length = max(lengths) + + self._stats.prompt_tokens += sum(lengths) + left_padding = [max_length - l for l in lengths] + inputs = _left_pad_prompts(inputs, max_length = max_length) + + if self.prompt_cache is not None: + prompt_cache = self.prompt_cache + elif len(uids) == 1 and max(left_padding) == 0: + # Single prompt with no padding: use standard caches to avoid + # numerical divergence from batch cache wrappers. + prompt_cache = cache.make_prompt_cache(self.model) + else: + prompt_cache = _make_cache(self.model, left_padding) + + # Slice batch data in kwargs to match current batch size + batch_size = len(uids) + for key, value in kwargs.items(): + if isinstance(value, mx.array) and value.ndim > 0: + kwargs[key] = value[:batch_size] + + inputs_embeds = kwargs.pop("inputs_embeds", None) + if inputs_embeds is None: + raise ValueError("inputs_embeds is required") + + if ( + self.prefill_step_size is not None + and inputs_embeds.shape[1] > self.prefill_step_size + ): + # Chunked prefill with embeddings + while inputs_embeds.shape[1] > 1: + n_to_process = min(self.prefill_step_size, inputs_embeds.shape[1] - 1) + self.model( + inputs[:, :n_to_process], + cache = prompt_cache, + inputs_embeds = inputs_embeds[:, :n_to_process], + n_to_process = n_to_process, + **kwargs, + ) + mx.eval([c.state for c in prompt_cache]) + inputs_embeds = inputs_embeds[:, n_to_process:] + inputs = inputs[:, n_to_process:] + mx.clear_cache() + + y, logprobs = self._step( + inputs, prompt_cache, inputs_embeds = inputs_embeds, **kwargs + ) + + mx.async_eval(y, logprobs) + mx.clear_cache() + return Batch( + list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache + ) + + def _step(self, input_tokens: mx.array, prompt_cache: List[Any], **kwargs): + output = self.model(input_tokens, cache = prompt_cache, **kwargs) + logits = output.logits[:, -1, :] + logprobs = logits - mx.logsumexp(logits, axis = -1, keepdims = True) + sampled = self.sampler(logprobs) + + # TODO: Add KV cache quantization if specified + return sampled, logprobs + + def stats(self): + self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time + self._stats.generation_tps = ( + self._stats.generation_tokens / self._stats.generation_time + ) + self._stats.peak_memory = mx.get_peak_memory() / 1e9 + return self._stats + + def _next(self, **kwargs): + tic = time.perf_counter() + + prompt_processing = False + batch = self.active_batch + num_active = len(batch) if batch else 0 + num_to_add = self.completion_batch_size - num_active + while num_to_add >= self.prefill_batch_size: + prompts = self.unprocessed_prompts[: self.prefill_batch_size] + # Finish processing the last examples of the last batch + if len(prompts) == 0 and num_active > 0: + break + # No more prompts and no more completions, all done + elif len(prompts) == 0: + self.active_batch = None + return [] + # Process prompts + if batch is not None and not prompt_processing: + # Finish any active completion tokens + mx.eval(batch.y, batch.logprobs) + self._stats.generation_time += time.perf_counter() - tic + tic = time.perf_counter() + + batch = self._process_prompts(prompts, **kwargs) + self.unprocessed_prompts = self.unprocessed_prompts[ + self.prefill_batch_size : + ] + prompt_processing = True + # If there was no active batch, set it + if self.active_batch is None: + self.active_batch = batch + else: + self.active_batch.extend(batch) + + num_active = len(self.active_batch) + num_to_add -= len(batch) + + batch = self.active_batch + y, logprobs = batch.y, batch.logprobs + batch.y, batch.logprobs = self._step(y[:, None], batch.cache) + mx.async_eval(batch.y, batch.logprobs) + + y = y.tolist() + toc = time.perf_counter() + if prompt_processing: + self._stats.prompt_time += toc - tic + else: + self._stats.generation_time += toc - tic + keep_idx = [] + end_idx = [] + responses = [] + + for e, (t, uid, num_tok, max_tok) in enumerate( + zip(y, batch.uids, batch.num_tokens, batch.max_tokens) + ): + num_tok += 1 + batch.num_tokens[e] = num_tok + if self.tokenizer.stopping_criteria(t): + finish_reason = "stop" + end_idx.append(e) + elif num_tok >= max_tok: + finish_reason = "length" + end_idx.append(e) + else: + finish_reason = None + keep_idx.append(e) + responses.append(self.Response(uid, t, logprobs[e], finish_reason)) + + # Remove any finished completions + if len(end_idx): + if len(keep_idx) > 0: + batch.filter(keep_idx) + else: + self.active_batch = None + + self._stats.generation_tokens += len(responses) + + if len(responses) > 0 and self._stats.generation_tokens % 100 == 0: + mx.clear_cache() + + return responses + + def next(self, **kwargs): + return self._next(**kwargs) + + +def batch_generate( + model, + processor, + images: Union[str, List[str]] = None, + audios: Union[str, List[str]] = None, + prompts: List[str] = None, + max_tokens: Union[int, List[int]] = 128, + verbose: bool = False, + group_by_shape: bool = True, + track_image_sizes: bool = True, + **kwargs, +): + """ + Generate responses for the given batch of prompts with variable-sized images. + + This function implements the transformers-style approach to batching: + 1. Group images with the same shape for efficient batch processing + 2. Process each group as a batch (no padding waste within groups) + 3. Track original image sizes for proper attention masking + 4. Restore results to original batch order + + Key insight: Instead of padding all images to the same spatial dimensions + (which wastes computation and may hurt accuracy), we group same-sized + images together so there's zero padding within each group. + + Args: + model (nn.Module): The language model. + processor (PreTrainedTokenizer): The tokenizer/processor. + images (Union[str, List[str]]): Images (paths, URLs, or PIL images). + audios (Union[str, List[str]]): Audio files (not yet supported for batching). + prompts (List[str]): The input prompts. + max_tokens (Union[int, List[int]]): Maximum number of output tokens. This + can be per prompt if a list is provided. + verbose (bool): If ``True``, print tokens and timing information. + group_by_shape (bool): If ``True``, group same-shaped images for efficient + batch processing. + track_image_sizes (bool): If ``True``, track and return original image sizes. + kwargs: The remaining options get passed to :obj:`BatchGenerator`. + See :obj:`BatchGenerator` for more details. + + Returns: + BatchResponse with generated texts, statistics, and optionally image_sizes. + """ + from PIL import Image + + from .utils import process_image + + processor.detokenizer.reset() + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Handle single image case + if isinstance(images, str): + images = [images] + + # Handle no images case + if images is None: + texts, stats = _generate_batch( + model, processor, prompts, None, max_tokens, verbose, **kwargs + ) + return BatchResponse(texts, stats) + + # Load and preprocess images + image_processor = ( + processor.image_processor if hasattr(processor, "image_processor") else None + ) + + processed_images = [] + image_sizes_original = [] + for img in images: + if isinstance(img, str): + pil_img = process_image(img, None, image_processor) + elif isinstance(img, Image.Image): + pil_img = img + else: + pil_img = img + processed_images.append(pil_img) + # Track original size + if hasattr(pil_img, "height"): + image_sizes_original.append((pil_img.height, pil_img.width)) + else: + image_sizes_original.append((0, 0)) + + # Group images by shape for efficient processing (no padding within groups) + if group_by_shape and len(processed_images) > 1: + grouped_images, grouped_indices = group_images_by_shape(processed_images) + + if verbose: + print(f"[batch_generate] Found {len(grouped_images)} unique image shapes") + else: + # Single image or grouping disabled - treat as one group + shape = ( + (processed_images[0].height, processed_images[0].width) + if processed_images + else (0, 0) + ) + grouped_images = {shape: processed_images} + grouped_indices = {shape: list(range(len(processed_images)))} + + # Process each shape group + all_texts = [None] * len(prompts) + all_image_sizes = [None] * len(prompts) + total_stats = BatchStats() + + for shape, indices in grouped_indices.items(): + # Get images and prompts for this shape group + group_images = [processed_images[i] for i in indices] + group_prompts = [prompts[i] for i in indices] + group_sizes = [image_sizes_original[i] for i in indices] + + # Handle per-sample max_tokens + if isinstance(max_tokens, list): + group_max_tokens = [max_tokens[i] for i in indices] + else: + group_max_tokens = max_tokens + + # Process the entire group at once (same shape = no padding needed) + chunk_texts, chunk_stats = _generate_batch( + model, + processor, + group_prompts, + group_images, + group_max_tokens, + **kwargs, + ) + + # Store results in original order + for j, orig_idx in enumerate(indices): + all_texts[orig_idx] = chunk_texts[j] + all_image_sizes[orig_idx] = group_sizes[j] + + # Accumulate stats + total_stats.prompt_tokens += chunk_stats.prompt_tokens + total_stats.prompt_time += chunk_stats.prompt_time + total_stats.generation_tokens += chunk_stats.generation_tokens + total_stats.generation_time += chunk_stats.generation_time + + mx.clear_cache() + + # Compute final stats + if total_stats.prompt_time > 0: + total_stats.prompt_tps = total_stats.prompt_tokens / total_stats.prompt_time + if total_stats.generation_time > 0: + total_stats.generation_tps = ( + total_stats.generation_tokens / total_stats.generation_time + ) + total_stats.peak_memory = mx.get_peak_memory() / 1e9 + + if verbose: + print(f"[batch_generate] Finished processing {len(prompts)} samples") + print( + f"[batch_generate] Prompt: {total_stats.prompt_tokens} tokens, {total_stats.prompt_tps:.3f} tokens-per-sec" + ) + print( + f"[batch_generate] Generation: {total_stats.generation_tokens} tokens, " + f"{total_stats.generation_tps:.3f} tokens-per-sec" + ) + print(f"[batch_generate] Peak memory: {total_stats.peak_memory:.3f} GB") + + response = BatchResponse(all_texts, total_stats) + if track_image_sizes: + response.image_sizes = all_image_sizes + return response + + +def _generate_batch( + model, + processor, + prompts: List[str], + images: List = None, + max_tokens: Union[int, List[int]] = 100, + verbose: bool = False, + **kwargs, +) -> Tuple[List[str], BatchStats]: + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + batch_size = len(prompts) + + num_images_list = [ + 1 if i < (len(images) if images is not None else 0) else 0 + for i in range(len(prompts)) + ] + formatted_prompts = [ + apply_chat_template( + processor, + model.config, + p, + num_images = num_images_list[i], + ) + for i, p in enumerate(prompts) + ] + + add_special_tokens = ( + getattr(processor, "chat_template", None) is None + if model.config.model_type in ["gemma3", "gemma3n", "gemma4"] + else True + ) + + resize_shape = normalize_resize_shape(kwargs.pop("resize_shape", None)) + image_token_index = getattr(model.config, "image_token_index", None) + + inputs = prepare_inputs( + processor, + images = images, + audio = None, + prompts = formatted_prompts, + image_token_index = image_token_index, + resize_shape = resize_shape, + add_special_tokens = add_special_tokens, + pad_to_uniform_size = False, # Since images are pre-grouped by shape, they're already uniform size + ) + input_ids = inputs.get("input_ids", None) + pixel_values = inputs.get("pixel_values", None) + mask = inputs.get("attention_mask", None) + + data_kwargs = { + k: v + for k, v in inputs.items() + if k not in ["input_ids", "pixel_values", "attention_mask"] + } + + if getattr(model, "no_chunked_prefill", False): + kwargs.pop("prefill_step_size", None) + kwargs["prefill_step_size"] = None + + # Use batch_size for prefill and completion to ensure consistent processing + gen = BatchGenerator( + model.language_model, + processor, + prefill_batch_size = batch_size, + completion_batch_size = batch_size, + **kwargs, + ) + + with wired_limit(model, [generation_stream]): + embedding_output = model.get_input_embeddings( + input_ids, pixel_values, mask = mask, **data_kwargs + ) + + gen_kwargs = {**data_kwargs, **embedding_output.to_dict()} + + uids = gen.insert(input_ids.tolist(), max_tokens) + results = {uid: [] for uid in uids} + while responses := gen.next(**gen_kwargs): + for r in responses: + if r.finish_reason != "stop": + results[r.uid].append(r.token) + + detokenizer = processor.detokenizer + texts = [] + for uid in uids: + detokenizer.reset() + for t in results[uid]: + detokenizer.add_token(t) + detokenizer.finalize() + texts.append(detokenizer.text) + return texts, gen.stats() + + +def main(): + args = parse_arguments() + if isinstance(args.image, str): + args.image = [args.image] + + model, processor = load( + args.model, + args.adapter_path, + revision = args.revision, + trust_remote_code = args.trust_remote_code, + quantize_activations = args.quantize_activations, + ) + config = model.config + + prompt = args.prompt + + num_images = len(args.image) if args.image is not None else 0 + num_audios = ( + 1 if args.audio is not None else 0 + ) # TODO: Support multiple audio files + + chat_template_kwargs = {"enable_thinking": args.enable_thinking} + + prompt = apply_chat_template( + processor, + config, + prompt, + num_images = num_images, + num_audios = num_audios, + **chat_template_kwargs, + ) + + kwargs = {} + + if args.eos_tokens is not None: + eos_tokens = [] + for token in args.eos_tokens: + try: + decoded_token = codecs.decode(token, "unicode_escape") + eos_tokens.append(decoded_token) + except (UnicodeDecodeError, UnicodeError): + eos_tokens.append(token) + kwargs["eos_tokens"] = eos_tokens + + if args.skip_special_tokens: + kwargs["skip_special_tokens"] = args.skip_special_tokens + + # Add processor kwargs from JSON + if args.processor_kwargs: + kwargs.update(args.processor_kwargs) + + # Add thinking kwargs + kwargs["enable_thinking"] = args.enable_thinking + if args.thinking_budget is not None: + kwargs["thinking_budget"] = args.thinking_budget + kwargs["thinking_end_token"] = args.thinking_end_token + if args.thinking_start_token is not None: + kwargs["thinking_start_token"] = args.thinking_start_token + + if args.chat: + from .vision_cache import VisionFeatureCache + + vision_cache = VisionFeatureCache() + chat = [] + if args.system: + chat.append({"role": "system", "content": args.system}) + while user := input("User:"): + chat.append({"role": "user", "content": user}) + prompt = apply_chat_template( + processor, config, chat, num_images = num_images, **chat_template_kwargs + ) + response = "" + print("Assistant:", end = "") + stream_kwargs = { + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "vision_cache": vision_cache, + **kwargs, + } + if args.resize_shape is not None: + stream_kwargs["resize_shape"] = args.resize_shape + if args.prefill_step_size is not None: + stream_kwargs["prefill_step_size"] = args.prefill_step_size + + for chunk in stream_generate( + model, + processor, + prompt, + args.image, + args.audio, + **stream_kwargs, + ): + response += chunk.text + print(chunk.text, end = "") + + chat.append({"role": "assistant", "content": response}) + print() + + else: + gen_kwargs = { + "image": args.image, + "audio": args.audio, + "temperature": args.temperature, + "max_tokens": args.max_tokens, + "verbose": args.verbose, + "max_kv_size": args.max_kv_size, + "kv_bits": args.kv_bits, + "kv_group_size": args.kv_group_size, + "kv_quant_scheme": getattr( + args, "kv_quant_scheme", DEFAULT_KV_QUANT_SCHEME + ), + "quantized_kv_start": args.quantized_kv_start, + **kwargs, + } + if args.resize_shape is not None: + gen_kwargs["resize_shape"] = args.resize_shape + if args.prefill_step_size is not None: + gen_kwargs["prefill_step_size"] = args.prefill_step_size + + result = generate( + model, + processor, + prompt, + **gen_kwargs, + ) + if not args.verbose: + print(result.text) + + +if __name__ == "__main__": + print( + "Calling `python -m mlx_vlm.generate ...` directly is deprecated." + " Use `mlx_vlm generate` or `python -m mlx_vlm generate` instead." + ) + main() diff --git a/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py b/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py new file mode 100644 index 0000000000..dcc85e0b54 --- /dev/null +++ b/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py @@ -0,0 +1,138 @@ +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn + +from ..base import InputEmbeddingsFeatures +from ..qwen3_vl import Model as Qwen3VLModel +from ..qwen3_vl import processing_qwen3_vl # noqa: F401 +from ..qwen3_vl.qwen3_vl import masked_scatter +from .config import ModelConfig +from .language import LanguageModel +from .vision import VisionModel + + +class Model(Qwen3VLModel): + def __init__(self, config: ModelConfig): + # only initialize nn.Module, skip the initialization of vision_tower and language_model in the parent class + nn.Module.__init__(self) + self.config = config + self.vision_tower = VisionModel(config.vision_config) + self.language_model = LanguageModel(config.text_config, config) + + def get_input_embeddings( + self, + input_ids: Optional[mx.array] = None, + pixel_values: Optional[mx.array] = None, + **kwargs, + ): + image_grid_thw = kwargs.get("image_grid_thw", None) + video_grid_thw = kwargs.get("video_grid_thw", None) + mask = kwargs.get("mask", None) + grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw + + if pixel_values is None: + return InputEmbeddingsFeatures( + inputs_embeds = self.language_model.model.embed_tokens(input_ids) + ) + + dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.astype(dtype) + + # Get the input embeddings from the language model + inputs_embeds = self.language_model.model.embed_tokens(input_ids) + + cached = kwargs.get("cached_image_features", None) + if cached is not None: + hidden_states = cached + else: + # Get the ouptut hidden states from the vision model + hidden_states, _ = self.vision_tower(pixel_values, grid_thw) + + # Insert special image tokens in the input_ids + inputs_embeds, _ = self.merge_input_ids_with_image_features( + hidden_states, + inputs_embeds, + input_ids, + self.config.image_token_index, + self.config.video_token_index, + ) + + # Pre-calculate position_ids for chunked prefill + if image_grid_thw is not None or video_grid_thw is not None: + position_ids, rope_deltas = self.language_model.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, mask + ) + self.language_model._position_ids = position_ids + self.language_model._rope_deltas = rope_deltas + + return InputEmbeddingsFeatures( + inputs_embeds = inputs_embeds, + ) + + @staticmethod + def merge_input_ids_with_image_features( + image_features, inputs_embeds, input_ids, image_token_index, video_token_index + ): + special_image_mask = input_ids == image_token_index + special_video_mask = input_ids == video_token_index + special_image_mask = special_image_mask | special_video_mask + n_image_tokens = special_image_mask.sum() + special_image_mask = special_image_mask[..., None] + special_image_mask = mx.broadcast_to(special_image_mask, inputs_embeds.shape) + + n_image_features = image_features.shape[0] + n_image_mask_elements = special_image_mask.sum() + if n_image_mask_elements != image_features.size: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + inputs_embeds = masked_scatter( + inputs_embeds, special_image_mask, image_features + ) + + return inputs_embeds, special_image_mask + + def sanitize(self, weights): + # ignore mtp weights + weights = {key: value for key, value in weights.items() if "mtp." not in key} + + if self.config.text_config.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + norm_keys = ( + ".input_layernorm.weight", + ".post_attention_layernorm.weight", + "model.norm.weight", + ".q_norm.weight", + ".k_norm.weight", + ) + + sanitized_weights = {} + for key, value in weights.items(): + if "model" in key: + if "model.language_model" in key: + key = key.replace("model.language_model", "language_model.model") + elif "model.visual" in key: + key = key.replace("model.visual", "vision_tower") + elif "lm_head" in key: + key = key.replace("lm_head", "language_model.lm_head") + + if "conv1d.weight" in key and value.shape[-1] != 1: + value = value.moveaxis(2, 1) + if any(key.endswith(sfx) for sfx in norm_keys): + if value.ndim == 1: + value += 1.0 + + sanitized_weights[key] = value + + return sanitized_weights + + @property + def quant_predicate(self): + return self.language_model.quant_predicate + + @property + def cast_predicate(self): + return self.language_model.cast_predicate