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

lalamo/model_import/decoder_configs/huggingface/mistral.py

@@ -7,24 +7,25 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
 )
 from lalamo.modules.activations import SiLU
 from lalamo.modules.normalization import UpcastMode
 
-from .common import HuggingFaceConfig
+from .common import HuggingFaceLMConfig
 
 __all__ = ["HFMistralConfig"]
 
 
 @dataclass(frozen=True)
-class HFMistralConfig(HuggingFaceConfig):
+class HFMistralConfig(HuggingFaceLMConfig):
     architectures: list[Literal["MistralForCausalLM"]]
     attention_dropout: float
     bos_token_id: int
@@ -42,7 +43,6 @@ class HFMistralConfig(HuggingFaceConfig):
     rope_theta: float
     sliding_window: int | None
     tie_word_embeddings: bool
-    torch_dtype: Literal["bfloat16", "float16", "float32"]
     transformers_version: str
     use_cache: bool
     vocab_size: int
@@ -74,12 +74,13 @@ class HFMistralConfig(HuggingFaceConfig):
             max_sequence_length=context_length or 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=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
 
         linear_config = FullPrecisionLinearConfig(
@@ -116,7 +117,7 @@ class HFMistralConfig(HuggingFaceConfig):
                 sliding_window_size=self.sliding_window,
             )
 
-            decoder_layer_config = DecoderLayerConfig(
+            transformer_layer_config = TransformerLayerConfig(
                 pre_mixer_norm_config=rmsnorm_config,
                 mixer_config=attention_config,
                 post_mixer_norm_config=None,
@@ -124,16 +125,20 @@ class HFMistralConfig(HuggingFaceConfig):
                 mlp_config=mlp_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/modern_bert.py

@@ -0,0 +1,241 @@
+from dataclasses import dataclass
+from typing import Literal
+
+import jax.numpy as jnp
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    Activation,
+    AttentionConfig,
+    ClassifierConfig,
+    DenseMLPConfig,
+    FullPrecisionLinearConfig,
+    NormalizationConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
+    UnscaledRoPEConfig,
+    UpcastMode,
+)
+from lalamo.modules.activations import GELU, SiLU
+from lalamo.modules.classifier import (
+    PoolingType,
+    PredictionHeadConfig,
+)
+from lalamo.modules.embedding import TiedEmbeddingConfig
+
+from .common import (
+    AWQQuantizationConfig,
+    GPTQQuantizationConfig,
+    HuggingFaceClassifierConfig,
+)
+
+__all__ = ["ModernBERTConfig"]
+
+
+def activation_from_str(activation: str) -> type[Activation]:
+    supported_activations = {
+        "silu": SiLU,
+        "gelu": GELU,
+    }
+    if activation in supported_activations:
+        return supported_activations[activation]
+
+    raise ValueError(
+        f"Only activations from the following list are supported by Classifier: {supported_activations.keys()}"
+    )
+
+
+@dataclass(frozen=True)
+class ModernBERTConfig(HuggingFaceClassifierConfig):
+    architectures: list[Literal["ModernBertForSequenceClassification"]]
+    attention_bias: bool
+    attention_dropout: float
+    bos_token_id: int | list[int]
+    classifier_activation: Literal["gelu"]
+    classifier_bias: bool
+    classifier_dropout: float
+    classifier_pooling: Literal["mean"]
+    cls_token_id: int | list[int]
+    decoder_bias: bool
+    deterministic_flash_attn: bool
+    embedding_dropout: float
+    eos_token_id: int | list[int]
+    global_attn_every_n_layers: int
+    global_rope_theta: float
+    gradient_checkpointing: bool
+    hidden_activation: Literal["gelu"]
+    hidden_size: int
+    initializer_cutoff_factor: float
+    initializer_range: float
+    intermediate_size: int
+    layer_norm_eps: float
+    local_attention: int
+    local_rope_theta: float
+    max_position_embeddings: int
+    mlp_bias: bool
+    mlp_dropout: float
+    model_type: Literal["modernbert"]
+    norm_bias: bool
+    norm_eps: float
+    num_attention_heads: int
+    num_hidden_layers: int
+    pad_token_id: int | list[int]
+    position_embedding_type: Literal["absolute"]
+    sep_token_id: int | list[int]
+    transformers_version: str
+    vocab_size: int
+    id2label: dict[int, str]
+    label2id: dict[str, int]
+
+    quantization_config: AWQQuantizationConfig | GPTQQuantizationConfig | None = None
+
+    def __post_init__(self) -> None:
+        if len(self.label2id) != len(self.id2label):
+            raise ValueError("Legnth of label2id and id2label is expected to be the same")
+
+    def calculate_sliding_windows(self, num_layers: int, global_attn_every_n_layers: int) -> tuple[None, ...]:
+        result = [None] * num_layers
+        for index in range(len(result)):
+            if index % global_attn_every_n_layers != 0:
+                result[index] = self.local_attention  # type: ignore
+            else:
+                pass
+        return tuple(result)
+
+    def to_classifier_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> ClassifierConfig:
+        embedding_config = TiedEmbeddingConfig(
+            input_scale=None,
+            logit_soft_cap=None,
+            precision=activation_precision,
+        )
+        embedding_norm_config = NormalizationConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=True,
+        )
+
+        global_rope_config = UnscaledRoPEConfig(
+            precision=activation_precision,
+            base=self.global_rope_theta,
+            max_sequence_length=context_length or self.max_position_embeddings,
+        )
+        local_rope_config = UnscaledRoPEConfig(
+            precision=activation_precision,
+            base=self.local_rope_theta,
+            max_sequence_length=context_length or self.max_position_embeddings,
+        )
+
+        sliding_window_sizes = self.calculate_sliding_windows(self.num_hidden_layers, self.global_attn_every_n_layers)
+
+        transformer_norm_config = NormalizationConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=True,
+        )
+        linear_config = FullPrecisionLinearConfig(
+            precision=activation_precision,
+        )
+        activation = activation_from_str(self.hidden_activation)
+        assert activation is SiLU or activation is GELU
+        mlp_config = DenseMLPConfig(
+            linear_config=linear_config,
+            activation=activation(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
+        )
+
+        # In ModernBERT architecture first Transformer layer has no pre-attention normalization
+        pre_attn_configs = [transformer_norm_config if i > 0 else None for i in range(self.num_hidden_layers)]
+
+        transformer_layer_configs = []
+        for sliding_window_size, pre_attn_config in zip(sliding_window_sizes, pre_attn_configs, strict=True):
+            attention_config = AttentionConfig(
+                qkv_projection_config=linear_config,
+                out_projection_config=linear_config,
+                query_norm_config=None,
+                key_norm_config=None,
+                logit_soft_cap=None,
+                has_sinks=False,
+                has_qkv_biases=self.attention_bias,
+                has_out_biases=False,
+                num_heads=self.num_attention_heads,
+                num_groups=self.num_attention_heads,
+                head_dim=self.hidden_size // self.num_attention_heads,
+                scale=None,
+                is_causal=False,
+                sliding_window_size=sliding_window_size,
+            )
+            layer_config = TransformerLayerConfig(
+                pre_mixer_norm_config=pre_attn_config,
+                mixer_config=attention_config,
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=transformer_norm_config,
+                mlp_config=mlp_config,
+                post_mlp_norm_config=None,
+            )
+            transformer_layer_configs.append(layer_config)
+
+        transformer_config = TransformerConfig(
+            global_rope_config=global_rope_config,
+            local_rope_config=local_rope_config,
+            layer_configs=tuple(transformer_layer_configs),
+            output_norm_config=transformer_norm_config,
+            model_dim=self.hidden_size,
+            hidden_dim=self.intermediate_size,
+            context_length=context_length or self.max_position_embeddings,
+        )
+
+        prediction_head_dense_config = FullPrecisionLinearConfig(
+            precision=activation_precision,
+        )
+        prediction_head_norm_config = NormalizationConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=jnp.float32,
+            epsilon=self.norm_eps,
+            scale_offset=0.0,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=True,
+        )
+        prediction_head_activation = activation_from_str(self.classifier_activation)
+        prediction_head_readout_config = FullPrecisionLinearConfig(
+            precision=activation_precision,
+        )
+        prediction_head_config = PredictionHeadConfig(
+            dense_config=prediction_head_dense_config,
+            activation=prediction_head_activation(),
+            normalization_config=prediction_head_norm_config,
+            readout_config=prediction_head_readout_config,
+            use_dense_bias=self.classifier_bias,
+        )
+
+        output_labels = [self.id2label[idx] for idx in range(len(self.id2label))]
+
+        return ClassifierConfig(
+            embedding_config=embedding_config,
+            embedding_norm_config=embedding_norm_config,
+            transformer_config=transformer_config,
+            prediction_head_config=prediction_head_config,
+            readout_config=prediction_head_readout_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.hidden_size,
+            hidden_dim=self.hidden_size,
+            attention_scale=None,
+            num_layers=self.num_hidden_layers,
+            context_length=self.max_position_embeddings,
+            num_labels=len(self.id2label),
+            classifier_pooling=PoolingType(self.classifier_pooling),
+            output_labels=tuple(output_labels),
+        )

lalamo/model_import/decoder_configs/huggingface/qwen2.py

@@ -7,12 +7,13 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
@@ -20,13 +21,13 @@ from lalamo.modules import (
 from lalamo.modules.activations import SiLU
 from lalamo.quantization import QuantizationMode
 
-from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
+from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceLMConfig
 
 __all__ = ["HFQwen2Config"]
 
 
 @dataclass(frozen=True)
-class HFQwen2Config(HuggingFaceConfig):
+class HFQwen2Config(HuggingFaceLMConfig):
     torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["Qwen2ForCausalLM"]]
     attention_dropout: float
@@ -89,12 +90,13 @@ class HFQwen2Config(HuggingFaceConfig):
             base=self.rope_theta,
             max_sequence_length=context_length or 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=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
         if self.quantization_config is None:
             linear_config = FullPrecisionLinearConfig(
@@ -136,7 +138,7 @@ class HFQwen2Config(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,
@@ -144,15 +146,20 @@ class HFQwen2Config(HuggingFaceConfig):
                 mlp_config=mlp_config,
                 post_mlp_norm_config=None,
             )
-            layer_configs.append(decoder_layer_config)
-        return DecoderConfig(
-            embedding_config=embedding_config,
+            layer_configs.append(transformer_layer_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/qwen3.py

@@ -7,12 +7,13 @@ from jaxtyping import DTypeLike
 from lalamo.modules import (
     AttentionConfig,
     DecoderConfig,
-    DecoderLayerConfig,
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
-    RMSNormConfig,
+    NormalizationConfig,
     TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
@@ -22,13 +23,13 @@ from lalamo.modules.embedding import MLXQuantizedTiedEmbeddingConfig
 from lalamo.modules.linear import MLXQuantizedLinearConfig
 from lalamo.quantization import QuantizationMode
 
-from .common import HuggingFaceConfig, MLXQuantizationConfig, QuantizationConfigType
+from .common import HuggingFaceLMConfig, MLXQuantizationConfig, QuantizationConfigType
 
 __all__ = ["HFQwen3Config"]
 
 
 @dataclass(frozen=True)
-class HFQwen3Config(HuggingFaceConfig):
+class HFQwen3Config(HuggingFaceLMConfig):
     eos_token_id: int | list[int]
     torch_dtype: Literal["bfloat16", "float16", "float32"]
     attention_bias: bool
@@ -100,12 +101,13 @@ class HFQwen3Config(HuggingFaceConfig):
             base=self.rope_theta,
             max_sequence_length=context_length or 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=None,
             upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
         )
         if self.quantization_config is None:
             linear_config = FullPrecisionLinearConfig(
@@ -153,7 +155,7 @@ class HFQwen3Config(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,
@@ -161,15 +163,18 @@ class HFQwen3Config(HuggingFaceConfig):
                 mlp_config=mlp_config,
                 post_mlp_norm_config=None,
             )
-            layer_configs.append(decoder_layer_config)
-        return DecoderConfig(
-            embedding_config=embedding_config,
+            layer_configs.append(transformer_layer_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/loaders/__init__.py

@@ -1,7 +1,8 @@
 # from .executorch import load_executorch
-from .huggingface import load_huggingface
+from .huggingface import load_huggingface_classifier, load_huggingface_decoder
 
 __all__ = [
+    "load_huggingface_classifier",
     # "load_executorch",
-    "load_huggingface",
+    "load_huggingface_decoder",
 ]

lalamo/model_import/loaders/executorch.py

@@ -6,7 +6,15 @@ from einops import rearrange
 from jaxtyping import Array, Float, Int
 
 from lalamo.common import ParameterPath
-from lalamo.modules import Attention, Decoder, DecoderLayer, DenseMLP, QLoRALinear, QuantizedTiedEmbedding, RMSNorm
+from lalamo.modules import (
+    Attention,
+    Decoder,
+    DenseMLP,
+    Normalization,
+    QLoRALinear,
+    QuantizedTiedEmbedding,
+    TransformerLayer,
+)
 
 from .common import load_parameters
 
@@ -95,7 +103,7 @@ def load_mlp(module: DenseMLP, weights_dict: Mapping[str, Array], path: Paramete
     )
 
 
-def load_rmsnorm(module: RMSNorm, weights_dict: Mapping[str, Array], path: ParameterPath) -> RMSNorm:
+def load_rmsnorm(module: Normalization, weights_dict: Mapping[str, Array], path: ParameterPath) -> Normalization:
     return load_parameters(lambda m: (m.scales,), module, (weights_dict[path / "weight"],))
 
 
@@ -175,18 +183,21 @@ def load_attention(
     )
 
 
-def load_decoder_layer(
-    module: DecoderLayer,
+def load_transformer_layer(
+    module: TransformerLayer,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
-) -> DecoderLayer:
+) -> TransformerLayer:
     if module.post_mixer_norm is not None:
         raise ValueError("Post attention normalization is not supported")
     if module.post_mlp_norm is not None:
         raise ValueError("Post MLP normalization is not supported")
-    attention_norm = load_rmsnorm(module.pre_mixer_norm, weights_dict, path / "attention_norm")
+    if module.pre_mixer_norm is not None:
+        attention_norm = load_rmsnorm(module.pre_mixer_norm, weights_dict, path / "attention_norm")
+    else:
+        attention_norm = None
     assert isinstance(module.mixer, Attention)
-    attention = load_attention(module.mixer, weights_dict, path / "attention")
+    attention = load_attention(module.mixer, weights_dict, path / "mixer")
     mlp_norm = load_rmsnorm(module.pre_mlp_norm, weights_dict, path / "ffn_norm")
     assert isinstance(module.mlp, DenseMLP)
     mlp = load_mlp(module.mlp, weights_dict, path / "feed_forward")
@@ -214,12 +225,13 @@ def load_executorch(module: Decoder, weights_dict: Mapping[str, Array]) -> Decod
         raise TypeError(f"Expected embedding to be QuantizedTiedEmbedding, got {type(module.embedding)}")
 
     embedding = load_embedding(module.embedding, weights_dict, root_path / "tok_embeddings")
-    decoder_layers = tuple(
-        load_decoder_layer(layer, weights_dict, root_path / f"layers.{i}") for i, layer in enumerate(module.layers)
+    transformer_layers = tuple(
+        load_transformer_layer(layer, weights_dict, root_path / f"layers.{i}")
+        for i, layer in enumerate(module.transformer.layers)
     )
-    output_norm = load_rmsnorm(module.output_norm, weights_dict, root_path / "norm")
+    output_norm = load_rmsnorm(module.transformer.output_norm, weights_dict, root_path / "norm")
     return load_parameters(
-        lambda m: (m.embedding, m.layers, m.output_norm),
+        lambda m: (m.embedding, m.transformer.layers, m.transformer.output_norm),
         module,
-        (embedding, decoder_layers, output_norm),
+        (embedding, transformer_layers, output_norm),
     )