lalamo 0.5.15__tar.gz → 0.5.16__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.15
+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.15 → lalamo-0.5.16}/lalamo/__init__.py

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

{lalamo-0.5.15 → 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.15 → lalamo-0.5.16}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.15
+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