lalamo 0.5.10__tar.gz → 0.5.11__tar.gz

{lalamo-0.5.10 → lalamo-0.5.11}/PKG-INFO

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

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

{lalamo-0.5.10 → lalamo-0.5.11}/lalamo/model_import/decoder_configs/huggingface/lfm2.py

@@ -9,6 +9,8 @@ from lalamo.modules import (
     DecoderConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
+    MLXQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
     NormalizationConfig,
     SeparableCausalConvConfig,
     ShortConvConfig,
@@ -20,14 +22,21 @@ from lalamo.modules import (
     UntiedEmbeddingConfig,
     UpcastMode,
 )
+from lalamo.quantization import QuantizationMode
 
 from .common import HuggingFaceLMConfig
 
 
+@dataclass(frozen=True)
+class QuantizationConfig:
+    group_size: int
+    bits: int
+
+
 @dataclass(frozen=True)
 class HFLFM2Config(HuggingFaceLMConfig):
     architectures: list[Literal["Lfm2ForCausalLM"]]
-    block_auto_adjust_ff_dim: Literal[False]
+    block_auto_adjust_ff_dim: bool
     block_dim: int
     block_ff_dim: int
     block_ffn_dim_multiplier: float
@@ -38,16 +47,14 @@ class HFLFM2Config(HuggingFaceLMConfig):
     block_use_swiglu: bool
     block_use_xavier_init: bool
     bos_token_id: int
-    conv_L_cache: int # noqa: N815
-    conv_bias: int
+    conv_L_cache: int  # noqa: N815
+    conv_bias: bool
     conv_dim: int
     conv_dim_out: int
     conv_use_xavier_init: bool
     eos_token_id: int
     hidden_size: int
     initializer_range: float
-    intermediate_size: int
-    layer_types: list[Literal["conv", "full_attention"]]
     max_position_embeddings: int
     model_type: Literal["lfm2"]
     norm_eps: float
@@ -57,14 +64,21 @@ class HFLFM2Config(HuggingFaceLMConfig):
     num_key_value_heads: int
     pad_token_id: int
     rope_theta: float
-    theta: float
-    tie_embedding: bool
     torch_dtype: Literal["bfloat16"]
     transformers_version: str
     use_cache: bool
     use_pos_enc: bool
     vocab_size: int
 
+    intermediate_size: int | None = None
+    layer_types: list[Literal["conv", "full_attention"]] | None = None
+    full_attn_idxs: list[int] | None = None
+    tie_embedding: bool = True
+    theta: float | None = None
+
+    quantization: QuantizationConfig | None = None
+    quantization_config: QuantizationConfig | None = None
+
     def to_decoder_config(
         self,
         context_length: int | None,
@@ -74,7 +88,18 @@ class HFLFM2Config(HuggingFaceLMConfig):
     ) -> DecoderConfig:
         assert self.num_attention_heads == self.num_heads
 
-        if self.tie_embedding:
+        if self.quantization_config is not None:
+            assert self.tie_embedding
+
+            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_embedding:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
                 logit_soft_cap=None,
@@ -93,7 +118,15 @@ class HFLFM2Config(HuggingFaceLMConfig):
             max_sequence_length=context_length or self.max_position_embeddings,
         )
 
-        linear_config = FullPrecisionLinearConfig(activation_precision)
+        if self.quantization_config is None:
+            linear_config = FullPrecisionLinearConfig(activation_precision)
+        else:
+            linear_config = MLXQuantizedLinearConfig(
+                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,
+            )
 
         block_norm_config = NormalizationConfig(
             scale_precision=activation_precision,
@@ -123,7 +156,7 @@ class HFLFM2Config(HuggingFaceLMConfig):
 
         short_conv_config = ShortConvConfig(
             in_projection_config=linear_config,
-            conv_config=SeparableCausalConvConfig(activation_precision, has_biases=False),
+            conv_config=SeparableCausalConvConfig(activation_precision, has_biases=self.conv_bias),
             out_projection_config=linear_config,
             kernel_size=self.conv_L_cache,
         )
@@ -137,6 +170,15 @@ class HFLFM2Config(HuggingFaceLMConfig):
             gate_clipping=None,
         )
 
+        if self.layer_types is not None:
+            layer_types = self.layer_types
+        elif self.full_attn_idxs is not None:
+            layer_types = [
+                "full_attention" if i in self.full_attn_idxs else "conv" for i in range(self.num_hidden_layers)
+            ]
+        else:
+            raise RuntimeError("Either layer_types or full_attn_idxs must be present.")
+
         layer_configs = [
             TransformerLayerConfig(
                 pre_mixer_norm_config=block_norm_config,
@@ -145,7 +187,8 @@ class HFLFM2Config(HuggingFaceLMConfig):
                 pre_mlp_norm_config=block_norm_config,
                 mlp_config=mlp_config,
                 post_mlp_norm_config=None,
-            ) for layer_type in self.layer_types
+            )
+            for layer_type in layer_types
         ]
 
         output_norm_config = NormalizationConfig(
@@ -157,13 +200,21 @@ class HFLFM2Config(HuggingFaceLMConfig):
             subtract_mean=False,
         )
 
+        if self.intermediate_size is not None:
+            hidden_dim = self.intermediate_size
+        else:
+            hidden_dim_adjusted = self.block_ff_dim * self.block_ffn_dim_multiplier * (2 / 3)
+            hidden_dim = int(
+                (hidden_dim_adjusted + self.block_multiple_of - 1) // self.block_multiple_of * self.block_multiple_of,
+            )
+
         transformer_config = TransformerConfig(
             global_rope_config=rope_config,
             local_rope_config=None,
             layer_configs=tuple(layer_configs),
             output_norm_config=output_norm_config,
             model_dim=self.hidden_size,
-            hidden_dim=self.intermediate_size,
+            hidden_dim=hidden_dim,
             context_length=context_length or self.max_position_embeddings,
         )
 

{lalamo-0.5.10 → lalamo-0.5.11}/lalamo/model_import/loaders/huggingface.py

@@ -18,6 +18,7 @@ from lalamo.modules import (
     Mamba2Config,
     MLXQuantizedLinear,
     MLXQuantizedTiedEmbedding,
+    MLXQuantizedTiedEmbeddingConfig,
     MLXSemiQuantizedUntiedEmbedding,
     Normalization,
     SeparableCausalConv,
@@ -411,6 +412,7 @@ def _load_conv(
     conv_module: SeparableCausalConv,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
+    permute_conv: bool,
 ) -> SeparableCausalConv:
     weight_path = path / "conv1d" / "weight"
     if weight_path not in weights_dict:
@@ -422,6 +424,8 @@ def _load_conv(
 
     if weight_path is not None:
         raw = weights_dict[weight_path]
+        if permute_conv:
+            raw = jnp.matrix_transpose(raw)
         conv_weight = raw.squeeze(1) if raw.ndim == 3 else raw
     else:
         conv_weight = conv_module.weights
@@ -450,10 +454,11 @@ def load_mamba2(
     module: Mamba2,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
+    permute_conv: bool,
 ) -> Mamba2:
     in_projection = load_linear(module.in_projection, weights_dict, path / "in_proj")
     out_projection = load_linear(module.out_projection, weights_dict, path / "out_proj")
-    conv = _load_conv(module.conv, weights_dict, path)
+    conv = _load_conv(module.conv, weights_dict, path, permute_conv)
 
     skip_connection_weight_path = path / "D"
     if skip_connection_weight_path in weights_dict:
@@ -484,10 +489,11 @@ def load_short_conv(
     module: ShortConv,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
+    permute_conv: bool,
 ) -> ShortConv:
     in_projection = load_linear(module.in_projection, weights_dict, path / "in_proj")
     out_projection = load_linear(module.out_projection, weights_dict, path / "out_proj")
-    conv = _load_conv(module.conv, weights_dict, path)
+    conv = _load_conv(module.conv, weights_dict, path, permute_conv)
 
     return load_parameters(
         lambda m: (m.in_projection, m.out_projection, m.conv),
@@ -508,6 +514,7 @@ def load_transformer_layer(
     up_proj_key: str,
     gate_proj_key: str,
     down_proj_key: str,
+    permute_conv: bool,
 ) -> TransformerLayer:
     if module.pre_mixer_norm is not None:
         pre_attention_norm = load_rmsnorm(
@@ -522,9 +529,9 @@ def load_transformer_layer(
     if isinstance(module.mixer, Attention):
         mixer = load_attention(module.mixer, weights_dict, mixer_path / mixer_key)
     elif isinstance(module.mixer, Mamba2):
-        mixer = load_mamba2(module.mixer, weights_dict, mixer_path / mixer_key)
+        mixer = load_mamba2(module.mixer, weights_dict, mixer_path / mixer_key, permute_conv)
     elif isinstance(module.mixer, ShortConv):
-        mixer = load_short_conv(module.mixer, weights_dict, mixer_path / mixer_key)
+        mixer = load_short_conv(module.mixer, weights_dict, mixer_path / mixer_key, permute_conv)
     else:
         mixer = module.mixer
 
@@ -678,6 +685,7 @@ def load_huggingface_decoder(
         embedding_path = decoder_path / "embedding"
         pre_mixer_norm_key = "input_layernorm"
         mixer_key = {Mamba2Config: "mixer"}
+        permute_conv = False
         pre_mlp_norm_key = "post_attention_layernorm"
         mlp_key = "mlp"
         up_proj_key = "up_proj"
@@ -691,6 +699,7 @@ def load_huggingface_decoder(
         embedding_path = base_path / "embedding.encoder"
         pre_mixer_norm_key = "norm"
         mixer_key = {Mamba2Config: "layer"}
+        permute_conv = False
         pre_mlp_norm_key = "norm"
         mlp_key = "layer"
         up_proj_key = "gate_proj"
@@ -704,6 +713,7 @@ def load_huggingface_decoder(
         embedding_path = decoder_path / "embed_tokens"
         pre_mixer_norm_key = "operator_norm"
         mixer_key = {ShortConvConfig: "conv", AttentionConfig: "self_attn"}
+        permute_conv = isinstance(module.config.embedding_config, MLXQuantizedTiedEmbeddingConfig)
         pre_mlp_norm_key = "ffn_norm"
         mlp_key = "feed_forward"
         up_proj_key = "w3"
@@ -717,6 +727,7 @@ def load_huggingface_decoder(
         embedding_path = decoder_path / "embed_tokens"
         pre_mixer_norm_key = "input_layernorm"
         mixer_key = {AttentionConfig: "self_attn"}
+        permute_conv = False
         pre_mlp_norm_key = "post_attention_layernorm"
         mlp_key = "mlp"
         up_proj_key = "up_proj"
@@ -755,6 +766,7 @@ def load_huggingface_decoder(
             up_proj_key,
             gate_proj_key,
             down_proj_key,
+            permute_conv,
         )
         for i, layer in enumerate(module.transformer.layers)
     )

lalamo-0.5.11/lalamo/model_import/model_specs/lfm2.py

@@ -0,0 +1,31 @@
+from lalamo.model_import.decoder_configs import HFLFM2Config
+from lalamo.quantization import QuantizationMode
+
+from .common import ConfigMap, FileSpec, ModelSpec
+
+__all__ = ["LFM2_MODELS"]
+
+
+def _lfm2_repo(size: str, quantization: QuantizationMode | None) -> tuple[str, str]:
+    organization = "LiquidAI" if quantization is None else "mlx-community"
+    name = f"LFM2-{size}{f'-{quantization.bits}bit' if quantization is not None else ''}"
+    return (organization, name)
+
+
+LFM2_MODELS = [
+    ModelSpec(
+        vendor="LiquidAI",
+        family="LFM2",
+        name=_lfm2_repo(size, quantization)[1],
+        size=size,
+        repo="/".join(_lfm2_repo(size, quantization)),
+        config_type=HFLFM2Config,
+        quantization=quantization,
+        configs=ConfigMap(
+            chat_template=FileSpec("chat_template.jinja"),
+        ),
+        use_cases=tuple(),
+    )
+    for size in ["350M", "700M", "1.2B", "2.6B"]
+    for quantization in [None, *([QuantizationMode.UINT4, QuantizationMode.UINT8] if size != "2.6B" else [])]
+]

{lalamo-0.5.10 → lalamo-0.5.11}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.10
+Version: 0.5.11
 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.10 → lalamo-0.5.11}/tests/test_huggingface_model_conversion.py

@@ -14,6 +14,7 @@ from safetensors.flax import save_file
 from lalamo.common import flatten_parameters
 from lalamo.model_import import REPO_TO_MODEL, ModelMetadata, import_model
 from lalamo.model_import.model_specs import ModelType
+from lalamo.model_import.model_specs.lfm2 import LFM2_MODELS
 from lalamo.models import ClassifierModelConfig, LanguageModelConfig
 from lalamo.modules import config_converter
 from tests.test_models import DType, ModelTestSpec
@@ -27,7 +28,8 @@ MODEL_LIST: list[ModelTestSpec] = [
     ModelTestSpec("meta-llama/Llama-3.2-1B-Instruct", DType.FLOAT32),
     ModelTestSpec("cartesia-ai/Llamba-1B", DType.FLOAT32),
     ModelTestSpec("cartesia-ai/Llamba-1B-4bit-mlx", DType.FLOAT32),
-]
+] + \
+[ModelTestSpec(model.repo, DType.FLOAT32) for model in LFM2_MODELS]
 
 MODEL_LIST += (
     [

{lalamo-0.5.10 → lalamo-0.5.11}/tests/test_lfm2_models.py

@@ -1,11 +1,10 @@
 import pytest
 
+from lalamo.model_import.model_specs.lfm2 import LFM2_MODELS
 from tests.lfm2_tracer import LFM2DecoderTracer
 from tests.test_models import DType, ModelTestSpec, _test_model
 
-MODEL_LIST = [
-    ModelTestSpec("LiquidAI/LFM2-2.6B", DType.FLOAT32),
-]
+MODEL_LIST = [ModelTestSpec(model.repo, DType.FLOAT32) for model in LFM2_MODELS if model.quantization is None]
 
 
 @pytest.mark.parametrize("test_spec", MODEL_LIST, ids=[m.model_repo for m in MODEL_LIST])

lalamo-0.5.10/lalamo/model_import/model_specs/lfm2.py

@@ -1,21 +0,0 @@
-from lalamo.model_import.decoder_configs import HFLFM2Config
-
-from .common import ConfigMap, FileSpec, ModelSpec
-
-__all__ = ["LFM2_MODELS"]
-
-LFM2_MODELS = [
-    ModelSpec(
-        vendor="LiquidAI",
-        family="LFM2",
-        name="LFM2-2.6B",
-        size="2.6B",
-        repo="LiquidAI/LFM2-2.6B",
-        config_type=HFLFM2Config,
-        quantization=None,
-        configs=ConfigMap(
-            chat_template=FileSpec("chat_template.jinja"),
-        ),
-        use_cases=tuple(),
-    ),
-]