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

lalamo/__init__.py

@@ -1,7 +1,7 @@
 from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
 from lalamo.modules import Decoder
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 __all__ = [
     "REPO_TO_MODEL",

lalamo/model_import/__init__.py

@@ -0,0 +1,8 @@
+from .common import REPO_TO_MODEL, ModelMetadata, ModelSpec, import_model
+
+__all__ = [
+    "REPO_TO_MODEL",
+    "ModelMetadata",
+    "ModelSpec",
+    "import_model",
+]

lalamo/model_import/common.py

@@ -0,0 +1,111 @@
+import importlib.metadata
+from dataclasses import dataclass
+from pathlib import Path
+from typing import NamedTuple
+
+import huggingface_hub
+import jax.numpy as jnp
+from jaxtyping import DTypeLike
+
+from lalamo.modules import Decoder, DecoderConfig
+from lalamo.quantization import QuantizationMode
+
+from .model_specs import REPO_TO_MODEL, ModelSpec, UseCase
+
+__all__ = [
+    "REPO_TO_MODEL",
+    "ModelMetadata",
+    "ModelSpec",
+    "import_model",
+]
+
+
+LALAMO_VERSION = importlib.metadata.version("lalamo")
+
+
+@dataclass(frozen=True)
+class ModelMetadata:
+    toolchain_version: str
+    vendor: str
+    family: str
+    name: str
+    size: str
+    quantization: QuantizationMode | None
+    repo: str
+    use_cases: tuple[UseCase, ...]
+    model_config: DecoderConfig
+    tokenizer_file_names: tuple[str, ...]
+
+
+def download_weights(model_spec: ModelSpec, output_dir: Path | str | None = None) -> list[Path]:
+    result = [
+        huggingface_hub.hf_hub_download(
+            repo_id=model_spec.repo,
+            local_dir=output_dir,
+            filename=filename,
+        )
+        for filename in model_spec.weights_file_names
+    ]
+    return [Path(path) for path in result]
+
+
+def download_config_file(model_spec: ModelSpec, output_dir: Path | str | None = None) -> Path:
+    result = huggingface_hub.hf_hub_download(
+        repo_id=model_spec.repo,
+        local_dir=output_dir,
+        filename=model_spec.config_file_name,
+    )
+    return Path(result)
+
+
+def download_tokenizer_files(model_spec: ModelSpec, output_dir: Path | str | None = None) -> tuple[Path, ...]:
+    result = [
+        huggingface_hub.hf_hub_download(
+            repo_id=tokenizer_file_spec.repo or model_spec.repo,
+            local_dir=output_dir,
+            filename=tokenizer_file_spec.filename,
+        )
+        for tokenizer_file_spec in model_spec.tokenizer_files
+    ]
+    return tuple(Path(path) for path in result)
+
+
+class ImportResults(NamedTuple):
+    model: Decoder
+    metadata: ModelMetadata
+    tokenizer_file_paths: tuple[Path, ...]
+
+
+def import_model(
+    model_spec: ModelSpec,
+    *,
+    context_length: int | None = None,
+    precision: DTypeLike | None = None,
+    accumulation_precision: DTypeLike = jnp.float32,
+) -> ImportResults:
+    foreign_config_file = download_config_file(model_spec)
+    foreign_config = model_spec.config_type.from_json(foreign_config_file)
+
+    tokenizer_file_paths = download_tokenizer_files(model_spec)
+    if precision is None:
+        precision = foreign_config.default_precision
+
+    weights_paths = download_weights(model_spec)
+    weights_dict = {}
+    for weights_path in weights_paths:
+        weights_dict.update(model_spec.weights_type.load(weights_path, precision))
+
+    model = foreign_config.load_model(context_length, precision, accumulation_precision, weights_dict)
+    metadata = ModelMetadata(
+        toolchain_version=LALAMO_VERSION,
+        vendor=model_spec.vendor,
+        family=model_spec.family,
+        name=model_spec.name,
+        size=model_spec.size,
+        quantization=model_spec.quantization,
+        repo=model_spec.repo,
+        use_cases=model_spec.use_cases,
+        model_config=model.config,
+        tokenizer_file_names=tuple(p.name for p in tokenizer_file_paths),
+    )
+    return ImportResults(model, metadata, tokenizer_file_paths)

lalamo/model_import/configs/__init__.py

@@ -0,0 +1,23 @@
+from .common import ForeignConfig
+from .executorch import ETLlamaConfig
+from .huggingface import (
+    HFGemma2Config,
+    HFGemma3Config,
+    HFGemma3TextConfig,
+    HFLlamaConfig,
+    HFMistralConfig,
+    HFQwen2Config,
+    HFQwen3Config,
+)
+
+__all__ = [
+    "ETLlamaConfig",
+    "ForeignConfig",
+    "HFGemma2Config",
+    "HFGemma3Config",
+    "HFGemma3TextConfig",
+    "HFLlamaConfig",
+    "HFMistralConfig",
+    "HFQwen2Config",
+    "HFQwen3Config",
+]

lalamo/model_import/configs/common.py

@@ -0,0 +1,62 @@
+import json
+from abc import abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar, Self
+
+import cattrs
+import jax
+from jaxtyping import Array, DTypeLike
+
+from lalamo.modules import Decoder, DecoderConfig
+
+__all__ = ["ForeignConfig"]
+
+
+@dataclass(frozen=True)
+class ForeignConfig:
+    _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
+    _converter.register_structure_hook(int | list[int], lambda v, _: v)
+
+    @property
+    @abstractmethod
+    def default_precision(self) -> DTypeLike: ...
+
+    @classmethod
+    def from_json(cls, json_path: Path | str) -> Self:
+        json_path = Path(json_path)
+        with open(json_path) as f:
+            config = json.load(f)
+        return cls._converter.structure(config, cls)
+
+    def to_json(self, json_path: Path | str) -> None:
+        json_path = Path(json_path)
+        with open(json_path, "w") as f:
+            json.dump(self._converter.unstructure(self), f, indent=2)
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> DecoderConfig:
+        raise NotImplementedError
+
+    @classmethod
+    def _load_weights(
+        cls,
+        model: Decoder,
+        weights_dict: dict[str, Array],
+    ) -> Decoder:
+        raise NotImplementedError
+
+    def load_model(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+        weights_dict: dict[str, Array],
+    ) -> Decoder:
+        config = self.to_decoder_config(context_length, activation_precision, accumulation_precision)
+        model = config.random_init(key=jax.random.PRNGKey(0))
+        return self._load_weights(model, weights_dict)

lalamo/model_import/configs/executorch.py

@@ -0,0 +1,166 @@
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from jaxtyping import Array, DTypeLike
+
+from lalamo.model_import.loaders import load_executorch
+from lalamo.modules import (
+    Activation,
+    AttentionConfig,
+    Decoder,
+    DecoderConfig,
+    DecoderLayerConfig,
+    LlamaRoPEConfig,
+    MLPConfig,
+    QLoRALinearConfig,
+    QuantizedTiedEmbeddingConfig,
+    RMSNormConfig,
+    UpcastMode,
+)
+from lalamo.quantization import QuantizationMode
+
+from .common import ForeignConfig
+
+__all__ = ["ETLlamaConfig"]
+
+
+# These parameters are not present in the config file, and are extracted from the executorch implementation
+LOW_FREQ_FACTOR = 1.0
+HIGH_FREQ_FACTOR = 4.0
+OLD_CONTEXT_LENGTH = 8192
+MAX_SEQUENCE_LENGTH = 8192 * 32
+
+ROPE_SCALING_FACTOR = 32.0
+
+EMBEDDING_QUANTIZATION_MODE = QuantizationMode.INT8
+ACTIVATION_QUANTIZATION_MODE = QuantizationMode.INT8
+WEIGHT_QUANTIZATION_MODE = QuantizationMode.UINT4
+
+
+@dataclass(frozen=True)
+class QuantizationConfig:
+    group_size: int
+
+
+@dataclass(frozen=True)
+class LoraConfig:
+    rank: int
+    scale: float
+
+
+@dataclass(frozen=True)
+class ExecutorchConfig(ForeignConfig):
+    @property
+    def default_precision(self) -> DTypeLike:
+        return jnp.bfloat16
+
+    @classmethod
+    def _load_weights(
+        cls,
+        model: Decoder,
+        weights_dict: dict[str, Array],
+    ) -> Decoder:
+        return load_executorch(model, weights_dict)
+
+
+@dataclass(frozen=True)
+class ETLlamaConfig(ExecutorchConfig):
+    dim: int
+    n_layers: int
+    n_heads: int
+    n_kv_heads: int
+    vocab_size: int
+    ffn_dim_multiplier: float
+    multiple_of: int
+    norm_eps: float
+    rope_theta: float
+    use_scaled_rope: bool
+    quantization_args: QuantizationConfig | None = None
+    lora_args: LoraConfig | None = None
+
+    def _find_hidden_size(self) -> int:
+        # Magic formula from executorch
+        size_candidate = int(8 / 3 * self.dim * self.ffn_dim_multiplier)
+        return size_candidate // self.multiple_of * self.multiple_of
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> DecoderConfig:
+        if self.lora_args is None:
+            raise ValueError("We only support QLoRA models for now.")
+
+        if self.quantization_args is None:
+            raise ValueError("Quantization arguments are required for QLoRA models.")
+
+        embedding_config = QuantizedTiedEmbeddingConfig(
+            input_scale=None,
+            logits_soft_cap=None,
+            embedding_quantization_mode=EMBEDDING_QUANTIZATION_MODE,
+            activation_quantization_mode=ACTIVATION_QUANTIZATION_MODE,
+            activation_precision=activation_precision,
+        )
+        rope_config = LlamaRoPEConfig(
+            precision=activation_precision,
+            base=self.rope_theta,
+            max_sequence_length=MAX_SEQUENCE_LENGTH,
+            scaling_factor=ROPE_SCALING_FACTOR,
+            original_context_length=OLD_CONTEXT_LENGTH,
+            low_frequency_factor=LOW_FREQ_FACTOR,
+            high_frequency_factor=HIGH_FREQ_FACTOR,
+        )
+        rmsnorm_config = RMSNormConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+        )
+        linear_config = QLoRALinearConfig(
+            group_size=self.quantization_args.group_size,
+            weight_quantization_mode=WEIGHT_QUANTIZATION_MODE,
+            activation_quantization_mode=ACTIVATION_QUANTIZATION_MODE,
+            activation_precision=activation_precision,
+            lora_rank=self.lora_args.rank,
+            lora_scale=self.lora_args.scale,
+        )
+        attention_config = AttentionConfig(
+            qkv_projection_config=linear_config,
+            out_projection_config=linear_config,
+            query_norm_config=None,
+            key_norm_config=None,
+            logit_soft_cap=None,
+            has_qkv_biases=False,
+            has_out_biases=False,
+        )
+        mlp_config = MLPConfig(
+            linear_config=linear_config,
+            activation=Activation.SILU,
+        )
+        decoder_layer_config = DecoderLayerConfig(
+            pre_attention_norm_config=rmsnorm_config,
+            attention_config=attention_config,
+            post_attention_norm_config=None,
+            pre_mlp_norm_config=rmsnorm_config,
+            mlp_config=mlp_config,
+            post_mlp_norm_config=None,
+        )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            global_rope_config=rope_config,
+            local_rope_config=None,
+            layer_config=decoder_layer_config,
+            output_norm_config=rmsnorm_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.dim,
+            hidden_dim=self._find_hidden_size(),
+            num_heads=self.n_heads,
+            num_groups=self.n_kv_heads,
+            head_dim=self.dim // self.n_heads,
+            attention_scale=None,
+            num_layers=self.n_layers,
+            sliding_window_sizes=None,
+            context_length=context_length or MAX_SEQUENCE_LENGTH,
+        )

lalamo/model_import/configs/huggingface/__init__.py

@@ -0,0 +1,18 @@
+from .common import HuggingFaceConfig
+from .gemma2 import HFGemma2Config
+from .gemma3 import HFGemma3Config, HFGemma3TextConfig
+from .llama import HFLlamaConfig
+from .mistral import HFMistralConfig
+from .qwen2 import HFQwen2Config
+from .qwen3 import HFQwen3Config
+
+__all__ = [
+    "HFGemma2Config",
+    "HFGemma3Config",
+    "HFGemma3TextConfig",
+    "HFLlamaConfig",
+    "HFMistralConfig",
+    "HFQwen2Config",
+    "HFQwen3Config",
+    "HuggingFaceConfig",
+]

lalamo/model_import/configs/huggingface/common.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from typing import Literal
+
+import jax.numpy as jnp
+from jaxtyping import Array, DTypeLike
+
+from lalamo.model_import.configs import ForeignConfig
+from lalamo.model_import.loaders import load_huggingface
+from lalamo.modules import Decoder
+
+__all__ = [
+    "HuggingFaceConfig",
+    "AWQQuantizationConfig",
+    "GPTQMetaConfig",
+    "GPTQQuantizationConfig"
+]
+
+
+@dataclass(frozen=True)
+class AWQQuantizationConfig:
+    backend: Literal["autoawq"] = "autoawq"
+    bits: Literal[4, 8] = 4
+    do_fuse: Literal[False] = False
+    exllama_config: None = None
+    fuse_max_seq_len: None = None
+    group_size: int = 128
+    modules_to_fuse: None = None
+    modules_to_not_convert: None = None
+    quant_method: Literal["awq"] = "awq"
+    version: Literal["gemm"] = "gemm"
+    zero_point: bool = True
+
+
+@dataclass(frozen=True)
+class GPTQMetaConfig:
+    damp_auto_increment: float
+    damp_percent: float
+    mse: float
+    quantizer: list[str]
+    static_groups: bool
+    true_sequential: bool
+    uri: str
+
+
+@dataclass(frozen=True)
+class GPTQQuantizationConfig:
+    bits: int
+    checkpoint_format: str
+    desc_act: bool
+    group_size: int
+    lm_head: bool
+    meta: GPTQMetaConfig
+    pack_dtype: str
+    quant_method: Literal["gptq"]
+    sym: bool
+
+
+@dataclass(frozen=True)
+class HuggingFaceConfig(ForeignConfig):
+    torch_dtype: Literal["bfloat16", "float16", "float32"]
+
+    @property
+    def default_precision(self) -> DTypeLike:
+        return jnp.dtype(self.torch_dtype)
+
+    @classmethod
+    def _load_weights(
+        cls,
+        model: Decoder,
+        weights_dict: dict[str, Array],
+    ) -> Decoder:
+        return load_huggingface(model, weights_dict)

lalamo/model_import/configs/huggingface/gemma2.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Literal
+
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    Activation,
+    AttentionConfig,
+    DecoderConfig,
+    DecoderLayerConfig,
+    FullPrecisionLinearConfig,
+    MLPConfig,
+    RMSNormConfig,
+    TiedEmbeddingConfig,
+    UnscaledRoPEConfig,
+    UpcastMode,
+)
+
+from .common import HuggingFaceConfig
+
+__all__ = ["HFGemma2Config"]
+
+
+@dataclass(frozen=True)
+class HFGemma2Config(HuggingFaceConfig):
+    architectures: list[Literal["Gemma2ForCausalLM"]]
+    attention_bias: bool
+    attention_dropout: float
+    attn_logit_softcapping: float
+    bos_token_id: int | list[int]
+    cache_implementation: Literal["hybrid"]
+    eos_token_id: int | list[int]
+    final_logit_softcapping: float
+    head_dim: int
+    hidden_act: Literal["gelu_pytorch_tanh"]
+    hidden_activation: Literal["gelu_pytorch_tanh"]
+    hidden_size: int
+    initializer_range: float
+    intermediate_size: int
+    max_position_embeddings: int
+    model_type: Literal["gemma2"]
+    num_attention_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    pad_token_id: int
+    query_pre_attn_scalar: float
+    rms_norm_eps: float
+    rope_theta: float
+    sliding_window: int
+    transformers_version: str
+    use_cache: bool
+    vocab_size: int
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+    ) -> DecoderConfig:
+        sliding_window_sizes = tuple(
+            self.sliding_window if not bool(i % 2) else None for i in range(self.num_hidden_layers)
+        )
+        embedding_input_scale = self.hidden_size**0.5
+        attention_scale = self.query_pre_attn_scalar**-0.5
+        embedding_config = TiedEmbeddingConfig(
+            input_scale=embedding_input_scale,
+            logits_soft_cap=self.final_logit_softcapping,
+            precision=activation_precision,
+        )
+        rope_config = UnscaledRoPEConfig(
+            precision=activation_precision,
+            base=self.rope_theta,
+            max_sequence_length=self.max_position_embeddings,
+        )
+        rmsnorm_config = RMSNormConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.rms_norm_eps,
+            scale_offset=1.0,
+            upcast_mode=UpcastMode.FULL_LAYER,
+        )
+        linear_config = FullPrecisionLinearConfig(
+            precision=activation_precision,
+        )
+        attention_config = AttentionConfig(
+            qkv_projection_config=linear_config,
+            out_projection_config=linear_config,
+            query_norm_config=None,
+            key_norm_config=None,
+            logit_soft_cap=self.attn_logit_softcapping,
+            has_qkv_biases=self.attention_bias,
+            has_out_biases=False,
+        )
+        mlp_config = MLPConfig(
+            linear_config=linear_config,
+            activation=Activation.GELU,
+        )
+        decoder_layer_config = DecoderLayerConfig(
+            pre_attention_norm_config=rmsnorm_config,
+            attention_config=attention_config,
+            post_attention_norm_config=rmsnorm_config,
+            pre_mlp_norm_config=rmsnorm_config,
+            mlp_config=mlp_config,
+            post_mlp_norm_config=rmsnorm_config,
+        )
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            global_rope_config=rope_config,
+            local_rope_config=None,
+            layer_config=decoder_layer_config,
+            output_norm_config=rmsnorm_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.hidden_size,
+            hidden_dim=self.intermediate_size,
+            num_heads=self.num_attention_heads,
+            num_groups=self.num_key_value_heads,
+            head_dim=self.head_dim,
+            attention_scale=attention_scale,
+            num_layers=self.num_hidden_layers,
+            sliding_window_sizes=sliding_window_sizes,
+            context_length=context_length or self.max_position_embeddings,
+        )