lalamo 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

lalamo/model_import/configs/huggingface/qwen3.py

@@ -0,0 +1,142 @@
+from dataclasses import dataclass
+from typing import Literal
+
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    Activation,
+    AttentionConfig,
+    DecoderConfig,
+    DecoderLayerConfig,
+    FullPrecisionLinearConfig,
+    GroupQuantizedLinearConfig,
+    MLPConfig,
+    RMSNormConfig,
+    TiedEmbeddingConfig,
+    UnscaledRoPEConfig,
+    UntiedEmbeddingConfig,
+    UpcastMode,
+)
+from lalamo.quantization import QuantizationMode
+
+from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
+
+__all__ = ["HFQwen3Config"]
+
+
+@dataclass(frozen=True)
+class HFQwen3Config(HuggingFaceConfig):
+    attention_bias: bool
+    hidden_act: Literal["silu"]
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    max_window_layers: int
+    model_type: Literal["qwen3"]
+    num_attention_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    sliding_window: int | None
+    tie_word_embeddings: bool
+    use_sliding_window: bool
+    vocab_size: int
+    head_dim: int
+
+    quantization_config: AWQQuantizationConfig | GPTQQuantizationConfig | None = None
+
+    def _get_sliding_window_sizes(self) -> tuple[int | None, ...]:
+        if not self.use_sliding_window:
+            return tuple([None] * self.num_hidden_layers)
+
+        # The HuggingFace Qwen3 implementation's comment states that bottom layers use SWA,
+        # but the code (`configuration_qwen3.py`) implements it for the top layers.
+        # We are following the code.
+        sliding_window_sizes = []
+        for i in range(self.num_hidden_layers):
+            if i >= self.max_window_layers:
+                sliding_window_sizes.append(self.sliding_window)
+            else:
+                sliding_window_sizes.append(None)
+        return tuple(sliding_window_sizes)
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> DecoderConfig:
+        if self.tie_word_embeddings:
+            embedding_config = TiedEmbeddingConfig(
+                input_scale=None,
+                logits_soft_cap=None,
+                precision=activation_precision,
+            )
+        else:
+            embedding_config = UntiedEmbeddingConfig(
+                input_scale=None,
+                logits_soft_cap=None,
+                precision=activation_precision,
+            )
+        rope_config = UnscaledRoPEConfig(
+            precision=activation_precision,
+            base=self.rope_theta,
+            max_sequence_length=self.max_position_embeddings,
+        )
+        rmsnorm_config = RMSNormConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.rms_norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+        )
+        if self.quantization_config is None:
+            linear_config = FullPrecisionLinearConfig(
+                precision=activation_precision,
+            )
+        else:
+            linear_config = GroupQuantizedLinearConfig(
+                group_size=self.quantization_config.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        attention_config = AttentionConfig(
+            qkv_projection_config=linear_config,
+            out_projection_config=linear_config,
+            query_norm_config=rmsnorm_config,
+            key_norm_config=rmsnorm_config,
+            logit_soft_cap=None,
+            has_qkv_biases=self.attention_bias,
+            has_out_biases=self.attention_bias,
+        )
+        mlp_config = MLPConfig(
+            linear_config=linear_config,
+            activation=Activation.SILU,
+        )
+        decoder_layer_config = DecoderLayerConfig(
+            pre_attention_norm_config=rmsnorm_config,
+            attention_config=attention_config,
+            post_attention_norm_config=None,
+            pre_mlp_norm_config=rmsnorm_config,
+            mlp_config=mlp_config,
+            post_mlp_norm_config=None,
+        )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            global_rope_config=rope_config,
+            local_rope_config=None,
+            layer_config=decoder_layer_config,
+            output_norm_config=rmsnorm_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.hidden_size,
+            hidden_dim=self.intermediate_size,
+            num_heads=self.num_attention_heads,
+            num_groups=self.num_key_value_heads,
+            head_dim=self.head_dim,
+            attention_scale=None,
+            num_layers=self.num_hidden_layers,
+            sliding_window_sizes=self._get_sliding_window_sizes(),
+            context_length=context_length or self.max_position_embeddings,
+        )

lalamo/model_import/loaders/__init__.py

@@ -0,0 +1,7 @@
+# from .executorch import load_executorch
+from .huggingface import load_huggingface
+
+__all__ = [
+    # "load_executorch",
+    "load_huggingface",
+]

lalamo/model_import/loaders/common.py

@@ -0,0 +1,45 @@
+from collections.abc import Callable, Iterable
+
+import equinox as eqx
+from jax.tree import leaves_with_path
+from jax.tree_util import keystr
+from jaxtyping import Array, PyTree
+
+__all__ = [
+    "load_parameters",
+]
+
+
+def _get_name(leaf: PyTree, tree: PyTree) -> str:
+    for path, value in leaves_with_path(tree):
+        if value is leaf:
+            return f"~{keystr(path)}"
+    raise ValueError(f"Leaf {leaf} not found in tree {tree}")
+
+
+def _check_compatible(old_value: PyTree, new_value: PyTree, module: eqx.Module) -> None:
+    if isinstance(old_value, Array) and isinstance(new_value, Array):
+        name = _get_name(old_value, module)
+        if old_value.shape != new_value.shape:
+            raise ValueError(f"Expected parameter {name} to have shape {old_value.shape}, got {new_value.shape}")
+        if old_value.dtype != new_value.dtype:
+            raise ValueError(f"Expected parameter {name} to have dtype {old_value.dtype}, got {new_value.dtype}")
+    elif type(old_value) is not type(new_value):
+        raise TypeError(f"Expected parameter of type {type(old_value)}, got {type(new_value)}")
+
+
+def load_parameters[M: eqx.Module](
+    selector: Callable[[M], Iterable[PyTree]],
+    module: M,
+    new_values: Iterable[PyTree],
+) -> M:
+    old_values = list(selector(module))
+    new_values = list(new_values)
+    casted_new_values = []
+    for old_value, new_value in zip(old_values, new_values, strict=True):
+        _check_compatible(old_value, new_value, module)
+        if isinstance(old_value, Array) and isinstance(new_value, Array):
+            casted_new_values.append(new_value.astype(old_value.dtype))
+        else:
+            casted_new_values.append(new_value)
+    return eqx.tree_at(selector, module, casted_new_values, is_leaf=lambda x: x is None)

lalamo/model_import/loaders/executorch.py

@@ -0,0 +1,223 @@
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass, replace
+
+import jax.numpy as jnp
+from einops import rearrange
+from jaxtyping import Array, Float, Int
+
+from lalamo.common import ParameterPath
+from lalamo.modules import MLP, Attention, Decoder, DecoderLayer, QLoRALinear, QuantizedTiedEmbedding, RMSNorm
+
+from .common import load_parameters
+
+__all__ = ["load_executorch"]
+
+
+@dataclass(frozen=True)
+class QLoRALinearParams:
+    weights: Int[Array, "out_channels in_channels"]
+    zero_points: Int[Array, "out_channels groups"]
+    scales: Float[Array, "out_channels groups"]
+    lora_down_weights: Float[Array, "total_lora_channels in_channels"]
+    lora_up_weights: tuple[Float[Array, "..."], ...]
+
+    def __iter__(self) -> Iterator[Array]:
+        yield self.weights
+        yield self.zero_points
+        yield self.scales
+        yield self.lora_down_weights
+        yield from self.lora_up_weights
+
+    def __len__(self) -> int:
+        return 3 + len(self.lora_up_weights)
+
+
+def params_selector(module: QLoRALinear) -> tuple:
+    return (
+        module.weights,
+        module.zero_points,
+        module.scales,
+        module.lora_down_weights,
+        *module.lora_up_weights,
+    )
+
+
+def get_qlora_linear_params(
+    weights_dict: dict[str, Array],
+    path: ParameterPath,
+    weights_dtype: jnp.dtype,
+) -> QLoRALinearParams:
+    shift_to_unsigned = 8
+
+    weights = weights_dict[path / "weight"].astype(weights_dtype)
+    scales = weights_dict[path / "scales"]
+
+    # We don't support signed int4 on the inference side, so we map int4 to uint4 and add zero-points.
+    weights = weights + shift_to_unsigned
+    zero_points = jnp.ones_like(scales) * shift_to_unsigned
+
+    lora_down_weights = weights_dict[path / "adaptor" / "A" / "weight"]
+    lora_up_weights = (weights_dict[path / "adaptor" / "B" / "weight"],)
+    return QLoRALinearParams(weights, scales, zero_points, lora_down_weights, lora_up_weights)
+
+
+def merge_linear_params(params_list: Iterable[QLoRALinearParams]) -> QLoRALinearParams:
+    params_list = list(params_list)
+    weights = jnp.concatenate([p.weights for p in params_list], axis=0)
+    scales = jnp.concatenate([p.scales for p in params_list], axis=0)
+    zero_points = jnp.concatenate([p.zero_points for p in params_list], axis=0)
+    lora_down_weights = jnp.concatenate([p.lora_down_weights for p in params_list], axis=0)
+    lora_up_weights = tuple(w for p in params_list for w in p.lora_up_weights)
+    return QLoRALinearParams(weights, scales, zero_points, lora_down_weights, lora_up_weights)
+
+
+def load_linear(module: QLoRALinear, weights_dict: dict[str, Array], path: ParameterPath) -> QLoRALinear:
+    params = get_qlora_linear_params(weights_dict, path, module.weights.dtype)
+    return load_parameters(params_selector, module, params)
+
+
+def load_mlp(module: MLP, weights_dict: dict[str, Array], path: ParameterPath) -> MLP:
+    if not isinstance(module.up_projection, QLoRALinear):
+        raise TypeError(f"Expected up_projection to be QLoRALinear, got {type(module.up_projection)}")
+    if not isinstance(module.down_projection, QLoRALinear):
+        raise TypeError(f"Expected down_projection to be QLoRALinear, got {type(module.down_projection)}")
+
+    up_proj_params = get_qlora_linear_params(weights_dict, path / "w3", module.up_projection.weights.dtype)
+    gate_proj_params = get_qlora_linear_params(weights_dict, path / "w1", module.down_projection.weights.dtype)
+    down_proj_params = get_qlora_linear_params(weights_dict, path / "w2", module.down_projection.weights.dtype)
+
+    fused_up_gate_params = merge_linear_params([up_proj_params, gate_proj_params])
+
+    return load_parameters(
+        lambda m: (*params_selector(m.up_projection), *params_selector(m.down_projection)),  # type: ignore
+        module,
+        (*fused_up_gate_params, *down_proj_params),
+    )
+
+
+def load_rmsnorm(module: RMSNorm, weights_dict: dict[str, Array], path: ParameterPath) -> RMSNorm:
+    return load_parameters(lambda m: (m.scales,), module, (weights_dict[path / "weight"],))
+
+
+def permute_qk_weights(weights: Array, input_dim: int, num_heads: int, head_dim: int) -> Array:
+    # Reference: https://github.com/huggingface/transformers/blob/15bd3e61f8d3680ca472c9314ad07584d20f7b81/src/transformers/models/llama/convert_llama_weights_to_hf.py#L222
+    return rearrange(
+        weights,
+        "(heads rotors reim) input_channels -> (heads reim rotors) input_channels",
+        heads=num_heads,
+        rotors=head_dim // 2,
+        reim=2,
+        input_channels=input_dim,
+    )
+
+
+def permute_qk_params(
+    *,
+    params: QLoRALinearParams,
+    model_dim: int,
+    num_heads: int,
+    head_dim: int,
+    quantization_group_size: int,
+    lora_rank: int,
+) -> QLoRALinearParams:
+    # Read https://github.com/huggingface/transformers/issues/25199 to understand WTF is going on here
+    return replace(
+        params,
+        weights=permute_qk_weights(params.weights, model_dim, num_heads, head_dim),
+        scales=permute_qk_weights(params.scales, model_dim // quantization_group_size, num_heads, head_dim),
+        lora_up_weights=tuple(permute_qk_weights(w, lora_rank, num_heads, head_dim) for w in params.lora_up_weights),
+    )
+
+
+def load_attention(
+    module: Attention,
+    weights_dict: dict[str, Array],
+    path: ParameterPath,
+) -> Attention:
+    if not isinstance(module.qkv_projection, QLoRALinear):
+        raise TypeError(f"Expected qkv_projection to be QLoRALinear, got {type(module.qkv_projection)}")
+
+    model_dim = module.model_dim
+    num_heads = module.num_heads
+    num_groups = module.num_groups
+    head_dim = module.head_dim
+    lora_rank = module.qkv_projection.config.lora_rank
+
+    q_params = get_qlora_linear_params(weights_dict, path / "wq", module.qkv_projection.weights.dtype)
+    q_params = permute_qk_params(
+        params=q_params,
+        model_dim=model_dim,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        quantization_group_size=module.qkv_projection.config.group_size,
+        lora_rank=lora_rank,
+    )
+
+    k_params = get_qlora_linear_params(weights_dict, path / "wk", module.qkv_projection.weights.dtype)
+    k_params = permute_qk_params(
+        params=k_params,
+        model_dim=model_dim,
+        num_heads=num_groups,
+        head_dim=head_dim,
+        quantization_group_size=module.qkv_projection.config.group_size,
+        lora_rank=lora_rank,
+    )
+
+    v_params = get_qlora_linear_params(weights_dict, path / "wv", module.qkv_projection.weights.dtype)
+
+    out_params = get_qlora_linear_params(weights_dict, path / "wo", module.qkv_projection.weights.dtype)
+
+    qkv_params = merge_linear_params([q_params, k_params, v_params])
+    return load_parameters(
+        lambda m: (*params_selector(m.qkv_projection), *params_selector(m.out_projection)),  # type: ignore
+        module,
+        (*qkv_params, *out_params),
+    )
+
+
+def load_decoder_layer(
+    module: DecoderLayer,
+    weights_dict: dict[str, Array],
+    path: ParameterPath,
+) -> DecoderLayer:
+    if module.post_attention_norm is not None:
+        raise ValueError("Post attention normalization is not supported")
+    if module.post_mlp_norm is not None:
+        raise ValueError("Post MLP normalization is not supported")
+    attention_norm = load_rmsnorm(module.pre_attention_norm, weights_dict, path / "attention_norm")
+    attention = load_attention(module.attention, weights_dict, path / "attention")
+    mlp_norm = load_rmsnorm(module.pre_mlp_norm, weights_dict, path / "ffn_norm")
+    mlp = load_mlp(module.mlp, weights_dict, path / "feed_forward")
+    return load_parameters(
+        lambda m: (m.pre_attention_norm, m.attention, m.pre_mlp_norm, m.mlp),
+        module,
+        (attention_norm, attention, mlp_norm, mlp),
+    )
+
+
+def load_embedding(
+    module: QuantizedTiedEmbedding,
+    weights_dict: dict[str, Array],
+    path: ParameterPath,
+) -> QuantizedTiedEmbedding:
+    weights = weights_dict[path / "weight"].astype(module.weights.dtype)
+    scales = weights_dict[path / "scales"].squeeze(1)
+
+    return load_parameters(lambda m: (m.weights, m.scales), module, (weights, scales))
+
+
+def load_executorch(module: Decoder, weights_dict: dict[str, Array]) -> Decoder:
+    root_path = ParameterPath()
+    if not isinstance(module.embedding, QuantizedTiedEmbedding):
+        raise TypeError(f"Expected embedding to be QuantizedTiedEmbedding, got {type(module.embedding)}")
+
+    embedding = load_embedding(module.embedding, weights_dict, root_path / "tok_embeddings")
+    decoder_layers = tuple(
+        load_decoder_layer(layer, weights_dict, root_path / f"layers.{i}") for i, layer in enumerate(module.layers)
+    )
+    output_norm = load_rmsnorm(module.output_norm, weights_dict, root_path / "norm")
+    return load_parameters(
+        lambda m: (m.embedding, m.layers, m.output_norm),
+        module,
+        (embedding, decoder_layers, output_norm),
+    )