(microgpt Mojo 版を MAX で高速化する)=
# microgpt Mojo 版を MAX で高速化する

## この章で学ぶこと

- Mojo 版 microgpt（Tape ベース）の学習はそのまま維持し、**推論だけ MAX に切り替える**方法
- Mojo の Tape ノード値を Python interop で numpy 配列に変換する方法
- MAX Graph で GPT 推論ループを再構築する方法
- `microgpt_mojo` と `microgpt_mojo_max` を比較して結果を確認する方法

---

## 方針：学習と推論を分離する

Mojo の Tape 実装はスカラーループのため推論が遅いですが、自動微分（学習）には最適です。
MAX Engine には学習（勾配計算）API がないため、学習は Tape のままにして推論だけ MAX に切り替えます。

```{mermaid}
flowchart LR
    A["学習<br/>（Mojo Tape + Adam）<br/>B1〜B6 変更なし"] -->|"t.node_data(id)<br/>→ numpy float32"| B["重み抽出<br/>b08_max_infer.mojo"]
    B -->|"Python interop<br/>max_infer_helper.py"| C["MAX Graph<br/>推論"]
    C --> D["生成結果"]
```

| フェーズ | 担当 | microgpt_mojo との違い |
|---------|------|-----------------------|
| 学習 | Mojo Tape + Adam（B1〜B6） | **なし** |
| 重みの抽出 | `b08_max_infer.mojo`（新規） | 新規 |
| 推論 | `max_infer_helper.py`（MAX Graph） | 新規（高速） |

---

## ディレクトリ構成

```
src/part3/microgpt_mojo_max/
├── b01_dataset.mojo      ← microgpt_mojo からコピー（変更なし）
├── b02_tokenizer.mojo    ← 同上
├── b03_value.mojo        ← 同上
├── b04_state_dict.mojo   ← 同上
├── b05_ops.mojo          ← 同上
├── b05_gpt.mojo          ← 同上
├── b06_train.mojo        ← 同上
├── b08_max_infer.mojo    ← 新規：重み抽出 + MAX 呼び出し
├── max_infer_helper.py   ← 新規：MAX Graph 定義・推論ループ
├── main.mojo             ← 変更：推論部分を b08 に差し替え
└── input.txt             ← microgpt_mojo からシンボリックリンク
```

---

## B8: 重みの抽出と MAX 呼び出し（b08_max_infer.mojo）

`b08_max_infer.mojo` は Mojo と MAX（Python）の橋渡しをする唯一の新規 Mojo ファイルです。

```{literalinclude} ../../../src/part3/microgpt_mojo_max/b08_max_infer.mojo
:language: mojo
```

**ポイント：**

- `_matrix_to_numpy()` では `t.node_data(mat[i][j])` で Tape ノードの値（Float64）を取り出し、
  `Python.list()` に積んで `np.array(...).reshape(rows, cols)` で ndarray に変換する
- キー名は `"attn_wq_0"`, `"mlp_fc1_0"` のように `名前_層番号` で統一する
- `sys.path.insert(0, ".")` で実行ディレクトリを Python パスに追加してから
  `max_infer_helper` をインポートする

---

## MAX Graph の定義（max_infer_helper.py）

MAX Graph のコード（`EmbedLayer`・`AttentionLayer`・`MLPLayer`・`build_gpt_graph`・`run_inference`）は
純粋な Python ファイルにまとめます。

```{literalinclude} ../../../src/part3/microgpt_mojo_max/max_infer_helper.py
:language: python
```

**b08_max_infer.mojo との対応：**

| Mojo 側（b08）の変数 | Python 側（helper）での使われ方 |
|---------------------|-------------------------------|
| `weights["wte"]` など | `build_gpt_graph(weights, hp)` に渡す |
| `hp_dict` | `hp["n_embd"]` などで参照 |
| `uchars_py` | `uchars[token_id]` で文字を逆引き |
| `bos` | 生成終端の判定に使用 |

---

## main.mojo の変更点

学習ループは `microgpt_mojo/main.mojo` と**まったく同じ**です。
違いは末尾の推論部分だけです。

```{literalinclude} ../../../src/part3/microgpt_mojo_max/main.mojo
:language: mojo
```

**差分まとめ：**

| 変更前（microgpt_mojo） | 変更後（microgpt_mojo_max） |
|------------------------|---------------------------|
| `from b07_infer import ...` | `from b08_max_infer import run_max_inference` |
| Tape ベースの推論ループ | `run_max_inference(t, sd, hp, ...)` 1行 |

---

## 比較実行

どちらも同じ `input.txt` と `seed(42)` を使うため、
学習の loss 推移は**一致**します。

```bash
# Tape ベース推論（従来版）
cd src/part3/microgpt_mojo && mojo run main.mojo

# MAX Graph 推論（新版）
cd src/part3/microgpt_mojo_max && mojo run main.mojo
```

| 比較項目 | 結果 |
|---------|------|
| 学習の loss 推移 | **同一**（同じ Tape 実装・同じ seed） |
| 推論結果 | 異なりうる（確率的サンプリングのため） |
| 推論速度 | MAX 版の方が速い（特に `n_embd` が大きい場合） |

---

## Tape 方式との速度比較

| 方式 | 推論 1 ステップの仕組み | 特徴 |
|------|----------------------|------|
| Mojo Tape | スカラー演算をループで追記・即時実行 | 遅い、自動微分可能 |
| MAX Graph | グラフをコンパイル後に行列演算として一括実行 | 速い、SIMD/GPU 対応 |

MAX Graph では行列演算をハードウェアの SIMD 命令や GPU カーネルにマッピングできるため、
`n_embd=16` 程度の小さなモデルでも推論のオーバーヘッドが大幅に減ります。
モデルが大きくなるほど（`n_embd=256`, `n_layer=6` など）差が顕著になります。

---

## まとめ

- 学習（Tape + Adam、B1〜B6）はそのまま活かし、**推論だけ MAX に切り替える**のが段階的な高速化の第一歩
- 重みの受け渡しは Mojo Python interop で `t.node_data()` → numpy に変換するだけ
- MAX Graph の定義（Python）は `max_infer_helper.py` に分離し、Mojo からは1行で呼ぶ
- `b08_max_infer.mojo` が Mojo と MAX の橋渡し役になる

本書で実装した microgpt の知識は、より大きなモデルへのスケールアップや MAX/Mojo による最適化の基礎になります。