NVIDIA-BioNeMo/bionemo-recipes

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

What it solves

BioNeMo Recipes helps the biological AI community scale the training of transformer-based foundation models efficiently. It removes the complexity of configuring high-performance parallelism and optimization settings by providing pre-optimized model checkpoints and reference training implementations.

How it works

The project provides three main components:

  • Models: Hugging Face-compatible model classes (like ESM-2, Geneformer, and AMPLIFY) that integrate NVIDIA TransformerEngine layers for hardware acceleration.
  • Recipes: Self-contained Docker containers containing training examples. These demonstrate how to use frameworks like PyTorch, HuggingFace Accelerate, and PyTorch Lightning with advanced scaling techniques such as Fully Sharded Data Parallel (FSDP), FP8 precision, and sequence packing (THD).
  • Interpretability Tools: Workflows for analyzing model behavior, including sparse autoencoder training and feature analysis.

Who it’s for

  • Foundation Model Developers: AI researchers and ML engineers building new biological foundation models who need to scale prototypes to multi-node clusters.
  • Foundation Model Customizers: Domain scientists fine-tuning existing models on proprietary biological data for drug discovery and research.

Highlights

  • High-Performance Precision: Supports BF16, FP8, MXFP8, and NVFP4 to increase training throughput.
  • Flexible Scaling: Supports scaling from a single GPU to multi-node training using megatron-FSDP and context parallelism.
  • Framework Interoperability: Compatible with vanilla PyTorch, PyTorch Lightning, and HuggingFace Accelerate.
  • Optimized Data Handling: Implements sequence packing (THD) to eliminate wasted compute on padding tokens.

Related

  • 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.

  • 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.

  • Dispatch
    Hugging Face Transformers and Ray Tune Integration

    Hugging Face Transformers 3.1 introduces an integration with Ray Tune, enabling users to easily implement advanced hyperparameter tuning algorithms like Population-Based Training and Bayesian Optimization.

  • Project
    whitecircle/halo

    Halo is an open‑source, Docker‑based framework from White Circle that trains Hugging Face LLM/VLM models with advanced parallelism (MoE, tensor, context), high‑throughput kernels (FlashAttention 4, Liger, Grouped GEMM) and asynchronous RL roll‑outs. It supports the full training lifecycle—pre‑training, SFT, preference learning, distillation, and multi‑turn RL—while keeping checkpoints compatible with standard Hugging Face formats.

  • Dispatch
    Hugging Face Transformers timm Integration

    Hugging Face has introduced the TimmWrapper, allowing any model from the PyTorch Image Models (timm) library to be used seamlessly within the transformers ecosystem for inference, quantization, and fine-tuning.