NVIDIA-BioNeMo/bionemo-recipes

何を解決するか

BioNeMo Recipes は、生物学的AIコミュニティがTransformerベースの基礎モデルのトレーニングを効率的にスケーリングできるように支援します。高パフォーマンスな並列処理と最適化設定の構成の複雑さを解消し、事前に最適化されたモデルチェックポイントと参照実装のトレーニングコードを提供します。

仕組み

このプロジェクトは以下の3つの主要なコンポーネントを提供します：

• モデル: NVIDIA TransformerEngineレイヤーを統合した、Hugging Face互換のモデルクラス（ESM-2、Geneformer、AMPLIFYなど）で、ハードウェア加速を実現します。
• レシピ: 自己完結型のDockerコンテナに含まれるトレーニング例。PyTorch、HuggingFace Accelerate、PyTorch Lightningなどのフレームワークと、Fully Sharded Data Parallel (FSDP)、FP8精度、シーケンスパッキング（THD）などの高度なスケーリング技術の使い方を示します。
• 解釈性ツール: スパースオートエンコーダーのトレーニングや特徴分析を含む、モデル挙動を分析するためのワークフロー。

対象ユーザー

• 基礎モデル開発者: 新しい生物学的基礎モデルを開発するAI研究者やMLエンジニアで、プロトタイプをマルチノードクラスタにスケーリングする必要がある人。
• 基礎モデルカスタマイザー: 薬物発見や研究のため、独自の生物学的データ上で既存モデルを微調整する分野の科学者。

主な特徴

• 高性能精度: BF16、FP8、MXFP8、NVFP4をサポートし、トレーニングスループットを向上します。
• 柔軟なスケーリング: megatron-FSDPとコンテキスト並列化を用いて、単一GPUからマルチノードトレーニングまでスケーリング可能。
• フレームワーク互換性: ベーシックなPyTorch、PyTorch Lightning、HuggingFace Accelerateと互換性があります。
• 最適化されたデータ処理: シーケンスパッキング（THD）を実装し、パディングトークンによる無駄な計算を排除します。