NVIDIA-BioNeMo/bionemo-recipes

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

What it solves

BioNeMo Recipes provides a streamlined path for the biological foundation model community to scale up transformer-based models efficiently. It addresses the need for tooling that balances high performance with extensibility and ease-of-use, allowing researchers to move from single-GPU prototyping to multi-node training without complex parallelism configurations.

How it works

Instead of a monolithic training framework, the project provides two primary components:

  1. Model Checkpoints: Hugging Face-compatible PreTrainedModel classes that integrate TransformerEngine (TE) layers for optimized performance (e.g., FP8 training and context parallelism).
  2. Training Recipes: Self-contained Docker containers that serve as reference implementations. These demonstrate how to achieve maximum throughput using popular frameworks like vanilla PyTorch, HuggingFace Accelerate, and PyTorch Lightning, leveraging technologies such as Fully Sharded Data Parallel (FSDP) and sequence packing (THD).

Additionally, the project includes research tools for model interpretability, such as sparse autoencoder training and feature analysis.

Who it’s for

  • Foundation Model Developers: AI researchers and ML engineers building novel biological foundation models who need efficient scaling strategies.
  • Foundation Model Customizers: Domain scientists fine-tuning existing models with proprietary biological data for drug discovery and biological research.

Highlights

  • Flexible Scaling: Supports scaling from a single GPU to multi-node clusters.
  • Framework Interoperability: Compatible with PyTorch, PyTorch Lightning, and HuggingFace Accelerate.
  • Performance Optimizations: Utilizes TransformerEngine, megatron-FSDP, and advanced precision formats (BF16, FP8, MXFP8, NVFP4).
  • Research-Friendly: Prioritizes readable, "hackable" code over complex abstractions to facilitate rapid experimentation.
  • Hugging Face Integration: Models are distributed via the Hugging Face Hub and are drop-in replacements for AutoModel.from_pretrained().

Related

  • Project
    NVIDIA-NeMo/Automodel

    NeMo AutoModel is a Python library from NVIDIA that provides ready‑to‑use, highly‑optimized training recipes for a huge range of LLMs, VLMs, and diffusion models. It supports full‑parameter and LoRA fine‑tuning, multi‑node tensor‑ and pipeline‑parallelism, and even speculative‑decoding drafters, all driven by simple YAML configs.

  • Dispatch
    Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel

    Hugging Face integrates PyTorch Fully Sharded Data Parallel (FSDP) into the Accelerate library, enabling the training of massive models by sharding optimizer states, gradients, and parameters across GPUs with optional CPU offloading.

  • Project
    NVIDIA-NeMo/Megatron-Bridge

    A PyTorch-native library that enables bidirectional checkpoint conversion and high-performance training between Hugging Face and NVIDIA Megatron Core.

  • Dispatch
    Training Language Models with Megatron-LM

    Hugging Face provides a guide on using NVIDIA's Megatron-LM framework to efficiently pre-train large language models on GPUs, including integration with the Transformers library.

  • Project
    PaddlePaddle/PaddleFormers

    A high-performance Transformers library for PaddlePaddle that enables efficient pre-training and post-training of LLMs and VLMs with extensive support for distributed strategies and domestic AI hardware.