③文系東大生がGPT型モデルをフルスクラッチで開発する話 Part3【モデル実装編】【株式会社ハイレゾ協力】

Last updated at 2026-06-24Posted at 2026-06-24

【機材協力について】
本記事（および本連載プロジェクト）は、株式会社ハイレゾ様よりGPUのクラウドサービス「GPUSOROBAN」の計算資源（NVIDIA A100 80GB）を無償でご提供いただき、開発・検証を行っています。

0．はじめに

こんにちは！東京大学文科3類2年で、東大AI研究会代表の青木です。
現在株式会社ハイレゾ様の協力のもと、25億パラメーターのGPT型モデルの実装と訓練を行っています。その名もSGT（Scratch Generative Transformer）プロジェクト！
この記事は、Part3【モデル実装編】として、実際に学習を行うGPT型AIモデルの実装過程をお送りします。
(Part1【企画背景編】、Part2【データ準備編】はリンクからどうぞ。)
　
　私たち東大AI研究会は東大で活動するサークルで、1年かけて0からGPT型モデルを開発するために、毎週水曜日に勉強会を開催しています。勉強会では初心者でも理解しやすいスライドと、充実したカリキュラムを用いて開発を進めています。

　株式会社ハイレゾ様は自社でGPUデータセンターを運営し「GPUSOROBAN」というGPUのクラウドサービスを提供する会社です。高性能GPUは非常に高価で、私たちのような学生サークルが簡単に手を出せるような代物ではありません。GPUSOROBANではクラウド上で必要な期間・量だけGPUを借りることができるため、初期投資を抑えてGPUを活用し開発を行うことができます。

1.モデル実装

前回、学習用のデータを約600億トークン分を用意しました。次にそのデータを学習するGPT型AIモデルを作成していきます。モデルは東大AI研究会のカリキュラムで作成した基本的なGPT型モデルをそのまま使用します。

2.モデルアーキテクチャ

モデルの概要は以下の図のようになっています。

今回実装するモデルはDecoder-only Transformerを採用した自己回帰型言語モデルです。つまり、過去の文から次の文字を予測していきます。

⓪Tokenize

コンピューターが文章をモデルへ入力できるように、文章を「トークン」と呼ばれる単位に区切り、それぞれを数字に変換します。この数字を「トークンID」と呼びます。
トークンは単語そのものとは限らず、単語の一部分や記号などになることもあります。
今回はモデル全体のトークナイズはモデルの学習を開始する前に行ったので、テキストのトークナイズはモデルに含まれていません。トークナイズの様子は技術記事Part2【データ準備編】（リンク）からご覧ください。

①Token Embedding

Token Embeddingでは、それぞれのトークンIDを、2,560個の数値からなるベクトルへ変換します。トークンIDは単なる番号であり、番号の大小そのものに意味はありません。
そこで、モデルが学習できる2,560次元のベクトルへ変換することで、トークンの特徴を数値として表現します。
このベクトルの値は、モデルの学習を通じて更新されます。ただし、この段階では同じトークンは基本的に同じベクトルへ変換されます。文章中の文脈に応じた意味は、この後のTransformer Blockによって加えられていきます。

②Transformer blocks構造・Attentionの方式

まず、各トークンの2,560次元のベクトルにLayerNormを適用します。
LayerNormは、それぞれのトークンベクトルについて、各次元の値の平均やばらつきを整える処理です。値の大きさが極端になることを抑え、学習を安定させる効果があります。
次に、Multi-Head Self-Attentionを使い、文章内の各トークンが他のトークンとどのような関係を持つかを計算します。
今回のモデルでは、Attentionを20個のAttention Headに分けて計算します。複数のHeadを使うことで、単語同士の関係を異なる観点から捉えられるようにします。
GPTは文章を左から右へ生成するため、あるトークンを処理するとき、それより後ろにある未来のトークンを参照できないようにします。この仕組みをCausal Maskと呼びます。
また、今回のモデルではRoPEを使用し、各トークンが文章内のどの位置にあるかという情報もAttentionの計算へ取り入れます。
Attentionの計算後には、処理前の入力ベクトルを計算結果へ足し戻します。これをResidual Connectionと呼び、深いモデルでも情報や勾配を伝わりやすくします。
ただし、Attention計算には、PyTorchのScaled Dot-Product Attentionを使用しました。 RoPEとScaled Dot-Product Attentionについては、3．実装方針で雰囲気だけ説明しています。

③FFNの構造

Attentionの後には、もう一度LayerNormを適用し、続いてFeed Forward Networkへ入力します。
Feed Forward Networkでは、それぞれのトークンを表す2,560次元のベクトルを、いったん10,240次元まで拡張します。
その後、ReLUと呼ばれる活性化関数を通し、再び2,560次元へ戻します。
処理の流れは次のとおりです。
2,560次元
→ 10,240次元
→ ReLU
→ 2,560次元
Attentionがトークン同士の関係を計算するのに対して、Feed Forward Networkは、それぞれのトークンの表現を個別に変換します。
Feed Forward Networkの計算後にも、処理前の入力を足し戻すResidual Connectionを適用します。
今回のモデルでは、AttentionとFeed Forward Networkを組み合わせたTransformer Blockを30回繰り返します。

④出力層

Transformer Blockを30回通過した後、最後にもう一度LayerNormを適用します。
その後、各トークンの2,560次元のベクトルを、語彙数と同じ50,257次元へ変換します。
この50,257個の数値は、語彙に含まれるそれぞれのトークンが「次に来そうな度合い」を表しています。この予測値をLogitsと呼びます。
文章生成時には、このLogitsをもとに次のトークンを1つ選びます。そして、選ばれたトークンを入力へ追加して同じ処理を繰り返すことで、文章を少しずつ生成していきます。

3.実装方針

今回はNVIDIA A100 80GB × 1枚を二か月間使用します。
大きなモデルを効率よく学習・推論するために、以下の工夫が重要になります。

・RoPE：位置情報をAttentionに取り入れるための仕組み
・SDPA / Flash Attention系の実装：Attention計算を高速化・省メモリ化するための仕組み
・KVキャッシュ：推論時に過去のKey/Valueを再利用して生成を高速化する仕組み
・勾配蓄積：メモリ制約下で実効バッチサイズを大きくするための工夫

・RoPE
・SDPA（Flash Attention）
・KVキャッシュ
・勾配蓄積
これらの処理について、私がどのように理解したかを示します。あくまで私の理解であり、内容の正確性は担保できません。

①RoPE(Rotary Positional Embedding)

RoPEは、トークンの位置情報をモデルに伝えるための仕組みです。
GPT型モデルは、単語の意味だけでなく、単語同士の位置関係も考える必要があります。RoPEを使用することで、Attentionの中に位置情報を自然に取り入れることができます。

②SDPA(Scaled Dot-Product Attention)

SDPAは、Attentionの計算を効率よく行うための仕組みです。
GPT型モデルの中核ともいえるAttentionは、文章が長くなるほど計算量やメモリ使用量が大きくなります。PyTorchのSDPAを使うことで、条件に応じて高速・省メモリなAttention計算が行われ、処理速度を向上させることができます。

③KVキャッシュ

推論の際に、AttentionにおいてKeyとValueの計算結果を一時的に保存しておくことで、計算回数を削減し、推論速度を向上させます。やってること自体はシンプルですが、コードに直すと100行以上になります。

④勾配蓄積

勾配を計算した後に毎回重みを更新するのではなく、計算した勾配を記録しておいて、数回分の勾配の平均を用いてまとめて重みを更新します。これにより、メモリを節約しながらバッチサイズを大きくすることができます。

4.パラメータ規模

今回実装したモデルは、30層のDecoder-only Transformerです。当初は25億パラメータの規模を目標としていましたが、実際の設定から計算した層パラメータ数は26億になりました。

項目	設定値	意味
Transfomer block数	30	AttentionとFFNを含むブロックの数
埋め込み次元	2,560	1トークンを表現するベクトルの次元数
Attention Head数	20	Attentionを20個のHeadに分割
1 Headあたりの次元数	128	2,560 ÷ 20
FFN中間次元	10,240	埋め込み次元の４倍
語彙数	50,257	GPT-2トークナイザーの語彙数
最大系列長	2,048	モデルが一度に処理するトークン数
RoPEの基数	1,000,000	位置情報の回転周波数を決める値
総パラメータ数	約26億

5.実装時に詰まった点

勾配蓄積の実装と、計算速度を考慮した実質バッチサイズの決定に苦労しました。今回はNVIDIA A100 80BG 1ノードを用いて計算するので、バッチサイズが大きすぎるとメモリ不足(Out Of Memory)を起こしてしまいます。約25億パラメーターでは、バッチサイズ４でメモリの限界が来てしまいました。これでは学習が安定せず、せっかくの計算資源が無駄になってしまいます。こういった経緯で、メモリの限界を超えたバッチサイズを再現できる勾配蓄積を実装しました。コード自体は単純なのですが、学習の速度に強く影響する実質バッチサイズを変えながら、学習と計算速度の確認を繰り返しました。バッチサイズの調整には結構時間がかかりました...
最終的に実質バッチサイズは512に設定しました。

6.小規模LLM開発で見えた課題

小規模LLM開発では『長期間、安定的に学習させる環境づくり』が大切になると実感しました。

「5.実装時に詰まった点」でも紹介した通り、汎化性能と向上だけを求めてパラメーター数を無駄に増やしたり、バッチサイズを大きくしてしまうと、メモリ不足を起こしてしまいます。いくら学習効率が良くても、学習が続かなければ意味はありません。逆にメモリを優先してバッチサイズを小さくすると、今度は学習が安定しなくなってしまいます。これでは学習効率が悪くなり、どれだけ学習を続けても、性能の向上は期待できません。

学習は「長く」そして「適度な効率を維持して」続けなければなりません。規模の差はあれども、この点はどのような開発環境でも意識しなければならないと考えました。「与えられた計算資源を最大限利用して、汎化性能をできるだけ高めたい」という思いを持って、モデルを調整していく経験をしたことで、このことを実感できました。

学習は現在進行中です！次回は 学習・チューニング の様子をお届けします！