Add SmolLM2 support (1.7B, 360M, 135M)

models/smollm2/README.md

@@ -0,0 +1,90 @@
+# SmolLM2
+
+HuggingFace's SmolLM2 for on-device inference via Core AI.
+
+## Supported Models
+
+| Model                 | Parameters | macOS | iOS |
+| --------------------- | ---------- | ----- | --- |
+| SmolLM2-1.7B-Instruct | 1.7B       | Yes   | Yes |
+| SmolLM2-360M-Instruct | 360M       | Yes   | Yes |
+| SmolLM2-135M-Instruct | 135M       | Yes   | Yes |
+
+## Setup to export models
+
+If you haven't installed `uv`, install it by
+```bash
+brew install uv
+```
+
+## Export models
+
+```bash
+# 1.7B (INT4 quantized with FP16 embedding)
+uv run coreai.llm.export HuggingFaceTB/SmolLM2-1.7B-Instruct
+
+# 360M (FP16)
+uv run coreai.llm.export HuggingFaceTB/SmolLM2-360M-Instruct
+
+# 135M (FP16)
+uv run coreai.llm.export HuggingFaceTB/SmolLM2-135M-Instruct
+```
+
+**Options:**
+
+```bash
+# Full precision
+uv run coreai.llm.export HuggingFaceTB/SmolLM2-1.7B-Instruct --compression none
+
+# iOS variant
+uv run coreai.llm.export HuggingFaceTB/SmolLM2-1.7B-Instruct --platform iOS
+```
+
+## Run a Core AI Language Model
+
+### In your iOS and macOS applications via Foundation Models
+
+```swift
+import FoundationModels
+import CoreAILanguageModels
+
+let model = try await CoreAILanguageModel(resourcesAt: modelURL)
+
+let session = LanguageModelSession(model: model)
+
+let response = try await session.respond(to: "What is quantum computing?")
+
+print(response)
+```
+
+### On your Mac using built-in Command Line Tool
+
+```bash
+swift run -c release llm-runner --model path/to/exported_model_folder --prompt "Hello"
+```
+
+## Benchmark a Core AI Language Model
+
+```bash
+swift run -c release llm-benchmark --model path/to/exported_model_folder
+```
+
+Defaults: 512 prompt tokens, 1024 generation tokens, 5 trials. Override with `-p`, `-g`, and `-n`.
+
+## Evaluation
+
+Perplexity score on the [`WikiText-2`](https://huggingface.co/datasets/EleutherAI/wikitext_document_level) dataset computed using the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/wikitext/README.md) with the Core AI PyTorch models.
+
+| Model | Compression                                | BPW   | Platform | Perplexity |
+| ----- | ------------------------------------------ | ----- | -------- | ---------- |
+| 1.7B  | none (`float16`)                           | 16.00 | macOS    | 12.27      |
+| 1.7B  | [INT4 with FP16 embedding][smollm2-yaml]   | 4.56  | macOS    | 14.17      |
+| 1.7B  | [6-bit palettized][smollm2-6bit-yaml]      | 6.00  | iOS      | 13.63      |
+| 360M  | none (`float16`)                           | 16.00 | macOS    | 18.23      |
+| 135M  | none (`float16`)                           | 16.00 | macOS    | 24.99      |
+
+The iOS 1.7B uses 6-bit palettization (64 centroids). 4-bit palettization (16 centroids)
+is too lossy for this model's weight distribution regardless of which layers are promoted.
+
+[smollm2-yaml]: smollm2_4bit_embedding_excluded.yaml
+[smollm2-6bit-yaml]: smollm2_1_7b_6bit.yaml

models/smollm2/smollm2_1_7b_6bit.yaml

@@ -0,0 +1,14 @@
+# 6-bit palettization for SmolLM2-1.7B iOS.
+# 64 centroids — middle ground between 4-bit (16) and 8-bit (256).
+kmeans_palettization_config:
+  global_config:
+    op_state_spec:
+      weight:
+        n_bits: 6
+        granularity:
+          type: per_grouped_channel
+          axis: 0
+          group_size: 8
+  module_type_configs:
+    torch.nn.modules.sparse.Embedding: null
+    coreai_models.primitives.ios.embedding.LoadEmbeddings: null

models/smollm2/smollm2_4bit_embedding_excluded.yaml

@@ -0,0 +1,21 @@
+# INT4 weight quantization with FP16 embedding/lm_head.
+# SmolLM2 ties these weights — INT4 on lm_head degrades generation quality.
+quantization_config:
+  execution_mode: eager
+  global_config:
+    op_state_spec:
+      weight:
+        dtype: int4
+        qscheme: symmetric_with_clipping
+        granularity:
+          type: per_block
+          block_size: 32
+          axis: 1
+    op_input_spec: null
+    op_output_spec: null
+  module_type_configs:
+    coreai_models.primitives.macos.sdpa.SDPA: null
+    coreai_models.primitives.macos.rope.RoPE: null
+    coreai_models.primitives.macos.rms_norm.RMSNorm: null
+    coreai_models.primitives.macos.rms_norm.RMSNormPlusOne: null
+    torch.nn.modules.sparse.Embedding: null

python/src/coreai_models/export/metadata.py

@@ -113,6 +113,33 @@ class AIModelMetadataFields:
             "Source: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1"
         ),
     ),
+    "HuggingFaceTB/SmolLM2-1.7B-Instruct": AIModelMetadataFields(
+        author="HuggingFace",
+        license="Apache-2.0",
+        model_description=(
+            "SmolLM2-1.7B-Instruct is a 1.7B-parameter instruction-tuned causal "
+            "language model from HuggingFace. "
+            "Source: https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct"
+        ),
+    ),
+    "HuggingFaceTB/SmolLM2-360M-Instruct": AIModelMetadataFields(
+        author="HuggingFace",
+        license="Apache-2.0",
+        model_description=(
+            "SmolLM2-360M-Instruct is a 360M-parameter instruction-tuned causal "
+            "language model from HuggingFace. "
+            "Source: https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct"
+        ),
+    ),
+    "HuggingFaceTB/SmolLM2-135M-Instruct": AIModelMetadataFields(
+        author="HuggingFace",
+        license="Apache-2.0",
+        model_description=(
+            "SmolLM2-135M-Instruct is a 135M-parameter instruction-tuned causal "
+            "language model from HuggingFace. "
+            "Source: https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct"
+        ),
+    ),
     "openai/gpt-oss-20b": AIModelMetadataFields(
         author="OpenAI",
         license="Apache-2.0",

python/src/coreai_models/export/pipeline.py

@@ -70,6 +70,8 @@ class ExportConfig:
     # QuantizerConfig) loaded from a user-provided YAML. When set, the pipeline
     # uses this directly and ignores `compression` for config resolution
     compression_config_object: Any = field(default=None, repr=False)
+    # Override HuggingFace model_type for registry lookup.
+    model_type_override: str | None = None
 
 
 def _generate_output_name(config: ExportConfig) -> str:
@@ -135,7 +137,7 @@ async def _async_export_model(config: ExportConfig) -> str:
 
     # ---- 1. Resolve model class ----
     hf_config = AutoConfig.from_pretrained(config.hf_model_id)
-    model_type = getattr(hf_config, "model_type", None)
+    model_type = config.model_type_override or getattr(hf_config, "model_type", None)
     if model_type is None:
         raise ValueError(
             f"Could not determine model_type from HuggingFace config for '{config.hf_model_id}'"

python/src/coreai_models/llm/export.py

@@ -26,7 +26,11 @@
     DEFAULT_IOS_COMPRESSION_PRESET as IOS_DEFAULT,
 )
 from coreai_models.export.presets import DEFAULT_MACOS_COMPRESSION_PRESET as MACOS_DEFAULT
-from coreai_models.model_registry import try_lookup_preset, try_lookup_preset_by_hf_id
+from coreai_models.model_registry import (
+    presets_for_type,
+    try_lookup_preset,
+    try_lookup_preset_by_hf_id,
+)
 from coreai_models.models.registry import list_models as list_llm_models
 
 
@@ -375,6 +379,7 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig:
         compression_config_object=compression_config_object,
         disable_embedding_quantization=args.disable_embedding_quantization_ios,
         include_debug_info=args.include_debug_info,
+        model_type_override=getattr(preset, "_model_type_override", None) if preset else None,
     )
 
 
@@ -408,7 +413,13 @@ def main() -> None:
         return
 
     if args.list_models:
-        print("LLM model types:")
+        print("LLM presets (use short name or HuggingFace ID):")
+        print()
+        for p in presets_for_type("llm"):
+            platform = p.variant or "macOS"
+            print(f"  {p.short_name:35s} {p.hf_id:45s} {platform}")
+        print()
+        print("Supported architectures:")
         print()
         for name in list_llm_models():
             print(f"  {name}")

python/src/coreai_models/model_registry.py

@@ -51,6 +51,10 @@ class ModelPreset:
     # Optional YAML config path, resolved relative to the repo root in the
     # source tree (e.g. `models/qwen3/qwen3_0_6b_mixed_4bit_8bit.yaml`).
     compression_config: str | None = None
+    # Override HuggingFace model_type for registry lookup. Use when a model's
+    # config.json reports a type we don't want to register globally (e.g.
+    # SmolLM2 reports "llama" but we use the qwen2 export path).
+    _model_type_override: str | None = None
 
 
 @dataclass(frozen=True)
@@ -131,6 +135,40 @@ class UtilityModel:
     ModelPreset(
         "gpt-oss-20b", "openai/gpt-oss-20b", "gpt-oss", "llm", "macOS", "none", "bfloat16", 32768
     ),
+    ModelPreset(
+        "smollm2-1.7b-instruct",
+        "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+        "smollm2",
+        "llm",
+        "macOS",
+        "4bit",
+        "float16",
+        8192,
+        compression_config="models/smollm2/smollm2_4bit_embedding_excluded.yaml",
+        _model_type_override="qwen2",
+    ),
+    ModelPreset(
+        "smollm2-360m-instruct",
+        "HuggingFaceTB/SmolLM2-360M-Instruct",
+        "smollm2",
+        "llm",
+        "macOS",
+        "none",
+        "float16",
+        8192,
+        _model_type_override="qwen2",
+    ),
+    ModelPreset(
+        "smollm2-135m-instruct",
+        "HuggingFaceTB/SmolLM2-135M-Instruct",
+        "smollm2",
+        "llm",
+        "macOS",
+        "none",
+        "float16",
+        8192,
+        _model_type_override="qwen2",
+    ),
     # --- iOS (compression = palettized) ---
     ModelPreset(
         "qwen3-0.6b",
@@ -164,6 +202,40 @@ class UtilityModel:
         IOS_DEFAULT_MAX_CONTEXT_LENGTH,
         compression_config="models/qwen3/qwen3_4b_mixed_4bit_8bit.yaml",
     ),
+    ModelPreset(
+        "smollm2-1.7b-instruct",
+        "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+        "smollm2",
+        "llm",
+        "iOS",
+        "none",
+        "float16",
+        4096,
+        compression_config="models/smollm2/smollm2_1_7b_6bit.yaml",
+        _model_type_override="qwen2",
+    ),
+    ModelPreset(
+        "smollm2-360m-instruct",
+        "HuggingFaceTB/SmolLM2-360M-Instruct",
+        "smollm2",
+        "llm",
+        "iOS",
+        "none",
+        "float16",
+        4096,
+        _model_type_override="qwen2",
+    ),
+    ModelPreset(
+        "smollm2-135m-instruct",
+        "HuggingFaceTB/SmolLM2-135M-Instruct",
+        "smollm2",
+        "llm",
+        "iOS",
+        "none",
+        "float16",
+        4096,
+        _model_type_override="qwen2",
+    ),
 ]
 
 # ---------------------------------------------------------------------------

python/src/coreai_models/models/ios/qwen2.py

@@ -34,10 +34,11 @@ def __init__(self, config: Qwen2Config, layer_idx: int) -> None:
         self.n_heads = n_heads = config.num_attention_heads
         self.n_kv_heads = n_kv_heads = config.num_key_value_heads
         self.head_dim = head_dim = getattr(config, "head_dim", dim // n_heads)
+        attention_bias = getattr(config, "attention_bias", True)
 
-        self.q_proj = nn.Conv2d(dim, n_heads * head_dim, kernel_size=1, bias=True)
-        self.k_proj = nn.Conv2d(dim, n_kv_heads * head_dim, kernel_size=1, bias=True)
-        self.v_proj = nn.Conv2d(dim, n_kv_heads * head_dim, kernel_size=1, bias=True)
+        self.q_proj = nn.Conv2d(dim, n_heads * head_dim, kernel_size=1, bias=attention_bias)
+        self.k_proj = nn.Conv2d(dim, n_kv_heads * head_dim, kernel_size=1, bias=attention_bias)
+        self.v_proj = nn.Conv2d(dim, n_kv_heads * head_dim, kernel_size=1, bias=attention_bias)
 
         self.o_proj = nn.Conv2d(n_heads * head_dim, dim, kernel_size=1, bias=False)
 
@@ -271,6 +272,13 @@ def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
             raise ValueError(err)
 
         for i in range(max_layer + 1):
+            # Remove spurious attention bias keys (HF may init them even when
+            # the original model has attention_bias=False)
+            if not getattr(self.config, "attention_bias", True):
+                for proj in ["q_proj", "k_proj", "v_proj"]:
+                    bias_key = f"model.layers.{i}.self_attn.{proj}.bias"
+                    state_dict.pop(bias_key, None)
+
             # Reshape attention weights for Conv2d
             for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
                 weight_key = f"model.layers.{i}.self_attn.{proj}.weight"

python/src/coreai_models/models/macos/qwen2.py

@@ -33,11 +33,12 @@ def __init__(self, config: Qwen2Config, layer_idx: int) -> None:
         self.n_heads = n_heads = config.num_attention_heads
         self.n_kv_heads = n_kv_heads = config.num_key_value_heads
         self.head_dim = head_dim = getattr(config, "head_dim", dim // n_heads)
+        self.attention_bias = getattr(config, "attention_bias", True)
 
         self.qkv_proj = nn.Linear(
             dim,
             n_heads * head_dim + n_kv_heads * head_dim + n_kv_heads * head_dim,
-            bias=True,
+            bias=self.attention_bias,
         )
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
 
@@ -177,27 +178,35 @@ def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None:
             err = "invalid state_dict"
             raise ValueError(err)
 
+        expects_bias = getattr(self.config, "attention_bias", True)
+
         for i in range(max_layer + 1):
             combined_weight = []
             combined_bias = []
-            need_to_fuse = True
+            has_weights = True
+            has_bias = True
             for proj in ["q_proj", "k_proj", "v_proj"]:
                 weight_key = f"model.layers.{i}.self_attn.{proj}.weight"
                 bias_key = f"model.layers.{i}.self_attn.{proj}.bias"
-                if weight_key not in state_dict or bias_key not in state_dict:
-                    need_to_fuse = False
+                if weight_key not in state_dict:
+                    has_weights = False
                     continue
                 combined_weight.append(state_dict[weight_key])
-                combined_bias.append(state_dict[bias_key])
                 del state_dict[weight_key]
-                del state_dict[bias_key]
-            if need_to_fuse:
+                if bias_key in state_dict:
+                    if expects_bias:
+                        combined_bias.append(state_dict[bias_key])
+                    del state_dict[bias_key]
+                else:
+                    has_bias = False
+            if has_weights and combined_weight:
                 state_dict[f"model.layers.{i}.self_attn.qkv_proj.weight"] = torch.concat(
                     combined_weight, axis=0
                 )
-                state_dict[f"model.layers.{i}.self_attn.qkv_proj.bias"] = torch.concat(
-                    combined_bias, axis=0
-                )
+                if expects_bias and has_bias and combined_bias:
+                    state_dict[f"model.layers.{i}.self_attn.qkv_proj.bias"] = torch.concat(
+                        combined_bias, axis=0
+                    )
 
     def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
         super().load_state_dict(state_dict, strict=strict, assign=assign)