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

lalamo/modules/normalization.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from enum import Enum
+
+import jax
+from jax import numpy as jnp
+from jaxtyping import Array, DTypeLike, Float
+
+from lalamo.common import ParameterDict
+
+from .common import LalamoModule, WeightLayout
+
+__all__ = [
+    "RMSNorm",
+    "RMSNormConfig",
+    "UpcastMode",
+]
+
+
+class UpcastMode(Enum):
+    ONLY_NORMALIZATION = "only_normalization"
+    FULL_LAYER = "full_layer"
+
+
+@dataclass(frozen=True)
+class RMSNormConfig:
+    scale_precision: DTypeLike
+    accumulation_precision: DTypeLike
+    epsilon: float
+    scale_offset: float | None
+    upcast_mode: UpcastMode
+
+    def init(self, channels: int) -> "RMSNorm":
+        scales = jnp.ones(channels, dtype=self.scale_precision)
+        return RMSNorm(self, scales=scales)
+
+
+class RMSNorm(LalamoModule[RMSNormConfig]):
+    scales: Float[Array, " channels"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.scale_precision
+
+    @property
+    def input_dim(self) -> int:
+        (result,) = self.scales.shape
+        return result
+
+    def __post_init__(self) -> None:
+        if self.config.scale_precision != self.scales.dtype:
+            raise ValueError(
+                f"Scales precision {self.scales.dtype} does not match the"
+                f" specified precision {self.config.scale_precision}",
+            )
+
+    def __call__(self, inputs: Float[Array, " channels"]) -> Float[Array, " channels"]:
+        upcasted_inputs = inputs.astype(self.config.accumulation_precision)
+
+        adjusted_variance = jnp.mean(jnp.square(upcasted_inputs)) + self.config.epsilon
+        normalized_x = upcasted_inputs * jax.lax.rsqrt(adjusted_variance)
+
+        if self.config.upcast_mode == UpcastMode.ONLY_NORMALIZATION:
+            normalized_x = normalized_x.astype(inputs.dtype)
+
+        if self.config.upcast_mode == UpcastMode.FULL_LAYER:
+            adjusted_scales = self.scales.astype(self.config.accumulation_precision)
+        else:
+            adjusted_scales = self.scales
+
+        if self.config.scale_offset is not None:
+            adjusted_scales = adjusted_scales + self.config.scale_offset
+
+        result = normalized_x * adjusted_scales
+        return result.astype(inputs.dtype)
+
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
+        return ParameterDict(scales=self.scales)

lalamo/modules/rope.py

@@ -0,0 +1,255 @@
+# Based on https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py
+# Original PyTorch code copyright notice:
+#
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+
+import equinox as eqx
+from jax import numpy as jnp
+from jaxtyping import Array, DTypeLike, Float, Int
+
+from lalamo.common import ParameterDict
+
+from .common import LalamoModule, WeightLayout, register_config_union
+
+__all__ = [
+    "LinearScalingRoPEConfig",
+    "LlamaRoPEConfig",
+    "PositionalEmbeddings",
+    "RoPE",
+    "RoPEConfigBase",
+    "UnscaledRoPEConfig",
+    "YARNRoPEConfig",
+]
+
+
+class PositionalEmbeddings(eqx.Module):
+    cosines: Float[Array, "tokens head_channels"]
+    sines: Float[Array, "tokens head_channels"]
+
+    @property
+    def head_dim(self) -> int:
+        return self.cosines.shape[-1]
+
+    def rotate_half(self, heads: Float[Array, "tokens head_channels"]) -> Float[Array, "tokens head_channels"]:
+        x1 = heads[..., : self.head_dim // 2]
+        x2 = heads[..., self.head_dim // 2 :]
+        return jnp.concatenate((-x2, x1), axis=-1)
+
+    def apply(self, heads: Float[Array, "tokens head_channels"]) -> Float[Array, "tokens head_channels"]:
+        return heads * self.cosines + self.rotate_half(heads) * self.sines
+
+    def export(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
+        return ParameterDict(
+            cosines=self.cosines,
+            sines=self.sines,
+        )
+
+
+@dataclass(frozen=True)
+class RoPEConfigBase:
+    precision: DTypeLike
+    base: float
+    max_sequence_length: int
+
+    @property
+    def _attention_scaling_factor(self) -> float:
+        return 1.0
+
+    def _scale_inverse_frequencies(
+        self,
+        inverse_frequencies: Float[Array, " tokens"],
+        head_dim: int,  # noqa: ARG002
+        max_sequence_length: int,  # noqa: ARG002
+    ) -> Float[Array, " tokens"]:
+        return inverse_frequencies
+
+    def init(
+        self,
+        head_dim: int,
+        num_timesteps: int,
+    ) -> "RoPE":
+        timesteps = jnp.arange(num_timesteps, dtype=jnp.float32)
+        channel_indices = jnp.arange(0, head_dim, 2, dtype=jnp.int32)
+        inverse_frequencies = 1.0 / (self.base ** (channel_indices.astype(jnp.float32) / head_dim))
+        inverse_frequencies = self._scale_inverse_frequencies(inverse_frequencies, head_dim, self.max_sequence_length)
+        outer_inverse_frequencies = jnp.outer(timesteps, inverse_frequencies)
+        embeddings = jnp.concatenate((outer_inverse_frequencies, outer_inverse_frequencies), axis=-1)
+        cosines = (jnp.cos(embeddings) * self._attention_scaling_factor).astype(self.precision)
+        sines = (jnp.sin(embeddings) * self._attention_scaling_factor).astype(self.precision)
+        return RoPE(config=self, cosines=cosines, sines=sines)
+
+
+class RoPE(LalamoModule[RoPEConfigBase]):
+    sines: Float[Array, "tokens head_channels"]
+    cosines: Float[Array, "tokens head_channels"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.precision
+
+    def __post_init__(self) -> None:
+        if self.cosines.dtype != self.config.precision:
+            raise ValueError(
+                f"Cosines dtype {self.cosines.dtype} does not match the specified precision {self.config.precision}",
+            )
+        if self.sines.dtype != self.config.precision:
+            raise ValueError(
+                f"Sines dtype {self.sines.dtype} does not match the specified precision {self.config.precision}",
+            )
+        if self.cosines.shape != self.sines.shape:
+            raise ValueError(
+                f"Cosines and sines shape mismatch: cosines have shape {self.cosines.shape},"
+                f" while sines have shape {self.sines.shape}",
+            )
+
+    @property
+    def head_dim(self) -> int:
+        _, result = self.sines.shape
+        return result
+
+    @property
+    def max_sequence_length(self) -> int:
+        result, _ = self.sines.shape
+        return result
+
+    def __call__(self, timesteps: Int[Array, " tokens"]) -> PositionalEmbeddings:
+        return PositionalEmbeddings(
+            cosines=self.cosines[timesteps],
+            sines=self.sines[timesteps],
+        )
+
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
+        return ParameterDict(cosines=self.cosines, sines=self.sines)
+
+
+class UnscaledRoPEConfig(RoPEConfigBase):
+    pass
+
+
+@dataclass(frozen=True)
+class LlamaRoPEConfig(RoPEConfigBase):
+    scaling_factor: float
+    original_context_length: int
+    low_frequency_factor: float
+    high_frequency_factor: float
+
+    def _scale_inverse_frequencies(
+        self,
+        inverse_frequencies: Float[Array, " tokens"],
+        head_dim: int,  # noqa: ARG002
+        max_sequence_length: int,  # noqa: ARG002
+    ) -> Float[Array, " tokens"]:
+        low_frequency_wavelength = self.original_context_length / self.low_frequency_factor
+        high_frequency_wavelength = self.original_context_length / self.high_frequency_factor
+
+        wavelengths = 2 * math.pi / inverse_frequencies
+
+        high_frequency_mask = wavelengths < high_frequency_wavelength
+        low_frequency_mask = wavelengths > low_frequency_wavelength
+        mid_frequency_mask = (~high_frequency_mask) & (~low_frequency_mask)
+
+        smoothing_factors = self.original_context_length / wavelengths - self.low_frequency_factor
+        smoothing_factors = smoothing_factors / (self.high_frequency_factor - self.low_frequency_factor)
+
+        scaled_frequencies = inverse_frequencies / self.scaling_factor
+        smoothly_scaled_frequencies = (
+            smoothing_factors * inverse_frequencies + (1 - smoothing_factors) * scaled_frequencies
+        )
+
+        result = inverse_frequencies * high_frequency_mask.astype(jnp.float32)
+        result = result + smoothly_scaled_frequencies * mid_frequency_mask.astype(jnp.float32)
+        result = result + scaled_frequencies * low_frequency_mask.astype(jnp.float32)
+
+        return result
+
+
+@dataclass(frozen=True)
+class YARNRoPEConfig(RoPEConfigBase):
+    scaling_factor: float
+    beta_fast: float
+    beta_slow: float
+
+    @classmethod
+    def _find_correction_dim(cls, num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float:
+        """Inverse dimension formula to find the dimension based on the number of rotations"""
+        return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+    @classmethod
+    def _find_correction_range(
+        cls,
+        low_rot: float,
+        high_rot: float,
+        dim: int,
+        base: float,
+        max_position_embeddings: int,
+    ) -> tuple[int, int]:
+        """Find dimension range bounds based on rotations"""
+        low = math.floor(cls._find_correction_dim(low_rot, dim, base, max_position_embeddings))
+        high = math.ceil(cls._find_correction_dim(high_rot, dim, base, max_position_embeddings))
+        return max(low, 0), min(high, dim - 1)
+
+    @classmethod
+    def _linear_ramp_factor(cls, min_value: float, max_value: float, dim: int) -> Float[Array, " head_dim"]:
+        if min_value == max_value:
+            max_value += 0.001  # Prevent singularity
+
+        linear_func = (jnp.arange(dim, dtype=jnp.float32) - min_value) / (max_value - min_value)
+        ramp_func = jnp.clip(linear_func, 0, 1)
+        return ramp_func
+
+    def _scale_inverse_frequencies(
+        self,
+        inverse_frequencies: Float[Array, " tokens"],
+        head_dim: int,
+        max_sequence_length: int,
+    ) -> Float[Array, " tokens"]:
+        scaled_frequencies = inverse_frequencies / self.scaling_factor
+
+        low, high = self._find_correction_range(
+            self.beta_fast,
+            self.beta_slow,
+            head_dim,
+            self.base,
+            max_sequence_length,
+        )
+
+        # Get n-dimensional rotational scaling corrected for extrapolation
+        smoothing_factor = 1 - self._linear_ramp_factor(low, high, head_dim // 2)
+        return scaled_frequencies * (1 - smoothing_factor) + inverse_frequencies * smoothing_factor
+
+    @property
+    def attention_scaling_factor(self) -> float:
+        return 0.1 * math.log(self.scaling_factor) + 1.0
+
+
+@dataclass(frozen=True)
+class LinearScalingRoPEConfig(RoPEConfigBase):
+    scaling_factor: float
+
+    def _scale_inverse_frequencies(
+        self,
+        inverse_frequencies: Float[Array, " tokens"],
+        head_dim: int,  # noqa: ARG002
+        max_sequence_length: int,  # noqa: ARG002
+    ) -> Float[Array, " tokens"]:
+        return inverse_frequencies / self.scaling_factor
+
+
+RoPEConfig = UnscaledRoPEConfig | LlamaRoPEConfig | YARNRoPEConfig | LinearScalingRoPEConfig
+
+register_config_union(RoPEConfig)

lalamo/modules/utils.py

@@ -0,0 +1,13 @@
+import jax
+from jaxtyping import Array, Float
+
+__all__ = [
+    "apply_soft_capping",
+]
+
+
+def apply_soft_capping(
+    values: Float[Array, "*"],
+    soft_cap: float,
+) -> Float[Array, "dst_tokens src_tokens"]:
+    return jax.nn.tanh(values / soft_cap) * soft_cap

{lalamo-0.2.1.dist-info → lalamo-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.1
+Version: 0.2.2
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

lalamo-0.2.2.dist-info/RECORD

@@ -0,0 +1,53 @@
+lalamo/__init__.py,sha256=1n--wwvN86epEr8MSH_-qSZPiHPDNRg45XRCMg6aV0o,217
+lalamo/common.py,sha256=uYLw68V4AF3zlENG3KAIKRpOFXVHv8xX_n0cc3qJnj4,1877
+lalamo/language_model.py,sha256=GiA_BDQuYCgVBFHljb_ltW_M7g3I1Siwm111M3Jc8MM,9286
+lalamo/main.py,sha256=K2RLyTcxvBCP0teSsminssj_oUkuQAQ5y9ixa1uOqas,9546
+lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
+lalamo/utils.py,sha256=QzkT0_82nd9pS5p0e7yOOdL_ZeKQr_Ftj4kFrWF35R8,1754
+lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
+lalamo/model_import/common.py,sha256=sHXEGQUtVb6TRT5FOGtJG9pz1Ohy5v_LtunubVxZKqQ,3303
+lalamo/model_import/configs/__init__.py,sha256=AbVkVT4tiofvQKym2maTv0dazEbalVrqFZjxqyVzk5o,456
+lalamo/model_import/configs/common.py,sha256=MKAinEL7WXkijS3IrfiTRgx2l6otpnIaJG_CajosMCU,1803
+lalamo/model_import/configs/executorch.py,sha256=yOa3mdM2FK3xahRclqhDbpste1DAW2kps2f-GgKnrfc,5255
+lalamo/model_import/configs/huggingface/__init__.py,sha256=kWHUnZDwGQCbA3Ucm-FEDr8zZ2yZ3yviPVftlNgMk30,460
+lalamo/model_import/configs/huggingface/common.py,sha256=p6oEKIT2Ezh_d8eDXYzHaJaqjPriQrAzz2bkEq_HkgY,1698
+lalamo/model_import/configs/huggingface/gemma2.py,sha256=oIefI_ad-7DtzXmisFczkKPuOQ-KkzMkKWTk9likaMs,4101
+lalamo/model_import/configs/huggingface/gemma3.py,sha256=1tkkmEs4pF0t0XlDS2Z5mWzcPGrRWb7FcLgazrFDJy8,6434
+lalamo/model_import/configs/huggingface/llama.py,sha256=_vOalgc24uhMcPyCqyxWOZk80hXxqN-dhMHvBGtbIlc,5444
+lalamo/model_import/configs/huggingface/mistral.py,sha256=39qsX_Twml8C0xz0CayVZse2uaHJtKS9-54B8nQw_5k,4148
+lalamo/model_import/configs/huggingface/qwen2.py,sha256=GnO1_DKDewiB4AW8lJu_x30lL-GgB9GYc64rl6XqfYI,4963
+lalamo/model_import/configs/huggingface/qwen3.py,sha256=UJ-EP0geHmGXnT_Ioy7Z7V4vns_dKz2YpPe-GLPQg20,5029
+lalamo/model_import/loaders/__init__.py,sha256=tocl2MQcMH5mLBkkGwdufDjscDvzbQ24Qz37_vKz1_o,144
+lalamo/model_import/loaders/common.py,sha256=2FigeDMUwlMPUebX8DAK2Yh9aLgVtsfTj0S431p7A0o,1782
+lalamo/model_import/loaders/executorch.py,sha256=nSvpylK8QL3nBk78P3FabLoyA87E3kv5CCpMfvuZe6Q,8886
+lalamo/model_import/loaders/huggingface.py,sha256=Ze_qB0fSxY8lH4ovH0t8jd5jiteasUWkS9HdgMZXCrs,10523
+lalamo/model_import/model_specs/__init__.py,sha256=_sJthAH1xXl5B9JPhRqMVP2t5KkhzqmKFHSRlOiFg8s,915
+lalamo/model_import/model_specs/common.py,sha256=ygfNjwVZBrjNkCVuv66R1vy5hXjgbAJyDc0QJfRfgik,3789
+lalamo/model_import/model_specs/deepseek.py,sha256=9l3pVyC-ZoIaFG4xWhPDCbKkD2TsND286o0KzO0uxKo,788
+lalamo/model_import/model_specs/gemma.py,sha256=y4aDeaGGl4JPIanAgPMOlyfD_cx3Q7rpTKgDgx5AsX0,2299
+lalamo/model_import/model_specs/huggingface.py,sha256=ktDJ_qZxSGmHREydrYQaWi71bXJZiHqzHDoZeORENno,784
+lalamo/model_import/model_specs/llama.py,sha256=oPnHw8qV2l_cfQcW6OPTfehatP-ovLMPppIZVJ8yOWI,3234
+lalamo/model_import/model_specs/mistral.py,sha256=xDX2SyTruGR7A8LI_Ypa6qAP5nVyYhxLffoxS2F6bmI,1649
+lalamo/model_import/model_specs/pleias.py,sha256=zLRjmT6PXFtykqSYpaRtVObP306urMjF2J6dTKdAbQM,747
+lalamo/model_import/model_specs/polaris.py,sha256=TiGlXI3j7HP9bs01jdcysBNFxvNKnxTF30wuv5Jg2mQ,768
+lalamo/model_import/model_specs/qwen.py,sha256=dsCo3uaSPtPPjGuWHerFUY27f5Pv_HFOao2lPjFFHJI,11302
+lalamo/model_import/model_specs/reka.py,sha256=YtAuM52ImgH24lVuICXDcS39mNNzG_b-ouoAy5uVYLk,766
+lalamo/modules/__init__.py,sha256=iNzQL_qIG10U157bWmblj9fZNewup0O0aB8IsMXuBPU,2164
+lalamo/modules/activations.py,sha256=ZgUd3E4VTAVgCZaj9HhYkXiJuiKrWBzK6so5JGnucOc,532
+lalamo/modules/attention.py,sha256=ZvhQPBpsgZc-RDM1EQZgG6a8CZGYvMthE7g3qDedUVM,12158
+lalamo/modules/common.py,sha256=6FOmvICxVJTxLll785WY7KVY7ixAvPuzaW_J1trYNj0,3171
+lalamo/modules/decoder.py,sha256=Erc8k_tjrCCZOMBJM1bL0b96zFjaKDGIuFmtEp-OJVk,9043
+lalamo/modules/decoder_layer.py,sha256=gVVE48hlkNTvWJIt5oKak7VhEIrYdCaErf7_6_A_9ys,9443
+lalamo/modules/embedding.py,sha256=6xnNFi_TrB0ymSAmCLwmwjAZW0pchgzjQDjDws22PPw,10684
+lalamo/modules/kv_cache.py,sha256=GLZ84VTl2QtJYdnADR3j0e4CUmnSAlSvnCb63E8AtIk,7225
+lalamo/modules/linear.py,sha256=loUGFu3wx-iGqDqGMphQorhqBm7b9lAqT4B0jAmoamk,24087
+lalamo/modules/mlp.py,sha256=bV8qJTjsQFGv-CA7d32UQFn6BX5zmCKWC5pgm29-W3U,2631
+lalamo/modules/normalization.py,sha256=BWCHv6ycFJ_qMGfxkusGfay9dWzUlbpuwmjbLy2rI68,2380
+lalamo/modules/rope.py,sha256=Vdt2J_W0MPDK52nHsroLVCfWMHyHW3AfrKZCZAE4VYs,9369
+lalamo/modules/utils.py,sha256=5QTdi34kEI5jix7TfTdB0mOYZbzZUul_T1y8eWCA6lQ,262
+lalamo-0.2.2.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.2.2.dist-info/METADATA,sha256=6vrCJSOr_hGJeCWIA46DbL6OnoEL3rdK3xYrxeqeVRo,2611
+lalamo-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.2.2.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.2.2.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.2.2.dist-info/RECORD,,

lalamo-0.2.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-lalamo/__init__.py,sha256=uKBR6vAH2AmdpPqz1q2zVVwQyCpWRWUHAfm-uQg8DAM,217
-lalamo/common.py,sha256=uYLw68V4AF3zlENG3KAIKRpOFXVHv8xX_n0cc3qJnj4,1877
-lalamo/language_model.py,sha256=GiA_BDQuYCgVBFHljb_ltW_M7g3I1Siwm111M3Jc8MM,9286
-lalamo/main.py,sha256=K2RLyTcxvBCP0teSsminssj_oUkuQAQ5y9ixa1uOqas,9546
-lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
-lalamo/utils.py,sha256=QzkT0_82nd9pS5p0e7yOOdL_ZeKQr_Ftj4kFrWF35R8,1754
-lalamo-0.2.1.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.2.1.dist-info/METADATA,sha256=1qDWPQiCYK_EIeff-oiaF7VeIksGNdZ4nCFikHXGJR4,2611
-lalamo-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.2.1.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.2.1.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.2.1.dist-info/RECORD,,