NVIDIA-BioNeMo/bionemo-recipes

解決的問題

BioNeMo Recipes 為生物基礎模型社群提供了一條高效擴展基於 Transformer 的模型的簡化途徑。它解決了在高效率、可擴展性與易用性之間取得平衡的工具需求，讓研究人員能從單一 GPU 原型開發，無需複雜的平行設定，直接過渡到多節點訓練。

工作原理

本專案不採用單一的訓練框架，而是提供兩個主要元件：

1. 模型檢查點：相容 Hugging Face 的 PreTrainedModel 類別，整合了 TransformerEngine (TE) 層以實現最佳化效能（例如 FP8 訓練與上下文平行）。
2. 訓練配方：自包含的 Docker 容器，作為參考實作。展示如何使用常見框架（如原生 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() 的即插即用替代品。