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

Gemma 3 1B Instruct — CoreML (seq=1024 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 Fixed [1, 512] — short-form enhancement / single-sentence translation. Smallest graph.
• grio-gemma-3-1b-coreml-anyLM-seq1024 (this repo) Fixed [1, 1024] — medium-context post-processing. Always pays for full 1024 positions per step.
• 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	1024 tokens, fixed shape (prompt + completion combined)
Inputs	inputIds, attentionMask
Input shape	Int32 [1, 1024]
Output	logits: Float16 [1, 1024, 262144] (rank-3, per-position)
Min OS	iOS 18 / macOS 15
Compute	.cpuAndGPU recommended; do not use .all on this shape
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 under .cpuAndGPU. Output is byte-identical to the matching seq512 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 since the weights and graph topology are identical.

Observed performance (M1 Pro, macOS)

Benchmark notes: full-pad 1024.

Compute	Predict	Load.	Result
.cpuAndGPU	~1781	~11s	✅ Clean. Production-recommended for this shape.
.cpuOnly	~2149	~13s	✅ Clean output on this fixed-shape variant (unlike the RangeDim variant).
.cpuAndNE	~5762	~219s	⚠️ Clean output but ANE compilation is extremely slow on this shape. Avoid.
.all	—	—	❌ RuntimeError at predict time. ANE compiler rejects this shape under the .all scheduler.

Smoke runs pad to the full 1024-position graph every step, so the ms/tok number reflects fixed-shape compute over 1024 positions — not the real-prompt-length cost you would see on a RangeDim build. For typical product prompts (<200 tokens), the seq2048 RangeDim variant in this collection is faster despite its larger maximum shape.

iPad device perf not benchmarked in this build run.

Runtime gotchas (please read before integrating)

1. .all is broken on this shape. RuntimeError at predict(). Use .cpuAndGPU. Different shapes in this same collection have different safe-set: seq512 works on .all, seq2048 works on .all but slowly. Per-shape verification is required.
2. Fixed context window — input + output ≤ 1024 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, 1024, 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. 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 even 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,  // do NOT use .all on this shape — see Runtime gotchas
    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, 1024), dtype=int),
        ct.TensorType(name="attentionMask", shape=(1, 1024), 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.