lalamo 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

lalamo/modules/decoder_layer.py

@@ -0,0 +1,240 @@
+from dataclasses import dataclass
+
+import equinox as eqx
+import jax
+from jax import vmap
+from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
+
+from lalamo.common import ParameterDict
+
+from .attention import Attention, AttentionConfig
+from .common import AttentionType, LalamoModule, WeightLayout
+from .kv_cache import KVCacheLayer, StaticKVCacheLayer
+from .mlp import MLP, MLPConfig
+from .normalization import RMSNorm, RMSNormConfig
+from .rope import PositionalEmbeddings
+
+__all__ = [
+    "DecoderLayer",
+    "DecoderLayerActivationTrace",
+    "DecoderLayerConfig",
+    "DecoderLayerResult",
+]
+
+
+class DecoderLayerActivationTrace(eqx.Module):
+    inputs: Float[Array, "suffix_tokens channels"]
+    positional_embeddings: PositionalEmbeddings
+    kv_cache: KVCacheLayer | None
+
+    mlp_inputs: Float[Array, "suffix_tokens channels"]
+    pre_attention_norm: Float[Array, "suffix_tokens channels"]
+    attention: Float[Array, "suffix_tokens channels"]
+    post_attention_norm: Float[Array, "suffix_tokens channels"] | None
+    pre_mlp_norm: Float[Array, "suffix_tokens channels"]
+    mlp: Float[Array, "suffix_tokens channels"]
+    post_mlp_norm: Float[Array, "suffix_tokens channels"] | None
+
+    def export(self) -> ParameterDict:
+        result = ParameterDict(
+            inputs=self.inputs,
+            positional_embeddings=self.positional_embeddings.export(),
+            mlp_inputs=self.mlp_inputs,
+            pre_attention_norm=self.pre_attention_norm,
+            attention=self.attention,
+            pre_mlp_norm=self.pre_mlp_norm,
+            mlp=self.mlp,
+        )
+        if self.kv_cache is not None:
+            result["kv_cache"] = self.kv_cache.export()
+        if self.post_attention_norm is not None:
+            result["post_attention_norm"] = self.post_attention_norm
+        if self.post_mlp_norm is not None:
+            result["post_mlp_norm"] = self.post_mlp_norm
+        return result
+
+
+class DecoderLayerResult(eqx.Module):
+    outputs: Float[Array, "suffix_tokens channels"]
+    updated_kv_cache: KVCacheLayer | None
+    activation_trace: DecoderLayerActivationTrace | None
+
+    def export(self) -> ParameterDict:
+        result = ParameterDict(
+            outputs=self.outputs,
+        )
+        if self.updated_kv_cache is not None:
+            result["updated_kv_cache"] = self.updated_kv_cache.export()
+        if self.activation_trace is not None:
+            result["activation_trace"] = self.activation_trace.export()
+        return result
+
+
+@dataclass(frozen=True)
+class DecoderLayerConfig:
+    pre_attention_norm_config: RMSNormConfig
+    attention_config: AttentionConfig
+    post_attention_norm_config: RMSNormConfig | None
+    pre_mlp_norm_config: RMSNormConfig
+    mlp_config: MLPConfig
+    post_mlp_norm_config: RMSNormConfig | None
+
+    def random_init(
+        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,
+        *,
+        key: PRNGKeyArray,
+    ) -> "DecoderLayer":
+        attention_key, mlp_key = jax.random.split(key)
+        pre_attention_norm = self.pre_attention_norm_config.init(model_dim)
+        attention = self.attention_config.random_init(
+            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,
+            key=attention_key,
+        )
+        if self.post_attention_norm_config is not None:
+            post_attention_norm = self.post_attention_norm_config.init(model_dim)
+        else:
+            post_attention_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(
+            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
+    attention: Attention
+    post_attention_norm: RMSNorm | None
+    pre_mlp_norm: RMSNorm
+    mlp: MLP
+    post_mlp_norm: RMSNorm | None
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.attention.activation_precision
+
+    @property
+    def attention_type(self) -> AttentionType:
+        return self.attention.attention_type
+
+    def __post_init__(self) -> None:
+        model_dim = self.pre_attention_norm.input_dim
+        if self.attention.model_dim != model_dim:
+            raise ValueError(
+                f"Attention model dim {self.attention.model_dim} does not match"
+                f" the first normalization layer dim {model_dim}",
+            )
+        if self.post_attention_norm is not None and self.post_attention_norm.input_dim != model_dim:
+            raise ValueError(
+                f"Post attention normalization dim {self.post_attention_norm.input_dim} does not match"
+                f" the first normalization layer dim {model_dim}",
+            )
+        if self.pre_mlp_norm.input_dim != model_dim:
+            raise ValueError(
+                f"Pre MLP normalization dim {self.pre_mlp_norm.input_dim} does not match"
+                f" the first normalization layer dim {model_dim}",
+            )
+        if self.mlp.model_dim != model_dim:
+            raise ValueError(
+                f"MLP up projection dim {self.mlp.up_projection.input_dim} does not match"
+                f" the first normalization layer dim {model_dim}",
+            )
+        if self.mlp.hidden_dim != self.mlp.down_projection.input_dim:
+            raise ValueError(
+                f"MLP down projection dim {self.mlp.down_projection.input_dim} does not match"
+                f" the up projection dim {self.mlp.hidden_dim}",
+            )
+
+    def __call__(
+        self,
+        inputs: Float[Array, "suffix_tokens channels"],
+        positional_embeddings: PositionalEmbeddings,
+        kv_cache: KVCacheLayer | None = None,
+        return_updated_kv_cache: bool = False,
+        return_activation_trace: bool = False,
+        length_without_padding: Int[Array, ""] | int | None = None,
+    ) -> DecoderLayerResult:
+        normalized_attention_inputs = vmap(self.pre_attention_norm, in_axes=0)(inputs)
+        attention_outputs, updated_kv_cache = self.attention(
+            normalized_attention_inputs,
+            positional_embeddings,
+            kv_cache=kv_cache,
+            return_updated_kv_cache=return_updated_kv_cache,
+            length_without_padding=length_without_padding,
+        )
+        if self.post_attention_norm is not None:
+            normalized_attention_outputs = vmap(self.post_attention_norm, in_axes=0)(attention_outputs)
+            mlp_inputs = inputs + normalized_attention_outputs
+        else:
+            normalized_attention_outputs = None
+            mlp_inputs = inputs + attention_outputs
+
+        normalized_mlp_inputs = vmap(self.pre_mlp_norm, in_axes=0)(mlp_inputs)
+        mlp_outputs = vmap(self.mlp, in_axes=0)(normalized_mlp_inputs)
+        if self.post_mlp_norm is not None:
+            normalized_mlp_outputs = vmap(self.post_mlp_norm, in_axes=0)(mlp_outputs)
+            outputs = mlp_inputs + normalized_mlp_outputs
+        else:
+            normalized_mlp_outputs = None
+            outputs = mlp_inputs + mlp_outputs
+
+        if return_activation_trace:
+            activation_trace = DecoderLayerActivationTrace(
+                inputs=inputs,
+                positional_embeddings=positional_embeddings,
+                kv_cache=kv_cache,
+                pre_attention_norm=normalized_attention_inputs,
+                attention=attention_outputs,
+                post_attention_norm=normalized_attention_outputs,
+                mlp_inputs=mlp_inputs,
+                pre_mlp_norm=normalized_mlp_inputs,
+                mlp=mlp_outputs,
+                post_mlp_norm=normalized_mlp_outputs,
+            )
+        else:
+            activation_trace = None
+
+        return DecoderLayerResult(
+            outputs=outputs,
+            updated_kv_cache=updated_kv_cache,
+            activation_trace=activation_trace,
+        )
+
+    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(
+            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),
+            mlp=self.mlp.export_weights(weight_layout),
+        )
+        if self.post_attention_norm is not None:
+            result["post_attention_norm"] = self.post_attention_norm.export_weights(weight_layout)
+        if self.post_mlp_norm is not None:
+            result["post_mlp_norm"] = self.post_mlp_norm.export_weights(weight_layout)
+        return result

lalamo/modules/embedding.py

@@ -0,0 +1,299 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+
+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.quantization import QuantizationMode, dynamically_quantize_activations, quantize_weights
+
+from .common import LalamoModule, WeightLayout, register_config_union
+from .utils import apply_soft_capping
+
+__all__ = [
+    "EmbeddingBase",
+    "EmbeddingConfig",
+    "QuantizedTiedEmbedding",
+    "QuantizedTiedEmbeddingConfig",
+    "TiedEmbedding",
+    "TiedEmbeddingConfig",
+    "UntiedEmbedding",
+    "UntiedEmbeddingConfig",
+]
+
+
+@dataclass(frozen=True)
+class EmbeddingConfigBase:
+    input_scale: float | None
+    logits_soft_cap: float | None
+
+    @abstractmethod
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "EmbeddingBase": ...
+
+
+class EmbeddingBase[ConfigT: EmbeddingConfigBase](LalamoModule[ConfigT]):
+    @abstractmethod
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]: ...
+
+    @abstractmethod
+    def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]: ...
+
+    @property
+    @abstractmethod
+    def vocab_size(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def model_dim(self) -> int: ...
+
+    @classmethod
+    def _default_weight_layout(cls) -> WeightLayout:
+        return WeightLayout.INPUT_OUTPUT
+
+    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
+
+    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:
+            logits = apply_soft_capping(logits, self.config.logits_soft_cap)
+        return logits
+
+
+@dataclass(frozen=True)
+class TiedEmbeddingConfig(EmbeddingConfigBase):
+    precision: DTypeLike
+
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "TiedEmbedding":
+        weights = jax.random.normal(key, (vocab_size, model_dim), dtype=self.precision)
+        return TiedEmbedding(config=self, weights=weights)
+
+
+class TiedEmbedding(EmbeddingBase[TiedEmbeddingConfig]):
+    weights: Float[Array, "vocabulary channels"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.precision
+
+    def __post_init__(self) -> None:
+        if self.config.precision != self.weights.dtype:
+            raise ValueError(
+                f"Embedding dtype {self.weights.dtype} does not match the specified precision {self.config.precision}",
+            )
+
+    @property
+    def model_dim(self) -> int:
+        _, model_dim = self.weights.shape
+        return model_dim
+
+    @property
+    def vocab_size(self) -> int:
+        vocab_size, _ = self.weights.shape
+        return vocab_size
+
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]:
+        return self.weights
+
+    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)
+
+
+@dataclass(frozen=True)
+class UntiedEmbeddingConfig(EmbeddingConfigBase):
+    precision: DTypeLike
+
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "UntiedEmbedding":
+        input_key, output_key = jax.random.split(key)
+        input_weights = jax.random.normal(input_key, (vocab_size, model_dim), dtype=self.precision)
+        output_weights = jax.random.normal(output_key, (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"]
+    output_weights: Float[Array, "vocabulary channels"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.precision
+
+    @property
+    def model_dim(self) -> int:
+        _, model_dim = self.input_weights.shape
+        return model_dim
+
+    @property
+    def vocab_size(self) -> int:
+        vocab_size, _ = self.input_weights.shape
+        return vocab_size
+
+    def __post_init__(self) -> None:
+        if self.config.precision != self.input_weights.dtype:
+            raise ValueError(
+                f"Embedding dtype {self.input_weights.dtype} does not match",
+                f" the specified precision {self.config.precision}",
+            )
+        if self.config.precision != self.output_weights.dtype:
+            raise ValueError(
+                f"Embedding dtype {self.output_weights.dtype} does not match"
+                f" the specified precision {self.config.precision}",
+            )
+        input_vocab_size, input_model_dim = self.input_weights.shape
+        output_vocab_size, output_model_dim = self.output_weights.shape
+        if input_vocab_size != output_vocab_size:
+            raise ValueError(
+                f"Input vocab size {input_vocab_size} does not match the output vocab size {output_vocab_size}",
+            )
+        if input_model_dim != output_model_dim:
+            raise ValueError(
+                f"Input model dim {input_model_dim} does not match the output model dim {output_model_dim}",
+            )
+
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]:
+        return self.input_weights
+
+    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()
+
+        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,
+        )
+
+
+@dataclass(frozen=True)
+class QuantizedTiedEmbeddingConfig(EmbeddingConfigBase):
+    embedding_quantization_mode: QuantizationMode
+    activation_quantization_mode: QuantizationMode | None
+    activation_precision: DTypeLike
+
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "QuantizedTiedEmbedding":
+        min_val, max_val = self.embedding_quantization_mode.range
+        min_abs_val = min(abs(min_val), abs(max_val))
+        scale = 1 / min_abs_val
+        scales = scale * jnp.ones(vocab_size, dtype=self.activation_precision)
+        weights = jax.random.normal(key, (vocab_size, model_dim), dtype=self.activation_precision)
+        weights = quantize_weights(weights * min_abs_val, self.embedding_quantization_mode)
+        return QuantizedTiedEmbedding(config=self, weights=weights, scales=scales)
+
+
+class QuantizedTiedEmbedding(EmbeddingBase[QuantizedTiedEmbeddingConfig]):
+    weights: Float[Array, "vocabulary channels"]
+    scales: Float[Array, " vocabulary"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.activation_precision
+
+    @property
+    def model_dim(self) -> int:
+        _, model_dim = self.weights.shape
+        return model_dim
+
+    @property
+    def vocab_size(self) -> int:
+        vocab_size, _ = self.weights.shape
+        return vocab_size
+
+    def __post_init__(self) -> None:
+        if self.weights.dtype != self.config.activation_precision:
+            raise ValueError(
+                f"Embedding dtype ({self.scales.dtype}) is not equal to specified activation precision"
+                f" ({self.config.activation_precision})."
+                " Quantized layers require parameter dtypes to be equal to the activation precision.",
+            )
+        if self.scales.dtype != self.config.activation_precision:
+            raise ValueError(
+                f"Scales dtype {self.scales.dtype} does not match the specified activation precision"
+                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
+        (scales_vocab_size,) = self.scales.shape
+        if weights_vocab_size != scales_vocab_size:
+            raise ValueError(
+                f"Embedding vocab size {weights_vocab_size} does not match"
+                f" the scales dimension size {scales_vocab_size}",
+            )
+
+    @property
+    def int_weights(self) -> Int[Array, "vocabulary channels"]:
+        result = quantize_weights(self.weights, self.config.embedding_quantization_mode)
+        return result.astype(self.config.embedding_quantization_mode.dtype)
+
+    def _prepare_weights(self) -> Float[Array, "vocabulary channels"]:
+        quantized_weights = quantize_weights(self.weights, self.config.embedding_quantization_mode)
+        quantized_weights = quantized_weights * self.scales.reshape(-1, 1)
+        return quantized_weights
+
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]:
+        return self._prepare_weights()
+
+    def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
+        return self._prepare_weights()
+
+    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,
+        )
+
+
+EmbeddingConfig = TiedEmbeddingConfig | UntiedEmbeddingConfig | QuantizedTiedEmbeddingConfig
+
+
+register_config_union(EmbeddingConfig)