Google ColabだけでLLMを作る！0.5BサイズのLLMをフルスクラッチで作る教材「EveryonesLLM」
Google Colabだけで、0.5BサイズのLLM/SLMをフルスクラッチで作る教材 EveryonesLLM を公開しました。
はじめまして、EveryonesLLMのメイン作者の本郷 颯人です。
学生団体「東大AI研究会」に所属しています。
GitHubはこちらです。
https://github.com/HayatoHongo/EveryonesLLM/tree/ja
この記事では、EveryonesLLM がどういう教材なのか、なぜ作ったのか、どこを大事にして設計したのかを書きます。
EveryonesLLMとは
EveryonesLLM は、Google ColabだけでLLMを0から作るための教材です。
単に「LLMを使う」のではなく、データローダー、Embedding、Attention、Transformer Block、学習ループ、事前学習、Instruction Tuning までを、順番に実装していきます。
最終的には、かなり小さいとはいえ、自分で作ったモデルが会話っぽい応答を返すところまで持っていくことを目標にしています。
この教材で大事にしたのは、次の3つです。
- Google Colabだけで完結すること
- なるべく環境構築で苦しまないこと
- LLMの「面白いところ」に集中できること
LLMを作る教材や記事は世の中にたくさんあります。ただ、いざ自分でやろうとすると、強いGPUが必要だったり、クラウドサーバーにSSHで接続する必要があったり、環境構築でかなり時間を持っていかれたりします。
もちろん環境構築も大事です。ただ、最初にLLMを学ぶ人にとっては、そこが一番面白い場所ではないと思っています。
僕が作りたかったのは、Colabを開いて、ノートブックを上から順番に読んで実行していけば、LLMが少しずつ形になっていく教材です。
なぜ作ったのか
1年前くらいの時点では、少なくとも「0.5BサイズくらいのLLMを、3B tokensくらいで学習させて、ある程度まともに会話できるようにする」というのは、個人がGoogle Colabだけでやるにはかなり厳しいと思っていました。
でも、将来的にはColabにもっと良いGPUが載るだろうし、計算環境も少しずつ手軽になるはずだと思っていました。
そうなったときに、
Google Colabだけで、ある程度まともに会話できるLLMを作れる教材があったら面白い
と思って、EveryonesLLMを作り始めました。
僕自身、プログラミングが全然できなかった頃にChatGPTに出会って、「これを自分でも作ってみたい」と思いました。
ただ、そこから実際に学ぼうとすると、かなり大変でした。
大変だった理由の多くは、LLMそのものというより、環境構築や周辺知識でした。GPU環境、クラウド、ライブラリ、設定、謎のエラー。加えて、AIという名前がついているけれど実態としては古典的な統計の講義にかなり時間を使ったりもしました。
もちろん、それらの知識も重要です。ただ、LLMに感動した人が、最初に体験したいのはそこではないはずです。
僕は、高校生でも、LLMを使って感動した人が、今度は「自分でLLMを作って感動する」体験を作りたかったです。
教材の対象者
現在の主な対象者は、PyTorchでMLPを実装したことがある方です。
そのため、PyTorch自体が完全に未経験の方は、まずPyTorchの使い方を身につける必要があります。
ただ、最終的にはプログラミング未経験の高校生でも進められる教材にしたいと思っています。
そのため、実装はなるべくシンプルにしています。数式も、悪い意味で退屈にならない範囲で、できるだけ簡単にしています。
行列の内積など最低限必要なものを除けば、なるべく高校数学で理解できるようにしている、という方針です。
モデル設計について
EveryonesLLM のモデルは、基本的には dense な decoder-only Transformer です。
最近の大規模言語モデルで使われているテクニックを全部盛りにすることはしていません。
むしろ、かなり削っています。
例えば、EveryonesLLM には以下のようなものは入れていません。
- KV Cache
- RoPE
- FlashAttention
- GQA
- SwiGLU
活性化関数も、かなり素朴に ReLU を使っています。GELUですらありません。
最近のLLM実装をよく見ている人からすると、かなり簡単な実装に見えると思います。
これは手抜きではなく、教材としての意図です。
LLMを作るうえで、最新のテクニックを全部理解してからでないと進めない、という構成にはしたくありませんでした。
例えばRoPEは非常に重要な技術ですが、ちゃんと理解しようとするとフーリエ級数などの知識も絡んできます。これは、高校生や初学者にとってはかなり重いです。
もちろん、性能を突き詰めるなら入れた方がよい技術はたくさんあります。ただ、最初にLLMを学ぶ段階では、わずかな性能改善のために複雑さを増やすよりも、「Transformerの本質的な流れを理解できること」の方が大事だと思っています。
一方で、相対位置エンベディングは入れています。
これは実装が比較的シンプルなわりに、性能への影響がかなり大きいと感じたからです。
つまり、EveryonesLLM の方針はこうです。
GPTに近い基本構造を保ちつつ、初学者にとって理解の負担が大きいものはできるだけ削る。
ただし、効果が大きく、教材としても扱いやすいものは入れる。
参考にしたものとしては、nanoGPT や MiniMind があります。ただし、それらをそのまま大きくするというより、教育用にかなり簡略化する方向で作っています。
どこまでできるのか
EveryonesLLM では、Chapter 00 から順番に進めることで、LLMの構成要素を少しずつ実装していきます。
大まかには、次のような流れです。
- データローダーを作る
- Token Embedding を作る
- Position Embedding を作る
- LayerNorm を作る
- Attention を作る
- Multi-Head Attention を作る
- Feed Forward を作る
- Transformer Block を組み立てる
- Trainer を作る
- GPUで学習する
- 学習率やScaling Lawを扱う
- TinyStoriesで学習する
- 相対位置エンベディングを入れる
- Checkpointを扱う
- 大規模WebデータでPretrainingする
- Instruction Tuningする
最終的には、Chapter 29 で Instruction Tuning まで進みます。
コストについて
コストはなんとおよそ3,000円です。
Colabのクレジットでいうと300クレジットです。
実際に会話できるモデルをこのコストで作れるのは
まさにGoogle Colabのおかげです。
この教材で体験してほしいこと
EveryonesLLMで一番体験してほしいのは、
LLMは完全なブラックボックスではなく、自分の手で少しずつ組み立てられるものなんだ
という感覚です。
ChatGPTのようなモデルを見ると、とても遠い存在に感じます。
でも、Token Embeddingを作り、Attentionを書き、Transformer Blockを積み、Trainerを作り、データを流していくと、少しずつ「モデル」が形になっていきます。
最初は意味不明な文字列しか出さなかったモデルが、学習が進むにつれて少しずつ文章らしいものを出すようになります。
さらにInstruction Tuningをすると、会話ができるようになります。
この体験は、かなり面白いです。
LLMを「使う」だけではなく、「作る」側に回ることで、AIを見る目がかなり変わると思います。
なぜここまで簡単にしたのか
最近のLLM実装には、性能を上げるための工夫がたくさん入っています。
それ自体は素晴らしいことです。
ただ、初学者がいきなりそれらを全部理解しようとすると、かなりしんどいです。
RoPE、FlashAttention、GQA、SwiGLU、KV Cache、分散学習、データセット整形、評価、推論高速化……。
全部大事です。でも、全部を最初から入れると、何が本質なのかが見えづらくなります。
EveryonesLLMでは、あえて多くのものを削りました。
その代わり、
- 入力トークンがEmbeddingになる
- Attentionで文脈を混ぜる
- Transformer Blockを積む
- ロスを下げる
- 文章生成が少しずつ変わる
という流れを、できるだけ見通しよくしました。
まずはここを体験してほしいです。
最新テクニックを学ぶのは、そのあとで十分だと思っています。
今後やりたいこと
今はPyTorch初心者を主な対象にしていますが、将来的にはプログラミング未経験の高校生でも進められるようにしたいです。
そのために、説明の追加、図の追加、エラー対応、より丁寧な導入などを進めていく予定です。
また、教材としてはまだ改善できるところがたくさんあります。
「ここがわかりにくい」
「このChapterで詰まった」
「この説明があると嬉しい」
などがあれば、ぜひIssuesで教えてください。
おわりに
EveryonesLLM は、LLMを使って感動した人が、今度はLLMを作って感動するための教材です。
僕自身、ChatGPTに出会って「これを作ってみたい」と思ったところから始まりました。
そこからかなり遠回りもしました。環境構築で苦しんだり、何を学べばいいのかわからなかったり、無駄な苦労もたくさんありました。
だからこそ、これから学ぶ人には、できるだけLLMの面白いところに早くたどり着いてほしいと思っています。
Google Colabを開くだけで、LLMを0から作る。
そんな体験を、できるだけ多くの人に届けられたら嬉しいです。
リポジトリはこちらです。
https://github.com/HayatoHongo/EveryonesLLM/tree/ja
出典
このチュートリアルはAndrej KarpathyさんのnanoGPTとjingyaogongさんのMinimindをベースにしています。Instruction Tuningについては Sebastian Raschka さんの著書Build a Large Language Model (From Scratch)を参考にしています。Vision LLMについては、LLaVAを参考にしています。
この場を借りて、心から感謝します。
謝辞文
本プロジェクト「EveryonesLLM」は、株式会社NTTデータMSE様のご支援のもと、学生団体「東大AI研究会」によるOSSプロジェクトとして完成しました。
本教材の開発にあたり、温かいご支援とご協力をいただいた株式会社NTTデータMSE様に、心より感謝申し上げます。