lalamo 0.5.14__tar.gz → 0.5.16__tar.gz

{lalamo-0.5.14 → lalamo-0.5.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.14
+Version: 0.5.16
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/__init__.py

@@ -15,7 +15,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.14"
+__version__ = "0.5.16"
 
 __all__ = [
     "AssistantMessage",

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -13,6 +13,7 @@ from lalamo.modules import (
     LlamaRoPEConfig,
     MLXQuantizedLinearConfig,
     MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbeddingConfig,
     NormalizationConfig,
     SiLU,
     TiedEmbeddingConfig,
@@ -89,27 +90,37 @@ class HFLlamaConfig(HuggingFaceLMConfig):
     ) -> DecoderConfig:
         quantization = self.quantization or self.quantization_config
         if isinstance(quantization, MLXQuantizationConfig):
-            assert self.tie_word_embeddings, "only tied embeddings are supported"
-            embedding_config = MLXQuantizedTiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                group_size=quantization.group_size,
-                embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
-                activation_quantization_mode=None,
-                activation_precision=activation_precision,
-            )
-        elif self.tie_word_embeddings:
-            embedding_config = TiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                precision=activation_precision,
-            )
-        else:
-            embedding_config = UntiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                precision=activation_precision,
-            )
+            if self.tie_word_embeddings:
+                embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=quantization.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+            else:
+                embedding_config = MLXQuantizedUntiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=quantization.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+        else: # noqa: PLR5501
+            if self.tie_word_embeddings:
+                embedding_config = TiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    precision=activation_precision,
+                )
+            else:
+                embedding_config = UntiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    precision=activation_precision,
+                )
         if self.rope_scaling is None:
             rope_config = UnscaledRoPEConfig(
                 precision=activation_precision,

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/model_import/decoder_configs/huggingface/qwen3.py

@@ -10,6 +10,8 @@ from lalamo.modules import (
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbeddingConfig,
     NormalizationConfig,
     TiedEmbeddingConfig,
     TransformerConfig,
@@ -19,7 +21,6 @@ from lalamo.modules import (
     UpcastMode,
 )
 from lalamo.modules.activations import SiLU
-from lalamo.modules.embedding import MLXQuantizedTiedEmbeddingConfig
 from lalamo.modules.linear import MLXQuantizedLinearConfig
 from lalamo.quantization import QuantizationMode
 
@@ -75,27 +76,37 @@ class HFQwen3Config(HuggingFaceLMConfig):
         metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
         if isinstance(self.quantization_config, MLXQuantizationConfig):
-            assert self.tie_word_embeddings, "only tied embeddings are supported"
-            embedding_config = MLXQuantizedTiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                group_size=self.quantization_config.group_size,
-                embedding_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
-                activation_quantization_mode=None,
-                activation_precision=activation_precision,
-            )
-        elif self.tie_word_embeddings:
-            embedding_config = TiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                precision=activation_precision,
-            )
-        else:
-            embedding_config = UntiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                precision=activation_precision,
-            )
+            if self.tie_word_embeddings:
+                embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=self.quantization_config.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+            else:
+                embedding_config = MLXQuantizedUntiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=self.quantization_config.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+        else: # noqa: PLR5501
+            if self.tie_word_embeddings:
+                embedding_config = TiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    precision=activation_precision,
+                )
+            else:
+                embedding_config = UntiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    precision=activation_precision,
+                )
         rope_config = UnscaledRoPEConfig(
             precision=activation_precision,
             base=self.rope_theta,

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/model_import/loaders/huggingface.py

@@ -29,6 +29,7 @@ from lalamo.modules import (
     UntiedEmbedding,
 )
 from lalamo.modules.classifier import Classifier
+from lalamo.modules.embedding import MLXQuantizedUntiedEmbedding
 from lalamo.modules.mlp import MixtureOfExperts, MLPBase
 from lalamo.quantization import QuantizationMode
 
@@ -625,6 +626,51 @@ def load_mlx_quantized_tied_embedding(
     return load_parameters(lambda m: (m.weights, m.scales, m.biases), module, (weights, scales, biases))
 
 
+def load_mlx_quantized_untied_embedding(
+    module: MLXQuantizedUntiedEmbedding,
+    weights_dict: Mapping[str, Array],
+    embedding_path: ParameterPath,
+    lm_head_path: ParameterPath,
+) -> MLXQuantizedUntiedEmbedding:
+    input_qweights = weights_dict[embedding_path / "weight"]
+    input_qscales = weights_dict[embedding_path / "scales"]
+    input_qbiases = weights_dict[embedding_path / "biases"]
+    output_qweights = weights_dict[lm_head_path / "weight"]
+    output_qscales = weights_dict[lm_head_path / "scales"]
+    output_qbiases = weights_dict[lm_head_path / "biases"]
+
+    input_weights = _process_quantized_tensor(
+        input_qweights,
+        module.config.embedding_quantization_mode,
+        module.activation_precision,
+        None,
+    )
+    input_scales = input_qscales.astype(module.activation_precision)
+    input_biases = input_qbiases.astype(module.activation_precision)
+
+    output_weights = _process_quantized_tensor(
+        output_qweights,
+        module.config.embedding_quantization_mode,
+        module.activation_precision,
+        None,
+    )
+    output_scales = output_qscales.astype(module.activation_precision)
+    output_biases = output_qbiases.astype(module.activation_precision)
+
+    return load_parameters(
+        lambda m: (
+            m.input_weights,
+            m.input_scales,
+            m.input_biases,
+            m.output_weights,
+            m.output_scales,
+            m.output_biases,
+        ),
+        module,
+        (input_weights, input_scales, input_biases, output_weights, output_scales, output_biases),
+    )
+
+
 def load_mlx_semi_quantized_untied_embedding(
     module: MLXSemiQuantizedUntiedEmbedding,
     weights_dict: Mapping[str, Array],
@@ -741,6 +787,8 @@ def load_huggingface_decoder(
         embedding = load_tied_embedding(module.embedding, weights_dict, embedding_path)
     elif isinstance(module.embedding, MLXQuantizedTiedEmbedding):
         embedding = load_mlx_quantized_tied_embedding(module.embedding, weights_dict, embedding_path)
+    elif isinstance(module.embedding, MLXQuantizedUntiedEmbedding):
+        embedding = load_mlx_quantized_untied_embedding(module.embedding, weights_dict, embedding_path, lm_head_path)
     elif isinstance(module.embedding, MLXSemiQuantizedUntiedEmbedding):
         embedding = load_mlx_semi_quantized_untied_embedding(
             module.embedding,
@@ -759,7 +807,7 @@ def load_huggingface_decoder(
             weights_dict,
             decoder_path / "layers" / ((i * 2) if alternating_layers else i),
             decoder_path / "layers" / ((i * 2 + 1) if alternating_layers else i),
-            mixer_key[type(layer.config.mixer_config)], # type: ignore
+            mixer_key[type(layer.config.mixer_config)],  # type: ignore
             mlp_key,
             pre_mixer_norm_key,
             pre_mlp_norm_key,

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/model_import/model_specs/qwen.py

@@ -223,6 +223,20 @@ QWEN3 = [
         repo="Qwen/Qwen3-8B-AWQ",
         config_type=HFQwen3Config,
     ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-8B-MLX-4bit",
+        size="8B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-8B-MLX-4bit",
+        config_type=HFQwen3Config,
+        configs=ConfigMap(
+            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-8B"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-8B"),
+            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-8B"),
+        ),
+    ),
     ModelSpec(
         vendor="Alibaba",
         family="Qwen3",

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/modules/__init__.py

@@ -18,6 +18,8 @@ from .embedding import (
     EmbeddingConfig,
     MLXQuantizedTiedEmbedding,
     MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbedding,
+    MLXQuantizedUntiedEmbeddingConfig,
     MLXSemiQuantizedUntiedEmbedding,
     MLXSemiQuantizedUntiedEmbeddingConfig,
     QuantizedTiedEmbedding,
@@ -120,6 +122,8 @@ __all__ = [
     "MLXQuantizedLinearConfig",
     "MLXQuantizedTiedEmbedding",
     "MLXQuantizedTiedEmbeddingConfig",
+    "MLXQuantizedUntiedEmbedding",
+    "MLXQuantizedUntiedEmbeddingConfig",
     "MLXSemiQuantizedUntiedEmbedding",
     "MLXSemiQuantizedUntiedEmbeddingConfig",
     "Mamba2",

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo/modules/embedding.py

@@ -24,6 +24,8 @@ __all__ = [
     "EmbeddingConfig",
     "MLXQuantizedTiedEmbedding",
     "MLXQuantizedTiedEmbeddingConfig",
+    "MLXQuantizedUntiedEmbedding",
+    "MLXQuantizedUntiedEmbeddingConfig",
     "MLXSemiQuantizedUntiedEmbedding",
     "MLXSemiQuantizedUntiedEmbeddingConfig",
     "QuantizedTiedEmbedding",
@@ -492,6 +494,172 @@ class MLXQuantizedTiedEmbedding(EmbeddingBase[MLXQuantizedTiedEmbeddingConfig]):
         )
 
 
+@dataclass(frozen=True)
+class MLXQuantizedUntiedEmbeddingConfig(EmbeddingConfigBase):
+    group_size: int
+    embedding_quantization_mode: QuantizationMode
+    activation_quantization_mode: QuantizationMode | None
+    activation_precision: DTypeLike
+
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "MLXQuantizedUntiedEmbedding":
+        raise NotImplementedError
+
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "MLXQuantizedUntiedEmbedding":
+        assert model_dim % self.group_size == 0
+        model_groups = model_dim // self.group_size
+        return MLXQuantizedUntiedEmbedding(
+            config=self,
+            input_weights=dummy_array((vocab_size, model_dim), dtype=self.activation_precision),
+            input_scales=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            input_biases=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            output_weights=dummy_array((vocab_size, model_dim), dtype=self.activation_precision),
+            output_scales=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            output_biases=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+        )
+
+
+class MLXQuantizedUntiedEmbedding(EmbeddingBase[MLXQuantizedUntiedEmbeddingConfig]):
+    input_weights: Float[Array, "vocabulary channels"]
+    input_scales: Float[Array, "vocabulary groups"]
+    input_biases: Float[Array, "vocabulary groups"]
+    output_weights: Float[Array, "vocabulary channels"]
+    output_scales: Float[Array, "vocabulary groups"]
+    output_biases: Float[Array, "vocabulary groups"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.activation_precision
+
+    @property
+    def model_dim(self) -> int:
+        _, model_dim = self.input_weights.shape
+        return model_dim
+
+    @property
+    def vocab_size(self) -> int:
+        vocab_size, _ = self.input_weights.shape
+        return vocab_size
+
+    @property
+    def int_input_weights(self) -> Int[Array, "vocabulary channels"]:
+        quantized = quantize_weights(self.input_weights, self.config.embedding_quantization_mode)
+        casted = quantized.astype(self.config.embedding_quantization_mode.dtype)
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            packed = jax_uint4_to_packed_uint8(casted)
+        else:
+            packed = casted
+
+        return packed
+
+    @property
+    def int_output_weights(self) -> Int[Array, "vocabulary channels"]:
+        quantized = quantize_weights(self.output_weights, self.config.embedding_quantization_mode)
+        casted = quantized.astype(self.config.embedding_quantization_mode.dtype)
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            packed = jax_uint4_to_packed_uint8(casted)
+        else:
+            packed = casted
+
+        return packed
+
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]:
+        quantized_weights = quantize_weights(self.input_weights, self.config.embedding_quantization_mode)
+        grouped_weights = rearrange(
+            quantized_weights,
+            "vocab (groups elements) -> vocab groups elements",
+            elements=self.config.group_size,
+        )
+
+        scales = rearrange(self.input_scales, "vocab groups -> vocab groups 1")
+
+        biases = rearrange(self.input_biases, "vocab groups -> vocab groups 1")
+
+        scaled_grouped_weights = grouped_weights * scales + biases
+
+        result = rearrange(
+            scaled_grouped_weights,
+            "vocab groups elements -> vocab (groups elements)",
+        )
+        return result
+
+    def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
+        quantized_weights = quantize_weights(self.output_weights, self.config.embedding_quantization_mode)
+        grouped_weights = rearrange(
+            quantized_weights,
+            "vocab (groups elements) -> vocab groups elements",
+            elements=self.config.group_size,
+        )
+
+        scales = rearrange(self.output_scales, "vocab groups -> vocab groups 1")
+
+        biases = rearrange(self.output_biases, "vocab groups -> vocab groups 1")
+
+        scaled_grouped_weights = grouped_weights * scales + biases
+
+        result = rearrange(
+            scaled_grouped_weights,
+            "vocab groups elements -> vocab (groups elements)",
+        )
+        return result
+
+    @eqx.filter_jit
+    def readout(self, x: Float[Array, " channels"]) -> Float[Array, " vocabulary"]:
+        if self.config.activation_quantization_mode is not None:
+            x = dynamically_quantize_activations(x, self.config.activation_quantization_mode)
+        return super().readout(x)
+
+    def export_weights(self) -> ParameterTree:
+        return {
+            "input_weights": self.int_input_weights,
+            "input_scales": self.input_scales,
+            "input_biases": self.input_biases,
+            "output_weights": self.int_output_weights,
+            "output_scales": self.output_scales,
+            "output_biases": self.output_biases,
+        }
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["input_weights"], Array)
+        assert isinstance(weights["input_scales"], Array)
+        assert isinstance(weights["input_biases"], Array)
+        assert isinstance(weights["output_weights"], Array)
+        assert isinstance(weights["output_scales"], Array)
+        assert isinstance(weights["output_biases"], Array)
+
+        unpacked_input_weights = weights["input_weights"]
+        unpacked_output_weights = weights["output_weights"]
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            unpacked_input_weights = jax_uint8_to_unpacked_uint4(weights["input_weights"])
+            unpacked_output_weights = jax_uint8_to_unpacked_uint4(weights["output_weights"])
+
+        return replace(
+            self,
+            input_weights=unpacked_input_weights.astype(self.input_weights.dtype),
+            input_scales=weights["input_scales"],
+            input_biases=weights["input_biases"],
+            output_weights=unpacked_output_weights.astype(self.output_weights.dtype),
+            output_scales=weights["output_scales"],
+            output_biases=weights["output_biases"],
+        )
+
+
 @dataclass(frozen=True)
 class MLXSemiQuantizedUntiedEmbeddingConfig(EmbeddingConfigBase):
     group_size: int
@@ -626,6 +794,7 @@ EmbeddingConfig = (
     | UntiedEmbeddingConfig
     | QuantizedTiedEmbeddingConfig
     | MLXQuantizedTiedEmbeddingConfig
+    | MLXQuantizedUntiedEmbeddingConfig
     | MLXSemiQuantizedUntiedEmbeddingConfig
 )
 

{lalamo-0.5.14 → lalamo-0.5.16}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.14
+Version: 0.5.16
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown