Skip to content

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

1LM-Blackwell — 1時間で作る自作ミニ言語モデル (Windows + RTX 50)

Windows 11 のネイティブ環境(WSL なし)と NVIDIA GPU 1枚で、ライブラリの学習済みモデルを 一切使わずに 言語モデルを作って会話するまでの一式です。ルールベースの応答ではなく、 コーパスから学習した Transformer が「次の1文字」を予測し続けることで会話が成立します。

  • フレームワーク: PyTorch 2.13 + CUDA 13.0(Blackwell / sm_120 対応 wheel)
  • モデル: 文字レベル ミニGPT / 6層 / 384次元 / 6ヘッド / 文脈256文字 / 11.53M パラメータ
  • データ: kunishou/oasst1-89k-ja(Apache-2.0)を整形した日本語会話 28,616 件
  • 学習時間: RTX 5090 で 61.3 秒(3,600ステップ / 最良 val loss 1.8677)

Apple Silicon / MLX 版の 1LM の移植です。 同じ学習を Windows で再現し、Mac の val loss 1.8571 に対して 1.8677(差 +0.0106) まで 一致させています。移植でどこがずれるかは 同値確認 に書きました。

チャットGUI

できること

インターフェース コマンド
環境診断(最初にこれ) python check_env.py
CLI チャット python src\chat_cli.py
Web GUI(Liquid Glass 風) python server.py --open
単発生成 python src\generate.py --prompt "こんにちは"

1時間の流れ

時間 やること
0:00 - 0:15 環境構築(uv + torch cu130)と check_env.py
0:15 - 0:25 コーパス作成、トークナイザとモデルの説明
0:25 - 0:30 学習(1分で終わるので、待ち時間ではなくコード解説に使う)
0:30 - 0:45 ベンチマーク(torch.compile / bf16 / SDPA / バッチサイズ)
0:45 - 0:56 CLI で会話、サンプリング設定で遊ぶ
0:56 - 1:00 GUI をブラウザで開いて完成

Mac 版で30分かかっていた学習が1分で終わるので、空いた時間を「なぜ速いのか」の実測に回す 構成にしています。

セットアップ

必要なもの: Windows 10/11 / NVIDIA GPU / uv。

# 1. 環境変数(ユーザースコープ。設定後にターミナルを開き直す)
setx PYTHONUTF8 1
setx PYTHONIOENCODING utf-8
setx HF_HOME E:\hf_cache             # 空きの大きいドライブへ
setx TORCHINDUCTOR_CACHE_DIR E:\ti_cache
setx CUDA_MODULE_LOADING LAZY

# 2. 仮想環境(3リポジトリで共用する想定なので親ディレクトリに置く)
uv venv C:\LLM\.venv-blackwell --python 3.13
C:\LLM\.venv-blackwell\Scripts\Activate.ps1

# 3. 依存(torch は PyPI ではなく NVIDIA のインデックスから)
uv pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu130
uv pip install -r requirements-dev.txt   # 作図ツールを使う場合のみ

# 4. 診断
python check_env.py

check_env.py は 13 項目を見ます。「GPU で実際に行列積を回す」まで通ってから学習に進んでください。 torch.cuda.is_available() が True でも、wheel の CUDA が古いと最初の行列積で no kernel image is available for execution on the device になります。

環境診断

RTX 50 番台(Blackwell / sm_120)を使う人向け

pip install torch で入る PyPI 版は sm_120 のカーネルを含まないことがあります。 必ず --extra-index-url https://download.pytorch.org/whl/cu130 を付けてください。 確認は次の1行です。

python -c "import torch; print(torch.__version__, torch.cuda.get_device_capability())"
# 2.13.0+cu130 (12, 0) と出れば正しい

「静かに遅くなる」対策(Windows 固有・重要)

Windows の NVIDIA ドライバは、VRAM が足りなくなるとエラーを出さずにシステムRAM (タスクマネージャの「共有GPUメモリ」)へ溢れさせます。落ちないので気づきません。 check_env.py はこれを実際に起こして測ります。

VRAM超過の実測

書き込み帯域が 3.3 倍遅くなり、OOM 例外は出ませんでした。 この倍率は実行ごとに 3〜5 倍の幅で動きます(システムRAM側の状況次第)。大事なのは倍率の正確さではなく、 例外が出ないまま遅くなるという事実です。

さらに、確保した時点では共有GPUメモリは増えません。実際にそのメモリに触った瞬間に 移るので、torch.empty() が通るかどうかを見る検査では絶対に見つかりません。

3つ全部やってください。

  1. NVIDIA コントロールパネル → 3D設定の管理 → CUDA - Sysmem Fallback Policy を Prefer No Sysmem Fallback(この項目は NVAPI からは設定できません。GUI のみ)
  2. runtime.configure(0.85) で PyTorch 側から上限を切る
  3. MemoryGuard で毎バッチ共有GPUメモリを見て、増えたら止める(src/train.py は既定で有効)

なお nvidia-smi にはこの溢れは映りません。runtime.shared_memory_gb() は Windows の パフォーマンスカウンタ \GPU Process Memory(*)\Shared Usage を直接読んでいます。

1. コーパスを作る

python data\prepare.py

data/corpus.txt に「1行1会話」のテキストができます。

<|user|>おすすめの本はありますか?<|assistant|>SFがお好きなら...<|end|>

主なオプション。

python data\prepare.py --max-a 150         # 短い返答だけ使う(学習が安定しやすい)
python data\prepare.py --min-char-freq 20  # 語彙をさらに絞る
python data\prepare.py --no-hf             # data/raw/ の自分のデータだけで作る

実測値。

項目 値
会話数 28,616 件
文字数 4,412,846 文字
文字種(語彙) 2,077

Mac 版と1文字も違わない数字が出ます。ずれたら改行を疑ってください。 Python の write_text は Windows だと \n を \r\n に変換するので、newline="\n" を 明示しないと \r が語彙に入り 2,078 になります。この状態で学習すると、 語彙2,077で作られた checkpoints/final/ と噛み合いません。

自分のデータで学習する

data/raw/ に次のいずれかの形式で置いて python data\prepare.py を実行するだけです。

// data/raw/mydata.jsonl
{"user": "調子はどう?", "assistant": "ばっちりです。"}
# data/raw/mydata.tsv
調子はどう?	ばっちりです。

2. 学習する

python src\train.py                      # 3,600ステップ(RTX 5090 で約1分)
python src\train.py --tokens 30_000_000  # トークン予算で止める
python src\train.py --resume             # 中断したところから再開
python src\train.py --equivalence-run    # Mac 版と同条件で回して差を出す
  • --resume は step / 重み / optimizer の状態 / 乱数の状態をまとめて復元します。 手元では 40ステップで中断して再開した場合の step 60 の val loss が 4.0196、 通しで回した場合が 4.0194 でした(ほぼ一致)。
  • --micro-bs を省略すると VRAM から自動で決めます。グローバルバッチのトークン数 (16,384)は固定で、VRAM が足りなければ grad_accum で分割します。バッチだけ変えると 実効学習率が変わり、別のモデルになってしまいます。
  • 検証損失が改善したときだけ checkpoints/final/ に保存します(os.replace による原子的保存)。
  • 250ステップごとに「こんにちは」への返答を出すので、賢くなっていく様子が見られます。

学習ログ

学習曲線

実測値(Windows 11 / RTX 5090 / torch 2.13.0+cu130)。

項目 値
学習時間 61.3 秒 / 3,600ステップ
最終 train loss 1.7831(直近50ステップ平均)
最良 val loss 1.8677
スループット 1,138k tok/s
専用VRAM ピーク 1.67 GB
共有GPUメモリ増分ピーク +0.01 GB
生成速度 400〜440 文字/秒

学習終了時のまとめ

学習時間は src/train.py が最後に出す合計値です。作図に使う runs/loss.csv の 最終ステップ時点は 60.8 秒で、差の 0.5 秒は最後のサンプル生成ぶんです。

同じシードでも val loss は実行ごとに動きます。 手元の3回は 1.8661 / 1.8677 / 1.8688 で、 幅は 0.003 でした。bf16 と torch.compile を使うと総和の順序が固定されないためです。 同値確認の許容を ±0.05 にしているのはこの揺れを含めるためで、 「1回走らせて小数第4位まで一致した」という比較は成立しません。

Mac (M1 Max / MLX) との比較

MacとWindowsの比較

項目 M1 Max (MLX) RTX 5090 (PyTorch)
3,600ステップ 1,907.0 秒 60.8 秒(31.4 倍)
最良 val loss 1.8571 1.8677(差 +0.0106)
スループット 31k tok/s 1,138k tok/s

秒数が上の表(61.3 秒)と違うのは、こちらが作図に使う runs/loss.csv の 最終ステップ時点だからです。train.py の合計値 61.3 秒で数えると 31.1 倍になります。

同じ図は自分のログから作れます。

python tools\compare_mac.py --mac-runs E:\ref\1LM\runs

ベンチマーク

python tools\bench.py

RTX 5090 での実測(40ステップ / ウォームアップ後)。

条件 スループット 初回コンパイル 実効TFLOPS VRAM
compile なし / fp32 681k tok/s 0.34 秒 45.8 2.95 GB
compile なし / bf16 1,159k tok/s 0.10 秒 77.9 1.98 GB
compile あり / fp32 726k tok/s 1.14 秒 48.8 2.81 GB
compile あり / bf16 1,365k tok/s 0.42 秒 91.8 1.67 GB

バッチサイズを上げると 256 で 1,512k tok/s まで伸びますが、グローバルバッチを変えると 別のモデルになるので学習では 64 に固定しています(速度のためにここを触ってはいけません)。

SDPA の backend 別(compile なし / bf16)。

backend スループット
flash 使えない(RuntimeError)
cudnn 1,113k tok/s
mem_efficient 1,123k tok/s
math 667k tok/s

Windows の PyTorch wheel には FlashAttention が入っていません。 Torch was not compiled with flash attention になります。cudnn か mem_efficient が 使えれば学習速度はほぼ変わらないので、実害はありません(math に落ちると 1.7 倍遅い)。

bf16 8192³ の行列積は 229.4 TFLOPS(公開スペックからの参考値 209.5 の 110%)。 ベンチ値は通しの実測より必ず速く出るので、記事に書くときは分けてください。

3. 会話する

CLI

python src\chat_cli.py

CLIチャット

チャット中のコマンド。

コマンド 意味
/temp 0.6 ランダムさを変える
/topk 40 候補を上位k文字に絞る
/penalty 1.2 繰り返しを抑える
/reset 会話履歴を消す
/exit 終了

上の画像は --seed 777 で撮ったものです。MLX 版と同じシードでも同じ文章は出ません。 乱数生成器が違うので、生成物を比べるときは複数回試して傾向で見てください。

Web GUI

python server.py --open

ウェルカム画面

右上のスライダーアイコンから、temperature などを触りながら挙動を比べられます。

生成設定

同値確認(移植が正しいかを数字で言う)

「動いた」ではなく「Mac 版と同じ学習になっている」を3段階で確認します。

python tools\equivalence.py            # 1段目と2段目
python src\train.py --equivalence-run  # 3段目
段 見るもの 判定基準 実測
1 fp32 の logits を NumPy の独立実装と比較 `max diff
2 step 0 の loss ln(2077) + 0.5 = 8.139 ± 0.3 通過
3 3,600ステップ後の val loss 1.8571 ± 0.05 1.8677

移植で実際にずれた場所を挙げておきます。ここを外すと3段目で落ちます。

  1. 埋め込みの初期化: MLX の nn.Embedding は N(0, 1/sqrt(dims))。PyTorch の既定は N(0, 1) で、標準偏差が約 19.6 倍大きい。明示的に上書きする。
  2. AdamW のバイアス補正: MLX の AdamW は既定で bias_correction=False、PyTorch は常に True。 これだけで val loss が 1.87 前後から 1.92 に悪化します(許容 ±0.05 を外れる)。 src/optim.py に MLX と同じ挙動の MlxAdamW を実装しました。
  3. GELU: F.gelu に approximate="tanh" を付けない(MLX の既定は erf 版)。
  4. bias の非対称: qkv/proj は bias なし、MLP は bias あり。原典どおりに揃える。
  5. step 0 の loss は ln(V) ではない: weight tying + 1/sqrt(n_embd) 初期化だと logits の分散が 1 になるため、期待値は ln(V) + σ²/2 = ln(V) + 0.5。 ln(V) を基準にすると「+0.5 ずれている」と誤診して原因を探し続けることになります。

仕組み

flowchart LR
    A["文字列<br/>こんにち"] --> B[トークナイザ<br/>1文字=1ID]
    B --> C[埋め込み<br/>+ 位置埋め込み]
    C --> D["Transformer Block x6<br/>因果マスク付き自己注意 + MLP"]
    D --> E["次の1文字の確率<br/>語彙2077次元"]
    E --> F[サンプリング<br/>temperature / top_k]
    F --> A
Loading

ファイルの役割。

ファイル 役割
check_env.py 環境診断13項目。最初に通す
runtime.py VRAM / 共有GPUメモリの監視、バッチ計画、MemoryGuard
src/tokenizer.py 文字レベルトークナイザ。<|user|> などのマーカーは1トークン扱い
src/model.py ミニGPT本体。因果マスク付き自己注意、Pre-LN、weight tying
src/optim.py MLX と同じ挙動の AdamW(バイアス補正なし)
src/train.py 学習ループ。torch.compile + bf16 autocast + fp32マスタ重み
src/generate.py サンプリング。temperature / top-k / 繰り返しペナルティ
src/chat_cli.py CLIチャット
server.py FastAPI。SSE でトークンを流す
web/ Liquid Glass 風のチャットGUI
data/prepare.py コーパス整形
tools/ 同値確認・ベンチ・作図・撮影(pip install -r requirements-dev.txt)

つまずきポイント

Windows で実際に踏んだものだけを挙げます。

  1. torch.cuda.is_available() は True なのに学習が始まらない → wheel の CUDA が古い。 sm_120 は cu130 以降。no kernel image is available は「GPU が無い」ではなく 「この GPU 向けのコードが wheel に入っていない」という意味。
  2. torch.compile が FileNotFoundError → キャッシュパスが 260 文字を超えている。 TORCHINDUCTOR_CACHE_DIR を E:\ti_cache のような短いパスにする。 根本対処は管理者権限で reg add HKLM\SYSTEM\CurrentControlSet\Control\FileSystem /v LongPathsEnabled /t REG_DWORD /d 1 /f。
  3. 学習が急に遅くなったのにエラーが出ない → 共有GPUメモリを見る。 nvidia-smi には映らない。上の「静かに遅くなる」対策へ。
  4. torch.compile(model) したのに速くならない → model.loss(...) のように メソッドを直接呼ぶとコンパイルを迂回します。torch.compile(model.loss) と 関数側をコンパイルしてください(実測で 1,159k → 1,365k tok/s の差)。
  5. 語彙が 2,078 になる / 学習済みモデルが読めない → \r が混入している。 ファイル書き出しは全部 encoding="utf-8", newline="\n" を明示する。
  6. PowerShell に流したテキストが文字化けする → PowerShell 5.1 の Get-Content は BOM なし UTF-8 を CP932 として読みます。-Encoding UTF8 を付ける。 逆に Out-File -Encoding utf8 は BOM を付けるので、読む側は utf-8-sig にする。
  7. 返答が同じ言葉を繰り返す → モデルではなくサンプリングを疑う。 repetition_penalty を 1.15 前後に。
  8. 返答が毎回崩れる → 推論前に model.eval() を呼んで Dropout を切る。

学習済みモデルについて

checkpoints/final/ に学習済みの重みを同梱しています(11.53M パラメータ fp32 で約44MB)。 クローンすればすぐ会話できます。

checkpoints/final/
├── model.safetensors   # 重み
├── config.json         # モデル構成
└── tokenizer.json      # 語彙(文字→ID)

自分で学習し直すと同じ場所が上書きされます。残しておきたい場合は --out checkpoints\myrun を指定してください。

ライセンス / クレジット

コードは MIT License です(LICENSE)。 Apple Silicon / MLX 版 1LM の移植で、原典も同じ作者・同じ MIT License です。

同梱の学習済みモデルは、以下の Apache-2.0 データセットから学習しています。 再配布・商用利用のいずれの場合も、この出典表示を残してください。

前処理で加えた変更点とライセンス全文は NOTICE と licenses/Apache-2.0.txt にあります。

生成される文章は、コーパスの統計から次の1文字を予測し続けた結果にすぎません。 事実性は一切保証されず、実在の人物や団体について誤った内容を出力することがあります。 出力を公開の場に掲載する場合は、機械生成物である旨を明記してください。

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages