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

lalamo/modules/attention.py

@@ -0,0 +1,326 @@
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import equinox as eqx
+import jax
+from einops import einsum, rearrange, repeat
+from jax import numpy as jnp
+from jax import vmap
+from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
+
+from lalamo.common import ParameterDict
+from lalamo.modules.normalization import RMSNorm, RMSNormConfig
+
+from .common import AttentionType, LalamoModule, WeightLayout
+from .kv_cache import DynamicKVCacheLayer, KVCacheLayer, StaticKVCacheLayer
+from .linear import LinearBase, LinearConfig
+from .rope import PositionalEmbeddings
+from .utils import apply_soft_capping
+
+__all__ = [
+    "Attention",
+    "AttentionConfig",
+]
+
+
+def _repeat_kv(
+    keys_or_values: Float[Array, "tokens groups channels"],
+    group_size: int,
+) -> Float[Array, "tokens groups*group_size channels"]:
+    return repeat(
+        keys_or_values,
+        "tokens groups channels -> tokens (groups group_size) channels",
+        group_size=group_size,
+    )
+
+
+def _soft_capped_attention_kernel(
+    queries: Float[Array, "dst_tokens heads head_channels"],
+    keys: Float[Array, "src_tokens groups head_channels"],
+    values: Float[Array, "src_tokens groups head_channels"],
+    mask: Bool[Array, "dst_tokens src_tokens"] | None,
+    scale: float | None,
+    logit_soft_cap: float,
+) -> Float[Array, "dst_tokens heads head_channels"]:
+    dst_length, num_heads, head_dim = queries.shape
+    src_length, num_groups, _ = keys.shape
+    if scale is None:
+        scale = head_dim**-0.5
+    group_size = num_heads // num_groups
+    keys = _repeat_kv(keys, group_size)
+    values = _repeat_kv(values, group_size)
+    queries_head_first = rearrange(queries, "dst_tokens heads channels -> heads dst_tokens channels")
+    keys_head_first = rearrange(keys, "src_tokens heads channels -> heads src_tokens channels")
+    attention_logits = einsum(
+        queries_head_first,
+        keys_head_first,
+        "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 = attention_logits * scale
+    attention_logits = apply_soft_capping(attention_logits, logit_soft_cap)
+    attention_weights = jax.nn.softmax(attention_logits, axis=-1)
+    return einsum(
+        attention_weights,
+        values,
+        "heads dst_tokens src_tokens, src_tokens heads channels -> dst_tokens heads channels",
+    )
+
+
+class AttentionResult(NamedTuple):
+    outputs: Float[Array, "suffix_tokens channels"]
+    kv_cache: KVCacheLayer | None = None
+
+
+@dataclass(frozen=True)
+class AttentionConfig:
+    qkv_projection_config: LinearConfig
+    out_projection_config: LinearConfig
+
+    query_norm_config: RMSNormConfig | None
+    key_norm_config: RMSNormConfig | None
+
+    logit_soft_cap: float | None
+    has_qkv_biases: bool
+    has_out_biases: bool
+
+    def random_init(
+        self,
+        model_dim: int,
+        num_heads: int,
+        num_groups: int,
+        head_dim: int,
+        is_causal: bool,
+        scale: float | None,
+        sliding_window_size: int | None,
+        *,
+        key: PRNGKeyArray,
+    ) -> "Attention":
+        qkv_key, out_key = jax.random.split(key)
+        qkv_projection = self.qkv_projection_config.random_init(
+            input_dim=model_dim,
+            output_dims=(
+                num_heads * head_dim,
+                num_groups * head_dim,
+                num_groups * head_dim,
+            ),
+            has_biases=self.has_qkv_biases,
+            key=qkv_key,
+        )
+        out_projection = self.out_projection_config.random_init(
+            num_heads * head_dim,
+            (model_dim,),
+            has_biases=self.has_out_biases,
+            key=out_key,
+        )
+
+        if self.query_norm_config is not None:
+            query_norm = self.query_norm_config.init(
+                channels=head_dim,
+            )
+        else:
+            query_norm = None
+
+        if self.key_norm_config is not None:
+            key_norm = self.key_norm_config.init(
+                channels=head_dim,
+            )
+        else:
+            key_norm = None
+
+        return Attention(
+            self,
+            qkv_projection=qkv_projection,
+            out_projection=out_projection,
+            query_norm=query_norm,
+            key_norm=key_norm,
+            num_heads=num_heads,
+            num_groups=num_groups,
+            head_dim=head_dim,
+            is_causal=is_causal,
+            scale=scale,
+            sliding_window_size=sliding_window_size,
+        )
+
+
+class Attention(LalamoModule[AttentionConfig]):
+    qkv_projection: LinearBase
+    out_projection: LinearBase
+
+    query_norm: RMSNorm | None
+    key_norm: RMSNorm | None
+
+    num_heads: int = eqx.field(static=True)
+    num_groups: int = eqx.field(static=True)
+    head_dim: int = eqx.field(static=True)
+
+    is_causal: bool = eqx.field(static=True)
+
+    scale: float | None = eqx.field(static=True)
+    sliding_window_size: int | None = eqx.field(static=True)
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.qkv_projection.activation_precision
+
+    @property
+    def model_dim(self) -> int:
+        return self.qkv_projection.input_dim
+
+    @property
+    def group_size(self) -> int:
+        return self.num_heads // self.num_groups
+
+    @property
+    def use_sliding_window(self) -> bool:
+        return self.sliding_window_size is not None
+
+    @property
+    def attention_type(self) -> AttentionType:
+        return AttentionType.SLIDING_WINDOW if self.sliding_window_size is not None else AttentionType.GLOBAL
+
+    def __post_init__(self) -> None:
+        if self.qkv_projection.has_biases != self.config.has_qkv_biases:
+            raise ValueError(
+                f"QKV projection has_biases {self.qkv_projection.has_biases} does not match"
+                f" the specified config has_qkv_biases {self.config.has_qkv_biases}",
+            )
+        if self.out_projection.has_biases != self.config.has_out_biases:
+            raise ValueError(
+                f"Output projection has_biases {self.out_projection.has_biases} does not match"
+                f" the specified config has_out_biases {self.config.has_out_biases}",
+            )
+        if self.query_norm is not None and self.query_norm.input_dim != self.head_dim:
+            raise ValueError(
+                f"Query normalization input dimension must match head_dim ({self.head_dim}),"
+                f" got {self.query_norm.input_dim}",
+            )
+        if self.key_norm is not None and self.key_norm.input_dim != self.head_dim:
+            raise ValueError(
+                f"Key normalization input dimension must match head_dim ({self.head_dim}),"
+                f" got {self.key_norm.input_dim}",
+            )
+        if self.num_heads % self.num_groups != 0:
+            raise ValueError(
+                "Number of heads must be divisible by the number of groups,"
+                f" got {self.num_heads} heads and {self.num_groups} groups",
+            )
+        if self.out_projection.input_dim != self.num_heads * self.head_dim:
+            raise ValueError(
+                f"Output projection input dimension must be num_heads * head_dim"
+                f" ({self.num_heads} * {self.head_dim} = {self.num_heads * self.head_dim}),"
+                f" got {self.out_projection.input_dim}",
+            )
+        q_output_dim, k_output_dim, v_output_dim = self.qkv_projection.output_dims
+        if q_output_dim != self.num_heads * self.head_dim:
+            raise ValueError(
+                f"Query projection output dimension must be num_heads * head_dim"
+                f" ({self.num_heads} * {self.head_dim} = {self.num_heads * self.head_dim}),"
+                f" got {q_output_dim}",
+            )
+        if k_output_dim != self.num_groups * self.head_dim:
+            raise ValueError(
+                f"Key projection output dimension must be num_groups * head_dim"
+                f" ({self.num_groups} * {self.head_dim} = {self.num_groups * self.head_dim}),"
+                f" got {k_output_dim}",
+            )
+        if v_output_dim != self.num_groups * self.head_dim:
+            raise ValueError(
+                f"Value projection output dimension must be num_groups * head_dim"
+                f" ({self.num_groups} * {self.head_dim} = {self.num_groups * self.head_dim}),"
+                f" got {v_output_dim}",
+            )
+
+    def __call__(
+        self,
+        inputs: Float[Array, "suffix_tokens channels"],
+        positional_embeddings: PositionalEmbeddings,
+        kv_cache: KVCacheLayer | None = None,
+        return_updated_kv_cache: bool = False,
+        length_without_padding: Int[Array, ""] | int | None = None,
+    ) -> AttentionResult:
+        queries, keys, values = vmap(self.qkv_projection, in_axes=0)(inputs)
+        queries = rearrange(
+            queries,
+            "tokens (heads head_channels) -> tokens heads head_channels",
+            heads=self.num_heads,
+            head_channels=self.head_dim,
+        )
+        keys = rearrange(
+            keys,
+            "tokens (groups head_channels) -> tokens groups head_channels",
+            groups=self.num_groups,
+            head_channels=self.head_dim,
+        )
+        values = rearrange(
+            values,
+            "tokens (groups head_channels) -> tokens groups head_channels",
+            groups=self.num_groups,
+            head_channels=self.head_dim,
+        )
+
+        if self.query_norm is not None:
+            queries = vmap(vmap(self.query_norm))(queries)
+        if self.key_norm is not None:
+            keys = vmap(vmap(self.key_norm))(keys)
+
+        apply_positional_embeddings = vmap(positional_embeddings.apply, in_axes=1, out_axes=1)
+        queries = apply_positional_embeddings(queries)
+        keys = apply_positional_embeddings(keys)
+
+        if kv_cache is None:
+            updated_kv_cache = DynamicKVCacheLayer.init(keys, values, length=length_without_padding)
+        else:
+            updated_kv_cache = kv_cache.extend(keys, values, added_length=length_without_padding)
+
+        num_suffix_tokens, _, _ = queries.shape
+        mask = updated_kv_cache.attention_mask(num_suffix_tokens, self.is_causal, self.sliding_window_size)
+
+        if self.config.logit_soft_cap is not None:
+            attention_output = _soft_capped_attention_kernel(
+                queries,
+                updated_kv_cache.keys,
+                updated_kv_cache.values,
+                mask=mask,
+                scale=self.scale,
+                logit_soft_cap=self.config.logit_soft_cap,
+            )
+        else:
+            attention_output = jax.nn.dot_product_attention(
+                queries,
+                updated_kv_cache.keys,
+                updated_kv_cache.values,
+                mask=mask,
+                scale=self.scale,
+            )
+        attention_output = rearrange(
+            attention_output,
+            "tokens heads head_channels -> tokens (heads head_channels)",
+            heads=self.num_heads,
+            head_channels=self.head_dim,
+        )
+        (result,) = vmap(self.out_projection, in_axes=0)(attention_output)
+
+        if not return_updated_kv_cache:
+            updated_kv_cache = None
+
+        return AttentionResult(
+            outputs=result,
+            kv_cache=updated_kv_cache,
+        )
+
+    def init_static_kv_cache(self, capacity: int) -> StaticKVCacheLayer:
+        return StaticKVCacheLayer.empty(capacity, self.num_groups, self.head_dim, self.activation_precision)
+
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
+        result = ParameterDict(
+            qkv_projection=self.qkv_projection.export_weights(weight_layout),
+            out_projection=self.out_projection.export_weights(weight_layout),
+        )
+        if self.query_norm is not None:
+            result["query_norm"] = self.query_norm.export_weights(weight_layout)
+        if self.key_norm is not None:
+            result["key_norm"] = self.key_norm.export_weights(weight_layout)
+        return result

lalamo/modules/common.py

@@ -0,0 +1,133 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from types import UnionType
+
+import equinox as eqx
+from cattrs import Converter
+from jax import numpy as jnp
+from jaxtyping import DTypeLike
+
+from lalamo.common import ParameterDict
+
+__all__ = [
+    "AttentionType",
+    "DummyUnionMember",
+    "LalamoModule",
+    "config_converter",
+    "register_config_union",
+]
+
+
+class WeightLayout(Enum):
+    AUTO = "auto"
+    INPUT_OUTPUT = "input_output"
+    OUTPUT_INPUT = "output_input"
+
+    def __str__(self) -> str:
+        match self:
+            case WeightLayout.AUTO:
+                return "auto"
+            case WeightLayout.INPUT_OUTPUT:
+                return "(input, output)"
+            case WeightLayout.OUTPUT_INPUT:
+                return "(output, input)"
+
+
+class AttentionType(Enum):
+    GLOBAL = "global"
+    SLIDING_WINDOW = "sliding_window"
+
+
+class LalamoModule[ConfigT](eqx.Module):
+    config: ConfigT = eqx.field(static=True)
+
+    @property
+    @abstractmethod
+    def activation_precision(self) -> DTypeLike: ...
+
+    @abstractmethod
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict: ...
+
+
+def _dtype_to_str(dtype: DTypeLike) -> str:
+    if dtype == jnp.bfloat16:
+        return "bfloat16"
+    try:
+        return str(dtype.dtype)  # type: ignore
+    except AttributeError:
+        return str(dtype)
+
+
+def _str_to_dtype(dtype_str: str) -> jnp.dtype:
+    return {
+        "int4": jnp.int4,
+        "int8": jnp.int8,
+        "int16": jnp.int16,
+        "int32": jnp.int32,
+        "int64": jnp.int64,
+        "bfloat16": jnp.bfloat16,
+        "float16": jnp.float16,
+        "float32": jnp.float32,
+        "float64": jnp.float64,
+    }[dtype_str]
+
+
+config_converter = Converter()
+
+
+config_converter.register_unstructure_hook_func(
+    lambda t: t in [jnp.dtype, DTypeLike],
+    _dtype_to_str,
+)
+
+config_converter.register_structure_hook_func(
+    lambda t: t in [jnp.dtype, DTypeLike],
+    lambda s, _: _str_to_dtype(s),
+)
+
+
+def register_config_union(union_type: UnionType) -> None:
+    union_members = union_type.__args__
+    name_to_type = {m.__name__: m for m in union_members}
+
+    def unstructure(obj: object) -> dict | None:
+        if obj is None:
+            return None
+        return {
+            "type": obj.__class__.__name__,
+            **config_converter.unstructure(obj),
+        }
+
+    config_converter.register_unstructure_hook(
+        union_type,
+        unstructure,
+    )
+
+    config_converter.register_unstructure_hook(
+        union_type | None,
+        unstructure,
+    )
+
+    def structure[T](config: dict | None, _: type[T]) -> T | None:
+        if config is None:
+            return None
+        new_config = dict(config)
+        type_name = new_config.pop("type")
+        target_type = name_to_type[type_name]
+        return name_to_type[type_name](**config_converter.structure(new_config, target_type))
+
+    config_converter.register_structure_hook(
+        union_type,
+        structure,
+    )
+
+    config_converter.register_structure_hook(
+        union_type | None,
+        structure,
+    )
+
+
+@dataclass
+class DummyUnionMember:
+    pass

lalamo/modules/decoder.py

@@ -0,0 +1,244 @@
+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 .common import AttentionType, LalamoModule, WeightLayout
+from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerResult
+from .embedding import EmbeddingBase, EmbeddingConfig
+from .kv_cache import KVCache
+from .normalization import RMSNorm, RMSNormConfig
+from .rope import PositionalEmbeddings, RoPE, RoPEConfig
+
+__all__ = [
+    "Decoder",
+    "DecoderActivationTrace",
+    "DecoderConfig",
+    "DecoderResult",
+]
+
+
+class DecoderActivationTrace(eqx.Module):
+    token_ids: Int[Array, " suffix_tokens"]
+    token_positions: Int[Array, " suffix_tokens"]
+    kv_cache: KVCache | None
+
+    local_positional_embeddings: PositionalEmbeddings
+    global_positional_embeddings: PositionalEmbeddings
+
+    layer_results: tuple[DecoderLayerResult, ...]
+
+    output_norm: Float[Array, "suffix_tokens channels"]
+
+    def export(self) -> ParameterDict:
+        result = ParameterDict(
+            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]
+        return result
+
+
+class DecoderResult(eqx.Module):
+    logits: Float[Array, "suffix_tokens channels"]
+    updated_kv_cache: KVCache | None = None
+    activation_trace: DecoderActivationTrace | None = None
+
+    def export(self) -> ParameterDict:
+        result = ParameterDict(
+            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.activation_trace is not None:
+            result["activation_trace"] = self.activation_trace.export()
+        return result
+
+
+@dataclass(frozen=True)
+class DecoderConfig:
+    embedding_config: EmbeddingConfig
+    global_rope_config: RoPEConfig
+    local_rope_config: RoPEConfig | None
+    layer_config: 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,
+        *,
+        key: PRNGKeyArray,
+    ) -> "Decoder":
+        embedding_key, layers_key = jax.random.split(key)
+        embedding = self.embedding_config.random_init(
+            vocab_size=self.vocab_size,
+            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,
+        )
+
+        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_keys = jax.random.split(layers_key, self.num_layers)
+        layers = tuple(
+            self.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)
+        )
+        output_norm = self.output_norm_config.init(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
+    global_rope: RoPE
+    local_rope: RoPE | None
+    layers: tuple[DecoderLayer, ...]
+    output_norm: RMSNorm
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.embedding.activation_precision
+
+    def __call__(
+        self,
+        token_ids: Int[Array, " suffix_tokens"],
+        token_positions: Int[Array, " suffix_tokens"],
+        kv_cache: KVCache | None = None,
+        return_updated_kv_cache: bool = False,
+        return_activation_trace: bool = False,
+        length_without_padding: Int[Array, ""] | int | None = None,
+    ) -> DecoderResult:
+        maybe_kv_cache = kv_cache or ([None] * len(self.layers))
+        inner_features = self.embedding.embed(token_ids)
+
+        global_positional_embeddings = self.global_rope(token_positions)
+        if self.local_rope is not None:
+            local_positional_embeddings = self.local_rope(token_positions)
+        else:
+            local_positional_embeddings = global_positional_embeddings
+
+        updated_kv_cache_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
+
+            layer_result = layer(
+                inner_features,
+                positional_embeddings_to_use,
+                kv_cache=kv_cache_slice,
+                return_updated_kv_cache=return_updated_kv_cache,
+                return_activation_trace=return_activation_trace,
+                length_without_padding=length_without_padding,
+            )
+            inner_features = layer_result.outputs
+            layer_results.append(layer_result)
+            updated_kv_cache_layers.append(layer_result.updated_kv_cache)
+
+        normalized_outputs = vmap(self.output_norm, in_axes=0)(inner_features)
+        logits = vmap(self.embedding.readout, in_axes=0)(normalized_outputs)
+
+        if return_activation_trace:
+            activation_trace = DecoderActivationTrace(
+                token_ids=token_ids,
+                token_positions=token_positions,
+                kv_cache=kv_cache,
+                global_positional_embeddings=global_positional_embeddings,
+                local_positional_embeddings=local_positional_embeddings,
+                layer_results=tuple(layer_results),
+                output_norm=normalized_outputs,
+            )
+        else:
+            activation_trace = None
+
+        if return_updated_kv_cache:
+            updated_kv_cache = KVCache(updated_kv_cache_layers)
+        else:
+            updated_kv_cache = None
+
+        return DecoderResult(
+            logits=logits,
+            updated_kv_cache=updated_kv_cache,
+            activation_trace=activation_trace,
+        )
+
+    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(
+            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],
+            output_norm=self.output_norm.export_weights(weight_layout),
+        )
+        if self.local_rope:
+            result["local_rope"] = self.local_rope.export_weights(weight_layout)
+        return result