Nanotron - ミニマリズム 3D 並列 LLM トレーニング チートシート
Nanotron(Hugging Face)は ミニマリズムフレームワーク — 大規模言語モデル プリトレーニング用 3D 並列化(データ・テンソル・パイプライン並列)を故意に読める Python で実装。ゴールは十分小さくハッキング可能で実際に理解できながら、多数 GPU にまたがるトレーニング分散可能。複雑さなく分散プリトレーニング機構を学ぶ・カスタマイズしたい場合の強力な選択肢。
プリトレーニングは計算集約的。スケーリング前に小さいモデルと数個 GPU で設定検証。
インストール
| 方法 | コマンド |
|---|
| ソースから | git clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e . |
| Flash attention | pip install flash-attn(推奨) |
| 要件 | PyTorch + CUDA GPU |
| 確認 | 1–2 GPU でサンプル設定実行 |
3D 並列化
| 次元 | 分割 | Config キー |
|---|
| Data Parallel(DP) | レプリカ間でバッチ | dp |
| Tensor Parallel(TP) | GPU 間でレイヤー | tp |
| Pipeline Parallel(PP) | ステージ間でモデル | pp |
積 dp × tp × pp は総 GPU 数と等しくする必要。
設定(YAML)
Nanotron トレーニングは モデル・データ・並列化・オプティマイザー記述 YAML 設定で定義。
parallelism:
dp: 2
tp: 2
pp: 2
pp_engine: 1f1b # パイプラインスケジュール
model:
model_config:
hidden_size: 2048
num_hidden_layers: 24
num_attention_heads: 16
optimizer:
learning_rate_scheduler:
learning_rate: 3.0e-4
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.95
tokens:
sequence_length: 2048
micro_batch_size: 4
batch_accumulation_per_replica: 8
トレーニング起動
# シングルノード、マルチ GPU
torchrun --nproc_per_node=8 run_train.py --config-file config.yaml
# スターター設定生成
python examples/config_tiny_llama.py
| キー | 意味 |
|---|
run_train.py | メイン トレーニング エントリポイント |
--config-file | YAML トレーニング設定 |
pp_engine | パイプラインスケジュール(例 1F1B、AFAB) |
micro_batch_size | ステップごとマイクロバッチ |
batch_accumulation_per_replica | 勾配累積 |
データ
| オプション | 注記 |
|---|
| プリトークン化データセット | トークン化データにポイント設定 |
| HF datasets | Hugging Face 経由でストリーム/ロード |
| Nanoset | Nanotron 効率的データセット形式 |
| Tokenizer | トークナイザー設定経由で設定 |
チェックポイント・Resume
| 項目 | 実行方法 |
|---|
| Save 頻度 | checkpoints.checkpoint_interval |
| Save パス | checkpoints.checkpoints_path |
| Resume | --config-file + 保存チェックポイント指定 |
| Sharded | チェックポイント並列化対応 |
監視
| メトリック | 意味 |
|---|
| loss / lm_loss | トレーニング loss |
| tokens/s | スループット |
| grad norm | 勾配大きさ(安定性) |
| Logging | W&B / TensorBoard 統合 |
Nanotron vs その他トレーニングフレームワーク
| 項目 | Nanotron | TorchTitan | Megatron-LM |
|---|
| 哲学 | ミニマル、ハッキング可能 | PyTorch ネイティブリファレンス | 機能豊富、max スケール |
| 並列化 | 3D(dp/tp/pp) | 4D(コンテキスト追加) | 3D+ 広範 |
| 読みやすさ | 高(から学ぶ) | 高 | 低(大規模コードベース) |
| 最適用途 | 3D 理解/カスタマイズ | PyTorch ネイティブスケーリング | フロンティアスケールトレーニング |
リソース