pythia-160m-stateful-coreml

English | 日本語

Model Summary

This is an unofficial Core ML conversion of EleutherAI/pythia-160m (GPT-NeoX architecture, one of EleutherAI's Pythia research model suite), using Core ML's stateful-model support (coremltools.StateType) for a real KV-cache implementation, rather than a KV-cache-free full-recompute design.

Different from earlier conversions

Earlier ASR/OCR conversions had shallow decoders, so a KV-cache-free design that recomputes the whole sequence every step was practical enough. But pythia-160m has a standard 12-layer autoregressive decoder, where the same design would be inefficient. This model instead uses coremltools' StateType (ct.StateType) to keep each layer's Key/Value cache as Core ML "state", decoding one token per step and updating that state in place.

GPT-NeoX-specific details (the parallel residual structure, partial rotary position embeddings with rotary_pct=0.25) are hand-reimplemented directly against the original weights, without going through transformers' Cache classes.

Usage (Python / coremltools)

import numpy as np
import coremltools as ct
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-160m")
mlmodel = ct.models.MLModel("pythia-160m-stateful_fp16.mlpackage")

prompt = "The history of artificial intelligence began"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids[0].tolist()

state = mlmodel.make_state()  # the KV-cache state object

pos = 0
logits = None
for tok in input_ids:  # prefill: feed the prompt one token at a time
    out = mlmodel.predict({
        "input_ids": np.array([[tok]], dtype=np.int32),
        "position": np.array([pos], dtype=np.int32),
    }, state=state)
    logits = out["logits"]
    pos += 1

generated = list(input_ids)
for _ in range(50):  # greedy decoding
    next_id = int(np.argmax(logits[0]))
    generated.append(next_id)
    out = mlmodel.predict({
        "input_ids": np.array([[next_id]], dtype=np.int32),
        "position": np.array([pos], dtype=np.int32),
    }, state=state)
    logits = out["logits"]
    pos += 1

print(tokenizer.decode(generated))

Traced with max_seq_len=256, so before position reaches 256 you need to create a fresh state (or provide a separate bucket for longer sequences).

Accuracy (an honest report)

The output token sequence does not exactly match the PyTorch fp32 reference (model.generate()). Isolating this to the first-step logits comparison: the cosine similarity between PyTorch fp32 and Core ML fp16 logits was 0.999999 (effectively identical). The divergence comes from the top-1 and top-2 candidates being extremely close (841.006 vs 840.990), so rounding to fp16 precision nearly ties them, and the argmax selection flips. Running the PyTorch model itself in fp16 makes the same choice, so this is not a Core ML conversion bug -- it's a known consequence of fp16 precision.

Item	Result
Logits cosine similarity (PyTorch fp32 vs Core ML fp16)	0.999999
Logits cosine similarity (PyTorch fp16 vs Core ML fp16)	0.9999998
Exact token-sequence match	✗ (path diverges at a near-tie)

For use cases that need strict reproducibility, consider re-converting with compute_precision=ct.precision.FLOAT32 (this release is fp16 only).

Specs

Item	Value
Base model	EleutherAI/pythia-160m (12 layers, 160M params)
Precision	float16
Max sequence length	256
Framework	Core ML (stateful KV-cache via coremltools.StateType)

Notes

• This is a community conversion, not an official release from EleutherAI.

Security

Audited against its upstream with model-audit-lite: weight format, bundled code, and a machine-readable lineage (ML-BOM). Details, checksums and how to reproduce: SECURITY.md.

---

モデルの概要

EleutherAI/pythia-160m（GPT-NeoXアーキテクチャ、EleutherAIのPythia研究用モデル群の一つ）を、KVキャッシュを持つstatefulなCore MLモデルとして変換したものです。元モデルの著作権はEleutherAIに帰属します。

これまでとの違い

これまでの音声認識・OCR系の変換ではデコーダーが浅かったため、KVキャッシュなしで 全系列を毎回再計算する設計でも十分実用的でした。しかしpythia-160mは標準的な 12層の自己回帰デコーダーを持つため、同じ設計は非効率です。本モデルでは coremltoolsのStateType（ct.StateType）を使い、各層のKey/Valueキャッシュを Core MLの「状態」として保持し、デコードステップごとに1トークンずつ推論・ インプレース更新する設計にしています。

GPT-NeoX特有の設計（parallel residual構造、部分的回転位置埋め込み rotary_pct=0.25）は、transformersのCacheクラスを介さず、重みをそのまま 流用しながら手動で再実装しています。

使い方（Python / coremltools）

import numpy as np
import coremltools as ct
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-160m")
mlmodel = ct.models.MLModel("pythia-160m-stateful_fp16.mlpackage")

prompt = "The history of artificial intelligence began"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids[0].tolist()

state = mlmodel.make_state()  # KVキャッシュの状態オブジェクト

pos = 0
logits = None
for tok in input_ids:  # プレフィル: プロンプトを1トークンずつ流す
    out = mlmodel.predict({
        "input_ids": np.array([[tok]], dtype=np.int32),
        "position": np.array([pos], dtype=np.int32),
    }, state=state)
    logits = out["logits"]
    pos += 1

generated = list(input_ids)
for _ in range(50):  # グリーディデコード
    next_id = int(np.argmax(logits[0]))
    generated.append(next_id)
    out = mlmodel.predict({
        "input_ids": np.array([[next_id]], dtype=np.int32),
        "position": np.array([pos], dtype=np.int32),
    }, state=state)
    logits = out["logits"]
    pos += 1

print(tokenizer.decode(generated))

max_seq_len=256でトレースしているため、positionが256に達する前に 新しいstateを作り直す（または長い系列用に別バケットを用意する）必要があります。

精度検証（重要：正直な報告）

PyTorchのfp32リファレンス（model.generate()）とのトークン列の完全一致は 確認できませんでした。原因を1ステップ目のロジット比較で特定したところ、 PyTorch fp32とCore ML fp16のロジットのコサイン類似度は0.999999（実質的に 同一）でした。分岐の原因は、1位と2位の候補が841.006対840.990という 僅差だったため、fp16精度に丸めるとほぼ同点になり、argmaxの選択が 入れ替わったことです。PyTorchモデル自体をfp16で実行しても同じ選択をするため、 これはCore ML変換のバグではなく、fp16精度に起因する既知の現象です。

項目	結果
ロジットのコサイン類似度（PyTorch fp32 vs Core ML fp16）	0.999999
ロジットのコサイン類似度（PyTorch fp16 vs Core ML fp16）	0.9999998
トークン列の完全一致	✗（僅差の分岐点で経路が変わる）

厳密な再現性が必要な用途では、compute_precision=ct.precision.FLOAT32で 再変換することを検討してください（本リリースはfp16のみ）。

Specs

Item	Value
ベースモデル	EleutherAI/pythia-160m（12層、160M params）
精度	float16
最大系列長	256
フレームワーク	Core ML（coremltools.StateTypeによるstateful KVキャッシュ）

備考

• 本変換は非公式のコミュニティ版です。EleutherAIによる公式リリースではありません。

セキュリティー

model-audit-lite で変換元と突き合わせて監査済みです（重みの形式、同梱コード、機械可読な系譜＝ML-BOM）。詳細・チェックサム・再現方法は SECURITY.md をご覧ください。