大規模言語モデル学習をPyTorchでネイティブに実現：TorchTitan入門ガイド

PyTorchネイティブなライブラリ「TorchTitan」で、大規模言語モデル（LLM）学習を始めませんか？ 設計思想からインストール、最新機能まで、わかりやすく解説します。

TorchTitanとは？ PyTorchネイティブなLLM学習ライブラリの概要

TorchTitanは、PyTorchの最新機能を活用して大規模LLM学習を支援するライブラリです。 他のLLM学習基盤（Megatron、DeepSpeedなど）を置き換えるものではなく、 PyTorchの機能をわかりやすく示すことを目的としています。 短時間で学習を開始できる手軽さが魅力です。

• 簡潔性: 理解、利用、拡張が容易な設計。
• 柔軟性: 1D, 2D Pa​​rallel（近日3D Parallelも）の適用がモデルコードへの最小限の変更で可能。
• モジュール性: 巨大なコードベースではなく、モジュール化されたコンポーネント。

TorchTitanで何ができる？ 主要機能を紹介

TorchTitanは、大規模言語モデルの学習を効率化するための様々な機能を提供します。

• FSDP2 (per param sharding): パラメータを分割してメモリ効率を向上。詳細はこちら
• Tensor Parallel: テンソルを分割して計算を分散。詳細はこちら
• Selective Layer/Operator Activation Checkpointing: メモリ消費を抑えつつ学習を高速化。
• Distributed Checkpointing: 学習途中のモデルを分散保存。
• 学習済みデータセット: 45K〜144Mトークンのデータセットが利用可能。
• TensorBoard連携: GPU使用率、MFU、トークン/秒などを可視化。
• 学習率スケジューラ: 最適な学習率を自動調整。
• Meta init: 効果的な初期化方法。
• Fused RMSNorm: (オプション) 高速な正規化手法。
• TOMLファイル: 設定を容易に変更。設定例
• Interoperable Checkpoints: TorchTuneで利用可能なチェックポイント。詳細はこちら

インストールから学習開始まで：TorchTitanを使い始める手順

以下の手順でTorchTitanをインストールし、学習を開始しましょう。

1. リポジトリをクローン:

   git clone https://github.com/pytorch/torchtitan
   cd torchtitan
   

2. 必要なパッケージをインストール:

   pip install -r requirements.txt
   pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 # または cu118
   pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly
   

3. トークナイザーをダウンロード: Llama 3またはLlama 2のトークナイザーをダウンロードします。 Hugging Faceのトークンが必要です。

   # Llama3 tokenizer.model
   python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Meta-Llama-3-8B --tokenizer_path "original" --hf_token=<YOUR_HF_TOKEN>
   
   # Llama2 tokenizer.model
   python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Llama-2-13b-hf --hf_token=<YOUR_HF_TOKEN>
   

4. 学習を開始: Llama 3 8Bモデルを8つのGPUでローカルに学習させるには、次のコマンドを実行します。

   CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh
   

TensorBoardで学習状況を可視化

TensorBoardを使って、学習中のモデルのメトリクスを可視化できます。

1. metrics.enable_tensorboardオプションを有効にする。
2. SSHトンネルを設定する。

   ssh -L 6006:127.0.0.1:6006 [username]@[hostname]
   

3. サーバー上でTensorBoardを起動する。

   tensorboard --logdir=./outputs/tb
   

4. ローカルのWebブラウザでhttp://localhost:6006/にアクセスする。

複数ノードでの分散学習

ParallelClusterやSlurmなどの環境で複数ノードでの分散学習を行うには、multinode_trainer.slurmファイルを使用します。

1. multinode_trainer.slurmファイルのノード数とGPU数を調整する。
2. sbatchジョブをサブミットする。
3. torchrunコマンドを実行する。

今後のアップデート予定

TorchTitanは今後もアップデートを予定しています。

• 非同期チェックポイント
• FP8サポート
• Context Parallel
• 3D Pipeline Parallel
• torch.compileサポート
• スケーラブルなデータローディング

まとめ：TorchTitanで大規模LLM学習を始めよう

TorchTitanは、PyTorchネイティブなLLM学習を支援する強力なツールです。 ぜひTorchTitanを活用して、大規模言語モデル学習に挑戦してみてください。