Nanotron

Nanotronは、テンソル並列、パイプライン並列、データ並列（3D並列）を備えた分散学習フレームワークです。数百枚のGPUにわたる大規模な学習ワークロード向けに設計されています。

convert_hf_to_nanotron.py スクリプトを使用して、任意の Transformers モデルを、事前学習用の最適化された Nanotron トランスフォーマーモデル実装に変換できます。

torchrun --nproc_per_node=1 examples/llama/convert_hf_to_nanotron.py \
    --checkpoint_path=meta-llama/Llama-2-7b-hf \
    --save_path=./llama-7b-nanotron

Transformers の統合

1. Llama などの対応する Transformers モデルを、from_pretrained() 関数で読み込みます。これにより、チェックポイントディレクトリから config.json ファイルを読み取り、LlamaConfig が作成されます。
2. Nanotron は LlamaConfig を独自の config 形式にマッピングし、Nanotron モデルを作成します。
3. Transformers の重みを Nanotron に変換します。重みのマッピングは、Nanotron のパラメータ名を Transformers のパラメータ名にどのように対応させるかを導きます。これには、QKV 射影や gate/up 射影の融合といった変換の処理が含まれます。

Nanotron は、前処理および生成時にテキストをトークン ID に変換するために AutoTokenizer に依存しています。

リソース

• Nanotronリポジトリ
• Ultrascale Playbookは、Nanotronを用いた効率的なトレーニングのスケーリング方法について説明しています