lalamo 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

lalamo/model_import/decoder_configs/huggingface/gemma3.py

@@ -8,11 +8,11 @@ from lalamo.modules import (
     DecoderConfig,
     TiedEmbeddingConfig,
 )
-from lalamo.modules.activations import Activation
+from lalamo.modules.activations import GELU
 from lalamo.modules.attention import AttentionConfig
 from lalamo.modules.decoder_layer import DecoderLayerConfig
 from lalamo.modules.linear import FullPrecisionLinearConfig
-from lalamo.modules.mlp import MLPConfig
+from lalamo.modules.mlp import DenseMLPConfig
 from lalamo.modules.normalization import RMSNormConfig, UpcastMode
 from lalamo.modules.rope import LinearScalingRoPEConfig, UnscaledRoPEConfig
 
@@ -75,7 +75,7 @@ class HFGemma3TextConfigRaw:
         attention_scale = self.query_pre_attn_scalar**-0.5
         embedding_config = TiedEmbeddingConfig(
             input_scale=input_scale,
-            logits_soft_cap=None,
+            logit_soft_cap=None,
             precision=activation_precision,
         )
         rms_norm_config = RMSNormConfig(
@@ -106,13 +106,21 @@ class HFGemma3TextConfigRaw:
         )
 
         linear_config = FullPrecisionLinearConfig(precision=activation_precision)
-        mlp_config = MLPConfig(linear_config=linear_config, activation=Activation.GELU)
+        mlp_config = DenseMLPConfig(
+            linear_config=linear_config,
+            activation=GELU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
+        )
         attention_config = AttentionConfig(
             qkv_projection_config=linear_config,
             out_projection_config=linear_config,
             query_norm_config=rms_norm_config,
             key_norm_config=rms_norm_config,
             logit_soft_cap=self.attn_logit_softcapping,
+            has_sinks=False,
             has_qkv_biases=self.attention_bias,
             has_out_biases=self.attention_bias,
         )
@@ -145,7 +153,7 @@ class HFGemma3TextConfigRaw:
 
 @dataclass(frozen=True)
 class HFGemma3TextConfig(HFGemma3TextConfigRaw, HuggingFaceConfig):
-    pass
+    torch_dtype: Literal["bfloat16", "float16", "float32"] = "bfloat16"
 
 
 @dataclass(frozen=True)
@@ -162,6 +170,7 @@ class HFGemma3VisionConfig:
 
 @dataclass(frozen=True)
 class HFGemma3Config(HuggingFaceConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["Gemma3ForConditionalGeneration"]]
     boi_token_index: int
     eoi_token_index: int

lalamo/model_import/decoder_configs/huggingface/gpt_oss.py

@@ -0,0 +1,195 @@
+from dataclasses import dataclass
+from typing import Literal
+
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    AttentionConfig,
+    DecoderConfig,
+    DecoderLayerConfig,
+    DenseMLPConfig,
+    FullPrecisionLinearConfig,
+    MixtureOfExpertsConfig,
+    RMSNormConfig,
+    SoftmaxRouting,
+    TiedEmbeddingConfig,
+    UntiedEmbeddingConfig,
+    UpcastMode,
+    YARNRoPEConfig,
+)
+from lalamo.modules.activations import SiLU
+
+from .common import HuggingFaceConfig
+
+__all__ = ["HFGPTOssConfig"]
+
+
+@dataclass(frozen=True)
+class YarnRopeScalingConfig:
+    factor: float
+    beta_fast: float
+    beta_slow: float
+    original_max_position_embeddings: int
+    rope_type: Literal["yarn"]
+    truncate: bool
+
+
+@dataclass(frozen=True)
+class HFGPTOssConfig(HuggingFaceConfig):
+    # Core HF fields
+    architectures: list[Literal["GptOssForCausalLM"]]
+    attention_bias: bool
+    attention_dropout: float
+    eos_token_id: int | list[int]
+    hidden_act: Literal["silu"]
+    hidden_size: int
+    initializer_range: float
+    intermediate_size: int
+    max_position_embeddings: int
+    model_type: Literal["gpt_oss"]
+    num_attention_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    pad_token_id: int
+    rms_norm_eps: float
+    rope_theta: float
+    tie_word_embeddings: bool
+    transformers_version: str
+    use_cache: bool
+    vocab_size: int
+
+    # GPT-OSS specifics
+    layer_types: list[Literal["sliding_attention", "full_attention"]] | None
+    sliding_window: int | None
+    swiglu_limit: float
+    head_dim: int | None
+    num_local_experts: int
+    num_experts_per_tok: int | None = None
+    experts_per_token: int | None = None  # some configs may use this alias
+    rope_scaling: YarnRopeScalingConfig | None = None
+    output_router_logits: bool | None = None
+    router_aux_loss_coef: float | None = None
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> DecoderConfig:
+        # Embedding
+        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 not None and self.rope_scaling.rope_type == "yarn":
+            rope_config = YARNRoPEConfig(
+                precision=activation_precision,
+                base=self.rope_theta,
+                max_sequence_length=context_length or self.max_position_embeddings,
+                scaling_factor=self.rope_scaling.factor,
+                original_context_length=self.rope_scaling.original_max_position_embeddings,
+                beta_fast=self.rope_scaling.beta_fast,
+                beta_slow=self.rope_scaling.beta_slow,
+                truncate=self.rope_scaling.truncate,
+            )
+        else:
+            rope_config = YARNRoPEConfig(
+                precision=activation_precision,
+                base=self.rope_theta,
+                max_sequence_length=context_length or self.max_position_embeddings,
+                scaling_factor=1.0,
+                original_context_length=self.max_position_embeddings,
+                beta_fast=32.0,
+                beta_slow=1.0,
+                truncate=True,
+            )
+
+        rmsnorm_config = RMSNormConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.rms_norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.FULL_LAYER,
+        )
+
+        # Linear layers
+        linear_config = FullPrecisionLinearConfig(precision=activation_precision)
+
+        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=True,
+            has_qkv_biases=self.attention_bias,
+            has_out_biases=self.attention_bias,
+        )
+
+        # Experts (MoE) scaffold
+        # Router: linear with bias; Experts: DenseMLP with SiLU(alpha=1.702) and value/gate clipping
+        experts_activation = SiLU(alpha=1.702)
+        experts_config = DenseMLPConfig(
+            linear_config=linear_config,
+            activation=experts_activation,
+            has_up_biases=True,
+            has_down_biases=True,
+            up_clipping=(-self.swiglu_limit + 1.0, self.swiglu_limit + 1.0),
+            gate_clipping=(None, self.swiglu_limit),
+        )
+        moe_config = MixtureOfExpertsConfig(
+            mixture_size=self.num_local_experts,
+            num_experts_per_token=(self.num_experts_per_tok or self.experts_per_token or 1),
+            routing_function=SoftmaxRouting(),
+            router_config=linear_config,
+            router_has_biases=True,
+            expert_config=experts_config,
+        )
+        decoder_layer_config = DecoderLayerConfig(
+            pre_attention_norm_config=rmsnorm_config,
+            attention_config=attention_config,
+            post_attention_norm_config=None,
+            pre_mlp_norm_config=rmsnorm_config,
+            mlp_config=moe_config,
+            post_mlp_norm_config=None,
+        )
+
+        # Per-layer sliding-window
+        if self.layer_types is not None and len(self.layer_types) == self.num_hidden_layers:
+            sliding_window_sizes = tuple(
+                self.sliding_window if layer_type == "sliding_attention" else None for layer_type in self.layer_types
+            )
+        else:
+            # Fallback: apply the same sliding window to all layers if provided
+            sliding_window_sizes = (
+                tuple([self.sliding_window] * self.num_hidden_layers) if self.sliding_window is not None else None
+            )
+
+        head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads
+
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            global_rope_config=rope_config,
+            local_rope_config=None,
+            layer_config=decoder_layer_config,
+            output_norm_config=rmsnorm_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.hidden_size,
+            hidden_dim=self.intermediate_size,
+            num_heads=self.num_attention_heads,
+            num_groups=self.num_key_value_heads,
+            head_dim=head_dim,
+            attention_scale=None,
+            num_layers=self.num_hidden_layers,
+            sliding_window_sizes=sliding_window_sizes,
+            context_length=context_length or self.max_position_embeddings,
+        )

lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -4,19 +4,20 @@ from typing import Literal
 from jaxtyping import DTypeLike
 
 from lalamo.modules import (
-    Activation,
     AttentionConfig,
     DecoderConfig,
     DecoderLayerConfig,
+    DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
     LlamaRoPEConfig,
-    MLPConfig,
     RMSNormConfig,
     TiedEmbeddingConfig,
     UnscaledRoPEConfig,
     UpcastMode,
+    YARNRoPEConfig,
 )
+from lalamo.modules.activations import SiLU
 from lalamo.modules.embedding import UntiedEmbeddingConfig
 from lalamo.quantization import QuantizationMode
 
@@ -34,8 +35,19 @@ class LlamaRopeScalingConfig:
     rope_type: Literal["llama3"]
 
 
+@dataclass(frozen=True)
+class YarnRopeScalingConfig:
+    factor: float
+    beta_fast: float
+    beta_slow: float
+    original_max_position_embeddings: int
+    rope_type: Literal["yarn"]
+    truncate: bool
+
+
 @dataclass(frozen=True)
 class HFLlamaConfig(HuggingFaceConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["LlamaForCausalLM"]]
     attention_bias: bool
     attention_dropout: float
@@ -53,7 +65,7 @@ class HFLlamaConfig(HuggingFaceConfig):
     num_key_value_heads: int
     pretraining_tp: int
     rms_norm_eps: float
-    rope_scaling: LlamaRopeScalingConfig | None
+    rope_scaling: LlamaRopeScalingConfig | YarnRopeScalingConfig | None
     rope_theta: float
     tie_word_embeddings: bool
     transformers_version: str
@@ -72,13 +84,13 @@ class HFLlamaConfig(HuggingFaceConfig):
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         else:
             embedding_config = UntiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         if self.rope_scaling is None:
@@ -87,7 +99,18 @@ class HFLlamaConfig(HuggingFaceConfig):
                 base=self.rope_theta,
                 max_sequence_length=context_length or self.max_position_embeddings,
             )
-        else:
+        elif isinstance(self.rope_scaling, YarnRopeScalingConfig):
+            rope_config = YARNRoPEConfig(
+                precision=activation_precision,
+                base=self.rope_theta,
+                max_sequence_length=context_length or self.max_position_embeddings,
+                scaling_factor=self.rope_scaling.factor,
+                original_context_length=self.rope_scaling.original_max_position_embeddings,
+                beta_fast=self.rope_scaling.beta_fast,
+                beta_slow=self.rope_scaling.beta_slow,
+                truncate=self.rope_scaling.truncate,
+            )
+        elif isinstance(self.rope_scaling, LlamaRopeScalingConfig):
             rope_config = LlamaRoPEConfig(
                 precision=activation_precision,
                 base=self.rope_theta,
@@ -97,6 +120,8 @@ class HFLlamaConfig(HuggingFaceConfig):
                 low_frequency_factor=self.rope_scaling.low_freq_factor,
                 high_frequency_factor=self.rope_scaling.high_freq_factor,
             )
+        else:
+            raise ValueError("Unsupported rope_scaling configuration")
         rmsnorm_config = RMSNormConfig(
             scale_precision=activation_precision,
             accumulation_precision=accumulation_precision,
@@ -121,12 +146,17 @@ class HFLlamaConfig(HuggingFaceConfig):
             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,
         )
-        mlp_config = MLPConfig(
+        mlp_config = DenseMLPConfig(
             linear_config=linear_config,
-            activation=Activation.SILU,
+            activation=SiLU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
         )
         decoder_layer_config = DecoderLayerConfig(
             pre_attention_norm_config=rmsnorm_config,

lalamo/model_import/decoder_configs/huggingface/mistral.py

@@ -4,17 +4,17 @@ from typing import Literal
 from jaxtyping import DTypeLike
 
 from lalamo.modules import (
-    Activation,
     AttentionConfig,
     DecoderConfig,
     DecoderLayerConfig,
+    DenseMLPConfig,
     FullPrecisionLinearConfig,
-    MLPConfig,
     RMSNormConfig,
     TiedEmbeddingConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
 )
+from lalamo.modules.activations import SiLU
 from lalamo.modules.normalization import UpcastMode
 
 from .common import HuggingFaceConfig
@@ -24,6 +24,7 @@ __all__ = ["HFMistralConfig"]
 
 @dataclass(frozen=True)
 class HFMistralConfig(HuggingFaceConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["MistralForCausalLM"]]
     attention_dropout: float
     bos_token_id: int
@@ -57,13 +58,13 @@ class HFMistralConfig(HuggingFaceConfig):
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         else:
             embedding_config = UntiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
 
@@ -91,13 +92,18 @@ class HFMistralConfig(HuggingFaceConfig):
             query_norm_config=None,
             key_norm_config=None,
             logit_soft_cap=None,
+            has_sinks=False,
             has_qkv_biases=False,
             has_out_biases=False,
         )
 
-        mlp_config = MLPConfig(
+        mlp_config = DenseMLPConfig(
             linear_config=linear_config,
-            activation=Activation.SILU,
+            activation=SiLU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
         )
 
         decoder_layer_config = DecoderLayerConfig(

lalamo/model_import/decoder_configs/huggingface/qwen2.py

@@ -4,19 +4,19 @@ from typing import Literal
 from jaxtyping import DTypeLike
 
 from lalamo.modules import (
-    Activation,
     AttentionConfig,
     DecoderConfig,
     DecoderLayerConfig,
+    DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
-    MLPConfig,
     RMSNormConfig,
     TiedEmbeddingConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
 )
+from lalamo.modules.activations import SiLU
 from lalamo.quantization import QuantizationMode
 
 from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
@@ -26,6 +26,7 @@ __all__ = ["HFQwen2Config"]
 
 @dataclass(frozen=True)
 class HFQwen2Config(HuggingFaceConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["Qwen2ForCausalLM"]]
     attention_dropout: float
     bos_token_id: int | list[int]
@@ -72,13 +73,13 @@ class HFQwen2Config(HuggingFaceConfig):
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         else:
             embedding_config = UntiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         rope_config = UnscaledRoPEConfig(
@@ -110,12 +111,17 @@ class HFQwen2Config(HuggingFaceConfig):
             query_norm_config=None,
             key_norm_config=None,
             logit_soft_cap=None,
+            has_sinks=False,
             has_qkv_biases=True,
             has_out_biases=False,
         )
-        mlp_config = MLPConfig(
+        mlp_config = DenseMLPConfig(
             linear_config=linear_config,
-            activation=Activation.SILU,
+            activation=SiLU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
         )
         decoder_layer_config = DecoderLayerConfig(
             pre_attention_norm_config=rmsnorm_config,

lalamo/model_import/decoder_configs/huggingface/qwen3.py

@@ -4,19 +4,19 @@ from typing import Literal
 from jaxtyping import DTypeLike
 
 from lalamo.modules import (
-    Activation,
     AttentionConfig,
     DecoderConfig,
     DecoderLayerConfig,
+    DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
-    MLPConfig,
     RMSNormConfig,
     TiedEmbeddingConfig,
     UnscaledRoPEConfig,
     UntiedEmbeddingConfig,
     UpcastMode,
 )
+from lalamo.modules.activations import SiLU
 from lalamo.quantization import QuantizationMode
 
 from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
@@ -26,6 +26,7 @@ __all__ = ["HFQwen3Config"]
 
 @dataclass(frozen=True)
 class HFQwen3Config(HuggingFaceConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
     attention_bias: bool
     hidden_act: Literal["silu"]
     hidden_size: int
@@ -70,13 +71,13 @@ class HFQwen3Config(HuggingFaceConfig):
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         else:
             embedding_config = UntiedEmbeddingConfig(
                 input_scale=None,
-                logits_soft_cap=None,
+                logit_soft_cap=None,
                 precision=activation_precision,
             )
         rope_config = UnscaledRoPEConfig(
@@ -108,12 +109,17 @@ class HFQwen3Config(HuggingFaceConfig):
             query_norm_config=rmsnorm_config,
             key_norm_config=rmsnorm_config,
             logit_soft_cap=None,
+            has_sinks=False,
             has_qkv_biases=self.attention_bias,
             has_out_biases=self.attention_bias,
         )
-        mlp_config = MLPConfig(
+        mlp_config = DenseMLPConfig(
             linear_config=linear_config,
-            activation=Activation.SILU,
+            activation=SiLU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
         )
         decoder_layer_config = DecoderLayerConfig(
             pre_attention_norm_config=rmsnorm_config,

lalamo/model_import/huggingface_tokenizer_config.py

@@ -72,10 +72,7 @@ class HFTokenizerConfig:
     def added_tokens(self) -> list[AddedToken]:
         if self.added_tokens_decoder is None:
             return []
-        return [
-            AddedToken(content=token.content, single_word=token.single_word, normalized=token.normalized)
-            for token in self.added_tokens_decoder.values()
-        ]
+        return [token.to_added_token() for token in self.added_tokens_decoder.values()]
 
     @classmethod
     def from_json(cls, json_path: Path | str) -> "HFTokenizerConfig":

lalamo/model_import/loaders/executorch.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable, Iterator
+from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import dataclass, replace
 
 import jax.numpy as jnp
@@ -6,7 +6,7 @@ from einops import rearrange
 from jaxtyping import Array, Float, Int
 
 from lalamo.common import ParameterPath
-from lalamo.modules import MLP, Attention, Decoder, DecoderLayer, QLoRALinear, QuantizedTiedEmbedding, RMSNorm
+from lalamo.modules import Attention, Decoder, DecoderLayer, DenseMLP, QLoRALinear, QuantizedTiedEmbedding, RMSNorm
 
 from .common import load_parameters
 
@@ -43,7 +43,7 @@ def params_selector(module: QLoRALinear) -> tuple:
 
 
 def get_qlora_linear_params(
-    weights_dict: dict[str, Array],
+    weights_dict: Mapping[str, Array],
     path: ParameterPath,
     weights_dtype: jnp.dtype,
 ) -> QLoRALinearParams:
@@ -76,7 +76,7 @@ def load_linear(module: QLoRALinear, weights_dict: dict[str, Array], path: Param
     return load_parameters(params_selector, module, params)
 
 
-def load_mlp(module: MLP, weights_dict: dict[str, Array], path: ParameterPath) -> MLP:
+def load_mlp(module: DenseMLP, weights_dict: Mapping[str, Array], path: ParameterPath) -> DenseMLP:
     if not isinstance(module.up_projection, QLoRALinear):
         raise TypeError(f"Expected up_projection to be QLoRALinear, got {type(module.up_projection)}")
     if not isinstance(module.down_projection, QLoRALinear):
@@ -95,7 +95,7 @@ def load_mlp(module: MLP, weights_dict: dict[str, Array], path: ParameterPath) -
     )
 
 
-def load_rmsnorm(module: RMSNorm, weights_dict: dict[str, Array], path: ParameterPath) -> RMSNorm:
+def load_rmsnorm(module: RMSNorm, weights_dict: Mapping[str, Array], path: ParameterPath) -> RMSNorm:
     return load_parameters(lambda m: (m.scales,), module, (weights_dict[path / "weight"],))
 
 
@@ -131,7 +131,7 @@ def permute_qk_params(
 
 def load_attention(
     module: Attention,
-    weights_dict: dict[str, Array],
+    weights_dict: Mapping[str, Array],
     path: ParameterPath,
 ) -> Attention:
     if not isinstance(module.qkv_projection, QLoRALinear):
@@ -177,7 +177,7 @@ def load_attention(
 
 def load_decoder_layer(
     module: DecoderLayer,
-    weights_dict: dict[str, Array],
+    weights_dict: Mapping[str, Array],
     path: ParameterPath,
 ) -> DecoderLayer:
     if module.post_attention_norm is not None:
@@ -187,6 +187,7 @@ def load_decoder_layer(
     attention_norm = load_rmsnorm(module.pre_attention_norm, weights_dict, path / "attention_norm")
     attention = load_attention(module.attention, weights_dict, path / "attention")
     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")
     return load_parameters(
         lambda m: (m.pre_attention_norm, m.attention, m.pre_mlp_norm, m.mlp),
@@ -197,7 +198,7 @@ def load_decoder_layer(
 
 def load_embedding(
     module: QuantizedTiedEmbedding,
-    weights_dict: dict[str, Array],
+    weights_dict: Mapping[str, Array],
     path: ParameterPath,
 ) -> QuantizedTiedEmbedding:
     weights = weights_dict[path / "weight"].astype(module.weights.dtype)
@@ -206,7 +207,7 @@ def load_embedding(
     return load_parameters(lambda m: (m.weights, m.scales), module, (weights, scales))
 
 
-def load_executorch(module: Decoder, weights_dict: dict[str, Array]) -> Decoder:
+def load_executorch(module: Decoder, weights_dict: Mapping[str, Array]) -> Decoder:
     root_path = ParameterPath()
     if not isinstance(module.embedding, QuantizedTiedEmbedding):
         raise TypeError(f"Expected embedding to be QuantizedTiedEmbedding, got {type(module.embedding)}")