lalamo 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

lalamo/model_import/decoder_configs/executorch.py

@@ -9,12 +9,13 @@ from lalamo.modules import (
     AttentionConfig,
     Decoder,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     LlamaRoPEConfig,
+    NormalizationConfig,
     QLoRALinearConfig,
     QuantizedTiedEmbeddingConfig,
-    RMSNormConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UpcastMode,
 )
 from lalamo.modules.activations import SiLU
@@ -62,7 +63,7 @@ class ExecutorchConfig(ForeignConfig):
         return jnp.bfloat16
 
     @classmethod
-    def _load_weights(
+    def _load_decoder_weights(
         cls,
         model: Decoder,
         weights_dict: Mapping[str, Array],
@@ -119,12 +120,13 @@ class ETLlamaConfig(ExecutorchConfig):
             low_frequency_factor=LOW_FREQ_FACTOR,
             high_frequency_factor=HIGH_FREQ_FACTOR,
         )
-        rmsnorm_config = RMSNormConfig(
+        rmsnorm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.norm_eps,
             scale_offset=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
         linear_config = QLoRALinearConfig(
             group_size=self.quantization_args.group_size,
@@ -158,7 +160,7 @@ class ETLlamaConfig(ExecutorchConfig):
             up_clipping=None,
             gate_clipping=None,
         )
-        decoder_layer_config = DecoderLayerConfig(
+        tranformer_layer_config = TransformerLayerConfig(
             pre_mixer_norm_config=rmsnorm_config,
             mixer_config=attention_config,
             post_mixer_norm_config=None,
@@ -166,14 +168,17 @@ class ETLlamaConfig(ExecutorchConfig):
             mlp_config=mlp_config,
             post_mlp_norm_config=None,
         )
-        return DecoderConfig(
-            embedding_config=embedding_config,
+        transformer_config = TransformerConfig(
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_configs=(decoder_layer_config,) * self.n_layers,
+            layer_configs=(tranformer_layer_config,) * self.n_layers,
             output_norm_config=rmsnorm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.dim,
             hidden_dim=self._find_hidden_size(),
             context_length=context_length or MAX_SEQUENCE_LENGTH,
         )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )

lalamo/model_import/decoder_configs/huggingface/__init__.py

@@ -1,10 +1,11 @@
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 from .gemma2 import HFGemma2Config
 from .gemma3 import HFGemma3Config, HFGemma3TextConfig
 from .gpt_oss import HFGPTOssConfig
 from .llama import HFLlamaConfig
 from .llamba import HFLlambaConfig
 from .mistral import HFMistralConfig
+from .modern_bert import ModernBERTConfig
 from .qwen2 import HFQwen2Config
 from .qwen3 import HFQwen3Config
 
@@ -18,5 +19,6 @@ __all__ = [
     "HFMistralConfig",
     "HFQwen2Config",
     "HFQwen3Config",
-    "HuggingFaceConfig",
+    "HuggingFaceLMConfig",
+    "ModernBERTConfig",
 ]

lalamo/model_import/decoder_configs/huggingface/common.py

@@ -6,15 +6,22 @@ import cattrs
 import jax.numpy as jnp
 from jaxtyping import Array, DTypeLike
 
-from lalamo.model_import.decoder_configs import ForeignConfig
-from lalamo.model_import.loaders import load_huggingface
+from lalamo.model_import.decoder_configs import ForeignLMConfig
+from lalamo.model_import.decoder_configs.common import ForeignClassifierConfig
+from lalamo.model_import.loaders import (
+    load_huggingface_classifier,
+    load_huggingface_decoder,
+)
 from lalamo.modules import Decoder
+from lalamo.modules.classifier import Classifier
+from lalamo.modules.common import LalamoModule
 
 __all__ = [
     "AWQQuantizationConfig",
     "GPTQMetaConfig",
     "GPTQQuantizationConfig",
-    "HuggingFaceConfig",
+    "HuggingFaceClassifierConfig",
+    "HuggingFaceLMConfig",
 ]
 
 
@@ -85,26 +92,45 @@ def _structure_quantization_config(v: object, _: object) -> QuantizationConfigTy
 
 
 @dataclass(frozen=True)
-class HuggingFaceConfig(ForeignConfig):
+class HuggingFaceLMConfig(ForeignLMConfig):
     _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
     _converter.register_structure_hook(int | list[int], lambda v, _: v)
     _converter.register_structure_hook(QuantizationConfigType, _structure_quantization_config)
 
     @property
     def eos_token_ids(self) -> list[int]:
-        if not hasattr(self, "eos_token_id"):
-            raise RuntimeError("model doesn't havve eos_token_id, override eos_token_ids in model config")
+        result = getattr(self, "eos_token_id", None)
+        if result is None:
+            raise RuntimeError("model doesn't have eos_token_id, override eos_token_ids in model config")
 
-        return [self.eos_token_id] if isinstance(self.eos_token_id, int) else self.eos_token_id  # type: ignore  (This is a bug in pyright)
+        if isinstance(result, int):
+            result = [result]
 
+        return result
+
+    @property
+    def default_precision(self) -> DTypeLike:
+        return jnp.dtype(getattr(self, "torch_dtype", "bfloat16"))
+
+    def _load_weights(
+        self,
+        model: LalamoModule,
+        weights_dict: Mapping[str, Array],
+    ) -> LalamoModule:
+        assert isinstance(model, Decoder)
+        return load_huggingface_decoder(model, weights_dict)
+
+
+@dataclass(frozen=True)
+class HuggingFaceClassifierConfig(ForeignClassifierConfig):
     @property
     def default_precision(self) -> DTypeLike:
         return jnp.dtype(getattr(self, "torch_dtype", "bfloat16"))
 
-    @classmethod
     def _load_weights(
-        cls,
-        model: Decoder,
+        self,
+        model: LalamoModule,
         weights_dict: Mapping[str, Array],
-    ) -> Decoder:
-        return load_huggingface(model, weights_dict)
+    ) -> LalamoModule:
+        assert isinstance(model, Classifier)
+        return load_huggingface_classifier(model, weights_dict)

lalamo/model_import/decoder_configs/huggingface/gemma2.py

@@ -7,23 +7,24 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UnscaledRoPEConfig,
     UpcastMode,
 )
 from lalamo.modules.activations import GELU
 
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 
 __all__ = ["HFGemma2Config"]
 
 
 @dataclass(frozen=True)
-class HFGemma2Config(HuggingFaceConfig):
+class HFGemma2Config(HuggingFaceLMConfig):
     architectures: list[Literal["Gemma2ForCausalLM"]]
     attention_bias: bool
     attention_dropout: float
@@ -72,12 +73,13 @@ class HFGemma2Config(HuggingFaceConfig):
             base=self.rope_theta,
             max_sequence_length=self.max_position_embeddings,
         )
-        rmsnorm_config = RMSNormConfig(
+        rmsnorm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.rms_norm_eps,
             scale_offset=1.0,
             upcast_mode=UpcastMode.FULL_LAYER,
+            subtract_mean=False,
         )
         linear_config = FullPrecisionLinearConfig(
             precision=activation_precision,
@@ -110,7 +112,7 @@ class HFGemma2Config(HuggingFaceConfig):
                 scale=attention_scale,
                 sliding_window_size=sliding_window_size,
             )
-            decoder_layer_config = DecoderLayerConfig(
+            transformer_layer_config = TransformerLayerConfig(
                 pre_mixer_norm_config=rmsnorm_config,
                 mixer_config=attention_config,
                 post_mixer_norm_config=rmsnorm_config,
@@ -118,16 +120,19 @@ class HFGemma2Config(HuggingFaceConfig):
                 mlp_config=mlp_config,
                 post_mlp_norm_config=rmsnorm_config,
             )
-            layer_configs.append(decoder_layer_config)
+            layer_configs.append(transformer_layer_config)
 
-        return DecoderConfig(
-            embedding_config=embedding_config,
+        transformer_config = TransformerConfig(
             global_rope_config=rope_config,
             local_rope_config=None,
             layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
             context_length=context_length or self.max_position_embeddings,
         )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )

lalamo/model_import/decoder_configs/huggingface/gemma3.py

@@ -1,23 +1,20 @@
 from collections.abc import Mapping
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Literal
 
 import jax.numpy as jnp
 from jaxtyping import DTypeLike
 
-from lalamo.modules import (
-    DecoderConfig,
-    TiedEmbeddingConfig,
-)
+from lalamo.modules import DecoderConfig, TiedEmbeddingConfig, TransformerConfig
 from lalamo.modules.activations import GELU
-from lalamo.modules.decoder_layer import DecoderLayerConfig
 from lalamo.modules.linear import FullPrecisionLinearConfig
 from lalamo.modules.mlp import DenseMLPConfig
-from lalamo.modules.normalization import RMSNormConfig, UpcastMode
+from lalamo.modules.normalization import NormalizationConfig, UpcastMode
 from lalamo.modules.rope import LinearScalingRoPEConfig, UnscaledRoPEConfig
-from lalamo.modules.token_mixers import AttentionConfig
+from lalamo.modules.token_mixers.attention import AttentionConfig
+from lalamo.modules.transformer_layer import TransformerLayerConfig
 
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 
 __all__ = ["HFGemma3Config", "HFGemma3TextConfig"]
 
@@ -80,12 +77,13 @@ class HFGemma3TextConfigRaw:
             logit_soft_cap=None,
             precision=activation_precision,
         )
-        rms_norm_config = RMSNormConfig(
+        rms_norm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.rms_norm_eps,
             scale_offset=1.0,
             upcast_mode=UpcastMode.FULL_LAYER,
+            subtract_mean=False,
         )
 
         if self.rope_scaling is not None:
@@ -134,7 +132,7 @@ class HFGemma3TextConfigRaw:
                 scale=attention_scale,
                 sliding_window_size=sliding_window_size,
             )
-            decoder_layer_config = DecoderLayerConfig(
+            transformer_layer_config = TransformerLayerConfig(
                 pre_mixer_norm_config=rms_norm_config,
                 mixer_config=attention_config,
                 post_mixer_norm_config=rms_norm_config,
@@ -142,23 +140,29 @@ class HFGemma3TextConfigRaw:
                 mlp_config=mlp_config,
                 post_mlp_norm_config=rms_norm_config,
             )
-            layer_configs.append(decoder_layer_config)
-        return DecoderConfig(
-            embedding_config=embedding_config,
+            layer_configs.append(transformer_layer_config)
+
+        transformer_config = TransformerConfig(
             global_rope_config=global_rope_config,
             local_rope_config=local_rope_config,
             layer_configs=tuple(layer_configs),
             output_norm_config=rms_norm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
             context_length=context_length or self.max_position_embeddings,
         )
 
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )
+
 
 @dataclass(frozen=True)
-class HFGemma3TextConfig(HFGemma3TextConfigRaw, HuggingFaceConfig):
+class HFGemma3TextConfig(HFGemma3TextConfigRaw, HuggingFaceLMConfig):
     torch_dtype: Literal["bfloat16", "float16", "float32"] = "bfloat16"
+    eos_token_id: int | list[int] = field(default_factory=list)
 
 
 @dataclass(frozen=True)
@@ -174,7 +178,7 @@ class HFGemma3VisionConfig:
 
 
 @dataclass(frozen=True)
-class HFGemma3Config(HuggingFaceConfig):
+class HFGemma3Config(HuggingFaceLMConfig):
     torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["Gemma3ForConditionalGeneration"]]
     boi_token_index: int

lalamo/model_import/decoder_configs/huggingface/gpt_oss.py

@@ -7,20 +7,21 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     MixtureOfExpertsConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     SoftmaxRouting,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
     YARNRoPEConfig,
 )
 from lalamo.modules.activations import SiLU
 
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 
 __all__ = ["HFGPTOssConfig"]
 
@@ -36,7 +37,7 @@ class YarnRopeScalingConfig:
 
 
 @dataclass(frozen=True)
-class HFGPTOssConfig(HuggingFaceConfig):
+class HFGPTOssConfig(HuggingFaceLMConfig):
     # Core HF fields
     architectures: list[Literal["GptOssForCausalLM"]]
     attention_bias: bool
@@ -115,12 +116,13 @@ class HFGPTOssConfig(HuggingFaceConfig):
                 truncate=True,
             )
 
-        rmsnorm_config = RMSNormConfig(
+        rmsnorm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.rms_norm_eps,
             scale_offset=None,
             upcast_mode=UpcastMode.FULL_LAYER,
+            subtract_mean=False,
         )
 
         # Linear layers
@@ -179,7 +181,7 @@ class HFGPTOssConfig(HuggingFaceConfig):
                 scale=None,
                 sliding_window_size=sliding_window_size,
             )
-            decoder_layer_config = DecoderLayerConfig(
+            transformer_layer_config = TransformerLayerConfig(
                 pre_mixer_norm_config=rmsnorm_config,
                 mixer_config=attention_config,
                 post_mixer_norm_config=None,
@@ -187,16 +189,20 @@ class HFGPTOssConfig(HuggingFaceConfig):
                 mlp_config=moe_config,
                 post_mlp_norm_config=None,
             )
-            layer_configs.append(decoder_layer_config)
+            layer_configs.append(transformer_layer_config)
 
-        return DecoderConfig(
-            embedding_config=embedding_config,
+        transformer_config = TransformerConfig(
             global_rope_config=rope_config,
             local_rope_config=None,
             layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
             context_length=context_length or self.max_position_embeddings,
         )
+
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )

lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -7,14 +7,15 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
     LlamaRoPEConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     SiLU,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
@@ -22,7 +23,7 @@ from lalamo.modules import (
 )
 from lalamo.quantization import QuantizationMode
 
-from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
+from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceLMConfig
 
 __all__ = ["HFLlamaConfig"]
 
@@ -47,7 +48,7 @@ class YarnRopeScalingConfig:
 
 
 @dataclass(frozen=True)
-class HFLlamaConfig(HuggingFaceConfig):
+class HFLlamaConfig(HuggingFaceLMConfig):
     torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["LlamaForCausalLM"]]
     attention_bias: bool
@@ -124,12 +125,13 @@ class HFLlamaConfig(HuggingFaceConfig):
             )
         else:
             raise ValueError("Unsupported rope_scaling configuration")
-        rmsnorm_config = RMSNormConfig(
+        rmsnorm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.rms_norm_eps,
             scale_offset=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
         if self.quantization_config is None:
             linear_config = FullPrecisionLinearConfig(
@@ -153,7 +155,7 @@ class HFLlamaConfig(HuggingFaceConfig):
             has_out_biases=False,
             num_heads=self.num_attention_heads,
             num_groups=self.num_key_value_heads,
-            head_dim=self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads,
+            head_dim=(self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads),
             is_causal=True,
             scale=None,
             sliding_window_size=None,
@@ -166,7 +168,7 @@ class HFLlamaConfig(HuggingFaceConfig):
             up_clipping=None,
             gate_clipping=None,
         )
-        decoder_layer_config = DecoderLayerConfig(
+        transformer_layer_config = TransformerLayerConfig(
             pre_mixer_norm_config=rmsnorm_config,
             mixer_config=attention_config,
             post_mixer_norm_config=None,
@@ -174,14 +176,17 @@ class HFLlamaConfig(HuggingFaceConfig):
             mlp_config=mlp_config,
             post_mlp_norm_config=None,
         )
-        return DecoderConfig(
-            embedding_config=embedding_config,
+        transformer_config = TransformerConfig(
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_configs=(decoder_layer_config,) * self.num_hidden_layers,
+            layer_configs=(transformer_layer_config,) * self.num_hidden_layers,
             output_norm_config=rmsnorm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
             context_length=context_length or self.max_position_embeddings,
         )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )

lalamo/model_import/decoder_configs/huggingface/llamba.py

@@ -6,23 +6,24 @@ from jaxtyping import DTypeLike
 
 from lalamo.modules import (
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     Identity,
     Mamba2Config,
     MLXQuantizedLinearConfig,
     MLXSemiQuantizedUntiedEmbeddingConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     SeparableCausalConvConfig,
     SiLU,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
 )
 from lalamo.quantization import QuantizationMode
 
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 
 
 @dataclass(frozen=True)
@@ -45,7 +46,7 @@ class HFLlambaSsmConfig:
 
 
 @dataclass(frozen=True)
-class HFLlambaConfig(HuggingFaceConfig):
+class HFLlambaConfig(HuggingFaceLMConfig):
     model_type: Literal["llamba"]
     vocab_size: int
     tie_embeddings: bool
@@ -74,7 +75,9 @@ class HFLlambaConfig(HuggingFaceConfig):
                 input_scale=None,
                 logit_soft_cap=None,
                 group_size=int(metadata_dict["quantization_kwargs.group_size"]),
-                embedding_quantization_mode=QuantizationMode.from_num_bits(int(metadata_dict["quantization_kwargs.bits"])),
+                embedding_quantization_mode=QuantizationMode.from_num_bits(
+                    int(metadata_dict["quantization_kwargs.bits"])
+                ),
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
             )
@@ -91,18 +94,21 @@ class HFLlambaConfig(HuggingFaceConfig):
                 precision=activation_precision,
             )
 
-        rmsnorm_config = RMSNormConfig(
+        rmsnorm_config = NormalizationConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
             epsilon=self.norm_epsilon,
             scale_offset=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
 
-        if "quantization_kwargs.group_size" in metadata_dict:
+        if metadata_dict and "quantization_kwargs.group_size" in metadata_dict:
             linear_config = MLXQuantizedLinearConfig(
                 group_size=int(metadata_dict["quantization_kwargs.group_size"]),
-                weight_quantization_mode=QuantizationMode.from_num_bits(int(metadata_dict["quantization_kwargs.bits"])),
+                weight_quantization_mode=QuantizationMode.from_num_bits(
+                    int(metadata_dict["quantization_kwargs.bits"])
+                ),
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
             )
@@ -148,7 +154,7 @@ class HFLlambaConfig(HuggingFaceConfig):
             has_out_biases=self.ssm_cfg.bias,
         )
 
-        decoder_layer_config = DecoderLayerConfig(
+        transformer_layer_config = TransformerLayerConfig(
             pre_mixer_norm_config=rmsnorm_config,
             mixer_config=mamba_config,
             post_mixer_norm_config=None,
@@ -156,15 +162,18 @@ class HFLlambaConfig(HuggingFaceConfig):
             mlp_config=mlp_config,
             post_mlp_norm_config=None,
         )
-
-        return DecoderConfig(
-            embedding_config=embedding_config,
+        transformer_config = TransformerConfig(
             global_rope_config=None,
             local_rope_config=None,
-            layer_configs=(decoder_layer_config,) * self.n_layer,
+            layer_configs=(transformer_layer_config,) * self.n_layer,
             output_norm_config=rmsnorm_config,
-            vocab_size=self.vocab_size,
             model_dim=self.d_model,
             hidden_dim=self.mlp_cfg.intermediate_size,
             context_length=context_length or 4096,
         )
+
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )