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

lalamo/modules/decoder.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from dataclasses import dataclass, replace
 from typing import Self
 
@@ -9,12 +9,16 @@ from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
 from lalamo.common import ParameterTree
 
-from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
-from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerForwardPassConfig, DecoderLayerResult
+from .common import ForwardPassMode, LalamoModule
 from .embedding import EmbeddingBase, EmbeddingConfig
-from .normalization import RMSNorm, RMSNormConfig
-from .rope import PositionalEmbeddings, RoPE, RoPEConfig
-from .token_mixers import AttentionConfig, State
+from .rope import PositionalEmbeddings
+from .token_mixers import State
+from .transformer import (
+    Transformer,
+    TransformerConfig,
+    TransformerForwardPassConfig,
+    TransformerLayerResult,
+)
 from .utils import vmap_twice
 
 __all__ = [
@@ -26,7 +30,7 @@ __all__ = [
 ]
 
 
-type DecoderForwardPassConfig = DecoderLayerForwardPassConfig
+type DecoderForwardPassConfig = TransformerForwardPassConfig
 
 
 class DecoderActivationTrace(eqx.Module):
@@ -37,7 +41,7 @@ class DecoderActivationTrace(eqx.Module):
     local_positional_embeddings: PositionalEmbeddings | None
     global_positional_embeddings: PositionalEmbeddings | None
 
-    layer_results: tuple[DecoderLayerResult, ...]
+    layer_results: tuple[TransformerLayerResult, ...]
 
     output_norm: Float[Array, "batch suffix_tokens channels"]
 
@@ -48,12 +52,12 @@ class DecoderActivationTrace(eqx.Module):
             layer_results=[layer_result.export() for layer_result in self.layer_results],
             output_norm=self.output_norm,
         )
+        if self.state is not None:
+            result["state"] = [state_layer.export() for state_layer in self.state]
         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
 
 
@@ -76,124 +80,46 @@ class DecoderResult(eqx.Module):
 @dataclass(frozen=True)
 class DecoderConfig:
     embedding_config: EmbeddingConfig
-    global_rope_config: RoPEConfig | None
-    local_rope_config: RoPEConfig | None
-    layer_configs: tuple[DecoderLayerConfig, ...]
-    output_norm_config: RMSNormConfig
+    transformer_config: TransformerConfig
 
     vocab_size: int
-    model_dim: int
-    hidden_dim: int
-    context_length: int
 
     def random_init(
         self,
         *,
         key: PRNGKeyArray,
     ) -> "Decoder":
-        embedding_key, layers_key = jax.random.split(key)
+        embedding_key, transformer_key = jax.random.split(key)
         embedding = self.embedding_config.random_init(
             vocab_size=self.vocab_size,
-            model_dim=self.model_dim,
+            model_dim=self.transformer_config.model_dim,
             key=embedding_key,
         )
+        transformer = self.transformer_config.random_init(key=transformer_key)
 
-        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:
-            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:
-            local_rope = None
-
-        layers_keys = jax.random.split(layers_key, len(self.layer_configs))
-        layers = tuple(
-            layer_config.random_init(
-                model_dim=self.model_dim,
-                hidden_dim=self.hidden_dim,
-                key=key,
-            )
-            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(
-            self,
+            config=self,
             embedding=embedding,
-            global_rope=global_rope,
-            local_rope=local_rope,
-            layers=layers,
-            output_norm=output_norm,
+            transformer=transformer,
         )
 
-    def empty(
-        self,
-    ) -> "Decoder":
+    def empty(self) -> "Decoder":
         embedding = self.embedding_config.empty(
             vocab_size=self.vocab_size,
-            model_dim=self.model_dim,
+            model_dim=self.transformer_config.model_dim,
         )
+        transformer = self.transformer_config.empty()
 
-        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:
-            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:
-            local_rope = None
-        layers = tuple(
-            layer_config.empty(
-                model_dim=self.model_dim,
-                hidden_dim=self.hidden_dim,
-            )
-            for layer_config in self.layer_configs
-        )
-        output_norm = self.output_norm_config.empty(self.model_dim)
         return Decoder(
-            self,
+            config=self,
             embedding=embedding,
-            global_rope=global_rope,
-            local_rope=local_rope,
-            layers=layers,
-            output_norm=output_norm,
+            transformer=transformer,
         )
 
 
 class Decoder(LalamoModule[DecoderConfig]):
     embedding: EmbeddingBase
-    global_rope: RoPE | None
-    local_rope: RoPE | None
-    layers: tuple[DecoderLayer, ...]
-    output_norm: RMSNorm
+    transformer: Transformer
 
     @property
     def activation_precision(self) -> DTypeLike:
@@ -213,93 +139,59 @@ class Decoder(LalamoModule[DecoderConfig]):
     ) -> DecoderResult:
         if token_ids.ndim != 2:
             raise ValueError(
-                f"token_ids must be a 2D arrays of size (batch_size, sequence_length), got {token_ids.shape}",
+                f"token_ids must be a 2D array of size (batch_size, sequence_length), got {token_ids.shape}",
             )
         if token_positions.ndim != 2:
             raise ValueError(
-                "token_positions must be a 2D arrays of size (batch_size, sequence_length),"
+                "token_positions must be a 2D array of size (batch_size, sequence_length),"
                 f" got {token_positions.shape}",
             )
 
-        maybe_state = state or ([None] * len(self.layers))
         inner_features = vmap(self.embedding.embed)(token_ids)
 
-        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_state_layers = []
-        layer_results = []
-        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,
-                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,
-                forward_pass_config=forward_pass_config,
-            )
-            inner_features = layer_result.outputs
-            layer_results.append(layer_result)
-            updated_state_layers.append(layer_result.updated_state)
+        transformer_result = self.transformer(
+            inner_features=inner_features,
+            token_positions=token_positions,
+            state=state,
+            return_updated_state=return_updated_state,
+            return_layer_results=return_activation_trace,
+            return_positional_embeddings=return_activation_trace,
+            lengths_without_padding=lengths_without_padding,
+            forward_pass_mode=forward_pass_mode,
+            forward_pass_config=forward_pass_config,
+        )
 
-        normalized_outputs = vmap_twice(self.output_norm)(inner_features)
-        logits = vmap_twice(self.embedding.readout)(normalized_outputs)
+        logits = vmap_twice(self.embedding.readout)(transformer_result.outputs)
 
         if return_activation_trace:
+            assert transformer_result.layer_results is not None
+
             activation_trace = DecoderActivationTrace(
                 token_ids=token_ids,
                 token_positions=token_positions,
                 state=state,
-                global_positional_embeddings=global_positional_embeddings,
-                local_positional_embeddings=local_positional_embeddings,
-                layer_results=tuple(layer_results),
-                output_norm=normalized_outputs,
+                global_positional_embeddings=transformer_result.global_positional_embeddings,
+                local_positional_embeddings=transformer_result.local_positional_embeddings,
+                layer_results=transformer_result.layer_results,
+                output_norm=transformer_result.outputs,
             )
         else:
             activation_trace = None
 
-        if return_updated_state:
-            updated_state = State(updated_state_layers)
-        else:
-            updated_state = None
-
         return DecoderResult(
             logits=logits,
-            updated_state=updated_state,
+            updated_state=transformer_result.updated_state,
             activation_trace=activation_trace,
         )
 
     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)
+        return self.transformer.init_static_state(batch_size, capacity)
 
     def export_weights(self) -> ParameterTree:
-        result = dict(
+        return dict(
             embedding=self.embedding.export_weights(),
-            layers=[layer.export_weights() for layer in self.layers],
-            output_norm=self.output_norm.export_weights(),
+            transformer=self.transformer.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
 
     def import_weights(
         self,
@@ -307,30 +199,10 @@ class Decoder(LalamoModule[DecoderConfig]):
     ) -> Self:
         assert isinstance(weights, Mapping)
         assert isinstance(weights["embedding"], 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
+        assert isinstance(weights["transformer"], Mapping)
 
-        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))
         return replace(
             self,
             embedding=self.embedding.import_weights(weights["embedding"]),
-            global_rope=global_rope,
-            layers=tuple(layers),
-            output_norm=self.output_norm.import_weights(weights["output_norm"]),
-            local_rope=local_rope,
+            transformer=self.transformer.import_weights(weights["transformer"]),
         )

lalamo/modules/mlp.py

@@ -16,7 +16,12 @@ from lalamo.common import ParameterTree
 from lalamo.modules.utils import vmap_twice
 
 from .activations import Activation
-from .common import DummyUnionMember, ForwardPassMode, LalamoModule, register_config_union
+from .common import (
+    DummyUnionMember,
+    ForwardPassMode,
+    LalamoModule,
+    register_config_union,
+)
 from .linear import LinearBase, LinearConfig
 
 __all__ = [
@@ -192,7 +197,10 @@ class DenseMLP(MLPBase[DenseMLPConfig]):
                 f" the gate output dimension {gate_output_dim}",
             )
         (down_output_dim,) = self.down_projection.output_dims
-        if (self.up_projection.input_dim, up_output_dim) != (down_output_dim, self.down_projection.input_dim):
+        if (self.up_projection.input_dim, up_output_dim) != (
+            down_output_dim,
+            self.down_projection.input_dim,
+        ):
             raise ValueError(
                 f"Down projection dimensions {self.down_projection.input_dim, down_output_dim} do not match"
                 f" the up projection output dimensions {self.up_projection.input_dim, up_output_dim}",
@@ -209,7 +217,10 @@ class DenseMLP(MLPBase[DenseMLPConfig]):
         return vmap_twice(self.call_unbatched)(inputs)
 
     @eqx.filter_jit
-    def call_unbatched(self, inputs: Float[Array, " channels"]) -> Float[Array, " channels"]:
+    def call_unbatched(
+        self,
+        inputs: Float[Array, " channels"],
+    ) -> Float[Array, " channels"]:
         if self.mixture_size is not None:
             raise ValueError(
                 "Mixtures of linear layers cannot be called directly."
@@ -222,6 +233,7 @@ class DenseMLP(MLPBase[DenseMLPConfig]):
             up_proj = jnp.clip(up_proj, *self.config.up_clipping)
         gate = self.config.activation(gate)
         (result,) = self.down_projection(up_proj * gate)
+
         return result
 
     def export_weights(self) -> ParameterTree:
@@ -450,10 +462,21 @@ class MixtureOfExperts(MLPBase[MixtureOfExpertsConfig]):
                     mode="drop",
                 )
 
-            return jax.lax.cond(jnp.any(token_indices_for_chunk != _SENTINEL), inner, lambda: accumulator), None
+            return (
+                jax.lax.cond(
+                    jnp.any(token_indices_for_chunk != _SENTINEL),
+                    inner,
+                    lambda: accumulator,
+                ),
+                None,
+            )
 
         result, _ = jax.lax.scan(loop_iteration, jnp.zeros_like(flattened_inputs), chunked_token_indices)
-        return rearrange(result, "(batch suffix_tokens) channels -> batch suffix_tokens channels", batch=batch_size)
+        return rearrange(
+            result,
+            "(batch suffix_tokens) channels -> batch suffix_tokens channels",
+            batch=batch_size,
+        )
 
     def export_weights(
         self,

lalamo/modules/normalization.py

@@ -13,8 +13,8 @@ from lalamo.common import ParameterTree, dummy_array
 from .common import LalamoModule
 
 __all__ = [
-    "RMSNorm",
-    "RMSNormConfig",
+    "Normalization",
+    "NormalizationConfig",
     "UpcastMode",
 ]
 
@@ -25,25 +25,26 @@ class UpcastMode(Enum):
 
 
 @dataclass(frozen=True)
-class RMSNormConfig:
+class NormalizationConfig:
     scale_precision: DTypeLike
     accumulation_precision: DTypeLike
     epsilon: float
     scale_offset: float | None
     upcast_mode: UpcastMode
+    subtract_mean: bool
 
-    def init(self, input_dim: int) -> "RMSNorm":
+    def init(self, input_dim: int) -> "Normalization":
         scales = jnp.ones(input_dim, dtype=self.scale_precision)
-        return RMSNorm(self, scales=scales)
+        return Normalization(self, scales=scales)
 
-    def empty(self, input_dim: int) -> "RMSNorm":
-        return RMSNorm(
+    def empty(self, input_dim: int) -> "Normalization":
+        return Normalization(
             config=self,
             scales=dummy_array(input_dim, dtype=self.scale_precision),
         )
 
 
-class RMSNorm(LalamoModule[RMSNormConfig]):
+class Normalization(LalamoModule[NormalizationConfig]):
     scales: Float[Array, " channels"]
 
     @property
@@ -66,6 +67,10 @@ class RMSNorm(LalamoModule[RMSNormConfig]):
     def __call__(self, inputs: Float[Array, " channels"]) -> Float[Array, " channels"]:
         upcasted_inputs = inputs.astype(self.config.accumulation_precision)
 
+        if self.config.subtract_mean:
+            mean = jnp.mean(upcasted_inputs)
+            upcasted_inputs = upcasted_inputs - mean
+
         adjusted_variance = jnp.mean(jnp.square(upcasted_inputs)) + self.config.epsilon
         normalized_x = upcasted_inputs * jax.lax.rsqrt(adjusted_variance)
 

lalamo/modules/token_mixers/attention.py

@@ -12,7 +12,7 @@ from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
 from lalamo.common import dummy_array
 from lalamo.modules.common import ParameterTree, PositionalEmbeddingSelector
 from lalamo.modules.linear import LinearBase, LinearConfig
-from lalamo.modules.normalization import RMSNorm, RMSNormConfig
+from lalamo.modules.normalization import Normalization, NormalizationConfig
 from lalamo.modules.rope import PositionalEmbeddings
 from lalamo.modules.utils import apply_soft_capping
 
@@ -58,7 +58,11 @@ def _soft_capped_attention_kernel(
         "heads dst_tokens channels, heads src_tokens channels -> heads dst_tokens src_tokens",
     )
     if mask is not None:
-        attention_logits = jnp.where(mask, attention_logits, jnp.array(float("-inf"), dtype=attention_logits.dtype))
+        attention_logits = jnp.where(
+            mask,
+            attention_logits,
+            jnp.array(float("-inf"), dtype=attention_logits.dtype),
+        )
 
     if scale is None:
         scale_val = head_dim**-0.5
@@ -82,8 +86,8 @@ class AttentionConfig(TokenMixerConfigBase):
     qkv_projection_config: LinearConfig
     out_projection_config: LinearConfig
 
-    query_norm_config: RMSNormConfig | None
-    key_norm_config: RMSNormConfig | None
+    query_norm_config: NormalizationConfig | None
+    key_norm_config: NormalizationConfig | None
 
     num_heads: int
     num_groups: int
@@ -217,8 +221,8 @@ class Attention(TokenMixerBase[AttentionConfig, KVCacheLayer]):
     qkv_projection: LinearBase
     out_projection: LinearBase
 
-    query_norm: RMSNorm | None
-    key_norm: RMSNorm | None
+    query_norm: Normalization | None
+    key_norm: Normalization | None
 
     sinks: Float[Array, " heads"] | None
 

lalamo/modules/token_mixers/state/kv_cache.py

@@ -89,7 +89,7 @@ class DynamicKVCacheLayer(KVCacheLayer):
         self,
         suffix_length: int,
         is_causal: bool,
-        suffix_length_without_padding: Int[Array, ""] | int | None = None,  # noqa: ARG002
+        suffix_length_without_padding: (Int[Array, ""] | int | None) = None,  # noqa: ARG002
         sliding_window_size: int | None = None,
     ) -> Bool[Array, "suffix_tokens tokens"]:
         self._raise_if_batched()
@@ -97,8 +97,11 @@ class DynamicKVCacheLayer(KVCacheLayer):
         result = jnp.ones((suffix_length, total_num_tokens), dtype=jnp.bool)
         if is_causal:
             result = jnp.tril(result, k=total_num_tokens - suffix_length)
-        if sliding_window_size is not None:
-            result = jnp.triu(result, k=1 - sliding_window_size)
+            if sliding_window_size is not None:
+                result = jnp.triu(result, k=1 - sliding_window_size)
+        elif sliding_window_size is not None:
+            top_zeroed = jnp.tril(result, k=sliding_window_size // 2)
+            result = jnp.triu(top_zeroed, k=-sliding_window_size // 2)
         if self.has_sinks:
             result = result.at[:, 0].set(True)
         if self.padding_mask is not None:
@@ -213,7 +216,14 @@ class StaticKVCacheLayer(KVCacheLayer):
         )
 
     @classmethod
-    def init(cls, has_sinks: bool, capacity: int, num_groups: int, head_dim: int, dtype: DTypeLike) -> Self:
+    def init(
+        cls,
+        has_sinks: bool,
+        capacity: int,
+        num_groups: int,
+        head_dim: int,
+        dtype: DTypeLike,
+    ) -> Self:
         return cls(
             has_sinks=has_sinks,
             keys=jnp.zeros((capacity, num_groups, head_dim), dtype=dtype),