lalamo 0.5.12__tar.gz → 0.5.13__tar.gz

{lalamo-0.5.12 → lalamo-0.5.13}/PKG-INFO

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

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

{lalamo-0.5.12 → lalamo-0.5.13}/lalamo/model_import/decoder_configs/huggingface/gemma3.py

@@ -5,7 +5,13 @@ from typing import Literal
 import jax.numpy as jnp
 from jaxtyping import DTypeLike
 
-from lalamo.modules import DecoderConfig, TiedEmbeddingConfig, TransformerConfig
+from lalamo.modules import (
+    DecoderConfig,
+    MLXQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
+    TiedEmbeddingConfig,
+    TransformerConfig,
+)
 from lalamo.modules.activations import GELU
 from lalamo.modules.linear import FullPrecisionLinearConfig
 from lalamo.modules.mlp import DenseMLPConfig
@@ -13,8 +19,9 @@ from lalamo.modules.normalization import NormalizationConfig, UpcastMode
 from lalamo.modules.rope import LinearScalingRoPEConfig, UnscaledRoPEConfig, YARNRoPEConfig
 from lalamo.modules.token_mixers.attention import AttentionConfig
 from lalamo.modules.transformer_layer import TransformerLayerConfig
+from lalamo.quantization import QuantizationMode
 
-from .common import HuggingFaceLMConfig
+from .common import HuggingFaceLMConfig, MLXQuantizationConfig, QuantizationConfigType
 
 __all__ = ["HFGemma3Config", "HFGemma3TextConfig"]
 
@@ -61,6 +68,9 @@ class HFGemma3TextConfigRaw:
     final_logit_softcapping: float | None = None
     vocab_size: int = 262208
 
+    quantization: QuantizationConfigType = None
+    quantization_config: QuantizationConfigType = None
+
     @property
     def sliding_window_sizes(self) -> list[int | None]:
         result = []
@@ -77,14 +87,28 @@ class HFGemma3TextConfigRaw:
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
         metadata_dict: Mapping[str, str],  # noqa: ARG002
+        fallback_quantization: QuantizationConfigType | None = None,
     ) -> DecoderConfig:
+        quantization = self.quantization or self.quantization_config or fallback_quantization
         input_scale = _round_to_bfloat16(self.hidden_size**0.5)
         attention_scale = self.query_pre_attn_scalar**-0.5
-        embedding_config = TiedEmbeddingConfig(
-            input_scale=input_scale,
-            logit_soft_cap=self.final_logit_softcapping,
-            precision=activation_precision,
-        )
+        if quantization is None:
+            embedding_config = TiedEmbeddingConfig(
+                input_scale=input_scale,
+                logit_soft_cap=self.final_logit_softcapping,
+                precision=activation_precision,
+            )
+        elif isinstance(quantization, MLXQuantizationConfig):
+            embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                input_scale=input_scale,
+                logit_soft_cap=self.final_logit_softcapping,
+                group_size=quantization.group_size,
+                embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        else:
+            raise RuntimeError(f"Unsupported quantization format: {type(quantization)}")
         rms_norm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
@@ -127,7 +151,17 @@ class HFGemma3TextConfigRaw:
             max_sequence_length=context_length or self.max_position_embeddings,
         )
 
-        linear_config = FullPrecisionLinearConfig(precision=activation_precision)
+        if quantization is None:
+            linear_config = FullPrecisionLinearConfig(precision=activation_precision)
+        elif isinstance(quantization, MLXQuantizationConfig):
+            linear_config = MLXQuantizedLinearConfig(
+                group_size=quantization.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        else:
+            raise RuntimeError(f"Unsupported quantization format: {type(quantization)}")
         mlp_config = DenseMLPConfig(
             linear_config=linear_config,
             activation=GELU(),
@@ -214,6 +248,9 @@ class HFGemma3Config(HuggingFaceLMConfig):
     transformers_version: str
     vision_config: HFGemma3VisionConfig
 
+    quantization: QuantizationConfigType = None
+    quantization_config: QuantizationConfigType = None
+
     def to_decoder_config(
         self,
         context_length: int | None,
@@ -221,9 +258,11 @@ class HFGemma3Config(HuggingFaceLMConfig):
         accumulation_precision: DTypeLike,
         metadata_dict: Mapping[str, str],
     ) -> DecoderConfig:
+        quantization = self.quantization or self.quantization_config
         return self.text_config.to_decoder_config(
             context_length=context_length,
             activation_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             metadata_dict=metadata_dict,
+            fallback_quantization=quantization,
         )

{lalamo-0.5.12 → lalamo-0.5.13}/lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -11,6 +11,8 @@ from lalamo.modules import (
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
     LlamaRoPEConfig,
+    MLXQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
     NormalizationConfig,
     SiLU,
     TiedEmbeddingConfig,
@@ -23,7 +25,7 @@ from lalamo.modules import (
 )
 from lalamo.quantization import QuantizationMode
 
-from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceLMConfig
+from .common import HuggingFaceLMConfig, MLXQuantizationConfig, QuantizationConfigType
 
 __all__ = ["HFLlamaConfig"]
 
@@ -75,7 +77,8 @@ class HFLlamaConfig(HuggingFaceLMConfig):
     vocab_size: int
     head_dim: int | None = None
 
-    quantization_config: AWQQuantizationConfig | GPTQQuantizationConfig | None = None
+    quantization: QuantizationConfigType = None
+    quantization_config: QuantizationConfigType = None
 
     def to_decoder_config(
         self,
@@ -84,7 +87,18 @@ class HFLlamaConfig(HuggingFaceLMConfig):
         accumulation_precision: DTypeLike,
         metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
-        if self.tie_word_embeddings:
+        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,
@@ -133,14 +147,21 @@ class HFLlamaConfig(HuggingFaceLMConfig):
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
             subtract_mean=False,
         )
-        if self.quantization_config is None:
+        if quantization is None:
             linear_config = FullPrecisionLinearConfig(
                 precision=activation_precision,
             )
+        elif isinstance(quantization, MLXQuantizationConfig):
+            linear_config = MLXQuantizedLinearConfig(
+                group_size=quantization.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                activation_quantization_mode=None,
+                activation_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),
+                group_size=quantization.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
             )

lalamo-0.5.13/lalamo/model_import/model_specs/gemma.py

@@ -0,0 +1,124 @@
+from lalamo.model_import.decoder_configs import (
+    HFGemma2Config,
+    HFGemma3Config,
+    HFGemma3TextConfig,
+)
+from lalamo.quantization import QuantizationMode
+
+from .common import ConfigMap, FileSpec, ModelSpec, WeightsType
+
+__all__ = ["GEMMA_MODELS"]
+
+GEMMA2 = [
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-2",
+        name="Gemma-2-2B-Instruct",
+        size="2B",
+        quantization=None,
+        repo="google/gemma-2-2b-it",
+        config_type=HFGemma2Config,
+    ),
+]
+
+GEMMA3 = [
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-1B-Instruct",
+        size="1B",
+        quantization=None,
+        repo="google/gemma-3-1b-it",
+        config_type=HFGemma3TextConfig,
+        weights_type=WeightsType.SAFETENSORS,
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-1B-Instruct-4bit",
+        size="1B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/gemma-3-1b-it-4bit",
+        config_type=HFGemma3TextConfig,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-1b-it")),
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-1B-Instruct-8bit",
+        size="1B",
+        quantization=QuantizationMode.UINT8,
+        repo="mlx-community/gemma-3-1b-it-8bit",
+        config_type=HFGemma3TextConfig,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-1b-it")),
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-4B-Instruct",
+        size="4B",
+        quantization=None,
+        repo="google/gemma-3-4b-it",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-4B-Instruct-4bit",
+        size="4B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/gemma-3-4b-it-4bit",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-4b-it")),
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-4B-Instruct-8bit",
+        size="4B",
+        quantization=QuantizationMode.UINT8,
+        repo="mlx-community/gemma-3-4b-it-8bit",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-4b-it")),
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-27B-Instruct",
+        size="27B",
+        quantization=None,
+        repo="google/gemma-3-27b-it",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-27B-Instruct-4bit",
+        size="27B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/gemma-3-27b-it-4bit",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-27b-it")),
+    ),
+    ModelSpec(
+        vendor="Google",
+        family="Gemma-3",
+        name="Gemma-3-27B-Instruct-8bit",
+        size="27B",
+        quantization=QuantizationMode.UINT8,
+        repo="mlx-community/gemma-3-27b-it-8bit",
+        config_type=HFGemma3Config,
+        weights_type=WeightsType.SAFETENSORS,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "google/gemma-3-27b-it")),
+    ),
+]
+
+
+GEMMA_MODELS = GEMMA2 + GEMMA3

lalamo-0.5.13/lalamo/model_import/model_specs/llama.py

@@ -0,0 +1,100 @@
+from lalamo.model_import.decoder_configs import HFLlamaConfig
+from lalamo.quantization import QuantizationMode
+
+from .common import ConfigMap, FileSpec, ModelSpec
+
+__all__ = ["LLAMA_MODELS"]
+
+LLAMA31 = [
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.1",
+        name="Llama-3.1-8B-Instruct",
+        size="8B",
+        quantization=None,
+        repo="meta-llama/Llama-3.1-8B-Instruct",
+        config_type=HFLlamaConfig,
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.1",
+        name="Llama-3.1-8B-Instruct-4bit",
+        size="8B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/Llama-3.1-8B-Instruct-4bit",
+        config_type=HFLlamaConfig,
+        use_cases=tuple(),
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.1-8B-Instruct")),
+    ),
+]
+
+
+LLAMA32 = [
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-1B-Instruct",
+        size="1B",
+        quantization=None,
+        repo="meta-llama/Llama-3.2-1B-Instruct",
+        config_type=HFLlamaConfig,
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-1B-Instruct-4bit",
+        size="1B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/Llama-3.2-1B-Instruct-4bit",
+        config_type=HFLlamaConfig,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-1B-Instruct")),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-1B-Instruct-8bit",
+        size="1B",
+        quantization=QuantizationMode.UINT8,
+        repo="mlx-community/Llama-3.2-1B-Instruct-8bit",
+        config_type=HFLlamaConfig,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-1B-Instruct")),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-3B-Instruct",
+        size="3B",
+        quantization=None,
+        repo="meta-llama/Llama-3.2-3B-Instruct",
+        config_type=HFLlamaConfig,
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-3B-Instruct-4bit",
+        size="3B",
+        quantization=QuantizationMode.UINT4,
+        repo="mlx-community/Llama-3.2-3B-Instruct-4bit",
+        config_type=HFLlamaConfig,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-3B-Instruct")),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-3B-Instruct-8bit",
+        size="3B",
+        quantization=QuantizationMode.UINT8,
+        repo="mlx-community/Llama-3.2-3B-Instruct-8bit",
+        config_type=HFLlamaConfig,
+        configs=ConfigMap(generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-3B-Instruct")),
+        use_cases=tuple(),
+    ),
+]
+
+LLAMA_MODELS = LLAMA31 + LLAMA32

{lalamo-0.5.12 → lalamo-0.5.13}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.12
+Version: 0.5.13
 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.12/lalamo/model_import/model_specs/gemma.py

@@ -1,57 +0,0 @@
-from lalamo.model_import.decoder_configs import (
-    HFGemma2Config,
-    HFGemma3Config,
-    HFGemma3TextConfig,
-)
-
-from .common import ModelSpec, WeightsType
-
-__all__ = ["GEMMA_MODELS"]
-
-GEMMA2 = [
-    ModelSpec(
-        vendor="Google",
-        family="Gemma-2",
-        name="Gemma-2-2B-Instruct",
-        size="2B",
-        quantization=None,
-        repo="google/gemma-2-2b-it",
-        config_type=HFGemma2Config,
-    ),
-]
-
-GEMMA3 = [
-    ModelSpec(
-        vendor="Google",
-        family="Gemma-3",
-        name="Gemma-3-1B-Instruct",
-        size="1B",
-        quantization=None,
-        repo="google/gemma-3-1b-it",
-        config_type=HFGemma3TextConfig,
-        weights_type=WeightsType.SAFETENSORS,
-    ),
-    ModelSpec(
-        vendor="Google",
-        family="Gemma-3",
-        name="Gemma-3-4B-Instruct",
-        size="4B",
-        quantization=None,
-        repo="google/gemma-3-4b-it",
-        config_type=HFGemma3Config,
-        weights_type=WeightsType.SAFETENSORS,
-    ),
-    ModelSpec(
-        vendor="Google",
-        family="Gemma-3",
-        name="Gemma-3-27B-Instruct",
-        size="27B",
-        quantization=None,
-        repo="google/gemma-3-27b-it",
-        config_type=HFGemma3Config,
-        weights_type=WeightsType.SAFETENSORS,
-    ),
-]
-
-
-GEMMA_MODELS = GEMMA2 + GEMMA3

lalamo-0.5.12/lalamo/model_import/model_specs/llama.py

@@ -1,44 +0,0 @@
-from lalamo.model_import.decoder_configs import HFLlamaConfig
-
-from .common import ModelSpec
-
-__all__ = ["LLAMA_MODELS"]
-
-LLAMA31 = [
-    ModelSpec(
-        vendor="Meta",
-        family="Llama-3.1",
-        name="Llama-3.1-8B-Instruct",
-        size="8B",
-        quantization=None,
-        repo="meta-llama/Llama-3.1-8B-Instruct",
-        config_type=HFLlamaConfig,
-        use_cases=tuple(),
-    ),
-]
-
-
-LLAMA32 = [
-    ModelSpec(
-        vendor="Meta",
-        family="Llama-3.2",
-        name="Llama-3.2-1B-Instruct",
-        size="1B",
-        quantization=None,
-        repo="meta-llama/Llama-3.2-1B-Instruct",
-        config_type=HFLlamaConfig,
-        use_cases=tuple(),
-    ),
-    ModelSpec(
-        vendor="Meta",
-        family="Llama-3.2",
-        name="Llama-3.2-3B-Instruct",
-        size="3B",
-        quantization=None,
-        repo="meta-llama/Llama-3.2-3B-Instruct",
-        config_type=HFLlamaConfig,
-        use_cases=tuple(),
-    ),
-]
-
-LLAMA_MODELS = LLAMA31 + LLAMA32