DOCSEARCHbyAlgolia
Submit Docs

BioNeMo Recipes

repository·main·Indexed 21 days ago

https://github.com/nvidia-bionemo/bionemo-recipes

Optimized model checkpoints and training recipes for biological foundation models. It features high-throughput training using NVIDIA TransformerEngine and FSDP, as well as mechanistic interpretability tools. The library includes a domain-agnostic Sparse Autoencoder (SAE) implementation (ReLUSAE, TopKSAE, MoESAE) with specific recipes for ESM-2 protein language models and CodonFM (Encodon) for DNA sequence analysis.

Tokens
164K
Snippets
473
Records
685
Agent score
74%

What's inside BioNeMo Recipes

  1. Overview of ESM-2 Models

    main

    ESM-2 is a pre-trained, bi-directional BERT-style encoder designed for amino acid sequences. It generates embeddings for amino acids used in downstream tasks like protein structure and function prediction. BioNeMo Recipes provides converted checkpoints for two specific variants:

    • 650M Parameter Model: 33 layers, 20 attention heads, 1280 hidden dimension.
    • 3B Parameter Model: 36 layers, 40 attention heads, 2560 hidden dimension.

    These models are optimized for commercial use and are compatible with NVIDIA Ampere, Hopper, and Volta architectures on Linux.

  3. Overview of available BioNeMo Recipes models

    main

    BioNeMo Recipes provides pre-trained models for various biological modalities, including Protein, DNA, and Single Cell data. These models are used for tasks such as Representation Learning and Generative AI. Detailed information for each model can be found in its specific Model Card.

    | **Model**                     | **Modality** | **Uses**                | **Training Location**                                                            |
| ----------------------------- | ------------ | ----------------------- | -------------------------------------------------------------------------------- |
| [AMPLIFY](./amplify.md)       | Protein      | Representation Learning | [BioNeMo Recipes](https://github.com/NVIDIA-BioNeMo/bionemo-framework/tree/main) |
| [ESM-2](./ESM-2/index.md)     | Protein      | Representation Learning | [BioNeMo Recipes](https://github.com/NVIDIA-BioNeMo/bionemo-framework/tree/main) |
| [Evo2](./evo2.md)             | DNA          | Generative AI           | [BioNeMo Recipes](https://github.com/NVIDIA-BioNeMo/bionemo-framework/tree/main) |
| [Geneformer](./geneformer.md) | Single Cell  | Representation Learning | [BioNeMo Recipes](https://github.com/NVIDIA-BioNeMo/bionemo-framework/tree/main) |

  4. Overview of BioNeMo Recipes

    main

    BioNeMo Recipes is a collection of model checkpoints and training recipes designed to help the biological foundation model community scale transformer-based models efficiently. It focuses on providing high-throughput training using TransformerEngine (TE) and supporting various scaling strategies like Fully Sharded Data Parallel (FSDP).

    Key features include:

    • Flexible scaling: Supports single-GPU to multi-node training.
    • Framework compatibility: Integrates with HuggingFace Accelerate, PyTorch Lightning, and vanilla PyTorch.
    • Performance optimization: Uses TransformerEngine and megatron-FSDP for state-of-the-art efficiency.
    • Research-friendly: Provides readable, hackable code for experimentation.

  5. Overview of the AMPLIFY model

    main

    AMPLIFY is a TransformerEngine-optimized protein language model variant of ESM-2. It features a modified layer structure and dataset construction designed for protein sequence understanding and prediction tasks. It is available in two parameter sizes: 120M and 350M.

    Key technical details:

    • Architecture Type: Transformer (ESM-2 variant).
    • Input: Protein sequences represented as strings of canonical amino acids.
    • Output: Numeric vectors (1D) representing amino acid and sequence-level embeddings.
    • Runtime Engines: BioNeMo, TransformerEngine.
    • Hardware Compatibility: NVIDIA Ampere and NVIDIA Hopper microarchitectures (tested on NVIDIA H100).
    • Operating System: Linux.
    • License: Apache 2.0.

  6. Overview of SAE Feature Dashboard capabilities

    main

    The SAE Feature Dashboard is an interactive tool for exploring Sparse Autoencoder (SAE) features. Key capabilities include:

    • UMAP Embedding View: An interactive scatter plot of feature embeddings with pan and zoom capabilities.
    • Crossfiltering: Selecting (brushing) areas on the UMAP plot or interacting with histograms filters the displayed feature list.
    • Feature Cards: Detailed views for each feature including description, activation stats, top positive/negative logits, and top activating examples with token highlighting.
    • Search: Text-based filtering of features by their description.
    • Coloring: Ability to color points in the UMAP view by categorical or sequential columns.

  7. Overview of BioNeMo Recipes

    main

    BioNeMo Recipes are self-contained training examples designed as educational reference implementations. They demonstrate best practices for scaling biological foundation models using TransformerEngine and megatron-FSDP).

    Each recipe is an isolated environment containing its own Dockerfile, pinned dependencies, and benchmarked training scripts, allowing users to learn from and adapt them for their own research without managing complex shared dependencies.

  8. Overview of Geneformer model

    main

    Geneformer is a foundation model designed to generate dense representations (embeddings) of single-cell RNA (scRNA) cells by learning co-expression patterns. It is a tabular count model trained on the Chan Zuckerberg CELLxGENE census. The model computes complete embeddings for each cell based on the top 1024 expressed genes, which can then be used as features for various predictive downstream tasks.

    Key Specifications:

    • Architecture Type: Bidirectional Encoder Representations from Transformers (BERT).
    • Runtime Engine: BioNeMo, TransformerEngine.
    • Supported Hardware: Ampere, Hopper, and Volta microarchitectures.
    • Supported OS: Linux.

  9. Overview of Codon FM Foundation Models

    main
    Codon FM is an open-source suite of foundation models trained on codon sequences to learn contextual representations for downstream codon-aware tasks. This specific recipe provides an NVIDIA TransformerEngine (TE) accelerated version of the original research code. It includes the EnCodon model family (80M, 600M, 1B parameters) using masked language modeling (MLM) and supports both PyTorch and TransformerEngine architectures.

  11. Overview of Mixtral Optimized with NVIDIA TransformerEngine

    main

    This implementation provides Mixtral-style Mixture of Experts (MoE) models that inherit from the Hugging Face transformers.PreTrainedModel class. It replaces standard attention layers with TransformerEngine (TE) MultiheadAttention and uses TE GroupedLinear for efficient parallel expert computation.

    Key Feature Support:

    • FP8: Supported on Hopper+ (compute capability 9.0+).
    • MXFP8: Supported on Blackwell (compute capability 10.0/10.3).
    • Sequence Packing / THD input format: Supported.
    • Expert Parallelism (EP): Supported via DTensor sharding, composable with FSDP2.
    • Checkpointing: Supports importing from and exporting to Hugging Face checkpoints.
    • Inference: Supports KV-cache inference.

  12. Overview of Geneformer (TransformerEngine-optimized)

    main

    Geneformer is a foundational transformer model pretrained on a large-scale corpus of single-cell transcriptomes across various human tissues. This specific implementation is optimized using NVIDIA's TransformerEngine.

    Key characteristics:

    • Compatibility: Within numerical precision, it has identical weights and outputs to the original Geneformer model.
    • Use Case: It is designed for fine-tuning on tasks that utilize gene expression data as input.
    • Library Dependency: It is built to work with the transformers library.