(microgpt.py を Mojo で書き直す)= # microgpt.py を Mojo で書き直す ## この章で学ぶこと - `microgpt.py` の**B1〜B7 ブロック**が Mojo でどう書かれているか、ソースコードを読みながら理解する - Mojo の**所有権・テープ（計算グラフ）・構造体**を、動くコードの上で確認する - Python の `Value` クラスを Mojo の `Tape`（インデックスベース）に置き換えた**設計上の理由**を知る - `main.mojo` で B1〜B7 を組み合わせた**学習・推論の全体像**を把握する --- ## 実行例 `src/part3/microgpt_mojo/` で次のように実行します。 ```bash cd src/part3/microgpt_mojo mojo run main.mojo ``` 出力（学習ステップは途中を省略）： ```text num docs: 32033 vocab size: 27 num params: 4192 step 1 / 1000 | loss 3.2589016467713185 step 2 / 1000 | loss 3.4757539661271393 step 3 / 1000 | loss 3.2850563256139123 step 4 / 1000 | loss 3.2385487413715050 step 5 / 1000 | loss 3.4347922579358574 ... step 996 / 1000 | loss 1.9243386681471330 step 997 / 1000 | loss 2.2576022798634730 step 998 / 1000 | loss 1.6500781354343423 step 999 / 1000 | loss 1.9544453781129398 step 1000 / 1000 | loss 2.5348443763481185 --- inference --- sample 1 : ala sample 2 : nayna sample 3 : anare sample 4 : alania sample 5 : odelen sample 6 : sona sample 7 : adila sample 8 : kais sample 9 : tayna sample 10 : adelel sample 11 : atena sample 12 : koncy sample 13 : maneza sample 14 : aleran sample 15 : amelin sample 16 : tirha sample 17 : saria sample 18 : tanin sample 19 : jaian sample 20 : astyan ``` loss が 3.2 前後から 1.6〜2.5 程度まで下がり、推論では名前らしい文字列（"ala", "sona", "adelel" など）が生成されていることが確認できます。 --- ## 全体構造 ### B1〜B7 の依存関係 各ブロックは以下のように依存しています。矢印は「使う→使われる」の方向です。 ```{mermaid} flowchart LR B1[b01_dataset] --> B2[b02_tokenizer] B2 --> B4[b04_state_dict] B3[b03_value] --> B4 B3 --> B5a[b05_ops] B3 --> B5b[b05_gpt] B4 --> B5b B5a --> B5b B5b --> B6[b06_train] B5b --> B7[b07_infer] B2 --> B6 B3 --> B6 B3 --> B7 ``` **B3（Tape）がすべての土台**です。演算のたびにノードをテープに追記し、逆伝播で勾配を計算します。B4〜B7 はテープ上のノード ID（`Int`）を受け渡しながら動きます。 ### ファイル構成 | ファイル | 対応する Python の概念 | 主な内容 | |---------|----------------------|---------| | `b01_dataset.mojo` | `docs` リスト構築 | テキストから文書リストを作る | | `b02_tokenizer.mojo` | `uchars` / `BOS` / トークン列 | 文字→IDの変換 | | `b03_value.mojo` | `Value` クラス / `backward` | 自動微分テープ | | `b04_state_dict.mojo` | `state_dict` / `params` | モデル重みの初期化・管理 | | `b05_ops.mojo` | `linear` / `rmsnorm` / `softmax` | テープ上のテンソル演算 | | `b05_gpt.mojo` | `gpt()` 関数 | GPT 1ステップの順伝播 | | `b06_train.mojo` | 学習ループの核 | 損失計算・Adam 更新 | | `b07_infer.mojo` | 推論ループの核 | 温度付きサンプリング | | `main.mojo` | スクリプト本体 | 全体を組み合わせた実行ファイル | ### データとコードの流れ ```{mermaid} flowchart TD A["input.txt
（32033件の名前）"] -->|load_docs_from_text| B["docs
List[String]"] B -->|build_tokenizer| C["Tokenizer
(uchars, bos)"] B -->|encode_doc| D["tokens
List[Int]"] C --> D E["Tape"] -->|init_state_dict| F["StateDict
（重み行列 = ノードID行列）"] F -->|flatten_params| G["params
List[NodeId]"] D -->|loss_on_document| H["loss_node
（テープ上のノード）"] E --> H F --> H H -->|backward| I["勾配 grad[i]"] I -->|adam_update| J["更新後の重み"] J -->|gpt_forward| K["logits
List[NodeId]"] K -->|"logits_with_temperature
sample_from_probs"| L["次のトークン ID"] ``` --- ## B1 — データセット（b01_dataset.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b01_dataset.mojo :language: mojo ``` ### 解説 `load_docs_from_text` は `input.txt` の中身（文字列）を受け取り、**空行を除いた行のリスト**を返します。 ```mojo var lines = text.splitlines() # 改行で分割 var line = String(String(lines[i]).strip()) # 前後の空白を除去 if len(line) > 0: # 空行をスキップ docs.append(line) ``` `String(String(lines[i]).strip())` と 2 重に `String()` が現れているのは、型変換が 2 段階必要なためです。 | 式 | 型 | 説明 | |----|-----|------| | `lines[i]` | `StringSlice` | `splitlines()` の要素（元テキストへの参照） | | `String(lines[i])` | `String` | 所有権のあるコピーに変換 | | `String(lines[i]).strip()` | `StringSlice` | 前後の空白を除いた部分への参照（所有権なし） | | `String(String(lines[i]).strip())` | `String` | 所有権のあるコピーに変換 ← これを `append` | `strip()` の戻り値が `StringSlice`（参照）であることが核心です。`List[String]` に `append` するには所有権のある `String` が必要なので、外側の `String(...)` でもう一度コピーを作っています。 関数の末尾で `return docs^` と書いているのは、**`docs` をコピーせず転送（move）して返す**ためです。`^` は「所有権をここで手放す」という演算子で、不要なコピーを避けます。 `load_sample_names` はテスト用の小さなデータセット（3件の名前）を固定文字列から作る便利関数です。 --- ## B2 — トークナイザー（b02_tokenizer.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b02_tokenizer.mojo :language: mojo ``` ### 解説 #### Tokenizer 構造体 ```mojo struct Tokenizer: var uchars: List[String] # 語彙: ソート済みユニーク文字の一覧 var bos: Int # BOS トークンの ID（= len(uchars)） ``` `uchars` は `['a', 'b', ..., 'z']` のような文字リストで、インデックスがそのままトークン ID になります。**BOS（Beginning of Sequence）** は文の先頭・末尾を示す特殊トークンで、文字数より 1 大きい ID を持ちます。 #### トークン化の例 文書 `"ada"` をトークン化すると次のようになります。 ```{mermaid} flowchart LR A["BOS
ID=26"] --> B["'a'
ID=0"] B --> C["'d'
ID=3"] C --> D["'a'
ID=0"] D --> E["BOS
ID=26"] style A fill:#f9c,stroke:#333 style E fill:#f9c,stroke:#333 ``` `encode_doc` は `[BOS, id('a'), id('d'), id('a'), BOS]` というリストを作ります。学習時は位置 `t` のトークンから位置 `t+1` のトークンを予測するタスクになります。 #### 文字の取り出し ```mojo def _char_at(doc: String, j: Int) -> String: return chr(Int(doc.as_bytes()[j])) ``` Python では `doc[j:j+1]` でスライスできましたが、Mojo 0.26 では `String` のスライスが直接 `String` を返さないため、`as_bytes()` でバイト列を取り出し `chr()` で文字に変換しています（ASCII 文字専用の教材向けの実装です）。 #### バブルソート Python の `sorted(set(...))` に相当する処理は `from std.builtin.sort import sort` でインポートした `sort()` で行います。`sort(chars)` と呼ぶだけで `List[String]` をインプレースにソートできます。 --- ## B3 — 自動微分テープ（b03_value.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b03_value.mojo :language: mojo ``` ### Tape の仕組み これが**最も重要なモジュール**です。GPT の学習には「どのパラメータをどの方向に動かすか」を知るための**勾配計算**が必要です。`Tape` はその計算グラフを管理します。 ```{note} **「テープ」という名前の由来** 「テープ」は自動微分の歴史的な用語です。1964 年に R. E. Wengert が提案した「演算を順番に書き留めたリスト（Wengert list）」が起源で、テープレコーダーが音声を順に記録するように、**順伝播中のすべての演算を記録していく**ことからこう呼ばれます。この方式を「テープベース自動微分（tape-based AD）」といいます。 TensorFlow には同じ概念の `tf.GradientTape` という API があり、PyTorch 内部でも "gradient tape" という用語が使われています。 Python 版の `Value` が「微分可能な値 1 個」を表す名前なのに対し、Mojo 版の `Tape` は「演算の記録装置全体」を表す名前です。Mojo の所有権制約で自己参照構造体が作れないため、全ノードを 1 つの構造体に集約する設計になり、その結果としてテープという概念と自然に一致しました。 ``` 演算（加算・乗算など）を呼ぶたびに、`Tape` は新しいノードを 4 つの配列に**追記**します。 | 配列 | 内容 | |------|------| | `data[i]` | ノード i の計算値 | | `grad[i]` | ノード i の勾配（逆伝播後） | | `children[i]` | 子ノードの ID 列 | | `local_grads[i]` | 各子ノードへの局所勾配 | #### 具体例: `c = a + b` ```mojo var a = t.leaf(2.0) # ノード 0: data=2.0 var b = t.leaf(3.0) # ノード 1: data=3.0 var c = t.add(a, b) # ノード 2: data=5.0, children=[0,1], local_grads=[1.0,1.0] ``` テープの内部状態はこうなります。 | index | data | grad | children | local_grads | |-------|------|------|----------|-------------| | 0 | 2.0 | 0.0 | `[]` | `[]` | | 1 | 3.0 | 0.0 | `[]` | `[]` | | 2 | 5.0 | 0.0 | `[0, 1]` | `[1.0, 1.0]` | 加算の局所勾配はどちらも 1.0 です（`∂(a+b)/∂a = 1`、`∂(a+b)/∂b = 1`）。 乗算 `c = a * b` の場合は `local_grads = [b.data, a.data]`（`∂(a*b)/∂a = b`、`∂(a*b)/∂b = a`）になります。 #### 逆伝播アルゴリズム `backward(tape, root)` はルートノードから勾配を伝播します。 ```{mermaid} flowchart TB subgraph "forward（順伝播）" n0["node0
leaf a=2.0"] --> n2["node2
add c=5.0"] n1["node1
leaf b=3.0"] --> n2 end subgraph "backward（逆伝播）" n2b["node2
grad=1.0
(ルートなので1)"] -->|"grad×1.0=1.0"| n0b["node0
grad+=1.0"] n2b -->|"grad×1.0=1.0"| n1b["node1
grad+=1.0"] end ``` 1. `clear_grads()` で全勾配をゼロにリセット 2. **トポロジカルソート**で「子より親が後」の順に並べる 3. ルートの勾配を 1.0 にセット（`∂loss/∂loss = 1`） 4. 逆順に辿りながら `child.grad += local_grad × node.grad` を累積 連鎖律（chain rule）の実装そのものです。 #### sum_nodes で sum() を使わない理由 `sum_nodes` は `xs: List[Int]` を受け取りますが、この `Int` は**数値ではなくノード ID**です。 ```{note} `xs` の中身はたとえば `[42, 57, 63]` のような整数ですが、これは「テープの 42 番目・57 番目・63 番目のノード」を指す**インデックス**です。 `sum(xs)` を使うと `42 + 57 + 63 = 162` という**インデックスの和**が計算されてしまい、まったく意味が違います。 `self.add(a, b)` を繰り返す理由は、単に数値を足すのではなく、**テープに加算ノードを追記する**ためです。 xs = [42, 57, 63] （ノード ID） step1: s = self.add(42, 57) → テープにノード 100 を追記（children=[42,57]） step2: s = self.add(100, 63) → テープにノード 101 を追記（children=[100,63]） return 101 （総和を表すノードの ID） このノードの連鎖があることで、`backward()` 時に勾配が各子ノードへ正しく伝播できます。`sum()` で単純に足すと計算グラフが作られず、逆伝播できません。 ``` --- ## B4 — ハイパーパラメータと重み（b04_state_dict.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b04_state_dict.mojo :language: mojo ``` ### 解説 #### HyperParams ```mojo var hp = HyperParams(1, 16, 16, 4) # n_layer=1, n_embd=16, block_size=16, n_head=4 ``` GPT の形状を決めるパラメータです。`head_dim()` メソッドは `n_embd ÷ n_head`（各アテンションヘッドの次元数）を返します。 #### StateDict — 重みの束 重み行列の各要素は**テープ上のノード ID**（`Int`）として保存されています。つまり `wte[0][3]` は「単語 0 の埋め込みベクトルの 3 番目の要素に対応するノードの ID」です。 ```{mermaid} graph TD subgraph "StateDict（ノードID の行列）" wte["wte
vocab_size × n_embd"] wpe["wpe
block_size × n_embd"] lm["lm_head
vocab_size × n_embd"] wq["attn_wq[layer]
n_embd × n_embd"] wk["attn_wk[layer]"] wv["attn_wv[layer]"] wo["attn_wo[layer]"] fc1["mlp_fc1[layer]
4*n_embd × n_embd"] fc2["mlp_fc2[layer]
n_embd × 4*n_embd"] end subgraph "Tape" nodes["data[], grad[]
children[], local_grads[]"] end wte --> nodes wpe --> nodes wq --> nodes ``` #### matrix_rand — 重みの初期化 ```mojo row.append(t.leaf(randn_float64(0.0, std))) ``` `from std.random import randn_float64` でインポートした `randn_float64(mean, std)` を使い、Python 版の `random.gauss(0, std)` と同等の正規分布乱数で重みを初期化しています。 #### flatten_params Adam 更新ループで全パラメータを一括処理するため、行列の入れ子をフラットな `List[NodeId]` に展開します。 --- ## B5a — 演算ユーティリティ（b05_ops.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b05_ops.mojo :language: mojo ``` ### 解説 すべての関数は `Tape` 上に演算ノードを追加しながら、**ノード ID のリスト**を入出力します。 #### linear（線形変換） $$\text{out}[r] = \sum_{c=0}^{C-1} W[r][c] \cdot x[c]$$ ``` out[r] = x[0]*w[r][0] + x[1]*w[r][1] + ... + x[C-1]*w[r][C-1] ``` 行列 `w`（形状 $n_\text{out} \times n_\text{in}$）の r 行目と入力ベクトル `x`（長さ $C = n_\text{in}$）の内積です。行列とベクトルの積 $\mathbf{out} = W\mathbf{x}$ を1行ずつ計算しています。Python の `sum(wi*xi ...)` に相当します。 #### rmsnorm（RMS 正規化） LayerNorm の簡略版です。各要素を「二乗平均の平方根」で割って正規化します。 $$\text{rms} = \sqrt{\frac{1}{n}\sum_{i=0}^{n-1} x[i]^2 + \varepsilon}, \qquad \text{out}[i] = \frac{x[i]}{\text{rms}}$$ ``` rms = sqrt( (x[0]² + x[1]² + ... + x[n-1]²) / n + ε ) out[i] = x[i] / rms ``` `pow_const(node, -0.5)` で $(\text{rms}^2)^{-0.5} = 1/\text{rms}$ を計算し、これを全要素に掛けています。 #### softmax logits を確率分布に変換します。数値安定化のため最大値 $m$ を引いてから exp を取ります（値が大きすぎて exp がオーバーフローするのを防ぐ）。 $$p[i] = \frac{\exp(\text{logits}[i] - m)}{\sum_j \exp(\text{logits}[j] - m)}, \qquad m = \max_j \text{logits}[j]$$ ``` m = max(logits[0], logits[1], ..., logits[V-1]) e[i] = exp(logits[i] - m) p[i] = e[i] / (e[0] + e[1] + ... + e[V-1]) ``` ```mojo var m = t.node_data(logits[0]) # 現在の値（.data）を取り出し for i in range(1, len(logits)): var v = t.node_data(logits[i]) m = m if m > v else v # 最大値を求める ``` ここで `t.node_data(...)` を呼んでいるのは、**最大値は定数扱い**で微分不要なためです（ソフトマックスの安定化はグラフに乗せない）。 #### scaled_attention_logits Query と Key の内積にスケール $1/\sqrt{d_k}$ を掛けます。 $$\text{logits}[t] = \frac{\mathbf{q} \cdot \mathbf{k}_t}{\sqrt{d_k}}$$ ``` logits[t] = (q[0]*k_t[0] + q[1]*k_t[1] + ... + q[d-1]*k_t[d-1]) / sqrt(d_k) ``` スケールしないと内積が $d_k$ に比例して大きくなり、softmax 後の分布が一点に集中してしまいます。 --- ## B5b — GPT 順伝播（b05_gpt.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b05_gpt.mojo :language: mojo ``` ### 解説 `gpt_forward` が1ステップの GPT 計算です。**1トークン分の入力**を受け取り、**次トークンの確率を示す logits ベクトル**を返します。 #### 埋め込み層 ```mojo var tok = row_embedding(sd.wte, token_id) # トークン埋め込み var pos = row_embedding(sd.wpe, pos_id) # 位置埋め込み var x = ... # tok + pos の要素和 x = rmsnorm(t, x^) # 初期正規化 ``` `row_embedding` は行列の 1 行（= ノード ID の列）をコピーして返します。`mat[idx]` を直接返すのではなく `mat[idx][i]` と**要素ごとに読む**のは、行リストのムーブによる行列破壊を防ぐためです（後述）。 #### Transformer 層 ```{mermaid} flowchart TD X["入力 x
（n_embd 次元）"] --> RN1["rmsnorm"] RN1 --> Q["Q = linear(attn_wq)"] RN1 --> K["K = linear(attn_wk)"] RN1 --> V["V = linear(attn_wv)"] K --> CACHE["KV キャッシュ
（過去のK/V）"] V --> CACHE Q --> ATT["scaled_attention_logits
softmax → 重み付き和"] CACHE --> ATT ATT --> WO["linear(attn_wo)"] WO --> ADD1["残差接続 (+x)"] X --> ADD1 ADD1 --> RN2["rmsnorm"] RN2 --> FC1["linear(mlp_fc1)"] FC1 --> RELU["relu"] RELU --> FC2["linear(mlp_fc2)"] FC2 --> ADD2["残差接続 (+x)"] ADD1 --> ADD2 ADD2 --> OUT["出力 x
（次層へ）"] ``` #### KV キャッシュ 推論時も学習時も、過去のトークンで計算した Key/Value をリストに蓄積します。位置 `t` での Attention は位置 0〜t の K/V すべてに注目できます（因果的 Attention）。 ```mojo append_kv_row(keys[li], k^) # この位置の Key をキャッシュに追加 append_kv_row(vals[li], v^) # この位置の Value をキャッシュに追加 for tix in range(len(keys[li])): k_h.append(head_slice_from_cache(keys[li], tix, hs, hd)) ... ``` `head_slice_from_cache` が `cache[tix]` を丸ごと渡さず `cache[tix][j + off]` と要素アクセスしているのは、行リストのムーブを避けるためです。 #### マルチヘッドアテンション ```mojo for h in range(hp.n_head): var hs = h * hd # このヘッドの開始インデックス var q_h = slice_vec(q, hs, hd) # Q をヘッドごとに切り出す # K/V も同様に head_slice_from_cache で切り出す var hpiece = attn_head(t, q_h, k_h, v_h, hd) heads.append(hpiece^) var merged = concat_heads(heads) # ヘッドを連結して n_embd 次元に戻す ``` `n_head` 個のヘッドを独立に計算して最後に連結します。 --- ## B6 — 学習ステップ（b06_train.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b06_train.mojo :language: mojo ``` ### 解説 #### loss_on_document — 1文書の損失 1文書分のトークン列に対して損失を計算します。 ```mojo for pos_id in range(n): var tid = tokens[pos_id] # 現在のトークン var target = tokens[pos_id + 1] # 次のトークン（正解ラベル） var logits = gpt_forward(...) var probs = softmax(t, logits) var nll = t.neg(t.log_(probs[target])) # 負の対数尤度 losses.append(nll) return t.scale(t.sum_nodes(losses), 1.0 / Float64(n)) # 平均損失 ``` 正解トークン `target` の確率 `probs[target]` を取り出し、その対数の符号反転（NLL）を損失とします。予測が正確なほど確率は 1 に近づき、NLL は 0 に近づきます。 #### adam_update — Adam 最適化 ```mojo var lr_t = lr * (1.0 - Float64(step_idx) / Float64(total_steps)) # 線形学習率減衰 m[i] = beta1 * m[i] + (1 - beta1) * g # 第1モーメント（勾配の移動平均） v[i] = beta2 * v[i] + (1 - beta2) * g * g # 第2モーメント（勾配二乗の移動平均） var m_hat = m[i] / (1 - b1_corr) # バイアス補正 var v_hat = v[i] / (1 - b2_corr) var newv = tape.node_data(pid) - lr_t * m_hat / (sqrt(v_hat) + eps) tape.set_data(pid, newv) # 重みの値を更新 tape.set_grad(pid, 0.0) # 勾配をリセット ``` Adam はモメンタム（第1モーメント）と RMSProp（第2モーメント）を組み合わせた最適化手法です。初期ステップのバイアスを $1 - beta^t$ で補正しています。 --- ## B7 — 推論ヘルパ（b07_infer.mojo） ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/b07_infer.mojo :language: mojo ``` ### 解説 #### logits_with_temperature — 温度スケーリング ```mojo var inv_t = 1.0 / temperature for i in range(len(logits)): scaled.append(t.scale(logits[i], inv_t)) return softmax(t, scaled) ``` `temperature` が低い（例 0.5）ほど logits の差が拡大し、**確率分布が尖って**確定的な出力になります。`temperature=1.0` では logits をそのまま使います。 :::: {container} mermaid-flow-half ```{mermaid} graph LR subgraph "t=0.5（確定的寄り）" L1["logits: [2, 1, 0]"] --> S1["probs: [0.84, 0.11, 0.04]"] end subgraph "t=1.0（フラット）" L2["logits: [2, 1, 0]"] --> S2["probs: [0.67, 0.24, 0.09]"] end ``` :::: #### sample_from_probs — 累積分布サンプリング ```mojo var r = random_float64() # [0, 1) の一様乱数 var cumsum = 0.0 for i in range(len(probs)): cumsum += tape.node_data(probs[i]) if r < cumsum: # 累積確率が乱数を超えた時点で選択 return i ``` 確率に比例した頻度でトークンを選びます。Python の `random.choices(range(n), weights=probs)` と同等です。 ```{mermaid} graph LR subgraph "probs = [0.6, 0.3, 0.1]" A["0.0"] -->|"0.6"| B["0.6"] B -->|"0.3"| C["0.9"] C -->|"0.1"| D["1.0"] end R["乱数 r = 0.75"] -->|"0.6 ≤ 0.75 < 0.9
→ index 1 を選択"| B ``` --- ## main.mojo — 全体を組み合わせる ### ソースコード ```{literalinclude} ../../../src/part3/microgpt_mojo/main.mojo :language: mojo ``` ### 解説 #### データ読み込みとシャッフル ```mojo var f = open("input.txt", "r") var text = f.read() f.close() var docs = load_docs_from_text(text) # Fisher-Yates シャッフル for i in range(len(docs) - 1, 0, -1): var j = Int(random_float64() * Float64(i + 1)) var tmp = docs[i] docs[i] = docs[j] docs[j] = tmp ``` Fisher-Yates シャッフルは末尾から先頭に向かって走査し、各要素をランダムな位置と入れ替えます。学習順序をランダム化することで、特定の文書の偏りを防ぎます。 #### 学習ループ ```mojo for step in range(num_steps): var doc = docs[step % len(docs)] # docs を順番に使い回す var keys = List[List[List[Int]]]() var vals = List[List[List[Int]]]() init_kv_cache(hp.n_layer, keys, vals) # KV キャッシュをリセット var tokens = encode_doc(tok, doc) var loss_node = loss_on_document(t, sd, hp, tokens, keys, vals) backward(t, loss_node) # 勾配を計算 adam_update(t, params, m_buf, v_buf, step, num_steps) # 重みを更新 ``` **各ステップで KV キャッシュを初期化する**のが重要です。KV キャッシュは1文書内の過去トークンを蓄積するもので、ステップをまたいで持ち越してはいけません。 #### 推論ループ ```mojo var token_id = tok.bos # BOS から生成スタート for pos_id in range(hp.block_size): var logits = gpt_forward(...) var probs = logits_with_temperature(t, logits^, 0.5) token_id = sample_from_probs(probs, t) if should_stop(token_id, tok.bos): # BOS が出たら終了 break result += tok.uchars[token_id] ``` BOS トークンを起点に1文字ずつ生成し、再び BOS が出たら（= 文の末尾を予測したら）停止します。`block_size` は生成の上限長になります。 --- ## Python 版との違いと実装上の工夫 ### なぜ Value 型ではなく Tape 型なのか Python 版の `Value` はクラスのインスタンスで、参照によって計算グラフを形成します。 ```python # Python: オブジェクト参照でグラフを作る class Value: def __init__(self, data, children=(), local_grads=()): self._children = children # 別の Value オブジェクトへの参照 ``` Mojo でこれを直接再現しようとすると問題が起きます。 ``` struct Value: var _children: List[Value] # ❌ Value が Value を含む → 無限サイズ ``` Mojo の struct は**固定サイズ**でなければならないため、自己参照型は定義できません。また、所有権ルールにより `List[Value]` 内の要素を `mut` 参照で変更することも困難です。 解決策が **テープ（Tape）パターン**です。 ```mojo struct Tape: var data: List[Float64] # 全ノードの値 var grad: List[Float64] # 全ノードの勾配 var children: List[List[Int]] # 各ノードの子インデックス var local_grads: List[List[Float64]] ``` ノード同士は**インデックス（Int = NodeId）で参照**します。`Tape` は複数の配列の束（SoA: Structure of Arrays）として表現されているので、サイズが固定でなくとも問題ありません。Python で言えば「グラフをオブジェクト参照で持つ」代わりに「グラフをインデックステーブルで持つ」イメージです。 ### `inout`/`owned` の廃止への対応 Mojo 0.26 では `inout`・`owned` キーワードが完全に削除されました。 | 旧構文 | 新構文 | 用途 | |--------|--------|------| | `fn __init__(inout self, ...)` | `def __init__(out self, ...)` | コンストラクタ | | `fn method(inout self, ...)` | `def method(mut self, ...)` | 変更ありメソッド | | `fn func(inout x: T, ...)` | `def func(mut x: T, ...)` | 変更あり引数 | | `fn func(owned x: T)` | `def func(x: T)` + 呼び側で `x^` | 所有権の転送 | `owned` の代替として、呼び出し側で `x^` を使いムーブ渡しします。 ```mojo return Tokenizer(chars^, bos) # chars の所有権をコンストラクタに渡す ``` ### List[T] の所有権 `List[T]` は**コピー不可能**（`ImplicitlyCopyable` でない）です。別の変数に代入すると所有権が移動します。 ```mojo var row = mat[0] # ❌ mat から行をコピーしようとしてエラー var row = mat[0][j] # ✅ 要素（Int）はコピー可能 ``` ネストしたリスト操作では**要素インデックスで個別アクセス**するのが基本方針です。`children` のような `List[List[Int]]` を取り出すときは `.copy()` が必要です。 ```mojo var ch = tape.children[v].copy() # .copy() で明示的にコピー ``` ### 転送演算子 `^` と変数の再代入 ```mojo x = rmsnorm(t, x^) # x を rmsnorm にムーブし、結果を x に代入 ``` `x^` は「x の所有権をここで放棄する」という宣言です。これ以降 x は未初期化になりますが、右辺の戻り値が即座に x に代入されるため問題ありません。 ただし**同じスコープで x を複数の関数に渡す場合**はムーブできないため、`rmsnorm` の結果を新しい変数名で受けています。 ```mojo var x_ln = rmsnorm(t, x^) # x はムーブ後に無効化 # x_res（別途保存した残差）を使う ``` ### 初期化方法の違い | 項目 | Python 版 | Mojo 版 | |------|----------|---------| | 重みの初期化 | `random.gauss(0, std)` | `randn_float64(0.0, std)` | | シード | `random.seed(42)` で固定 | `seed(42)` で固定 | | docs シャッフル | `random.shuffle(docs)` | Fisher-Yates + `random_float64` | | input.txt の取得 | 自動ダウンロード | 手動で配置（要事前準備） | `seed(42)` を `main.mojo` の冒頭で呼ぶことで、重みの初期化・シャッフル・推論サンプリングすべてに同じシードが適用され、Python 版と同様に再現性が確保されます。 --- ## まとめ - B1〜B7 は `src/part3/microgpt_mojo/b0*.mojo` に対応し、`b*_test.mojo` でブロック単位に検証できます - **Tape パターン**は Mojo の所有権制約のもとで計算グラフを安全に扱うための鍵です - Python の参照ベースの設計を、Mojo ではインデックスベースの SoA に置き換えることで、型安全性を保ちながら同等の自動微分を実現しています - `inout`/`owned` の廃止、`List[T]` の所有権制約、`^` 転送演算子など、Mojo 固有のルールに沿った書き方が全体を通して現れています ## 次に読む章 {numref}`MAX を導入する意味`（{ref}`MAX を導入する意味`）へ進みます。