NVIDIA-BioNeMo/bionemo-recipes

BioNeMo Recipes: For building and adapting AI models in drug discovery at scale

解决的问题

BioNeMo Recipes 帮助生物AI社区高效地扩展基于Transformer的基础模型训练。它通过提供预优化的模型检查点和参考训练实现，消除了配置高性能并行和优化设置的复杂性。

工作原理

该项目提供三个主要组件：

  • 模型：兼容Hugging Face的模型类（如ESM-2、Geneformer和AMPLIFY），集成了NVIDIA TransformerEngine层以实现硬件加速。
  • 训练配方：包含训练示例的自包含Docker容器。这些示例展示了如何使用PyTorch、HuggingFace Accelerate和PyTorch Lightning等框架，结合Fully Sharded Data Parallel (FSDP)、FP8精度和序列打包（THD）等高级扩展技术。
  • 可解释性工具：用于分析模型行为的工作流，包括稀疏自编码器训练和特征分析。

适用人群

  • 基础模型开发者：需要将原型扩展到多节点集群的AI研究人员和ML工程师，正在构建新的生物基础模型。
  • 基础模型定制者：在专有生物数据上对现有模型进行微调，用于药物发现和研究的领域科学家。

主要亮点

  • 高性能精度：支持BF16、FP8、MXFP8和NVFP4，以提高训练吞吐量。
  • 灵活扩展：支持从单个GPU扩展到使用megatron-FSDP和上下文并行的多节点训练。
  • 框架互操作性：与原生PyTorch、PyTorch Lightning和HuggingFace Accelerate兼容。
  • 优化的数据处理：实现序列打包（THD），消除填充标记带来的计算浪费。

相关

  • Dispatch
    使用 PyTorch Fully Sharded Data Parallel 加速大模型训练

    Hugging Face 将 PyTorch Fully Sharded Data Parallel (FSDP) 集成到了 Accelerate 库中，通过在 GPU 之间分片优化器状态、梯度和参数，并提供可选的 CPU offloading，从而实现了海量模型的训练。

  • Dispatch
    使用 Megatron-LM 训练语言模型

    Hugging Face 提供了一份指南，介绍如何使用 NVIDIA 的 Megatron-LM 框架在 GPU 上高效预训练大型语言模型，并与 Transformers 库进行集成。

  • Dispatch
    Hugging Face Transformers 与 Ray Tune 集成

    Hugging Face Transformers 3.1 引入了与 Ray Tune 的集成，使用户能够轻松实现诸如基于种群的训练 (Population-Based Training) 和贝叶斯优化 (Bayesian Optimization) 等先进的超参数调优算法。

  • 项目
    whitecircle/halo

    Halo 是 White Circle 推出的一个基于 Docker 的开源框架，用于训练 Hugging Face 的 LLM/VLM 模型，支持高级并行（MoE、张量、上下文）、高吞吐内核（FlashAttention 4、Liger、Grouped GEMM）和异步 RL 回滚。它覆盖完整的训练生命周期——预训练、SFT、偏好学习、蒸馏和多轮 RL——同时保持检查点与标准 Hugging Face 格式兼容。

  • Dispatch
    Hugging Face Transformers timm 集成

    Hugging Face 已经引入了 TimmWrapper，使得 PyTorch Image Models (timm) 库中的任何模型都能够在 transformers 生态系统中无缝用于推理、量化和微调。