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

Gemma 3 1B Instruct — CoreML (seq=512 fixed, 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 (this repo) Fixed [1, 512] — short-form text enhancement / single-sentence translation. Smallest graph. Multiple safe backends.
• grio-gemma-3-1b-coreml-anyLM-seq1024 Fixed [1, 1024] — medium-context post-processing.
• grio-gemma-3-1b-coreml-anyLM-seq2048 RangeDim [1, 1..2048] — variable-length / long-context. Per-step cost scales with actual prompt length; fastest for typical product prompts. Must use .cpuAndGPU on the RangeDim variant.

Model details

Spec	Value
Base	google/gemma-3-1b-it
Precision	Float16 (mlprogram)
Context	512 tokens, fixed shape (prompt + completion combined)
Inputs	inputIds, attentionMask
Input shape	Int32 [1, 512]
Output	logits: Float16 [1, 512, 262144] (rank-3, per-position)
Min OS	iOS 18 / macOS 15
Compute	.cpuAndGPU fastest; .cpuOnly and .all also work
Format	.mlpackage (compiled on first load)
Toolchain	coremltools 9.0 + torch 2.7 + transformers 5.8.1

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

Verified output (greedy, deterministic)

Actual output from this .mlpackage. Output is byte-identical to the matching seq1024 and seq2048 variants in this collection (same weights, same graph topology; only the static seq dimension differs).

• 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)

See the seq2048 variant card for the full translation + post-processing verified-output table — outputs are byte-identical between variants.

Observed performance (M1 Pro, macOS)

Benchmark notes: full-pad 512.

Compute	Predict	Load	Result
.cpuAndGPU	~1106	~12s	✅ Clean. Fastest.
.cpuOnly	~1106	~9s	✅ Clean output (fixed-shape variant; FP16 BNNS path is healthy at this scale).
.all	~1869	~197s	✅ Clean output but ANE compilation fails internally; CoreML silently falls back to GPU/CPU. Very slow load.
.cpuAndNE	~2883	~156s	⚠️ Clean output but ANE compilation slow.

Smoke runs pad to the full 512-position graph every step, so the ms/tok number reflects fixed-shape compute over 512 positions. For typical product prompts (<200 tokens), the seq2048 RangeDim variant is faster.

iPad device perf not benchmarked in this build run.

Runtime gotchas (please read before integrating)

1. All four compute units produce clean output on this fixed-shape variant. Different than the seq2048 RangeDim variant in this collection (where .cpuOnly silently produces wrong output and .cpuAndNE errors). Choose by perf, not correctness — .cpuAndGPU is fastest.
2. Fixed context window — input + output ≤ 512 tokens. Pad shorter prompts with zeros and set attentionMask = 0 on padded positions. The conversion script's smoke harness does this automatically.
3. logits is per-position, rank-3 [1, 512, 262144]. Pick the row at the last real prompt token for greedy/sampled decode.
4. 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. swift-transformers ≥ 1.3 handles this correctly.
5. The chat template lives in chat_template.jinja (a separate file from tokenizer_config.json). swift-transformers ≥ 1.3 reads it correctly. Older callers that only look in tokenizer_config.json will not find an inline template.

Conversion notes (for the CoreML community)

Findings worth flagging for others converting Gemma 3 to CoreML:

1. Use attn_implementation="sdpa" — gives a cleaner, fewer-op MIL graph that lowers reliably.
2. Use torch.export.default_decompositions() — the all-decompositions mode ({}) can SIGSEGV at 1B+ scale.
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. Don't reflexively reconvert on seeing it — verify with a PyTorch comparison instead.
4. FP16 / greedy decoding is not byte-deterministic across backends. Outputs are semantically equivalent to PyTorch CPU FP16 reference but may differ on near-tied tokens. Expected behavior, not a conversion bug.
5. Compute-unit safety differs per shape within this collection — see the perf table above. .all works on seq512, errors on seq1024. Always sweep your build before recommending a backend.
6. chat_template.jinja must be bundled alongside the .mlpackage. So must generation_config.json, special_tokens_map.json, added_tokens.json, tokenizer.model. All included in this repo.

Usage (Swift)

import AnyLanguageModel

let modelURL: URL = // path to this .mlpackage on disk
let lm = try await CoreMLLanguageModel(
    url: modelURL,
    computeUnits: .cpuAndGPU,
    chatTemplateHandler: { instructions, prompt in
        // Gemma 3 uses its own <start_of_turn>...<end_of_turn> chat template;
        // swift-transformers loads it from chat_template.jinja at runtime.
        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 helpful assistant.")
let response = try await session.respond(to: "Improve this text: …")
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",
)
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()
ep = torch.export.export(
    wrapper,
    (torch.randint(0, 262144, (1, 128), dtype=torch.int32),
     torch.ones((1, 128), dtype=torch.int32)),
).run_decompositions(torch.export.default_decompositions())

ct.convert(
    ep,
    inputs=[
        ct.TensorType(name="inputIds",      shape=(1, 512), dtype=int),
        ct.TensorType(name="attentionMask", shape=(1, 512), 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.