ゼロから作る日本語 LLM
GPT-2 の推論・学習の可視化から Modal での事前学習まで

推論の全体像 — 1 回の forward で 1 トークン

GPT-2 が入力文字列から次の 1 トークンを出すまでにやっていることは以下のとおりです。先に流れを見ておきます。

1. 分割: 文字列を整数の列 (トークン ID) に切る。ニューラルネットは文字を直接食えないため
2. ベクトル化: 各 ID を 768 次元ベクトルに変換し、位置情報も足す。語の意味と語順を 1 本のベクトルに乗せるため
3. 文脈を取り込む: 各位置が他の位置の情報を集めながら、同じ処理 (Attention + MLP) を 12 回重ねる。次のトークンを当てるには自分だけでなく前後の文脈が必要なため
4. 確率化: 最終位置のベクトルを語彙サイズ (50,257) に戻し、softmax で確率分布にしてから 1 トークン選ぶ
5. ループ: 出たトークンを入力の末尾に足してまた 1 に戻る。長文を生成するには、これを何十〜何百回と繰り返す

① Tokenize — 文字列をトークン列に

テキストはまず BPE (Byte Pair Encoding) で分割されます。 GPT-2 の語彙は 50,257 個、各トークンは整数 ID を持ちます。

クリックすると、そのトークンの ID・バイト列・先頭スペースの扱いが下に出ます。

Copy   256   ← 1 バイトで表現できる全パターン (0x00 〜 0xFF)
 50,000   ← 学習で得た BPE merges の数
      1   ← 特殊トークン <|endoftext|>
─────────
 50,257

② Embed — ID をベクトルに

学習済みの 埋め込み が 2 枚あります:

• WE: 50,257 × 768 — 各語彙 ID に対する 768 次元ベクトル
• WP: 1,024 × 768 — 各位置 (0〜1023) に対する 768 次元ベクトル

各トークンについて、この 2 枚から ID に対応する行と位置に対応する行をそれぞれ引いてきて、足し合わせます。 これだけで「その語が」「その位置にある」という 2 つの情報が 1 本のベクトルに乗ります。

縦は位置 (0〜N)、横は次元 (768 を 64 にブロック平均して表示) を示します。 色は 赤が正, 青が負。 行をクリックすると、その位置の足し算を下で詳しく見られます。

③ Attention — 過去を振り返る

ここからが Transformer の心臓部です。② で作ったベクトル列が入ってきて、同じ形のベクトル列が出ていきます。ただし中身は書き換わっていて、各位置のベクトルに「過去の位置の情報」が混ざり込んだ状態になります。

各位置でやるのは 1 つだけです。自分と自分より前の位置のベクトルを、相性のいい位置ほど大きな重みで、合計が 1 になるように加重平均する。それだけです。未来は見ません (これを causal と呼びます)。具体的な重み分布は、このあとヒートマップで本物のデータを動かせます。

どうやって重みを決めるか: Q/K/V

相性を測るために、各位置のベクトルから Query (Q)、Key (K)、Value (V) という 3 つのベクトルを作ります。役割はこんなイメージです:

• Q (Query): いま自分は何を探しているか
• K (Key): 自分は何を提供できるか (看板)
• V (Value): 実際に渡す情報の中身

「位置 5 が位置 2 をどれだけ見るか」は、5 の Q と 2 の K の内積で決まります。その値を softmax で合計 1 の重みに直して、その重みで全位置の V を加重平均する。これが Attention の正味の中身です。

なお各位置での Q/K/V は、元のベクトルに別々の行列を掛けるだけで作ります。その行列自体もデータから学ばれるので、「何を探すか」「何を看板に出すか」「何を渡すか」まで人間が決めずに済むのが Attention のご利益です。式 1 行の正体や $\sqrt{d_k}$ で割るスケーリングの話は付録「Attention の式とスケーリングの意味」で扱います。

マスクで未来を隠す

位置 i が位置 i+1 以降を見てしまうと、答え合わせが見えてしまってカンニングになります。そこで attention 行列の右上三角を softmax 前に −∞ で埋めて、softmax 後の重みを 0 にします。これが causal mask で、下のヒートマップの右上が真っ白なのはその効果です。

1 層で 12 並列 (Multi-Head)

GPT-2 small は、この Q/K/V の仕組みを 1 層あたり 12 個並列で持ちます。12 個のヘッドがそれぞれ別の「注目の切り口」を学んでくれるので、あるヘッドは直前の単語を、別のヘッドは主語と動詞の対応を、といった分業が自然に発生します (Elhage+ 2021)。Head ごとの具体的な次元の割り方や、全位置を 1 回の行列積でまとめて処理できる並列性の話は付録に寄せています。

では実データで眺めていきます。プロンプト全体の N × N ヒートマップと、その下で行 1 つを棒グラフに拡大したもの、2 つを並べます。Layer / Head を動かすと、ヘッドごとに注目パターンの性格がはっきり変わるのが見えます:

行 i 列 j のマスは「位置 i が位置 j をどれだけ見るか」を表します。行の合計は 1、右上三角は causal マスクで全部 0 です。

④ 12 層を積む — ブロックは何をしているのか

③ で眺めた Attention は、実は 1 ブロックの中にあるパーツの 1 つでした。ヒートマップで Layer を 0 〜 11 で切り替えられたのも、この 12 個のブロックの中身を切り替えていたからです。ここからはカメラを引いて、1 ブロックの構造と、それを 12 個積むことで GPT-2 の本体になる全体像を眺めます。

1 ブロックのしていること

1 ブロックは 2 ステップだけです:

1. Attention: ③ で見たとおり、各位置が過去の位置のベクトルを加重平均して情報を集めます
2. MLP: 各位置で独立にベクトルを変換します。位置同士のやりとりはなく、ここでは「集めた情報を加工する」ことに専念します

ひとことで言うと、Attention が位置横断の情報収集、MLP が各位置での変換です。この 2 ステップを交互に繰り返すのが Transformer の基本形になります。

そしてもう 1 つ大事なのが、どちらのステップも結果を上書きしない、という点です。出力を元の入力に足し込むだけで済ませます。これを残差接続 (residual connection, He+ 2016) と呼びます。上書きしないおかげで深く積んでも信号が壊れにくくなります。

ひとことで言うと、GPT-2 は 768 次元のベクトル列を、Attention (位置横断で情報を集約) と MLP (各位置で変換) の 2 ステップで 12 回編集するだけの装置です。これで文章が書けているのは、改めて驚異的だなと思います。ブロックが「上書きせず足し込むだけ」という性質の深い含意は付録「残差ストリームという見方」で扱います。

⑤ Logits → 確率 → サンプリング

12 ブロックを抜けた 768 次元のベクトルを、最後に語彙サイズ (50,257) に戻します。これを ② の埋め込み WE の転置で行うのが GPT-2 の流儀で、出てくるのが語彙ごとのロジット (確率に変える前の生の値) です。softmax を通せば確率分布になり、そこから 1 語を選ぶ、という流れです。

ただ生の確率のままサンプリングすると稀な語が顔を出して文が崩れます。そこで分布を整えるフィルタが 3 つあり、下のスライダーで挙動を動かせます。中身の計算例や使い分けは付録に寄せています。

温度を下げると尖って greedy に近づき、上げると一様分布に近づきます。 top-k/top-p は「上位だけ残して再正規化」するフィルターです。

⑥ 生成ループ — 一語ずつ伸ばす

一語生成したら、それを入力の末尾にくっつけてもう一度モデルに食わせます。 これを繰り返すのが自己回帰生成 (autoregressive generation) です。 下は greedy (各ステップで最確率を採用) のトレースになります。

各ステップで候補 20 個の確率が出ており、一番濃いのが採用された語です。

Copy  (no cache)                       (with KV cache)
  step 1: recompute K₁V₁           step 1: compute K₁V₁, save
  step 2: recompute K₁V₁,K₂V₂      step 2: compute only K₂V₂, append
  step 3: recompute K₁V₁,K₂V₂,K₃V₃ step 3: compute only K₃V₃, append

  per-step attn cost: O(t²)        per-step attn cost: O(t)
  N トークン生成の合計  : O(N³)       N トークン生成の合計  : O(N²)

学習の全体像 — forward → loss → backward → update の繰り返し

推論が「入力 → forward → 次のトークン」で終わるのに対して、学習は 1 ステップが 4 段階に分かれていて、それをひたすら回し続けるループです。

1. forward: 推論と同じ forward pass で、各位置の次トークン確率分布を出す
2. loss: 出た分布と「本当の次のトークン」を比較して、cross-entropy でズレを数値化する。正解に尖った分布ほど loss が小さい
3. backward: loss を下げるには各重みをどう動かせばいいかを、勾配として逆伝播で計算する
4. update: 勾配を使って AdamW optimizer が重みを少しだけ動かす。これを何千〜何億ステップ繰り返す

1 ステップ単位では「確率分布を正解に近づくよう微調整する」だけの地味な操作です。でも、これを何兆トークンぶん繰り返した結果として、モデルが文法や世界知識を持っているように振る舞うようになる、というのが LLM 事前学習の骨格です。

⑦ 学習タスク — 次のトークンを当てるだけ

GPT-2 がやっていることは、たった一つです。文の途中までを見て、次に来るトークンを当てる。これを膨大なテキストで何千億回と繰り返すと、勝手に文法も世界知識も論理推論も身についてくる、というのが 2019 年当時の驚きでした (Radford+ 2019)。具体的にこの「当てる」をどう数値化するか (cross-entropy) は ⑨ で扱います。

データ: WebText という独自コーパス

GPT-2 の学習データは OpenAI が独自に集めた WebText です。Reddit で 3 karma 以上の投稿からリンクされた外部ページだけを集める、という作り方になっていて、「人間がある程度は読む価値があると判断した文章」を自動でフィルタした格好になっています。規模は 8M ドキュメント / 40GB / 約 10B トークンで、2026 年の LLM と比べると相当小さいです (LLaMA 3 で 15T、LLaMA 4 Scout は約 40T トークン規模)。

それでも「自己回帰予測だけ」でここまで文章を書けるモデルが作れると示したのが GPT-2 の功績でした。データ規模やパラメータ数の細かい数字、必要な FLOPs などは付録「学習スケールを数字で」にまとめています。

⑧ Teacher forcing — 並列化の鍵

次トークン予測を素直に実装すると、「1 語予測 → それを入力に戻して次を予測 → …」という逐次計算になりそうです。 でもそれでは GPU が暇になってしまいます。ここで使う工夫が teacher forcing です:

• 学習中は、モデルが生成した語ではなく「正解の語」 xt を次の入力として使う
• こうすると、全位置の予測を 1 回の forward pass で同時に出せる。長さ T のサンプルなら T−1 個の予測が並列化される

入力列を x1..T-1、教師 (ラベル) 列を x2..T と 1 トークンずらすと、 各位置 t が「次は何か?」を独立に答える形になります:

「未来を覗いてカンニング」できないのは、③ で出てきた masked self-attention のおかげです。 位置 t の attention は x1..t までにしか重みを置けないので、x_t+1 以降の情報はモデルに見えません。 だから teacher forcing で正解列を流し込んでも、各位置の予測は正当に保たれます。

ちなみに推論時には teacher forcing は使えません (正解がそもそも無いため)。 推論では 1 語ずつ生成し、それを入力の末尾に追加して再 forward、を繰り返します (⑥ の生成ループ)。 この非対称性が、「学習は全位置並列でサクサク、推論は逐次でじわじわ」という挙動の理由になっています。

⑨ Loss — Cross-Entropy を最小化する

⑧ で「各位置で予測を並列に出す」ところまで来ました。次はその予測と正解のズレを、1 つの数値 (loss) に落とし込みます。ここで使うのが cross-entropy で、やっていることは「正解トークンの予測確率を対数にして、マイナスを付ける」だけです。

確率が 100% 振れていれば −log 1 = 0、50% なら約 0.69、1% なら約 4.6 になります。自信をもって外したときほどペナルティが重くなるのが特徴で、学習の初期に大きな信号を与えてくれます。

これを全位置・全バッチで平均したものが最終的な loss です。式と実数で追いたい方は下の付録をどうぞ。

⑩ 勾配と更新 — loss を下げる方向に重みを動かす

⑨ で「正解とのズレ」が数値になりました。あとは loss が下がる方向に重みを動かすだけです。1 ステップは 3 段階です:

1. forward: プロンプトを流して loss を計算する (⑦⑨ でやった内容)
2. backward: 各重みについて「少し動かしたら loss がどう変わるか」(= 勾配) を求める
3. update: 勾配と逆方向に、重みをほんの少し動かす

勾配は「こっちに動かせば loss が下がる」の矢印

勾配が +0.3 なら「重みを増やすと loss が増える」ので減らす方向へ、−0.5 なら逆に増やす方向へ。GPT-2 small では約 1 億 2 千万個ぶんの重み全部について、これを一気にやります。これで 1 ステップだけ「いまより少しマシ」な重みになります。

backward pass: 上から順に掛け算で降ろすだけ

素朴に「重みを微小に動かして loss の差分を見る」をやると、1 億個ぶんで forward を 2 億回走らせることになって現実的ではありません。そこで使うのが連鎖律 (chain rule) で、これを使うと全層の勾配が forward と同じオーダー (実装上はだいたい 2 倍) の計算量で一気に求まります。

中身を直感だけで言うと、loss から始めて、上から順に「ローカルな微分」を 1 回ずつ掛けながら下の層に流すだけです。各層は「ひとつ上から渡された勾配」を受け取り、重みの勾配をその場に溜めつつ、入力側の勾配だけを下の層に渡します。数式で追いたい方は付録「backward の数式: 連鎖律と residual」をどうぞ。

図で見るとこんな感じです。上が loss、下が入力側 (embed) で、▼ が backward の勾配の流れになります:

residual が勾配を守る

深いモデルでは、連鎖律で勾配が「0 に近い数字の何層もの掛け算」になり、下の層まで届かず消えてしまう (vanishing gradient) 問題があります。ここで効くのが ④ で見た residual connection です。残差接続があると勾配にも「+1 の通り道」がくっつくので、ブロックが何層あっても信号が embed まで届きます。12 層積んでも学習できるのはこれのおかげで、Transformer を深くできる直接の根拠になっています (数式は付録参照)。

重みを更新する — AdamW

勾配が手に入ったら、あとは引き算するだけです。一番素朴な SGD なら「勾配に学習率を掛けて引く」で終わりです。実運用ではその改良版の AdamW (Loshchilov & Hutter 2019) が使われます。AdamW は 3 つの工夫を載せています:

• モメンタム: 過去の勾配の移動平均を使う。ガタつく方向は打ち消され、一貫した方向には慣性がつく
• 適応学習率: パラメータごとに、過去の勾配の大きさで割って歩幅を調整する
• 重み減衰 (weight decay): 毎ステップ重みを 0 にごくわずかに引き寄せる。過学習を抑える

さらに学習率自体もステップで動かすのが定番で、warmup で徐々に上げてから cosine でゆっくり下げます。AdamW の更新式や学習率の具体的な数値例は付録「AdamW の更新式」と付録「学習率スケジュール」にまとめています。

以上が学習ループの 1 ステップで、forward → loss → backward → AdamW をひたすら繰り返します。これを数十万回回すうちに、初期はランダムだった重みが「文章を書ける」状態に収束していきます。GPT-2 XL (1.5B) でも当時 V100 数百台 × 数週間で済んだというのが 2019 年当時の衝撃でした (2026 年のフロンティア LLM は数千 GPU × 数ヶ月、LLaMA 4 Scout で約 40T トークン、Gemma 4 や Qwen3.6 系も trillions of tokens 規模)。

Copystep       0:   η = 0
step   1,000:   η = 1.25e-4      (warmup 途中, 50%)
step   2,000:   η = 2.5e-4       (warmup 完了, 最大)
step  50,000:   η ≈ 2.47e-4      (cosine でゆっくり下降)
step 300,000:  η ≈ 1.37e-4      (中盤, 約 55%)
step 600,000:  η = 2.5e-5        (最小, 10%)

⑪ 実践編とは — 自分で GPT-2 を作る

Part 1 と Part 2 の話は、結局のところ全部「人の作ったモデルを覗き見る」話でした。Attention のヒートマップも、cross-entropy の数値例も、HuggingFace から落としてきた gpt2 small (英語) を後ろから観察したものです。これはこれで仕組みを理解する近道なのですが、自分で動かして初めて分かることが結構あります。たとえば次のようなところ:

• iter 0 のとき loss が ln(vocab_size) ぴったりから始まること (一様分布の cross-entropy そのもの)
• warmup の 500 iter を抜けたあたりで loss が一段下がる気持ちよさ
• iter 500 で生成すると「お風呂にする。 16年、朝起きたとき」みたいな崩壊文だったのに、iter 4000 では「歯周病で歯槽膿漏が起きました」とちゃんと文が成立してくる段階変化
• 110M というサイズで Chinchilla 最適に学習しても、固有名詞や事実関係はほぼ出鱈目になること (これがスケールが効く理由)

なので実践編では、Part 1 / Part 2 の各セクションがコードのどこに対応するかを示しながら、4 つのステップ (トークナイザ学習 → データ準備 → 学習 → 公開) を一通り走らせます。読み終わったあと、自分の手元で modal run --detach scripts/03_train.py と打てば同じものが回せる、という状態を目標にします。

まず最終成果物を先に出しておきます。本稿で作ったものはそれぞれ次の場所に公開してあります:

• コード: github.com/nyosegawa/gpt2-jp-from-scratch
• 学習済みモデル: huggingface.co/sakasegawa/gpt2-jp-small
• W&B プロジェクト: wandb.ai/nyosegawa-mushoku/gpt2-jp-from-scratch
• W&B run (loss curve やサンプル生成のテーブル): runs/cpn804ot

パイプラインの全体像と実測値はこんな感じです。Modal に投げる関数 4 本だけで、ローカルマシンは原則「コマンドを叩く」だけの役割になります。

合計でだいたい 5 時間、コストは Modal の A100 80GB ($2.50/hr) を中心に 10〜13 ドルくらいに収まります。

⑫ インフラ — Modal でクラウド GPU を関数単位で借りる

さて、5 時間で 110M を回す、と言いましたが、これは A100 80 GB を 1 枚使った場合の話です。手元の RTX 3090 (24 GB) で同じ設定を動かすと、bfloat16 が出ないとか TF32 が遅いとかを考慮しても素直に 7〜10 倍くらい遅くなります。1 週間 GPU を専有することになって、その間 LLM が叩けなくなるのは結構つらいです。

なので今回は Modal で A100 を借りました。Modal は「Python 関数にデコレータを付けるだけでクラウド GPU 上で実行されて、終わったら課金が止まる」というサービスで、事前学習みたいな「数時間〜数日だけ走らせて終わり」のワークロードと相性が良いです。料金もシンプルで、現在 (2026-04) は次のようになっています:

嬉しいのが、Starter プラン (個人用、$0/月) に、月 $30 ぶんの無料コンピュートクレジットがついてくる点です (Modal Pricing)。A100 で換算すると約 12 時間ぶんあるので、smoke test (5 分 ≈ $0.21) を何十回やっても月内では無料で終わります。本番の 4 時間学習 ($10) も、月あたり 1 回なら実質 $0 で収まります。なお公式ページには繰り越しに関する記述はありませんが、実運用では月ごとに $30 がリセットされる挙動と理解しておいてください。

では実際にコードを見ます。今回の学習スクリプトはだいたいこんな形でした:

Copyimport modal

app = modal.App("gpt2-jp-train", image=modal.Image.debian_slim().pip_install(
    "torch==2.5.1", "wandb", "sentencepiece", "huggingface-hub", "hf-transfer",
).add_local_dir("src/gpt2_jp", remote_path="/root/gpt2_jp"))

data_volume = modal.Volume.from_name("gpt2-jp-data", create_if_missing=True)
ckpt_volume = modal.Volume.from_name("gpt2-jp-ckpt", create_if_missing=True)

@app.function(
    gpu="A100-80GB",
    volumes={"/data": data_volume, "/ckpt": ckpt_volume},
    secrets=[modal.Secret.from_name("huggingface"),
             modal.Secret.from_name("wandb")],
    timeout=60 * 60 * 10,
)
def train_fn(...):
    from gpt2_jp.trainer import TrainConfig, train
    cfg = TrainConfig.from_yaml("/root/configs/small.yaml")
    cfg.train_bin = "/data/train.bin"
    cfg.out_dir = "/ckpt/out"
    return train(cfg, ...)

要点は次の 4 つだけです:

• modal.Image: 必要な pip パッケージとローカルディレクトリを指定すると、Modal 側でコンテナ image をビルドしてくれる
• gpu="A100-80GB": 関数を実行する GPU の指定 (H100 や L40S も指定できる)
• modal.Volume: tokenizer 学習で作った train.bin を、別の関数 (学習) からも読めるように共有する永続ストレージ
• modal.Secret: HuggingFace の token と W&B の API key を環境変数として安全に注入する

ローカルからは modal run --detach scripts/03_train.py と打つだけで、Modal が image を build → A100 を確保 → 関数を実行 → checkpoint を Volume に保存、まで全部やってくれます。--detach を付けると、ローカルの ssh が切れても Modal 側のジョブは生き続けるので、寝ているあいだに学習が終わります。

学習ログの可視化と実験管理は Weights & Biases (W&B) に流します。これは wandb.init(project=..., name=...) で run を作ると、loss も lr も生成サンプルも全部 web 上のダッシュボードで見られるサービスです。今回は run の URL を public にしてあるので、後半の 結果のセクション で実物を貼ります。

重みの配布は HuggingFace Hub です。学習が終わったら state_dict を Hub に push しておけば、誰でも GPT2LMHeadModel.from_pretrained("sakasegawa/gpt2-jp-small") で読めます。

⑬ 日本語トークナイザを学習する

インフラの話が終わったので、ここからはパイプラインを順に作っていきます。最初は (Part 1 の最初のセクションと同じく) トークナイザです。

結論から言うと、GPT-2 のオリジナル BPE をそのまま日本語に使うのは無理筋です。理由は単純で、英語前提の語彙だと日本語が 3〜4 倍に水増しされてしまうからです。Part 1 の付録でちらっと触れた byte-level BPE を思い出してください。日本語の漢字 1 文字は UTF-8 で 3 バイトに分かれるので、語彙の中に「日」「本」みたいな漢字パーツが入っていないと、漢字 1 文字 = トークン 3〜4 個になります。

実測してみると、GPT-2 BPE と自前 SentencePiece (32k vocab) の差はだいたい次のようになります。比較対象に英語の文も 1 つ混ぜています:

この差はそのままコンテキスト長と学習コストに効いてきます。英語の GPT-2 が 1024 トークンで原稿用紙 2 枚弱を扱うのと同じ感覚で日本語を扱うには、日本語側のトークナイザを別に用意したほうが今回の学習設計では効率的です (同じ 1024 トークンのコンテキストに 2〜3 倍の日本語を詰め込めるうえ、同じ学習予算でカバーできるコーパス量も比例して増えます)。

アルゴリズムの選択: SentencePiece unigram

日本語向けに自前のサブワードを切る選択肢はいくつかあります (BPE, unigram, WordPiece など)。今回は SentencePiece の unigram モデルを選びました。rinna さんの japanese-gpt2-medium など、過去の日本語 GPT 系でも SentencePiece が採用されてきた流れに沿った選択でもあります。理由は:

• 日本語は単語境界が明示されないので、空白前トークン化を前提とする手法 (普通の BPE) より、生テキストをそのまま食わせられる SentencePiece の方が素直
• unigram は確率的に「もっともらしい分割」を選ぶので、形態素境界に近い切れ目になりやすい (BPE は頻度ベースなので「いです」のような半端な切れ目が混ざることがある)
• byte_fallback=True を立てておけば、語彙にない文字 (絵文字など) はバイト単位にフォールバックして、未知語問題が原理的に起こらない

vocab size は 32,000 にしました。GPT-2 オリジナルの 50,257 よりだいぶ小さいですが、日本語と英語と数字とバイトフォールバック (256 個) が入って語彙の頻度分布が滑らかなあたりが、110M モデルの埋め込み行列とのバランスでちょうど良かったです。

学習コーパスと実装

トークナイザ学習用のコーパスは Wikipedia 日本語版 (2023-11-01 ダンプ) と 青空文庫 (clean 版) を組み合わせました。ジャンルが偏らないこと、現代語と文語の両方をカバーすること、を狙って選んでいます。コードの中身はだいたい次の通りです:

Copyimport sentencepiece as spm

spm.SentencePieceTrainer.Train(
    input="/tmp/corpus_sample.txt",
    model_prefix="/tmp/tokenizer",
    vocab_size=32000,
    model_type="unigram",
    character_coverage=0.9995,
    input_sentence_size=500_000,
    shuffle_input_sentence=True,
    num_threads=16,
    normalization_rule_name="nmt_nfkc",
    user_defined_symbols=["<|endoftext|>", "<|pad|>"],
    byte_fallback=True,
    split_digits=True,
    max_sentencepiece_length=16,
)

character_coverage=0.9995 は「上位 99.95% の文字までは語彙でカバーする (残りは byte fallback)」という意味で、漢字の長い裾野を持つ日本語ではこの値が標準です。split_digits=True は数字を 1 文字ずつにバラす設定で、これを入れておくと 2026 が 1 トークンに固有化されたりせず、数字の汎化が効きやすくなります。

Modal で 64 GB マシンを借りて回しました。実時間で 17 分くらいです。生成された tokenizer.model は 800 KB の小さなファイルなので、後段の学習スクリプトでも HuggingFace push の際にも、これをそのまま同梱します。

⑭ 学習データを準備する — 2.3B トークン

トークナイザができたので、次は事前学習用のコーパスを tokenize して、ディスクに連続した整数列として並べます。ここで「何トークン使うか」を決める必要があるのですが、これは 1 つ前のサブトピックとしてすぐ片付くので、Chinchilla の話を短く挟みます。

何トークン食わせるか: Chinchilla 則

大雑把に言うと、2022 年までの事前学習では「同じ計算予算なら、パラメータをなるべく大きくすれば勝ち」という空気でした。ここに待ったをかけたのが DeepMind の Hoffmann et al. (2022) の Chinchilla 論文です。論文のメッセージはシンプルで、「計算予算 $C$ が与えられたとき、パラメータ数 $N$ と学習トークン数 $D$ は、ほぼ等しい割合で一緒に増やすのが最適」というもの。具体的には次の対応が経験的に成り立つ、と示しました:

言い換えると、$N = 110\text{M}$ のパラメータを持つモデルを「与えられた計算量で最大限賢く」するには、だいたい $110\text{M} \times 20 = 2.2\text{B}$ トークンを食わせるのが筋、という指針です。これより少ないと underfit (データが足りず賢くなりきらない)、多いと overfit ではなく単に計算効率が悪くなる (同じ計算を使って、$N$ を大きくした方が loss が下がった)、というのがこの論文の主張の核でした。

なので今回も Chinchilla 則を素直に採用します。目標は $110\text{M} \times 20 = 2.2\text{B}$ トークン。キリの良いところで 2.3B トークンに切りました。

コーパス: FineWeb-2-edu-japanese

2.3B トークンを何から取るか、が次の問題です。ここは hotchpotch/fineweb-2-edu-japanese の sample_10BT サブセットをそのまま使わせてもらいました。このデータセットは セコンさん (@hotchpotch) が 2025-02-20 の記事で公開したもので、HuggingFace 公式の多言語 web コーパス FineWeb-2 の日本語版 (376M 件) を、DeepSeek で作った「教育的価値スコア」の教師データで分類器 (fineweb-2-edu-japanese-classifier) を学習させ、その分類器でスコア 2.5 以上の文書だけを残して作られています (結果およそ 120M 件 / 89.3B トークン)。

英語には元々同系列の FineWeb-Edu があって、「Common Crawl をそのまま食わせるより、教育文書っぽく抽出したサブセットの方が同じトークン数でも性能が伸びる」というのは先行研究で知られていました。日本語には該当するものが長らく無かったので、今回の 110M スケールだと一番ありがたい選択肢になっています。

実装上は sample_10BT を streaming でロードして、先頭 500 万トークンを val にとってから、残りを train として 2.3B トークンに到達するまで詰め込みます。tokenize は SentencePiece の Encode(batch, num_threads=16) でバッチ並列化してあるので、Modal の 16 CPU で 22 分程度です。

Copyfrom datasets import load_dataset
import numpy as np
import sentencepiece as spm

sp = spm.SentencePieceProcessor()
sp.Load("/root/tokenizer/tokenizer.model")
eos_id = sp.PieceToId("<|endoftext|>")

ds = load_dataset(
    "hotchpotch/fineweb-2-edu-japanese",
    name="sample_10BT", split="train", streaming=True,
)

# 文書ごとに <|endoftext|> を挟みながら uint16 で連続書き込み
with open("/vol/train.bin", "wb") as f:
    for batch in chunks(ds, 2000):
        ids_list = sp.Encode([row["text"] for row in batch],
                             num_threads=16, out_type=int)
        for ids in ids_list:
            ids.append(eos_id)
            f.write(np.asarray(ids, dtype=np.uint16).tobytes())
        if total_tokens >= 2_300_000_000:
            break

トークン ID を uint16 (2 バイト) で書いているのは、語彙サイズが 32,000 で 16 ビットに収まるからです。これで 2.3B トークン = 4.6 GB に収まり、A100 のホストメモリにも余裕で乗ります。学習側はこのファイルを numpy.memmap で開き、ランダムなオフセットから block_size=1024 の窓を切り出してミニバッチを作ります。Part 2 の x = batch[:, :-1], y = batch[:, 1:] の話そのままです。

⑮ モデル実装 — nanoGPT 流の最小構成

データができたので、ようやくモデル本体です。実装は karpathy/nanoGPT をベースにしました。nanoGPT は GPT-2 の最小実装として有名で、モデル定義本体が約 300 行に収まっています (README 曰く「~300-line GPT model definition」)。Part 1 の図がそのままコードに対応していて読みやすいです。

まず config と Block の構造はこんな感じです (Part 1 で何度も見た「LayerNorm → Attention → 残差 → LayerNorm → MLP → 残差」がそのまま出てきます):

Copy@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 32000
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = False  # PaLM 以降の慣習で bias を外す。原典 GPT-2 (TF) は bias 込みで学習されているので、HF GPT-2 形式に export するときは zero bias を埋める

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln_1 = LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.attn = CausalSelfAttention(cfg)
        self.ln_2 = LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.mlp  = MLP(cfg)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # 残差 + attention
        x = x + self.mlp(self.ln_2(x))    # 残差 + MLP
        return x

Attention は F.scaled_dot_product_attention(..., is_causal=True) を呼ぶだけです。PyTorch 2 系では入力 dtype (fp16/bf16)・head_dim・GPU 世代 (SM 8.0+) などの条件を満たしていれば、自動で FlashAttention-2 や cuDNN attention の融合カーネルに dispatch されます (PyTorch SDPA docs)。条件が合わない場合 (CPU 実行や fp32) は math backend に fallback します。Flash Attention は大雑把に言うと「$QK^\top / \sqrt{d_k}$ の巨大な行列をわざわざ HBM (GPU のグローバルメモリ) に書き出さず、SRAM (GPU 上の高速な小さいキャッシュ) に載る小さなブロック単位でそのまま softmax まで済ませる」実装です (Dao et al. 2022)。普通に q @ k.transpose(-2, -1) → mask → softmax → ... @ v と書くと、シーケンス長 $N \times N$ の巨大な中間テンソルを HBM に往復させることになって、計算そのものよりメモリ帯域がボトルネックになります。Flash Attention はその往復を消してしまうので、概ね 1.5〜3 倍速くなり、使うメモリも減る、というのが仕組みです (公式 docs は backend ごとに浮動小数点演算の融合順序が変わるため bit-exact ではないと明記しているので、結果は fp16/bf16 の許容誤差内で一致する程度に思っておくのが安全です)。今回の 結果のセクション の 156K tok/s という数字は、これを踏んだうえでの値です。

実装としてはこの 1 行の裏に、Part 1 で散々眺めた causal mask 付きの $\text{softmax}(QK^\top / \sqrt{d_k})\,V$ が全部圧縮されている格好です。

小さな工夫: weight tying と c_proj 補正

モデル全体の組み立てで、地味だけど重要な工夫が 2 つあります:

Copyself.transformer.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
self.transformer.wte.weight = self.lm_head.weight  # ← weight tying

for pn, p in self.named_parameters():
    if pn.endswith("c_proj.weight"):
        nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))

1 つ目の weight tying は、入口の embedding (wte) と出口の lm_head の重みを同じテンソルとして共有する話です。「ID → ベクトル」と「ベクトル → ID」が逆向きの操作だから、同じ行列を使い回せばいい、という発想で、これでパラメータが $|V| \times d_\text{model} = 32000 \times 768 \approx 25\text{M}$ 個も節約できます。32k vocab で全 110M のうちの 22% に相当するので、これがあるかないかで非埋め込みパラメータ数が大きく変わります。

2 つ目の c_proj 補正は、各 Block の出口の重み (attention の c_proj と MLP の c_proj) だけ初期化の標準偏差を $0.02 / \sqrt{2L}$ に下げる、という GPT-2 流のトリックです。Block を素直に積むと残差ストリームの分散が層ごとに増えていって学習初期に発散しやすいのですが、出口を細めに初期化しておくと層数 $L$ が大きくても入口での分散がほぼ保たれます。これが入っていないと iter 100 あたりで loss が NaN に飛ぶことが結構あります。

パラメータ数を集計すると、最終的にこんな内訳になります:

Chinchilla 換算で「$110\text{M} \times 20 \approx 2.2\text{B}$ トークン」と言ったのは、この非埋め込みパラメータを基準にした数字です。

⑯ 学習ループ — 4 時間で 4,386 iters

モデルもデータも揃ったので、いよいよ学習ループです。Part 2 (学習タスク から 勾配と更新 まで) で見てきた「forward → loss → backward → AdamW」を素直にコードに落とすだけなのですが、A100 の VRAM (80 GB) を使い切るためのちょっとした工夫が入ります。

バッチ構成と gradient accumulation

まずバッチ構成です。A100 80 GB に batch_size=16, block_size=1024 のミニバッチを bfloat16 で乗せると、forward / backward 込みでだいたい 60 GB くらい使います。これで 1 ステップあたり $16 \times 1024 = 16{,}384$ トークン。残念ながらこれだと「512 sequences で 1 ステップ」という GPT-2 オリジナルのバッチに届きません。

そこで gradient accumulation を 32 回に設定します。1 つの optimizer step を打つ前に、バッチ 16 を 32 回ぶん forward+backward して勾配をまとめて貯める手法です。これで:

コードはこんな感じです:

Copyscaler = torch.amp.GradScaler(enabled=False)  # bfloat16 では不要
ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)

for it in range(max_iters):
    lr = get_lr(it, warmup=500, decay_iters=4386,
                max_lr=6e-4, min_lr=6e-5)
    for pg in optimizer.param_groups:
        pg["lr"] = lr

    for micro in range(grad_accum_steps):  # 32 回
        x, y = train_ds.get_batch(batch_size, block_size, device)
        with ctx:
            _, loss = model(x, y)
            loss = loss / grad_accum_steps
        loss.backward()                      # 勾配を貯めていく

    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

ポイントを 4 つ並べておきます:

• autocast(dtype=bfloat16) で、行列積など重い演算だけ bf16 に落とす (LayerNorm や softmax は fp32 のまま)。A100 / H100 / RTX 4090 で速い
• bfloat16 は表現範囲が float32 と同じ (指数 8 ビット) なので、勾配スケーリング (GradScaler) が要らない。fp16 と違ってここがすごく楽
• loss = loss / grad_accum_steps で割っておかないと、貯まる勾配が 32 倍になる。意外と忘れがち
• clip_grad_norm_(., 1.0) は GPT-2 原論文では明示されていないものの、GPT-3 の付録で明記され、以降 LLaMA や nanoGPT を含め「GPT-2 系の学習デファクト」として受け継がれています。Attention の確率分布が一時的に偏ったときの暴走を防ぐのが目的

学習率スケジュール

学習率は Part 2 の 勾配と更新 — backward と AdamW の付録でやった warmup + cosine decay をそのまま使います。今回のスケールに合わせて:

max_lr が 6e-4 なのは、有効バッチサイズ 524k tokens (約 0.5M) というスケールに合わせた経験則の選択です (バッチを大きくすると最適 lr もある程度大きくできる)。warmup 500 は全体の約 11% で、これも GPT-2 系の標準的な値です。

eval とサンプル生成

250 step ごとに「val 100 batch で平均 loss を計算」する eval を回しています。本当はもっと長い split を全部舐めるのが綺麗ですが、4,386 step に対して eval だけで 30 分以上かかると目も当てられないので、100 batch (約 1.6M トークン) のサンプリング推定にしました。

500 step ごとには 4 つのプロンプト (むかしむかし、あるところに, 日本の首都は, 人工知能とは, 吾輩は猫である。) で実際に 80 トークン生成して、jsonl と W&B Table に書き出します。これがあるおかげで「数値の loss」と「実際の生成」を後で並べて見ることができて、次の 結果のセクション で示すように、学習の段階性がすごくよく見えます。

⑰ 結果 — loss curve と生成の段階変化

ここまで仕込んだら、あとは modal run --detach scripts/03_train.py を叩いて寝ます。本番の前に --smoke-test で 30 iter だけ回してエラーがないことを確認しておくのがおすすめです (5 分で終わります)。本番は実時間で 4 時間 10 分でした。

Loss curve

iter 0 で loss は 10.5532 から始まりました。これは $\ln(32000) \approx 10.3735$ のすぐ近くで、Part 2 の Loss — Cross-Entropy を最小化する で見た通り「初期化直後のモデルが語彙上ほぼ一様な分布を吐いている」ことを意味します (差 0.18 は init 標準偏差や label smoothing なしのバラつきの範囲です)。逆に言うと、ここで ln(V) から桁単位でずれていたら、初期化・loss の reduction (mean/sum)・label shift のどこかがおかしいので、生成を始める前にチェックを挟むのが安心です。

ここから 4,386 iter のあいだ loss がどう動いたかを下に貼ります (train 線は同梱の training_log.jsonl を 10 iter 刻みでそのまま、val 線は EVAL_POINTS として 250 iter ごとに 100 batch ぶんを平均した値を app.js に同梱しています。下表の数値も EVAL_POINTS 由来です):

数値で書くと、節目はだいたいこうなりました:

best val loss は 3.0386 (iter 4000) でした。ここで $\text{perplexity} = \exp(3.04) \approx 20.9$ なので、平均すると次のトークンを「20 個くらいの候補のなかから 1 つ選ぶくらいの自信」で出している状態です。32k vocab に対しての 20 強なので、まあまあ言語モデルとして機能しはじめている雰囲気を感じます。

実際の W&B run はこちらに置いてあります (上のグラフは jsonl から自前で描いたものですが、W&B 側だと train/lr や train/tokens_per_sec も同時に見られます): wandb.ai/.../runs/cpn804ot。

生成サンプルの段階変化

loss の数字よりも面白いのは、500 step ごとに走らせた生成サンプルがどう変わっていくか、です。同じプロンプト (むかしむかし、あるところに や 日本の首都は) に対する出力を、iter ごとに並べて見てみます:

iter 500 の段階では「同じトークンの繰り返し」と「文として成立していない助詞並び」が支配的で、これは先ほどの表で iter 500 の eval loss が 4.68 だったこと ($\text{perplexity} = \exp(4.68) \approx 108$) と辻褄があいます。「次のトークンが 100 個の候補から 1 つ」というレベルだと、まだ単語の連鎖がまともに学習されていないわけです。

iter 1500〜2500 になると、文単位での日本語の「形」は出てきます。ただし内容は支離滅裂で、「広島 = 首都」とか「東京の人口比率は 3,835.147%」とか、明らかに事実と違うことを自信たっぷりに書きます。これは 110M というサイズで「事実」を覚えるには容量が足りないことを直接見ているところで、より大きいモデルが必要な理由でもあります。

iter 4000〜4386 だと、最初のフレーズに引きずられて文法的に自然に続きが書けるようになっています。「日本の首都は東京です!」までは合っていて、その後の数字は適当ですが、文として読めなくはない、というレベルになりました。これくらいまで行ければ、110M の事前学習モデルとしてはまずまず合格だと思います。

⑱ 公開 — HuggingFace と W&B

モデルが学習できたので、ここから公開作業です。「自分の checkpoint をローカルに置いて満足する」だと再現性も拡張性もないので、最終的には HuggingFace Hub に置いて誰でも from_pretrained で読める状態にします。

HuggingFace に push する

nanoGPT 風の state_dict をそのまま置くと transformers 側からは形が合いません。理由は、nanoGPT の attention・MLP 射影が nn.Linear (weight.shape = (out, in)) なのに対し、HF の GPT2LMHeadModel は内部で Conv1D (weight.shape = (in, out)) を使っていて、向きが逆だからです (huggingface/transformers/.../modeling_gpt2.py)。さらに bias=False で学習しているので Conv1D / LayerNorm の bias 用にゼロを埋める必要があります。tokenizer 側も tokenizer.model (SentencePiece) を 1 ファイル置いただけだと AutoTokenizer が tokenizer class を解決できないので、最低限 tokenizer_config.json を 1 つ足してやる必要があります。

まとめると、HF 互換に push する前に必要な手続きはこの 3 つです (学習レポ scripts/04_push_to_hub.py に Modal ジョブとしてまとめてあります):

CopyTRANSPOSE = ("attn.c_attn.weight", "attn.c_proj.weight",
             "mlp.c_fc.weight",  "mlp.c_proj.weight")

sd = torch.load("ckpt_best.pt", map_location="cpu", weights_only=False)["model"]
sd = {k.removeprefix("_orig_mod."): v for k, v in sd.items()}
sd = {k: v for k, v in sd.items()
      if not k.endswith(("attn.bias", "attn.masked_bias"))}
for k in list(sd):
    if any(k.endswith(s) for s in TRANSPOSE):
        sd[k] = sd[k].t().contiguous()  # ← Linear → Conv1D の向きに直す

cfg = GPT2Config(vocab_size=32000, n_positions=1024, n_embd=768,
                 n_layer=12, n_head=12,
                 bos_token_id=1, eos_token_id=1,  # SP の <|endoftext|> = 1
                 tie_word_embeddings=True)
model = GPT2LMHeadModel(cfg)
model.load_state_dict(sd, strict=False)  # bias 系の missing は zero 埋めで OK
model.save_pretrained(out_dir)            # ← model.safetensors と config.json

# SentencePiece を AutoTokenizer から読めるように LlamaTokenizer でラップする。
# 特殊トークンは SP 側に最初から入っている ID をそのまま使う
# (id 0: , id 1: <|endoftext|>, id 2: <|pad|>)
tok = LlamaTokenizer(vocab_file="tokenizer.model",
                     bos_token="<|endoftext|>", eos_token="<|endoftext|>",
                     unk_token="", pad_token="<|pad|>",
                     add_bos_token=False, add_eos_token=False, legacy=False)
tok.save_pretrained(out_dir)  # ← tokenizer_config.json などを書き出す

HfApi().upload_folder(folder_path=out_dir,
                      repo_id="sakasegawa/gpt2-jp-small")

ポイントは 3 つです:

• nanoGPT の Linear 重みを HF Conv1D の向きに転置する。対象は attn.c_attn, attn.c_proj, mlp.c_fc, mlp.c_proj の 4 種類だけ (nanoGPT 自身の from_pretrained でも逆方向で同じ転置が走っています)
• tokenizer.model 単体だと AutoTokenizer が tokenizer class を解決できないので、LlamaTokenizer (SentencePiece backend) でラップして save_pretrained する。これで tokenizer_config.json と special_tokens_map.json が生成され、利用側は AutoTokenizer.from_pretrained 1 行で読めるようになる。特殊トークンは SP の語彙に最初から入っている ID (0: <unk>, 1: <|endoftext|>, 2: <|pad|>) を再利用するのが大事で、新しいトークンを足してしまうと埋め込み行列のサイズを超えて IndexError になります
• README.md に「学習データ」「ハイパーパラメータ」「使い方」「val loss」を書いておく。これがそのままモデルカードになる

公開した repo はこちらです: huggingface.co/sakasegawa/gpt2-jp-small。読み込みは普通の from_pretrained で済みます:

Copyimport torch
from transformers import GPT2LMHeadModel, AutoTokenizer

model = GPT2LMHeadModel.from_pretrained("sakasegawa/gpt2-jp-small")
tok   = AutoTokenizer.from_pretrained("sakasegawa/gpt2-jp-small")

ids = tok("日本の首都は", return_tensors="pt").input_ids
y = model.generate(ids, max_new_tokens=50, top_k=40, do_sample=True)
print(tok.decode(y[0], skip_special_tokens=True))

W&B の学習ログ

loss curve と生成サンプルの Table は W&B に上げていて、誰でも読めるように public にしてあります。学習中の挙動をそのまま眺めたい方はどうぞ:

• Project: wandb.ai/nyosegawa-mushoku/gpt2-jp-from-scratch
• Run (本番): wandb.ai/.../runs/cpn804ot

⑲ 振り返りと次の一歩

4 時間と 17 ドル弱で、val loss 3.04 / perplexity 20.9 の日本語 GPT-2 small ができあがりました。これが何を意味するかというと:

• Part 1 で図にした構造は、本当に 200 行ちょっとの PyTorch コードに落ちて学習が回る
• Part 2 で見た「次トークン予測 + cross-entropy + AdamW + warmup+cosine」は、110M スケールでも素直に動いて loss を 10.55 → 3.04 に下げる
• Chinchilla 最適 (110M × 20 = 2.2B) だと、文法的に成立する日本語は出るが、事実を「書く」レベルには届かない

最後の点について、少しだけ補足します。今回の生成サンプルは「文として成立してきた」けれど、「日本の首都は広島」みたいな事実誤りが平気で出てきました。これは Chinchilla 最適でも、110M というサイズでは知識としての密度が不足している、という直接の証拠です。100B 以上の事前学習データを 1B モデルに食わせる、いまの「Chinchilla 超え」のレシピは、ここを腕力で押し切っている格好です。

ここから先に進むなら、たぶん次の順で広げていくと面白いと思います:

• 同じ 110M で Chinchilla の 5〜10 倍 (10B〜20B トークン) まで学習を伸ばしてみる。事実の密度がどう変わるかを観察する
• medium (350M) や large (760M) にスケールアップする (A100 80 GB の 1 枚でも一応収まる)
• 事前学習後に SFT (instruction tuning) を 1 〜 2 epoch 入れる。生成サンプルの「役立ち感」が一気に変わる
• vLLM や SGLang に gpt2-jp-small を載せて、KV cache 込みで推論ベンチマークを取ってみる (Part 1 の 生成ループ — 一語ずつ伸ばす で触れた話の実測)

ここまで読んでくださってありがとうございました。記事や実装に間違いを見つけたら、リポジトリの issue か @gyakuse までこっそり教えてください。