lalamo 0.2.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

lalamo/modules/decoder.py

@@ -1,11 +1,13 @@
-from dataclasses import dataclass
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, replace
+from typing import Self
 
 import equinox as eqx
 import jax
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree
 
 from .common import AttentionType, LalamoModule, WeightLayout
 from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerResult
@@ -34,8 +36,8 @@ class DecoderActivationTrace(eqx.Module):
 
     output_norm: Float[Array, "suffix_tokens channels"]
 
-    def export(self) -> ParameterDict:
-        result = ParameterDict(
+    def export(self) -> ParameterTree:
+        result = dict(
             token_ids=self.token_ids,
             token_positions=self.token_positions,
             local_positional_embeddings=self.local_positional_embeddings.export(),
@@ -53,8 +55,8 @@ class DecoderResult(eqx.Module):
     updated_kv_cache: KVCache | None = None
     activation_trace: DecoderActivationTrace | None = None
 
-    def export(self) -> ParameterDict:
-        result = ParameterDict(
+    def export(self) -> ParameterTree:
+        result: dict[str, ParameterTree | Array] = dict(
             logits=self.logits,
         )
         if self.updated_kv_cache is not None:
@@ -152,6 +154,56 @@ class DecoderConfig:
             output_norm=output_norm,
         )
 
+    def empty(
+        self,
+    ) -> "Decoder":
+        embedding = self.embedding_config.empty(
+            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:
+            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
+            )
+            local_rope = self.local_rope_config.init(
+                head_dim=self.head_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 = tuple(
+            self.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
+        )
+        output_norm = self.output_norm_config.empty(self.model_dim)
+        return Decoder(
+            self,
+            embedding=embedding,
+            global_rope=global_rope,
+            local_rope=local_rope,
+            layers=layers,
+            output_norm=output_norm,
+        )
+
 
 class Decoder(LalamoModule[DecoderConfig]):
     embedding: EmbeddingBase
@@ -164,6 +216,7 @@ class Decoder(LalamoModule[DecoderConfig]):
     def activation_precision(self) -> DTypeLike:
         return self.embedding.activation_precision
 
+    @eqx.filter_jit
     def __call__(
         self,
         token_ids: Int[Array, " suffix_tokens"],
@@ -232,8 +285,8 @@ class Decoder(LalamoModule[DecoderConfig]):
     def init_static_kv_cache(self, capacity: int) -> KVCache:
         return KVCache(layer.init_static_kv_cache(capacity) for layer in self.layers)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        result = ParameterDict(
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        result = dict(
             embedding=self.embedding.export_weights(weight_layout),
             global_rope=self.global_rope.export_weights(weight_layout),
             layers=[layer.export_weights(weight_layout) for layer in self.layers],
@@ -242,3 +295,32 @@ class Decoder(LalamoModule[DecoderConfig]):
         if self.local_rope:
             result["local_rope"] = self.local_rope.export_weights(weight_layout)
         return result
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> 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"], weight_layout)
+        else:
+            local_rope = None
+
+        layers = []
+        for layer, layer_weights in zip(self.layers, weights["layers"], strict=True):
+            assert isinstance(layer_weights, Mapping)
+            layers.append(layer.import_weights(layer_weights, weight_layout))
+        return replace(
+            self,
+            embedding=self.embedding.import_weights(weights["embedding"], weight_layout),
+            global_rope=self.global_rope.import_weights(weights["global_rope"], weight_layout),
+            layers=tuple(layers),
+            output_norm=self.output_norm.import_weights(weights["output_norm"], weight_layout),
+            local_rope=local_rope,
+        )

lalamo/modules/decoder_layer.py

@@ -1,11 +1,13 @@
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from typing import Self
 
 import equinox as eqx
 import jax
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree
 
 from .attention import Attention, AttentionConfig
 from .common import AttentionType, LalamoModule, WeightLayout
@@ -35,8 +37,8 @@ class DecoderLayerActivationTrace(eqx.Module):
     mlp: Float[Array, "suffix_tokens channels"]
     post_mlp_norm: Float[Array, "suffix_tokens channels"] | None
 
-    def export(self) -> ParameterDict:
-        result = ParameterDict(
+    def export(self) -> ParameterTree:
+        result = dict(
             inputs=self.inputs,
             positional_embeddings=self.positional_embeddings.export(),
             mlp_inputs=self.mlp_inputs,
@@ -59,8 +61,8 @@ class DecoderLayerResult(eqx.Module):
     updated_kv_cache: KVCacheLayer | None
     activation_trace: DecoderLayerActivationTrace | None
 
-    def export(self) -> ParameterDict:
-        result = ParameterDict(
+    def export(self) -> ParameterTree:
+        result: dict[str, ParameterTree | Array] = dict(
             outputs=self.outputs,
         )
         if self.updated_kv_cache is not None:
@@ -123,6 +125,46 @@ class DecoderLayerConfig:
             post_mlp_norm=post_mlp_norm,
         )
 
+    def empty(
+        self,
+        model_dim: int,
+        hidden_dim: int,
+        num_heads: int,
+        num_groups: int,
+        head_dim: int,
+        attention_scale: float | None,
+        sliding_window_size: int | None,
+    ) -> "DecoderLayer":
+        pre_attention_norm = self.pre_attention_norm_config.empty(model_dim)
+        attention = self.attention_config.empty(
+            model_dim=model_dim,
+            num_heads=num_heads,
+            num_groups=num_groups,
+            head_dim=head_dim,
+            is_causal=True,
+            scale=attention_scale,
+            sliding_window_size=sliding_window_size,
+        )
+        if self.post_attention_norm_config is not None:
+            post_attention_norm = self.post_attention_norm_config.empty(model_dim)
+        else:
+            post_attention_norm = None
+        pre_mlp_norm = self.pre_mlp_norm_config.empty(model_dim)
+        mlp = self.mlp_config.empty(model_dim, hidden_dim)
+        if self.post_mlp_norm_config is not None:
+            post_mlp_norm = self.post_mlp_norm_config.empty(model_dim)
+        else:
+            post_mlp_norm = None
+        return DecoderLayer(
+            config=self,
+            pre_attention_norm=pre_attention_norm,
+            attention=attention,
+            post_attention_norm=post_attention_norm,
+            pre_mlp_norm=pre_mlp_norm,
+            mlp=mlp,
+            post_mlp_norm=post_mlp_norm,
+        )
+
 
 class DecoderLayer(LalamoModule[DecoderLayerConfig]):
     pre_attention_norm: RMSNorm
@@ -168,6 +210,7 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
                 f" the up projection dim {self.mlp.hidden_dim}",
             )
 
+    @eqx.filter_jit
     def __call__(
         self,
         inputs: Float[Array, "suffix_tokens channels"],
@@ -226,8 +269,8 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
     def init_static_kv_cache(self, capacity: int) -> StaticKVCacheLayer:
         return self.attention.init_static_kv_cache(capacity)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        result = ParameterDict(
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        result = dict(
             pre_attention_norm=self.pre_attention_norm.export_weights(weight_layout),
             attention=self.attention.export_weights(weight_layout),
             pre_mlp_norm=self.pre_mlp_norm.export_weights(weight_layout),
@@ -238,3 +281,37 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
         if self.post_mlp_norm is not None:
             result["post_mlp_norm"] = self.post_mlp_norm.export_weights(weight_layout)
         return result
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["pre_attention_norm"], Mapping)
+        assert isinstance(weights["attention"], Mapping)
+        assert isinstance(weights["mlp"], Mapping)
+        assert isinstance(weights["pre_mlp_norm"], Mapping)
+
+        if self.post_attention_norm is not None:
+            assert isinstance(weights["post_attention_norm"], Mapping)
+            post_attention_norm = self.post_attention_norm.import_weights(
+                weights["post_attention_norm"],
+                weight_layout,
+            )
+        else:
+            post_attention_norm = None
+        if self.post_mlp_norm is not None:
+            assert isinstance(weights["post_mlp_norm"], Mapping)
+            post_mlp_norm = self.post_mlp_norm.import_weights(weights["post_mlp_norm"], weight_layout)
+        else:
+            post_mlp_norm = None
+        return replace(
+            self,
+            pre_attention_norm=self.pre_attention_norm.import_weights(weights["pre_attention_norm"], weight_layout),
+            attention=self.attention.import_weights(weights["attention"], weight_layout),
+            post_attention_norm=post_attention_norm,
+            pre_mlp_norm=self.pre_mlp_norm.import_weights(weights["pre_mlp_norm"], weight_layout),
+            mlp=self.mlp.import_weights(weights["mlp"], weight_layout),
+            post_mlp_norm=post_mlp_norm,
+        )

lalamo/modules/embedding.py

@@ -1,15 +1,23 @@
 from abc import abstractmethod
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from typing import Self
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
-from einops import rearrange
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree, dummy_array
 from lalamo.quantization import QuantizationMode, dynamically_quantize_activations, quantize_weights
 
-from .common import LalamoModule, WeightLayout, register_config_union
+from .common import (
+    LalamoModule,
+    WeightLayout,
+    from_layout,
+    into_layout,
+    register_config_union,
+)
 from .utils import apply_soft_capping
 
 __all__ = [
@@ -38,6 +46,13 @@ class EmbeddingConfigBase:
         key: PRNGKeyArray,
     ) -> "EmbeddingBase": ...
 
+    @abstractmethod
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "EmbeddingBase": ...
+
 
 class EmbeddingBase[ConfigT: EmbeddingConfigBase](LalamoModule[ConfigT]):
     @abstractmethod
@@ -54,16 +69,14 @@ class EmbeddingBase[ConfigT: EmbeddingConfigBase](LalamoModule[ConfigT]):
     @abstractmethod
     def model_dim(self) -> int: ...
 
-    @classmethod
-    def _default_weight_layout(cls) -> WeightLayout:
-        return WeightLayout.INPUT_OUTPUT
-
+    @eqx.filter_jit
     def embed(self, x: Int[Array, " tokens"]) -> Float[Array, "tokens channels"]:
         result = self._prepare_input_weights()[x]
         if self.config.input_scale is not None:
             result = result * jnp.array(self.config.input_scale, dtype=result.dtype)
         return result
 
+    @eqx.filter_jit
     def readout(self, x: Float[Array, " channels"]) -> Float[Array, " vocabulary"]:
         logits = self._prepare_output_weights() @ x
         if self.config.logits_soft_cap is not None:
@@ -85,6 +98,14 @@ class TiedEmbeddingConfig(EmbeddingConfigBase):
         weights = jax.random.normal(key, (vocab_size, model_dim), dtype=self.precision)
         return TiedEmbedding(config=self, weights=weights)
 
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "TiedEmbedding":
+        weights = dummy_array((vocab_size, model_dim), dtype=self.precision)
+        return TiedEmbedding(config=self, weights=weights)
+
 
 class TiedEmbedding(EmbeddingBase[TiedEmbeddingConfig]):
     weights: Float[Array, "vocabulary channels"]
@@ -115,8 +136,16 @@ class TiedEmbedding(EmbeddingBase[TiedEmbeddingConfig]):
     def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
         return self.weights
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
-        return ParameterDict(weights=self.weights)
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:  # noqa: ARG002
+        return {"weights": self.weights}
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,  # noqa: ARG002
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        return replace(self, weights=weights["weights"])
 
 
 @dataclass(frozen=True)
@@ -139,6 +168,19 @@ class UntiedEmbeddingConfig(EmbeddingConfigBase):
             output_weights=output_weights,
         )
 
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "UntiedEmbedding":
+        input_weights = dummy_array((vocab_size, model_dim), dtype=self.precision)
+        output_weights = dummy_array((vocab_size, model_dim), dtype=self.precision)
+        return UntiedEmbedding(
+            config=self,
+            input_weights=input_weights,
+            output_weights=output_weights,
+        )
+
 
 class UntiedEmbedding(EmbeddingBase[UntiedEmbeddingConfig]):
     input_weights: Float[Array, "vocabulary channels"]
@@ -186,21 +228,22 @@ class UntiedEmbedding(EmbeddingBase[UntiedEmbeddingConfig]):
     def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
         return self.output_weights
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        if weight_layout == WeightLayout.AUTO:
-            weight_layout = self._default_weight_layout()
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        return {
+            "input_weights": self.input_weights,
+            "output_weights": into_layout(self.output_weights, weight_layout),
+        }
 
-        match weight_layout:
-            case WeightLayout.OUTPUT_INPUT:
-                output_weights = self.output_weights
-            case WeightLayout.INPUT_OUTPUT:
-                output_weights = rearrange(self.output_weights, "token_ids channels -> channels token_ids")
-            case _:
-                raise ValueError(f"Unsupported weight layout: {weight_layout}")
-
-        return ParameterDict(
-            input_weights=self.input_weights,
-            output_weights=output_weights,
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        return replace(
+            self,
+            input_weights=weights["input_weights"],
+            output_weights=from_layout(weights["output_weights"], weight_layout),
         )
 
 
@@ -225,6 +268,15 @@ class QuantizedTiedEmbeddingConfig(EmbeddingConfigBase):
         weights = quantize_weights(weights * min_abs_val, self.embedding_quantization_mode)
         return QuantizedTiedEmbedding(config=self, weights=weights, scales=scales)
 
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "QuantizedTiedEmbedding":
+        scales = dummy_array(vocab_size, dtype=self.activation_precision)
+        weights = dummy_array((vocab_size, model_dim), dtype=self.activation_precision)
+        return QuantizedTiedEmbedding(config=self, weights=weights, scales=scales)
+
 
 class QuantizedTiedEmbedding(EmbeddingBase[QuantizedTiedEmbeddingConfig]):
     weights: Float[Array, "vocabulary channels"]
@@ -257,7 +309,7 @@ class QuantizedTiedEmbedding(EmbeddingBase[QuantizedTiedEmbeddingConfig]):
                 f" {self.config.activation_precision}"
                 " Quantized layers require parameter dtypes to be equal to the activation precision.",
             )
-        weights_vocab_size, weights_model_dim = self.weights.shape
+        weights_vocab_size, _ = self.weights.shape
         (scales_vocab_size,) = self.scales.shape
         if weights_vocab_size != scales_vocab_size:
             raise ValueError(
@@ -281,15 +333,29 @@ class QuantizedTiedEmbedding(EmbeddingBase[QuantizedTiedEmbeddingConfig]):
     def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
         return self._prepare_weights()
 
+    @eqx.filter_jit
     def readout(self, x: Float[Array, " channels"]) -> Float[Array, " vocabulary"]:
         if self.config.activation_quantization_mode is not None:
             x = dynamically_quantize_activations(x, self.config.activation_quantization_mode)
         return super().readout(x)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
-        return ParameterDict(
-            weights=self.int_weights,
-            scales=self.scales,
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        return {
+            "weights": into_layout(self.int_weights, weight_layout),
+            "scales": into_layout(self.scales, weight_layout),
+        }
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["weights"], Array)
+        return replace(
+            self,
+            weights=from_layout(weights["weights"].astype(self.weights.dtype), weight_layout),
+            scales=from_layout(weights["scales"], weight_layout),
         )
 
 

lalamo/modules/kv_cache.py

@@ -7,7 +7,7 @@ from jax.lax import dynamic_update_slice_in_dim
 from jax.tree_util import register_pytree_node_class
 from jaxtyping import Array, Bool, DTypeLike, Float, Int
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree
 
 __all__ = ["DynamicKVCacheLayer", "KVCache", "KVCacheLayer", "StaticKVCacheLayer"]
 
@@ -43,8 +43,8 @@ class KVCacheLayer(eqx.Module):
         added_length: Int[Array, ""] | int | None = None,
     ) -> Self: ...
 
-    def export(self) -> ParameterDict:
-        return ParameterDict(
+    def export(self) -> ParameterTree:
+        return dict(
             keys=self.keys,
             values=self.values,
         )