コンテンツにスキップ

Nanotron - ミニマリズム 3D 並列 LLM トレーニング チートシート

Nanotron - ミニマリズム 3D 並列 LLM トレーニング チートシート

Nanotron(Hugging Face)は ミニマリズムフレームワーク — 大規模言語モデル プリトレーニング用 3D 並列化(データ・テンソル・パイプライン並列)を故意に読める Python で実装。ゴールは十分小さくハッキング可能で実際に理解できながら、多数 GPU にまたがるトレーニング分散可能。複雑さなく分散プリトレーニング機構を学ぶ・カスタマイズしたい場合の強力な選択肢。

プリトレーニングは計算集約的。スケーリング前に小さいモデルと数個 GPU で設定検証。

インストール

方法コマンド
ソースからgit clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e .
Flash attentionpip 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-fileYAML トレーニング設定
pp_engineパイプラインスケジュール(例 1F1B、AFAB)
micro_batch_sizeステップごとマイクロバッチ
batch_accumulation_per_replica勾配累積

データ

オプション注記
プリトークン化データセットトークン化データにポイント設定
HF datasetsHugging Face 経由でストリーム/ロード
NanosetNanotron 効率的データセット形式
Tokenizerトークナイザー設定経由で設定

チェックポイント・Resume

項目実行方法
Save 頻度checkpoints.checkpoint_interval
Save パスcheckpoints.checkpoints_path
Resume--config-file + 保存チェックポイント指定
Shardedチェックポイント並列化対応

監視

メトリック意味
loss / lm_lossトレーニング loss
tokens/sスループット
grad norm勾配大きさ(安定性)
LoggingW&B / TensorBoard 統合

Nanotron vs その他トレーニングフレームワーク

項目NanotronTorchTitanMegatron-LM
哲学ミニマル、ハッキング可能PyTorch ネイティブリファレンス機能豊富、max スケール
並列化3D(dp/tp/pp)4D(コンテキスト追加)3D+ 広範
読みやすさ高(から学ぶ)低(大規模コードベース)
最適用途3D 理解/カスタマイズPyTorch ネイティブスケーリングフロンティアスケールトレーニング

リソース