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

lalamo/modules/transformer.py

@@ -0,0 +1,273 @@
+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 ParameterTree
+from lalamo.modules.token_mixers import AttentionConfig
+from lalamo.modules.utils import vmap_twice
+
+from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
+from .normalization import Normalization, NormalizationConfig
+from .rope import PositionalEmbeddings, RoPE, RoPEConfig
+from .token_mixers import State
+from .transformer_layer import (
+    TransformerLayer,
+    TransformerLayerConfig,
+    TransformerLayerForwardPassConfig,
+    TransformerLayerResult,
+)
+
+__all__ = [
+    "Transformer",
+    "TransformerConfig",
+    "TransformerResult",
+]
+
+
+type TransformerForwardPassConfig = TransformerLayerForwardPassConfig
+
+
+class TransformerResult(eqx.Module):
+    outputs: Float[Array, "batch suffix_tokens channels"]
+    updated_state: State | None = None
+    layer_results: tuple[TransformerLayerResult, ...] | None = None
+    global_positional_embeddings: PositionalEmbeddings | None = None
+    local_positional_embeddings: PositionalEmbeddings | None = None
+
+    def export(self) -> ParameterTree:
+        result: dict[str, ParameterTree | Array] = dict(
+            outputs=self.outputs,
+        )
+        if self.updated_state is not None:
+            result["updated_state"] = [state_layer.export() for state_layer in self.updated_state]
+        if self.layer_results is not None:
+            result["layer_results"] = [layer_result.export() for layer_result in self.layer_results]
+        if self.global_positional_embeddings is not None:
+            result["global_positional_embeddings"] = self.global_positional_embeddings.export()
+        if self.local_positional_embeddings is not None:
+            result["local_positional_embeddings"] = self.local_positional_embeddings.export()
+        return result
+
+
+@dataclass(frozen=True)
+class TransformerConfig:
+    global_rope_config: RoPEConfig | None
+    local_rope_config: RoPEConfig | None
+    layer_configs: tuple[TransformerLayerConfig, ...]
+    output_norm_config: NormalizationConfig
+    model_dim: int
+    hidden_dim: int
+    context_length: int
+
+    def random_init(self, *, key: PRNGKeyArray) -> "Transformer":
+        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(key, num=len(self.layer_configs))
+        layers = tuple(
+            layer_config.random_init(
+                model_dim=self.model_dim,
+                hidden_dim=self.hidden_dim,
+                key=layer_key,
+            )
+            for layer_key, layer_config in zip(layers_keys, self.layer_configs, strict=True)
+        )
+        output_norm = self.output_norm_config.init(self.model_dim)
+
+        return Transformer(
+            config=self,
+            global_rope=global_rope,
+            local_rope=local_rope,
+            layers=layers,
+            output_norm=output_norm,
+        )
+
+    def empty(self) -> "Transformer":
+        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:
+            local_rope = self.local_rope_config.init(
+                head_dim=first_layer_config.rope_dim,
+                num_timesteps=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 Transformer(
+            config=self,
+            global_rope=global_rope,
+            local_rope=local_rope,
+            layers=layers,
+            output_norm=output_norm,
+        )
+
+
+class Transformer(LalamoModule[TransformerConfig]):
+    global_rope: RoPE | None
+    local_rope: RoPE | None
+    layers: tuple[TransformerLayer, ...]
+    output_norm: Normalization
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.layers[0].activation_precision
+
+    @eqx.filter_jit
+    def __call__(
+        self,
+        inner_features: Float[Array, "batch suffix_tokens channels"],
+        token_positions: Int[Array, "batch suffix_tokens"],
+        state: State | None,
+        return_updated_state: bool,
+        return_layer_results: bool,
+        return_positional_embeddings: bool,
+        lengths_without_padding: Int[Array, " batch"] | None,
+        forward_pass_mode: ForwardPassMode,
+        forward_pass_config: TransformerForwardPassConfig | None,
+    ) -> TransformerResult:
+        if inner_features.ndim != 3:
+            raise ValueError(
+                f"inner_features must be a 3D array of size (batch_size, sequence_length, hidden_dim), got {inner_features.shape}",
+            )
+        if token_positions.ndim != 2:
+            raise ValueError(
+                "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))
+
+        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_layer_results,
+                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)
+
+        normalized_outputs = vmap_twice(self.output_norm)(inner_features)
+
+        return TransformerResult(
+            outputs=normalized_outputs,
+            updated_state=(State(updated_state_layers) if return_updated_state else None),
+            layer_results=tuple(layer_results) if return_layer_results else None,
+            global_positional_embeddings=(global_positional_embeddings if return_positional_embeddings else None),
+            local_positional_embeddings=(local_positional_embeddings if return_positional_embeddings else None),
+        )
+
+    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(
+            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
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["layers"], Sequence)
+        assert isinstance(weights["output_norm"], Mapping)
+
+        if self.global_rope:
+            assert isinstance(weights["global_rope"], Mapping)
+            global_rope = self.global_rope.import_weights(weights["global_rope"])
+        else:
+            global_rope = None
+
+        if self.local_rope:
+            assert isinstance(weights["local_rope"], Mapping)
+            local_rope = self.local_rope.import_weights(weights["local_rope"])
+        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))
+
+        return replace(
+            self,
+            global_rope=global_rope,
+            layers=tuple(layers),
+            output_norm=self.output_norm.import_weights(weights["output_norm"]),
+            local_rope=local_rope,
+        )

lalamo/modules/{decoder_layer.py → transformer_layer.py}

@@ -13,24 +13,24 @@ from lalamo.common import ParameterTree
 
 from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
 from .mlp import MLPBase, MLPConfig, MLPForwardPassConfig
-from .normalization import RMSNorm, RMSNormConfig
+from .normalization import Normalization, NormalizationConfig
 from .rope import PositionalEmbeddings
 from .token_mixers import KVCacheLayer, StateLayerBase, StaticKVCacheLayer, TokenMixerBase, TokenMixerConfig
 from .utils import vmap_twice
 
 __all__ = [
-    "DecoderLayer",
-    "DecoderLayerActivationTrace",
-    "DecoderLayerConfig",
-    "DecoderLayerForwardPassConfig",
-    "DecoderLayerResult",
+    "TransformerLayer",
+    "TransformerLayerActivationTrace",
+    "TransformerLayerConfig",
+    "TransformerLayerForwardPassConfig",
+    "TransformerLayerResult",
 ]
 
 
-type DecoderLayerForwardPassConfig = MLPForwardPassConfig
+type TransformerLayerForwardPassConfig = MLPForwardPassConfig
 
 
-class DecoderLayerActivationTrace(eqx.Module):
+class TransformerLayerActivationTrace(eqx.Module):
     inputs: Float[Array, "batch suffix_tokens channels"]
     positional_embeddings: PositionalEmbeddings | None
     state: StateLayerBase | None
@@ -63,10 +63,10 @@ class DecoderLayerActivationTrace(eqx.Module):
         return result
 
 
-class DecoderLayerResult(eqx.Module):
-    outputs: Float[Array, "suffix_tokens channels"]
+class TransformerLayerResult(eqx.Module):
+    outputs: Float[Array, "batch tokens channels"]
     updated_state: KVCacheLayer | None
-    activation_trace: DecoderLayerActivationTrace | None
+    activation_trace: TransformerLayerActivationTrace | None
 
     def export(self) -> ParameterTree:
         result: dict[str, ParameterTree | Array] = dict(
@@ -80,13 +80,13 @@ class DecoderLayerResult(eqx.Module):
 
 
 @dataclass(frozen=True)
-class DecoderLayerConfig:
-    pre_mixer_norm_config: RMSNormConfig
+class TransformerLayerConfig:
+    pre_mixer_norm_config: NormalizationConfig | None
     mixer_config: TokenMixerConfig
-    post_mixer_norm_config: RMSNormConfig | None
-    pre_mlp_norm_config: RMSNormConfig
+    post_mixer_norm_config: NormalizationConfig | None
+    pre_mlp_norm_config: NormalizationConfig
     mlp_config: MLPConfig
-    post_mlp_norm_config: RMSNormConfig | None
+    post_mlp_norm_config: NormalizationConfig | None
 
     @property
     def rope_dim(self) -> int:
@@ -98,28 +98,31 @@ class DecoderLayerConfig:
         hidden_dim: int,
         *,
         key: PRNGKeyArray,
-    ) -> "DecoderLayer":
+    ) -> "TransformerLayer":
         attention_key, mlp_key = jax.random.split(key)
-        pre_attention_norm = self.pre_mixer_norm_config.init(model_dim)
+        if self.pre_mixer_norm_config is not None:
+            pre_mixer_norm = self.pre_mixer_norm_config.init(model_dim)
+        else:
+            pre_mixer_norm = None
         mixer = self.mixer_config.random_init(
             model_dim=model_dim,
             key=attention_key,
         )
         if self.post_mixer_norm_config is not None:
-            post_attention_norm = self.post_mixer_norm_config.init(model_dim)
+            post_mixer_norm = self.post_mixer_norm_config.init(model_dim)
         else:
-            post_attention_norm = None
+            post_mixer_norm = None
         pre_mlp_norm = self.pre_mlp_norm_config.init(model_dim)
         mlp = self.mlp_config.random_init(model_dim, hidden_dim, key=mlp_key)
         if self.post_mlp_norm_config is not None:
             post_mlp_norm = self.post_mlp_norm_config.init(model_dim)
         else:
             post_mlp_norm = None
-        return DecoderLayer(
+        return TransformerLayer(
             config=self,
-            pre_mixer_norm=pre_attention_norm,
+            pre_mixer_norm=pre_mixer_norm,
             mixer=mixer,
-            post_mixer_norm=post_attention_norm,
+            post_mixer_norm=post_mixer_norm,
             pre_mlp_norm=pre_mlp_norm,
             mlp=mlp,
             post_mlp_norm=post_mlp_norm,
@@ -129,39 +132,42 @@ class DecoderLayerConfig:
         self,
         model_dim: int,
         hidden_dim: int,
-    ) -> "DecoderLayer":
-        pre_attention_norm = self.pre_mixer_norm_config.empty(model_dim)
+    ) -> "TransformerLayer":
+        if self.pre_mixer_norm_config is not None:
+            pre_mixer_norm = self.pre_mixer_norm_config.empty(model_dim)
+        else:
+            pre_mixer_norm = None
         attention = self.mixer_config.empty(
             model_dim=model_dim,
         )
         if self.post_mixer_norm_config is not None:
-            post_attention_norm = self.post_mixer_norm_config.empty(model_dim)
+            post_mixer_norm = self.post_mixer_norm_config.empty(model_dim)
         else:
-            post_attention_norm = None
+            post_mixer_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(
+        return TransformerLayer(
             config=self,
-            pre_mixer_norm=pre_attention_norm,
+            pre_mixer_norm=pre_mixer_norm,
             mixer=attention,
-            post_mixer_norm=post_attention_norm,
+            post_mixer_norm=post_mixer_norm,
             pre_mlp_norm=pre_mlp_norm,
             mlp=mlp,
             post_mlp_norm=post_mlp_norm,
         )
 
 
-class DecoderLayer(LalamoModule[DecoderLayerConfig]):
-    pre_mixer_norm: RMSNorm
+class TransformerLayer(LalamoModule[TransformerLayerConfig]):
+    pre_mixer_norm: Normalization | None
     mixer: TokenMixerBase
-    post_mixer_norm: RMSNorm | None
-    pre_mlp_norm: RMSNorm
+    post_mixer_norm: Normalization | None
+    pre_mlp_norm: Normalization
     mlp: MLPBase
-    post_mlp_norm: RMSNorm | None
+    post_mlp_norm: Normalization | None
 
     @property
     def activation_precision(self) -> DTypeLike:
@@ -172,7 +178,7 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
         return self.mixer.positional_embedding_selector
 
     def __post_init__(self) -> None:
-        model_dim = self.pre_mixer_norm.input_dim
+        model_dim = self.pre_mixer_norm.input_dim if self.pre_mixer_norm is not None else self.mixer.model_dim
         if self.mixer.model_dim != model_dim:
             raise ValueError(
                 f"Attention model dim {self.mixer.model_dim} does not match"
@@ -204,15 +210,21 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
         return_activation_trace: bool = False,
         lengths_without_padding: Int[Array, " batch"] | None = None,
         forward_pass_mode: ForwardPassMode = ForwardPassMode.MULTI_TOKEN,
-        forward_pass_config: DecoderLayerForwardPassConfig | None = None,
-    ) -> DecoderLayerResult:
+        forward_pass_config: TransformerLayerForwardPassConfig | None = None,
+    ) -> TransformerLayerResult:
         if inputs.ndim != 3:
             raise ValueError(
                 f"Inputs to decoder layers must be a 3D arrays of size (batch_size, sequence_length, hidden_dim),"
                 f" got {inputs.shape}",
             )
-        normalized_mixer_inputs = vmap_twice(self.pre_mixer_norm)(inputs)
-        batched_mixer_fn = vmap(partial(self.mixer, return_updated_state=return_updated_state))
+        if self.pre_mixer_norm is not None:
+            normalized_mixer_inputs = vmap_twice(self.pre_mixer_norm)(inputs)
+        else:
+            normalized_mixer_inputs = inputs
+
+        batched_mixer_fn = vmap(
+            partial(self.mixer, return_updated_state=return_updated_state or return_activation_trace),
+        )
         mixer_outputs, updated_state = batched_mixer_fn(
             normalized_mixer_inputs,
             positional_embeddings,
@@ -240,7 +252,7 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
             outputs = mlp_inputs + mlp_outputs
 
         if return_activation_trace:
-            activation_trace = DecoderLayerActivationTrace(
+            activation_trace = TransformerLayerActivationTrace(
                 inputs=inputs,
                 positional_embeddings=positional_embeddings,
                 state=state,
@@ -255,7 +267,7 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
         else:
             activation_trace = None
 
-        return DecoderLayerResult(
+        return TransformerLayerResult(
             outputs=outputs,
             updated_state=updated_state,
             activation_trace=activation_trace,
@@ -269,11 +281,12 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
 
     def export_weights(self) -> ParameterTree:
         result = dict(
-            pre_mixer_norm=self.pre_mixer_norm.export_weights(),
             mixer=self.mixer.export_weights(),
             pre_mlp_norm=self.pre_mlp_norm.export_weights(),
             mlp=self.mlp.export_weights(),
         )
+        if self.pre_mixer_norm is not None:
+            result["pre_mixer_norm"] = self.pre_mixer_norm.export_weights()
         if self.post_mixer_norm is not None:
             result["post_mixer_norm"] = self.post_mixer_norm.export_weights()
         if self.post_mlp_norm is not None:
@@ -285,7 +298,6 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
         weights: ParameterTree[Array],
     ) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["pre_mixer_norm"], Mapping)
         assert isinstance(weights["mixer"], Mapping)
         assert isinstance(weights["mlp"], Mapping)
         assert isinstance(weights["pre_mlp_norm"], Mapping)
@@ -302,9 +314,14 @@ class DecoderLayer(LalamoModule[DecoderLayerConfig]):
             post_mlp_norm = self.post_mlp_norm.import_weights(weights["post_mlp_norm"])
         else:
             post_mlp_norm = None
+        if self.pre_mixer_norm is not None:
+            assert isinstance(weights["pre_mixer_norm"], Mapping)
+            pre_mixer_norm = self.pre_mixer_norm.import_weights(weights["pre_mixer_norm"])
+        else:
+            pre_mixer_norm = None
         return replace(
             self,
-            pre_mixer_norm=self.pre_mixer_norm.import_weights(weights["pre_mixer_norm"]),
+            pre_mixer_norm=pre_mixer_norm,
             mixer=self.mixer.import_weights(weights["mixer"]),
             post_mixer_norm=post_mixer_norm,
             pre_mlp_norm=self.pre_mlp_norm.import_weights(weights["pre_mlp_norm"]),

lalamo/speculator/__init__.py

@@ -1,11 +1,15 @@
 from .common import Speculator
-from .inference import inference_collect_traces
+from .estimator import estimate_batchsize_from_memory
+from .inference import CollectTracesEvent, inference_collect_traces
 from .ngram import NGramSpeculator
-from .utils import train_speculator
+from .utils import SpeculatorTrainingEvent, train_speculator
 
 __all__ = [
+    "CollectTracesEvent",
     "NGramSpeculator",
     "Speculator",
+    "SpeculatorTrainingEvent",
+    "estimate_batchsize_from_memory",
     "inference_collect_traces",
     "train_speculator",
 ]

lalamo/speculator/estimator.py

@@ -0,0 +1,91 @@
+import functools
+import itertools
+from collections.abc import Callable
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+
+from lalamo.models import LanguageModel
+
+
+def estimate_memory_from_batchsize(
+    model: LanguageModel,
+    max_input_length: int,
+    max_output_length: int,
+    num_logits_per_token: int,
+    batch_size: int,
+) -> int:
+    memory_analysis = (
+        jax.jit(
+            functools.partial(
+                model.generate_tokens,
+                max_output_length=max_output_length,
+                num_top_logits_to_return=num_logits_per_token,
+            ),
+            backend="cpu", # cuda backend tries to allocate in .compile() and ooms
+        )
+        .lower(
+            prompt_token_ids=jax.ShapeDtypeStruct((batch_size, max_input_length), jnp.int32),
+            prompt_lengths_without_padding=jax.ShapeDtypeStruct((batch_size,), jnp.int32),
+        )
+        .compile()
+        .memory_analysis()
+    )
+
+    assert hasattr(memory_analysis, "argument_size_in_bytes")
+    assert hasattr(memory_analysis, "output_size_in_bytes")
+    assert hasattr(memory_analysis, "temp_size_in_bytes")
+
+    return (
+        memory_analysis.argument_size_in_bytes  # type: ignore (pyright bug)
+        + memory_analysis.output_size_in_bytes  # type: ignore (pyright bug)
+        + memory_analysis.temp_size_in_bytes    # type: ignore (pyright bug)
+    )
+
+
+class EstimateBatchsizeFromMemoryEvent(NamedTuple):
+    lo: int
+    hi: int | None
+
+
+def estimate_batchsize_from_memory(
+    model: LanguageModel,
+    max_input_length: int,
+    max_output_length: int,
+    num_logits_per_token: int,
+    target_mem: int,
+    progress: Callable[[EstimateBatchsizeFromMemoryEvent], None] | None = None,
+) -> int:
+    mem_for_bs = functools.cache(
+        functools.partial(
+            estimate_memory_from_batchsize,
+            model,
+            max_input_length,
+            max_output_length,
+            num_logits_per_token,
+        ),
+    )
+
+    lo = 0
+    hi = 0
+    for candidate_exp in itertools.count():
+        lo = hi
+        hi = 2**candidate_exp
+
+        if progress is not None:
+            progress(EstimateBatchsizeFromMemoryEvent(lo, None))
+        if target_mem < mem_for_bs(hi):
+            break
+
+    while hi - lo > 1:
+        mid = (lo + hi) // 2
+
+        if progress is not None:
+            progress(EstimateBatchsizeFromMemoryEvent(lo, hi))
+        if target_mem < mem_for_bs(mid):
+            hi = mid
+        else:
+            lo = mid
+
+    return lo