anziank/grio-gemma-3-1b-coreml-anyLM-seq2048

Gemma 3 1B Instruct — CoreML (seq≤2048, AnyLanguageModel-compatible)

On-device CoreML .mlpackage converted from `google/gemma-3-1b-it` for use with HuggingFace's AnyLanguageModel Swift framework and `swift-transformers` ≥ 1.0.

Input/output tensor names match the inputIds / attentionMask / logits convention required by swift-transformers 1.x LanguageModel. Drop-in compatible with CoreMLLanguageModel(url:computeUnits:chatTemplateHandler:).

Variants in this collection

• —`grio-gemma-3-1b-coreml-anyLM-seq512` Fixed [1, 512] — short-form enhancement / single-sentence translation. Always pays for full 512 positions per step.
• —`grio-gemma-3-1b-coreml-anyLM-seq1024` Fixed [1, 1024] — medium-context post-processing. Always pays for full 1024 positions per step.
• —`grio-gemma-3-1b-coreml-anyLM-seq2048` (this repo) RangeDim [1, 1..2048] — recommended production variant. Variable-length translation / long-context. Per-step cost scales with real prompt length, so for typical product prompts (<200 tokens) this is the fastest of the three. Must use `.cpuAndGPU`; .cpuAndNE rejects, .cpuOnly produces wrong output, and .all works but is slow.

Model details

Architecture note: sliding-window attention (window=512, every 6th layer is global), 4 heads / 1 KV head (extreme GQA).

Verified output (greedy, deterministic)

Actual outputs from this .mlpackage under .cpuAndGPU. Token IDs are byte-identical to PyTorch FP16 reference (model.generate(do_sample=False, use_cache=True)).

• —ChatML System: "You are a helpful assistant. Answer concisely." User: "What is the capital of France?" Output: "Paris." (stops at <end_of_turn> after 3 tokens)
• —ChatML translation System: "You are a French translator. Translate the user message to French. Output only the translation." User: "The quick brown fox jumps over the lazy dog." Output: "Le rapide renard brun saute par-dessus le chien paresseux."
• —ChatML translation System: "You are a German translator. Translate the user message to German. Output only the translation." User: "I would like to order a coffee, please." Output: "Ich möchte einen Kaffee bestellen."
• —ChatML rewrite System: "Rewrite the user's text in clear formal English. Output only the rewrite." User: "ok so the guy was like really mad cuz his package didnt show up" Output: "The individual expressed considerable frustration as his package had not arrived."
• —ChatML transcription edit System: "You are a transcription editor. Improve the grammar and punctuation of the user's text. Output only the improved text." User: "so um like i was thinking we should go to the store maybe tomorrow if its not raining" Output: "So, um, I was thinking we should go to the store maybe tomorrow if it's not raining."

Observed performance (M1 Pro, macOS)

Benchmark notes: real prompt about 27 tokens.

iPad device perf not benchmarked in this build run.

Runtime gotchas (please read before integrating)

1. 1.`.cpuOnly` silently produces wrong output on this RangeDim variant. Not NaN — wrong. The model returns valid-looking but incorrect token IDs. Always use .cpuAndGPU for this variant. The fixed-shape seq512/seq1024 variants in this collection do not have this issue (.cpuOnly is clean on those).
2. 2.`.cpuAndNE` is rejected with an explicit error about data-dependent shapes. Gemma 3's combination of sliding-window attention and RangeDim input shapes produces ops the ANE compiler refuses.
3. 3.Gemma 3 uses a two-id chat-EOS list: <eos> (id 1) for general end-of-sequence and <end_of_turn> (id 106) for chat-turn termination. generation_config.json lists both. If your runtime stops only on tokenizer.eos_token_id (which returns <eos>), the model will decode past <end_of_turn> into out-of-distribution territory and produce multilingual gibberish that looks like graph corruption but isn't. Read the full eos_token_id list from generation_config.json and stop on any of them. swift-transformers ≥ 1.3 handles this correctly.
4. 4.The chat template lives in `chat_template.jinja` (a separate file from tokenizer_config.json). swift-transformers ≥ 1.3 reads it correctly via Hub.swift. Older callers that look for an inline chat_template key in tokenizer_config.json will not find one.
5. 5.`logits` shape metadata is empty in the .mlpackage description (a coremltools artifact for RangeDim outputs). Verify at runtime with a real predict() — the actual output is rank-3 [1, seq_len, 262144] and works with swift-transformers' assert(scores.rank == 3).

Conversion notes (for the CoreML community)

This model was produced from PyTorch source via torch.export + coremltools.convert. Findings worth flagging for others converting Gemma 3 (or similar HF causal LMs) to CoreML:

1. 1.Use `attn_implementation="sdpa"`. The HF default "eager" for Gemma 3 may produce a corrupted graph under RangeDim (we saw this on Qwen 2.5 RangeDim builds). SDPA gives a cleaner, fewer-op MIL graph that lowers reliably.
2. 2.Use `torch.export.default_decompositions()` for run_decompositions. The all-decompositions mode ({}) can SIGSEGV inside optimize_repeat_ops at 1B+ scale on RangeDim graphs.
3. 3.The `optimize_repeat_ops.py:433 RuntimeWarning: overflow encountered in cast` fires during conversion of this model. Despite the warning, output is byte-identical to PyTorch FP16 reference under .cpuAndGPU. The warning indicates internal range arithmetic overflow but does not corrupt the produced graph for Gemma 3 + SDPA. Don't reflexively reconvert on seeing it — verify with a PyTorch comparison instead.
4. 4.FP16 / greedy decoding is not byte-deterministic across backends. Outputs are semantically equivalent to PyTorch CPU FP16 reference but may differ on tokens where the model has near-tied top-1 candidates. This is expected behavior, not a conversion bug.
5. 5.`chat_template.jinja` must be bundled alongside the .mlpackage. So must generation_config.json (for the multi-id EOS list), special_tokens_map.json, added_tokens.json, and tokenizer.model. This repo includes all of them.

Usage (Swift)

import AnyLanguageModel

let modelURL: URL = // path to this .mlpackage on disk
let lm = try await CoreMLLanguageModel(
    url: modelURL,
    computeUnits: .cpuAndGPU,  // REQUIRED for this RangeDim variant — see Runtime gotchas
    chatTemplateHandler: { instructions, prompt in
        // Gemma 3 uses <start_of_turn>...<end_of_turn> chat template; tokenizer.json's
        // Jinja template (loaded by swift-transformers from chat_template.jinja) applies it.
        var messages: [Message] = []
        if let system = instructions?.description, !system.isEmpty {
            messages.append(["role": "system", "content": system])
        }
        messages.append(["role": "user", "content": prompt.description])
        return messages
    }
)
let session = LanguageModelSession(model: lm, instructions: "You are a French translator. Output only the translation.")
let response = try await session.respond(to: "The capital of France is Paris.")
print(response.content)

Keep tokenizer.json, tokenizer_config.json, config.json, chat_template.jinja, generation_config.json, special_tokens_map.json, added_tokens.json, and tokenizer.model (all bundled in this repo) as siblings of the .mlpackage on disk.

Reproducibility

Conversion done with coremltools==9.0, torch==2.7.0, transformers==5.8.1. Approximate single-call recipe:

import coremltools as ct, torch, torch.nn as nn
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-3-1b-it",
    torch_dtype=torch.float16,
    attn_implementation="sdpa",         # see Conversion notes
)
model.eval()

class Wrapper(nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, inputIds, attentionMask):
        return self.m(input_ids=inputIds, attention_mask=attentionMask, use_cache=False).logits

wrapper = Wrapper(model).eval()
seq = torch.export.Dim("sequence_length", min=1, max=2048)
ep = torch.export.export(
    wrapper,
    (torch.randint(0, 262144, (1, 128), dtype=torch.int32),
     torch.ones((1, 128), dtype=torch.int32)),
    dynamic_shapes={"inputIds": {1: seq}, "attentionMask": {1: seq}},
).run_decompositions(torch.export.default_decompositions())  # NOT {}

ct.convert(
    ep,
    inputs=[
        ct.TensorType(name="inputIds",      shape=(1, ct.RangeDim(1, 2048)), dtype=int),
        ct.TensorType(name="attentionMask", shape=(1, ct.RangeDim(1, 2048)), dtype=int),
    ],
    outputs=[ct.TensorType(name="logits")],
    minimum_deployment_target=ct.target.iOS18,
    compute_precision=ct.precision.FLOAT16,
    convert_to="mlprogram",
)

License

Gemma Terms of Use. Weights from `google/gemma-3-1b-it` by Google DeepMind. Re-uploaded as a CoreML port; original model card terms apply. By using this model you agree to Google's Gemma Terms of Use and the Prohibited Use Policy.