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

lalamo/model_import/model_specs/llamba.py

@@ -0,0 +1,40 @@
+from lalamo.model_import.decoder_configs import HFLlambaConfig
+from lalamo.quantization import QuantizationMode
+
+from .common import ConfigMap, FileSpec, ModelSpec
+
+__all__ = ["LLAMBA_MODELS"]
+
+LLAMBA_MODELS = [
+    ModelSpec(
+        vendor="Cartesia",
+        family="Llamba",
+        name="Llamba-1B",
+        size="1B",
+        quantization=None,
+        repo="cartesia-ai/Llamba-1B",
+        config_type=HFLlambaConfig,
+        configs=ConfigMap(
+            tokenizer=FileSpec("tokenizer.json", "meta-llama/Llama-3.2-1B-Instruct"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "meta-llama/Llama-3.2-1B-Instruct"),
+            generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-1B-Instruct"),
+        ),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Cartesia",
+        family="Llamba",
+        name="Llamba-1B-4bit-mlx",
+        size="1B",
+        quantization=QuantizationMode.UINT4,
+        repo="cartesia-ai/Llamba-1B-4bit-mlx",
+        config_type=HFLlambaConfig,
+        configs=ConfigMap(
+            model_config=FileSpec("config.json", "cartesia-ai/Llamba-1B"),
+            tokenizer=FileSpec("tokenizer.json", "meta-llama/Llama-3.2-1B-Instruct"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "meta-llama/Llama-3.2-1B-Instruct"),
+            generation_config=FileSpec("generation_config.json", "meta-llama/Llama-3.2-1B-Instruct"),
+        ),
+        use_cases=tuple(),
+    ),
+]

lalamo/model_import/model_specs/qwen.py

@@ -1,7 +1,7 @@
 from lalamo.model_import.decoder_configs import HFQwen2Config, HFQwen3Config
 from lalamo.quantization import QuantizationMode
 
-from .common import ModelSpec, UseCase, WeightsType
+from .common import ConfigMap, FileSpec, ModelSpec, UseCase, WeightsType
 
 __all__ = ["QWEN_MODELS"]
 
@@ -148,6 +148,20 @@ QWEN3 = [
         repo="Qwen/Qwen3-0.6B",
         config_type=HFQwen3Config,
     ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-0.6B-MLX-4bit",
+        size="0.6B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-0.6B-MLX-4bit",
+        config_type=HFQwen3Config,
+        configs=ConfigMap(
+            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-0.6B"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-0.6B"),
+            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-0.6B"),
+        ),
+    ),
     ModelSpec(
         vendor="Alibaba",
         family="Qwen3",
@@ -177,6 +191,20 @@ QWEN3 = [
         repo="Qwen/Qwen3-4B-AWQ",
         config_type=HFQwen3Config,
     ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-4B-MLX-4bit",
+        size="4B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-4B-MLX-4bit",
+        config_type=HFQwen3Config,
+        configs=ConfigMap(
+            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-4B"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-4B"),
+            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-4B"),
+        ),
+    ),
     ModelSpec(
         vendor="Alibaba",
         family="Qwen3",

lalamo/modules/__init__.py

@@ -1,6 +1,5 @@
-from .activations import GELU, Activation, SiLU
-from .attention import Attention, AttentionConfig
-from .common import AttentionType, ForwardPassMode, LalamoModule, config_converter
+from .activations import GELU, Activation, Identity, SiLU
+from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector, config_converter
 from .decoder import Decoder, DecoderActivationTrace, DecoderConfig, DecoderForwardPassConfig, DecoderResult
 from .decoder_layer import (
     DecoderLayer,
@@ -12,6 +11,10 @@ from .decoder_layer import (
 from .embedding import (
     EmbeddingBase,
     EmbeddingConfig,
+    MLXQuantizedTiedEmbedding,
+    MLXQuantizedTiedEmbeddingConfig,
+    MLXSemiQuantizedUntiedEmbedding,
+    MLXSemiQuantizedUntiedEmbeddingConfig,
     QuantizedTiedEmbedding,
     QuantizedTiedEmbeddingConfig,
     TiedEmbedding,
@@ -19,7 +22,6 @@ from .embedding import (
     UntiedEmbedding,
     UntiedEmbeddingConfig,
 )
-from .kv_cache import DynamicKVCacheLayer, KVCache, KVCacheLayer, StaticKVCacheLayer
 from .linear import (
     FullPrecisionLinear,
     FullPrecisionLinearConfig,
@@ -27,6 +29,8 @@ from .linear import (
     GroupQuantizedLinearConfig,
     LinearBase,
     LinearConfig,
+    MLXQuantizedLinear,
+    MLXQuantizedLinearConfig,
     QLoRALinear,
     QLoRALinearConfig,
 )
@@ -51,13 +55,24 @@ from .rope import (
     UnscaledRoPEConfig,
     YARNRoPEConfig,
 )
+from .token_mixers import (
+    Attention,
+    AttentionConfig,
+    DynamicKVCacheLayer,
+    KVCacheLayer,
+    Mamba2,
+    Mamba2Config,
+    SeparableCausalConv,
+    SeparableCausalConvConfig,
+    State,
+    StaticKVCacheLayer,
+)
 
 __all__ = [
     "GELU",
     "Activation",
     "Attention",
     "AttentionConfig",
-    "AttentionType",
     "Decoder",
     "DecoderActivationTrace",
     "DecoderConfig",
@@ -78,7 +93,7 @@ __all__ = [
     "FullPrecisionLinearConfig",
     "GroupQuantizedLinear",
     "GroupQuantizedLinearConfig",
-    "KVCache",
+    "Identity",
     "KVCacheLayer",
     "LalamoModule",
     "LinearBase",
@@ -88,8 +103,17 @@ __all__ = [
     "MLPBase",
     "MLPConfig",
     "MLPForwardPassConfig",
+    "MLXQuantizedLinear",
+    "MLXQuantizedLinearConfig",
+    "MLXQuantizedTiedEmbedding",
+    "MLXQuantizedTiedEmbeddingConfig",
+    "MLXSemiQuantizedUntiedEmbedding",
+    "MLXSemiQuantizedUntiedEmbeddingConfig",
+    "Mamba2",
+    "Mamba2Config",
     "MixtureOfExperts",
     "MixtureOfExpertsConfig",
+    "PositionalEmbeddingSelector",
     "PositionalEmbeddings",
     "QLoRALinear",
     "QLoRALinearConfig",
@@ -100,8 +124,11 @@ __all__ = [
     "RoPE",
     "RoPEConfig",
     "RoutingFunction",
+    "SeparableCausalConv",
+    "SeparableCausalConvConfig",
     "SiLU",
     "SoftmaxRouting",
+    "State",
     "StaticKVCacheLayer",
     "TiedEmbedding",
     "TiedEmbeddingConfig",

lalamo/modules/activations.py

@@ -10,6 +10,7 @@ from lalamo.modules.common import register_config_union
 __all__ = [
     "GELU",
     "Activation",
+    "Identity",
     "SiLU",
 ]
 
@@ -34,7 +35,13 @@ class GELU(ActivationBase):
         return jax.nn.gelu(x)
 
 
-Activation = SiLU | GELU
+@dataclass(frozen=True)
+class Identity(ActivationBase):
+    def __call__(self, x: Float[Array, "*dims"]) -> Float[Array, "*dims"]:
+        return x
+
+
+Activation = SiLU | GELU | Identity
 
 
-register_config_union(Activation)
+register_config_union(Activation)  # type: ignore (pyright bug)

lalamo/modules/common.py

@@ -2,7 +2,7 @@ from abc import abstractmethod
 from dataclasses import dataclass
 from enum import Enum
 from types import UnionType
-from typing import Self
+from typing import Any, Self
 
 import equinox as eqx
 from cattrs import Converter
@@ -12,18 +12,19 @@ from jaxtyping import Array, DTypeLike
 from lalamo.common import ParameterTree
 
 __all__ = [
-    "AttentionType",
     "DummyUnionMember",
     "ForwardPassMode",
     "LalamoModule",
+    "PositionalEmbeddingSelector",
     "config_converter",
     "register_config_union",
 ]
 
 
-class AttentionType(Enum):
+class PositionalEmbeddingSelector(Enum):
     GLOBAL = "global"
-    SLIDING_WINDOW = "sliding_window"
+    LOCAL = "sliding_window"
+    NONE = "none"
 
 
 class ForwardPassMode(Enum):
@@ -128,4 +129,8 @@ def register_config_union(union_type: UnionType) -> None:
 
 @dataclass
 class DummyUnionMember:
-    pass
+    def __getattribute__(self, name: str, /) -> Any:  # noqa: ANN401
+        raise NotImplementedError
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        raise NotImplementedError

lalamo/modules/decoder.py

@@ -8,14 +8,14 @@ from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
 from lalamo.common import ParameterTree
-from lalamo.modules.utils import vmap_twice
 
-from .common import AttentionType, ForwardPassMode, LalamoModule
+from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
 from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerForwardPassConfig, DecoderLayerResult
 from .embedding import EmbeddingBase, EmbeddingConfig
-from .kv_cache import KVCache
 from .normalization import RMSNorm, RMSNormConfig
 from .rope import PositionalEmbeddings, RoPE, RoPEConfig
+from .token_mixers import AttentionConfig, State
+from .utils import vmap_twice
 
 __all__ = [
     "Decoder",
@@ -32,42 +32,42 @@ type DecoderForwardPassConfig = DecoderLayerForwardPassConfig
 class DecoderActivationTrace(eqx.Module):
     token_ids: Int[Array, "batch suffix_tokens"]
     token_positions: Int[Array, "batch suffix_tokens"]
-    kv_cache: KVCache | None
+    state: State | None
 
-    local_positional_embeddings: PositionalEmbeddings
-    global_positional_embeddings: PositionalEmbeddings
+    local_positional_embeddings: PositionalEmbeddings | None
+    global_positional_embeddings: PositionalEmbeddings | None
 
     layer_results: tuple[DecoderLayerResult, ...]
 
     output_norm: Float[Array, "batch suffix_tokens channels"]
 
     def export(self) -> ParameterTree:
-        result = dict(
+        result: dict[str, ParameterTree | Array] = dict(
             token_ids=self.token_ids,
             token_positions=self.token_positions,
-            local_positional_embeddings=self.local_positional_embeddings.export(),
-            global_positional_embeddings=self.global_positional_embeddings.export(),
             layer_results=[layer_result.export() for layer_result in self.layer_results],
             output_norm=self.output_norm,
         )
-        if self.kv_cache is not None:
-            result["kv_cache"] = [kv_cache_layer_slice.export() for kv_cache_layer_slice in self.kv_cache]
+        if self.local_positional_embeddings is not None:
+            result["local_positional_embeddings"] = self.local_positional_embeddings.export()
+        if self.global_positional_embeddings is not None:
+            result["global_positional_embeddings"] = self.global_positional_embeddings.export()
+        if self.state is not None:
+            result["state"] = [state_layer.export() for state_layer in self.state]
         return result
 
 
 class DecoderResult(eqx.Module):
     logits: Float[Array, "batch suffix_tokens channels"]
-    updated_kv_cache: KVCache | None = None
+    updated_state: State | None = None
     activation_trace: DecoderActivationTrace | None = None
 
     def export(self) -> ParameterTree:
         result: dict[str, ParameterTree | Array] = dict(
             logits=self.logits,
         )
-        if self.updated_kv_cache is not None:
-            result["updated_kv_cache"] = [
-                kv_cache_layer_slice.export() for kv_cache_layer_slice in self.updated_kv_cache
-            ]
+        if self.updated_state is not None:
+            result["updated_state"] = [state_layer.export() for state_layer in self.updated_state]
         if self.activation_trace is not None:
             result["activation_trace"] = self.activation_trace.export()
         return result
@@ -76,33 +76,16 @@ class DecoderResult(eqx.Module):
 @dataclass(frozen=True)
 class DecoderConfig:
     embedding_config: EmbeddingConfig
-    global_rope_config: RoPEConfig
+    global_rope_config: RoPEConfig | None
     local_rope_config: RoPEConfig | None
-    layer_config: DecoderLayerConfig
+    layer_configs: tuple[DecoderLayerConfig, ...]
     output_norm_config: RMSNormConfig
 
     vocab_size: int
     model_dim: int
     hidden_dim: int
-    num_heads: int
-    num_groups: int
-    head_dim: int
-    attention_scale: float | None
-    num_layers: int
-    sliding_window_sizes: tuple[int | None, ...] | None
     context_length: int
 
-    def __post_init__(self) -> None:
-        if self.local_rope_config is not None and self.sliding_window_sizes is None:
-            raise ValueError("Sliding window sizes must be provided when using local RoPE")
-        if self.sliding_window_sizes is None:
-            return
-        if len(self.sliding_window_sizes) != self.num_layers:
-            raise ValueError(
-                f"Number of sliding window sizes {len(self.sliding_window_sizes)} does not match"
-                f" the number of layers {self.num_layers}",
-            )
-
     def random_init(
         self,
         *,
@@ -114,40 +97,38 @@ class DecoderConfig:
             model_dim=self.model_dim,
             key=embedding_key,
         )
-        global_rope = self.global_rope_config.init(
-            head_dim=self.head_dim,
-            num_timesteps=self.context_length,
-        )
+
+        first_layer_config, *_ = self.layer_configs
+
+        if self.global_rope_config:
+            global_rope = self.global_rope_config.init(
+                head_dim=first_layer_config.rope_dim,
+                num_timesteps=self.context_length,
+            )
+        else:
+            global_rope = None
 
         if self.local_rope_config:
-            assert self.sliding_window_sizes is not None
             max_sliding_window_size = max(
-                window_size for window_size in self.sliding_window_sizes if window_size is not None
+                layer_config.mixer_config.sliding_window_size or 0
+                for layer_config in self.layer_configs
+                if isinstance(layer_config.mixer_config, AttentionConfig)
             )
             local_rope = self.local_rope_config.init(
-                head_dim=self.head_dim,
+                head_dim=first_layer_config.rope_dim,
                 num_timesteps=max(max_sliding_window_size, self.context_length),
             )
         else:
             local_rope = None
 
-        if self.sliding_window_sizes is None:
-            sliding_window_sizes = [None] * self.num_layers
-        else:
-            sliding_window_sizes = self.sliding_window_sizes
-        layers_keys = jax.random.split(layers_key, self.num_layers)
+        layers_keys = jax.random.split(layers_key, len(self.layer_configs))
         layers = tuple(
-            self.layer_config.random_init(
+            layer_config.random_init(
                 model_dim=self.model_dim,
                 hidden_dim=self.hidden_dim,
-                num_heads=self.num_heads,
-                num_groups=self.num_groups,
-                head_dim=self.head_dim,
-                attention_scale=self.attention_scale,
-                sliding_window_size=sliding_window_size,
                 key=key,
             )
-            for sliding_window_size, key in zip(sliding_window_sizes, layers_keys, strict=True)
+            for layer_config, key in zip(self.layer_configs, layers_keys, strict=False)
         )
         output_norm = self.output_norm_config.init(self.model_dim)
         return Decoder(
@@ -166,34 +147,35 @@ class DecoderConfig:
             vocab_size=self.vocab_size,
             model_dim=self.model_dim,
         )
-        global_rope = self.global_rope_config.init(
-            head_dim=self.head_dim,
-            num_timesteps=self.context_length,
-        )
 
-        if self.local_rope_config:
-            local_rope = self.local_rope_config.init(
-                head_dim=self.head_dim,
+        first_layer_config, *_ = self.layer_configs
+
+        if self.global_rope_config:
+            global_rope = self.global_rope_config.init(
+                head_dim=first_layer_config.rope_dim,
                 num_timesteps=self.context_length,
             )
         else:
-            local_rope = None
+            global_rope = None
 
-        if self.sliding_window_sizes is None:
-            sliding_window_sizes = [None] * self.num_layers
+        if self.local_rope_config:
+            max_sliding_window_size = max(
+                layer_config.mixer_config.sliding_window_size or 0
+                for layer_config in self.layer_configs
+                if isinstance(layer_config.mixer_config, AttentionConfig)
+            )
+            local_rope = self.local_rope_config.init(
+                head_dim=first_layer_config.rope_dim,
+                num_timesteps=max(max_sliding_window_size, self.context_length),
+            )
         else:
-            sliding_window_sizes = self.sliding_window_sizes
+            local_rope = None
         layers = tuple(
-            self.layer_config.empty(
+            layer_config.empty(
                 model_dim=self.model_dim,
                 hidden_dim=self.hidden_dim,
-                num_heads=self.num_heads,
-                num_groups=self.num_groups,
-                head_dim=self.head_dim,
-                attention_scale=self.attention_scale,
-                sliding_window_size=sliding_window_size,
             )
-            for sliding_window_size in sliding_window_sizes
+            for layer_config in self.layer_configs
         )
         output_norm = self.output_norm_config.empty(self.model_dim)
         return Decoder(
@@ -208,7 +190,7 @@ class DecoderConfig:
 
 class Decoder(LalamoModule[DecoderConfig]):
     embedding: EmbeddingBase
-    global_rope: RoPE
+    global_rope: RoPE | None
     local_rope: RoPE | None
     layers: tuple[DecoderLayer, ...]
     output_norm: RMSNorm
@@ -218,12 +200,12 @@ class Decoder(LalamoModule[DecoderConfig]):
         return self.embedding.activation_precision
 
     @eqx.filter_jit
-    def __call__(
+    def __call__(  # noqa: PLR0912
         self,
         token_ids: Int[Array, "batch suffix_tokens"],
         token_positions: Int[Array, "batch suffix_tokens"],
-        kv_cache: KVCache | None = None,
-        return_updated_kv_cache: bool = False,
+        state: State | None = None,
+        return_updated_state: bool = False,
         return_activation_trace: bool = False,
         lengths_without_padding: Int[Array, " batch"] | None = None,
         forward_pass_mode: ForwardPassMode = ForwardPassMode.MULTI_TOKEN,
@@ -239,28 +221,35 @@ class Decoder(LalamoModule[DecoderConfig]):
                 f" got {token_positions.shape}",
             )
 
-        maybe_kv_cache = kv_cache or ([None] * len(self.layers))
+        maybe_state = state or ([None] * len(self.layers))
         inner_features = vmap(self.embedding.embed)(token_ids)
 
-        global_positional_embeddings = vmap(self.global_rope)(token_positions)
+        if self.global_rope is not None:
+            global_positional_embeddings = vmap(self.global_rope)(token_positions)
+        else:
+            global_positional_embeddings = None
+
         if self.local_rope is not None:
             local_positional_embeddings = vmap(self.local_rope)(token_positions)
         else:
             local_positional_embeddings = global_positional_embeddings
 
-        updated_kv_cache_layers = []
+        updated_state_layers = []
         layer_results = []
-        for layer, kv_cache_slice in zip(self.layers, maybe_kv_cache, strict=True):
-            if layer.attention_type == AttentionType.SLIDING_WINDOW:
-                positional_embeddings_to_use = local_positional_embeddings
-            else:
-                positional_embeddings_to_use = global_positional_embeddings
+        for layer, state_layer in zip(self.layers, maybe_state, strict=True):
+            match layer.positional_embedding_selector:
+                case PositionalEmbeddingSelector.LOCAL:
+                    positional_embeddings_to_use = local_positional_embeddings
+                case PositionalEmbeddingSelector.GLOBAL:
+                    positional_embeddings_to_use = global_positional_embeddings
+                case PositionalEmbeddingSelector.NONE:
+                    positional_embeddings_to_use = None
 
             layer_result = layer(
                 inner_features,
                 positional_embeddings_to_use,
-                kv_cache=kv_cache_slice,
-                return_updated_kv_cache=return_updated_kv_cache,
+                state=state_layer,
+                return_updated_state=return_updated_state,
                 return_activation_trace=return_activation_trace,
                 lengths_without_padding=lengths_without_padding,
                 forward_pass_mode=forward_pass_mode,
@@ -268,7 +257,7 @@ class Decoder(LalamoModule[DecoderConfig]):
             )
             inner_features = layer_result.outputs
             layer_results.append(layer_result)
-            updated_kv_cache_layers.append(layer_result.updated_kv_cache)
+            updated_state_layers.append(layer_result.updated_state)
 
         normalized_outputs = vmap_twice(self.output_norm)(inner_features)
         logits = vmap_twice(self.embedding.readout)(normalized_outputs)
@@ -277,7 +266,7 @@ class Decoder(LalamoModule[DecoderConfig]):
             activation_trace = DecoderActivationTrace(
                 token_ids=token_ids,
                 token_positions=token_positions,
-                kv_cache=kv_cache,
+                state=state,
                 global_positional_embeddings=global_positional_embeddings,
                 local_positional_embeddings=local_positional_embeddings,
                 layer_results=tuple(layer_results),
@@ -286,27 +275,28 @@ class Decoder(LalamoModule[DecoderConfig]):
         else:
             activation_trace = None
 
-        if return_updated_kv_cache:
-            updated_kv_cache = KVCache(updated_kv_cache_layers)
+        if return_updated_state:
+            updated_state = State(updated_state_layers)
         else:
-            updated_kv_cache = None
+            updated_state = None
 
         return DecoderResult(
             logits=logits,
-            updated_kv_cache=updated_kv_cache,
+            updated_state=updated_state,
             activation_trace=activation_trace,
         )
 
-    def init_static_kv_cache(self, batch_size: int, capacity: int) -> KVCache:
-        return KVCache(layer.init_static_kv_cache(batch_size, capacity) for layer in self.layers)
+    def init_static_state(self, batch_size: int, capacity: int) -> State:
+        return State(layer.init_static_state(batch_size, capacity) for layer in self.layers)
 
     def export_weights(self) -> ParameterTree:
         result = dict(
             embedding=self.embedding.export_weights(),
-            global_rope=self.global_rope.export_weights(),
             layers=[layer.export_weights() for layer in self.layers],
             output_norm=self.output_norm.export_weights(),
         )
+        if self.global_rope:
+            result["global_rope"] = self.global_rope.export_weights()
         if self.local_rope:
             result["local_rope"] = self.local_rope.export_weights()
         return result
@@ -317,15 +307,21 @@ class Decoder(LalamoModule[DecoderConfig]):
     ) -> Self:
         assert isinstance(weights, Mapping)
         assert isinstance(weights["embedding"], Mapping)
-        assert isinstance(weights["global_rope"], Mapping)
         assert isinstance(weights["layers"], Sequence)
         assert isinstance(weights["output_norm"], Mapping)
+
         if self.local_rope:
             assert isinstance(weights["local_rope"], Mapping)
             local_rope = self.local_rope.import_weights(weights["local_rope"])
         else:
             local_rope = None
 
+        if self.global_rope:
+            assert isinstance(weights["global_rope"], Mapping)
+            global_rope = self.global_rope.import_weights(weights["global_rope"])
+        else:
+            global_rope = None
+
         layers = []
         for layer, layer_weights in zip(self.layers, weights["layers"], strict=True):
             assert isinstance(layer_weights, Mapping)
@@ -333,7 +329,7 @@ class Decoder(LalamoModule[DecoderConfig]):
         return replace(
             self,
             embedding=self.embedding.import_weights(weights["embedding"]),
-            global_rope=self.global_rope.import_weights(weights["global_rope"]),
+            global_rope=global_rope,
             layers=tuple(layers),
             output_norm=self.output_norm.import_weights(weights["output_norm"]),
             local_rope=local_rope,