Apple Silicon の Mac 1台で、学習済みモデルを一切使わずに Transformer を書き、
日本語のコーパスで学習させて、チャットできるところまでやりました。
- モデル: 文字レベル ミニGPT / 6層 / 384次元 / 6ヘッド / 文脈256文字 / 11.53Mパラメータ
- フレームワーク: MLX（Apple純正。
pip install mlxだけでGPUが動く） - データ:
kunishou/oasst1-89k-ja（Apache-2.0）を整形した日本語会話 28,616件 / 441万文字 - 学習: MacBook Pro (M1 Max, 32コアGPU) で 32分 / 3,600ステップ
コードは全部 MIT で公開しています。
この記事は MLX を使ったときにだけ踏む落とし穴 に絞って書きました。
PyTorch から来ると「同じ書き味なのに挙動が違う」箇所がいくつかあり、
そのどれもが エラーを出さずに静かに壊れる タイプだったからです。
記事中の数値はすべて手元の M1 Max での実測値です。
前提: MLXは本当に「入れるだけ」
pip install mlx numpy
python -c "import mlx.core as mx; print(mx.default_device())"
# Device(gpu, 0)
これだけでGPUが使えます。PyTorch + MPS のときに必要だった
PYTORCH_ENABLE_MPS_FALLBACK=1 の指定も、「この演算はMPS未対応です」との格闘も、
今回は一度もありませんでした。Apple Silicon でゼロから書くなら MLX が圧倒的に楽です。
ただし conda で環境を作る場合、CONDA_SUBDIR の指定だけは忘れないでください。
CONDA_SUBDIR=osx-arm64 conda create -n 1lm python=3.11 -y
conda activate 1lm
conda config --env --set subdir osx-arm64
付けないと x86_64 のパッケージが混ざり、実行時にこう化けます。
ImportError: dlopen(...): mach-o, but wrong architecture
conda config --env --set subdir osx-arm64 を打っておくと、
後から pip install や conda install を追加したときも巻き添えを防げます。
実装の核: 「次の1文字」を予測するだけ
言語モデルがやっていることは1つだけです。
これまでの文字列を見て、次に来る1文字の確率分布を出す
会話に見えるのは、学習データをこの形にしておくからです。
<|user|>質問<|assistant|>返答<|end|>
<|assistant|> の続きを予測させれば、それが「返答」になります。
ルールも分岐も if 文も書きません。フォーマットを学習させるだけです。
トークナイザ: マーカーは1トークンにする
文字レベル（1文字=1トークン）は日本語と相性が良く、SentencePiece の学習も要りません。
ただし <|user|> を素直に8文字へ分解すると、モデルは「開始記号」を覚えるだけで
貴重な容量を使い、生成時に <|use のような壊れた記号を吐きます。
USER, ASSISTANT, END, UNK = "<|user|>", "<|assistant|>", "<|end|>", "<|unk|>"
_MARKER_RE = re.compile("(" + "|".join(re.escape(t) for t in (USER, ASSISTANT, END)) + ")")
def encode(self, text):
ids = []
for chunk in _MARKER_RE.split(text): # マーカーで分割してから
if chunk in SPECIAL_TOKENS:
ids.append(self.stoi[chunk]) # マーカーは1個のID
else:
ids.extend(self.stoi.get(c, self.unk_id) for c in chunk) # 残りは1文字1ID
return ids
<|unk|> も語彙に入れておきます。コーパスには出てきませんが、
推論時にユーザーが未知の文字を打ってくる可能性があるためです。
なお文字レベルは語彙が数千種類に膨らみます。今回は「出現10回未満の文字を含む会話は丸ごと捨てる」
フィルタで 文字種 2,072 まで落としました。65,536未満なら
トークン列を uint16 で持てるので、441万文字が9MBに収まり全部メモリに載ります。
モデル: MLXはPyTorchとほぼ同じ書き味
class CausalSelfAttention(nn.Module):
def __init__(self, cfg):
super().__init__()
self.n_head = cfg.n_head
self.head_dim = cfg.n_embd // cfg.n_head
self.scale = self.head_dim ** -0.5
self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
def __call__(self, x):
B, T, C = x.shape
q, k, v = mx.split(self.qkv(x), 3, axis=-1)
shape = (B, T, self.n_head, self.head_dim)
q = q.reshape(shape).transpose(0, 2, 1, 3)
k = k.reshape(shape).transpose(0, 2, 1, 3)
v = v.reshape(shape).transpose(0, 2, 1, 3)
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask="causal")
return self.proj(out.transpose(0, 2, 1, 3).reshape(B, T, C))
MLX には mx.fast.scaled_dot_product_attention があり、mask="causal" と
文字列で渡せます。自分で下三角行列を作って -inf を埋める必要はありません。
ブロックは Pre-LN + 残差接続です。
def __call__(self, x):
x = x + self.attn(self.ln1(x))
return x + self.mlp(self.ln2(x))
出力層は埋め込み行列を転用します（weight tying）。
def __call__(self, idx):
_, T = idx.shape
x = self.drop(self.tok_emb(idx) + self.pos_emb(mx.arange(T)))
for block in self.blocks:
x = block(x)
return self.tok_emb.as_linear(self.ln_f(x)) # 埋め込みを出力層として使い回す
nn.Embedding.as_linear() が最初から用意されているのが MLX の気持ちいいところです。
語彙2,077 × 384次元ぶんのパラメータを節約できて、小さいモデルではむしろ精度が上がります。
これで 11.53M パラメータです。
ハマりどころ 1: mx.eval を呼ばないと計算が走らない
MLX は遅延評価です。演算はグラフを積むだけで、結果が必要になるまで実行されません。
学習ループで mx.eval を忘れると、グラフが延々と伸びてメモリを食い続け、
「異常に速いのに loss が出てこない」 という状態になります。
state = [model.state, optimizer.state, mx.random.state]
@partial(mx.compile, inputs=state, outputs=state)
def step(x, y):
loss, grads = loss_and_grad(model, x, y)
grads, _ = optim.clip_grad_norm(grads, 1.0)
optimizer.update(model, grads)
return loss
loss = step(x, y)
mx.eval(state) # ここで初めてGPUが動く
mx.eval に渡すのは loss だけでは足りません。オプティマイザの状態も含めた
state 全体を評価しないと、更新のグラフが残ったままになります。
PyTorch の .item() に相当する感覚ですが、あちらは呼び忘れても同期が入るのに対し、
MLX は本当に何も起きません。ここが一番の違いです。
ハマりどころ 2: mx.compile の inputs に mx.random.state を入れ忘れる
これが今回いちばん気づきにくかった罠です。
mx.compile はグラフを固定します。inputs / outputs に mx.random.state を
含めないと、Dropout の乱数が固定されます。
毎ステップまったく同じ位置のニューロンを落とすので、正則化として機能しません。
そして エラーも警告も一切出ません。loss はそれらしく下がっていきます。
# NG: 乱数状態が固定される
@partial(mx.compile)
def step(x, y): ...
# OK
state = [model.state, optimizer.state, mx.random.state]
@partial(mx.compile, inputs=state, outputs=state)
def step(x, y): ...
「mx.compile を使うときは、関数の外側で変化するものを全部 inputs/outputs に列挙する」
と覚えるのが安全です。乱数状態はその筆頭で、忘れやすいわりに影響が大きい。
ハマりどころ 3: mask="causal" を外すと、lossは下がるのに生成が壊れる
これは MLX 固有というより Transformer 全般の話ですが、
MLX だと mask="causal" という引数1つで済むぶん、消しても動いてしまうのが厄介です。
因果マスクは「未来の文字を見てはいけない」という制約で、言語モデルの心臓部です。
外すと学習時にモデルが答えを覗けるので、train loss も val loss も綺麗に下がります。
ところが生成時には未来の文字が存在しないため、出力は完全に破綻します。
「損失は良いのに生成がおかしい」という症状に出会ったら、まずここを疑ってください。
損失グラフだけ見ていると絶対に気づけません。
対策はシンプルで、250ステップごとに固定プロンプトで実際に生成させてログに残すことです。
step 500 | val 2.50 | sample: こんにちは！私の情報を助けることができます。
step 1000 | val 2.18 | sample: 私はオープンソースです、あなたの答えを助けることができません。
損失という数字だけでなく、生成という「出力そのもの」を監視対象に入れる。
これは自作モデル全般に効く習慣だと思います。
ハマりどころ 4: 推論前の model.eval() 忘れ
MLX の nn.Module にも学習モードと推論モードがあります。
model.eval() を呼ばないと Dropout が有効なままです。
model = MiniGPT(cfg)
model.load_weights(str(ckpt / "model.safetensors"))
model.eval() # Dropout を切る
症状は「同じ入力なのに返答が毎回ぶれる」「なんとなく壊れている」。
温度やシード固定で再現性を取ろうとしても取れないので、そこで気づけます。
PyTorch でも定番の事故ですが、MLX は情報が少ないぶん
「MLX 特有の問題では？」と疑って遠回りしがちです。まずここを確認してください。
ハマりどころ 5: mx.topk は並び順を保証しない
top-k サンプリングを書くときの罠です。
mx.topk(logits, k) は「大きい方から k 個の値」を返しますが、
それがソート済みである保証はありません。
なので「k番目の値」を取るつもりで [..., -1] を取ると、しきい値がずれます。
返ってきた k 個の 最小値 を取るのが正解です。
def _sample(logits, temperature, top_k):
if temperature <= 0:
return mx.argmax(logits)
logits = logits * (1.0 / temperature)
if top_k and 0 < top_k < logits.size:
threshold = mx.min(mx.topk(logits, top_k)) # k番目に大きい値
logits = mx.where(logits < threshold, -float("inf"), logits)
return mx.random.categorical(logits)
ついでに、小さいモデルが同じ言葉を繰り返すときは
モデルを疑う前にサンプリング設定を疑ってください。今回効いた値はこのあたりです。
-
temperature0.8前後 ── 低いと堅くなるが同じ言い回しに固まる。高いと崩れる -
top_k40前後 ── 語彙2,077のうち確率が極端に低い文字を切る -
repetition_penalty1.15前後 ── 直近に出した文字のロジットを割り引く。ループから抜ける効果が一番大きい
学習をやり直す前に、この3つを振ってみる価値は十分あります。
実測値: どれくらい出るのか
バッチサイズを速度で選ぶ意味はなかった
定常状態のスループットを測り直した結果です（M1 Max 32コアGPU / 11.53M / 文脈256）。
batch_size 16 : 38.0k tok/s
batch_size 32 : 36.7k tok/s
batch_size 64 : 37.0k tok/s
batch_size 128 : 39.3k tok/s
batch_size 192 : 36.4k tok/s
意外なことに ほぼ横ばいでした。差は±4%程度で、測り直すと順位が入れ替わる範囲です。
このサイズのモデルでは GPU がすでに飽和していて、バッチを増やしても
1トークンあたりのコストが変わらない、ということだと思います。
つまりバッチサイズは速度ではなく 学習の安定性 で選んでください。今回は64にしました。
ベンチマーク値で総ステップ数を逆算すると必ず足りない
cosine_decay は「総ステップ数」を前提に学習率を下げるので、
時間切れで打ち切ると学習率が高いまま終わって性能を取り逃がします。
そこで先にスループットを測って逆算したのですが、ここでも外しました。
37,000 × 60 × 32 ÷ 16,384 ≒ 4,300ステップ と見積もって回したところ、
32分で到達したのは3,600ステップでした。
原因は、この計算に 検証とサンプル生成の時間が入っていない ことです。
250ステップごとに検証を20バッチ回し、さらに固定プロンプトへの返答まで生成しています。
通しの実測スループットは 31k tok/s で、ベンチマーク値の84%しかありませんでした。
逆算にはベンチマーク値ではなく通しの実測値を使ってください。
それが分からないうちは、逆算値を1〜2割少なめに見積もっておくのが安全です。
最終結果
- 学習時間 ── 32分 / 3,600ステップ
- 最終 train loss ── 1.751
- 最良 val loss ── 1.857
- スループット ── 31k tok/s（通し）
- 生成速度 ── 200〜300 文字/秒（KVキャッシュなし）
ランダムに答えたときの損失は ln(2077) = 7.64 なので、そこから1.86まで下がっています。
で、会話はどうなったか
あなた> こんにちは
1LM > こんにちは、私はあなたを助けることができますか？
(24 文字 / 0.1秒 / 234 文字毎秒)
あなた> AIとは何ですか？
1LM > AI言語モデルとして、私は意識や情報を提供する能力に基づいた人工知能です。
どちらの人間は、人間が、最も人類を持つ人々によって生成されたコミュニティです。
(107 文字 / 0.3秒 / 306 文字毎秒)
内容はもちろんデタラメです。11.53Mパラメータに知識は入りません。
それでも助詞の使い方、句読点の位置、「〜として、私は〜です」という言い回しは
ちゃんと日本語になっています。文法も会話の作法も一切教えていません。
「次の1文字を当てろ」という問題を3,600回解かせただけです。
なぜ賢くならないのかも、学習曲線を見るとはっきり分かります。
train と val の線が 1,250ステップあたりから離れ始め、
そこから差は一度も縮まりませんでした。
step 2000 train 1.9168 val 1.9660 差 +0.049
step 2500 train 1.8562 val 1.9116 差 +0.055
step 3000 train 1.7886 val 1.8778 差 +0.089
step 3500 train 1.7570 val 1.8571 差 +0.100
これは「データを覚えきった」サインです。441万文字を13エポック回しているので、
足りないのは時間ではなくデータでした。長く回しても伸びません。
伸ばすならこの順番
- BPE / SentencePiece トークナイザにする。1トークンあたりの情報量が増え、同じ計算量で扱える文脈が数倍になります
- KVキャッシュを実装する。今は毎回フル系列を再計算しています
-
損失を
<|assistant|>の部分だけに絞る（instruction masking）。返答の質が上がります - 素の日本語で事前学習してから対話データで仕上げる。今回は対話2万8千件だけから日本語そのものを学んでいるので、語彙も言い回しも痩せています。1億文字を3周でも 31k tok/s なら約2.7時間、寝ている間に終わります
まとめ
MLX で踏んだ罠は、どれも エラーを出さずに静かに壊れる タイプでした。
-
mx.eval忘れ → メモリが増え続け loss が出ない -
mx.compileのinputsにmx.random.state忘れ → Dropout が効かない（無警告） -
mask="causal"を外す → loss は下がるのに生成が破綻する -
model.eval()忘れ → 返答が毎回ぶれる -
mx.topkの並び順を仮定する → top-k のしきい値がずれる
共通の教訓は、損失という数字だけを監視対象にしないことだと思います。
特に3番目と4番目は、250ステップごとに固定プロンプトで生成させてログに残していなければ、
最後まで気づけなかったはずです。
コード一式（データ整形・学習・CLI・Web GUI）はこちらです。
data/raw/*.jsonl に自分のデータを置けば差し替えられるようにしてあります。
質問や「うちのMacだとこうなった」があればコメントで教えてください。
ライセンスについて
- コード: MIT
- 学習データ: kunishou/oasst1-89k-ja（Apache-2.0、OpenAssistant/oasst1 の日本語訳）
- 学習済み重みは Apache-2.0 データからの派生物として扱い、リポジトリの
NOTICEに帰属表示をまとめています
補足
この題材については、GUI（Liquid Glass風）の作り方やコーパス設計の判断、
今回省いた分も含めた全13件の詰まりどころを、別途 note に有料記事としてまとめています。
この記事の内容だけで手元で再現できるようにはしてあるので、
興味があれば覗いてみてください。
お盆休みの3日間で「自分だけの言語モデル」を作る (有料)