NVIDIA-BioNeMo/bionemo-recipes

何を解決するか

BioNeMo Recipes は、生物学的基礎モデルのコミュニティがトランスフォーマーに基づくモデルを効率的にスケーリングするためのスムーズな道筋を提供します。高性能と拡張性、使いやすさのバランスを取ったツールの必要性に対応し、単一GPUでのプロトタイピングから複数ノードでのトレーニングへ移行する際に、複雑な並列化設定を必要とせずに済むようにします。

仕組み

モノリシックなトレーニングフレームワークではなく、本プロジェクトは以下の2つの主要なコンポーネントを提供します：

1. モデルチェックポイント：TransformerEngine (TE) レイヤーを統合した Hugging Face互換の PreTrainedModel クラス。最適化されたパフォーマンス（例：FP8トレーニング、コンテキスト並列化）を実現します。
2. トレーニングレシピ：自己完結型の Docker コンテナ。vanilla PyTorch、HuggingFace Accelerate、PyTorch Lightning などの人気フレームワークで最大スループットを達成する方法を参考実装として示します。Fully Sharded Data Parallel (FSDP) やシーケンスパッキング（THD）などの技術を活用します。

さらに、モデルの解釈性を高めるための研究ツールも含まれており、スパースオートエンコーダーのトレーニングや特徴分析が可能です。

対象ユーザー

• 基礎モデル開発者：新しい生物学的基礎モデルを構築するAI研究者やMLエンジニア。効率的なスケーリング戦略が必要な方。
• 基礎モデルカスタマイザー：独自の生物学的データで既存モデルをファインチューニングし、創薬や生物学的研究に活用する分野の科学者。

特徴

• 柔軟なスケーリング：単一GPUからマルチノードクラスタまでスケーリング可能。
• フレームワークの相互運用性：PyTorch、PyTorch Lightning、HuggingFace Accelerate と互換性あり。
• パフォーマンス最適化：TransformerEngine、megatron-FSDP、および高度な精度形式（BF16、FP8、MXFP8、NVFP4）を活用。
• 研究に適した設計：複雑な抽象化よりも読みやすく「ハッカブル」なコードを優先し、迅速な実験を可能にします。
• Hugging Face統合：モデルは Hugging Face Hub を通じて配布され、AutoModel.from_pretrained() の即時置き換えとして利用可能。