Fine-Tuning Biological Foundation Models with LoRA Using NVIDIA BioNeMo Recipes

Foundation models are reshaping computational biology. Pretrained on massive corpora of protein or genomic sequences, models such as ESM2 (a protein language model) and Evo 2 (a DNA language model) capture statistical regularities of biological sequences. These transfer well to a wide range of downstream tasks, including structure prediction, variant effect, and functional annotation.

Yet adapting these models to a specific task is nontrivial: at billions of parameters, full fine-tuning quickly becomes impractical, both in compute and storage of optimizer state and checkpoints.

Low-Rank Adaptation (LoRA) directly addresses this challenge. By keeping the pretrained backbone frozen and training only a small set of low-rank adapter matrices, LoRA can match full fine-tuning quality on many tasks while training ~1% of the parameters, fitting a single billion-scale model and its adapter state on a single workstation GPU.

To reduce the difficulty of building these workflows, NVIDIA BioNeMo Recipes provide step-by-step training recipes built on familiar PyTorch, Hugging Face, Megatron-Bridge patterns. Performance-oriented components such as NVIDIA Transformer Engine (TE) and scale-out strategies are integrated where they pay off, but the recipes themselves stay readable.

All the source code to customize or reproduce these results are available in NVIDIA BioNeMo Recipes.

Before diving into the case studies, a quick refresher on the method. Full fine-tuning is resource-heavy because it requires storing and updating all model parameters and their optimizer states, which quickly becomes impractical as models scale.

LoRA is a practical method to fine-tune large pretrained transformers without updating or storing optimizer state for all model parameters. The core idea behind LoRA is that instead of updating a dense model's weight matrix \(W\), LoRA adds a new trainable low-rank matrix \(W=BA\) in parallel and keeps \(W\) frozen. This dramatically reduces the number of trainable parameters and the optimizer/memory footprint.

Figure 1. LoRA reparametrization, where \(W\) (green) denotes the original model parameters; only the \(A\) and \(B\) matrices (black) are trained

LoRA is parameterized by a small set of hyperparameters that trade off capacity, stability, and cost. The rank \(r\) controls the size of the added low-rank matrices and therefore the number of trainable parameters. The target modules specify which layers receive adapters, with common choices such as attention and MLP projections. For small datasets, LoRA dropout can be enabled as an additional form of regularization.

While the following two case studies differ in modality (protein versus DNA), task type (token classification versus sequence classification), and underlying architecture (transformer versus striped Hyena), both use the same LoRA recipe pattern.

PSSP is the task of assigning a structural label to each amino acid in a protein sequence. Secondary structure labels describe local backbone conformations—helices and strands—without requiring a full 3D structure prediction. For many proteins, these local patterns correlate with functional motifs and global fold organization.

PSSP is a core building block for many downstream applications in biology. Because local structure is strongly correlated with protein function, PSSP can provide useful functional context. In addition, these predictions can inform tertiary structure prediction, solvent accessibility prediction, protein-protein interaction prediction, and structural class- or domain-related prediction.

At a modeling level, PSSP is a token classification problem: the input is an amino-acid sequence, and the output is a structural label for each residue.

ESM2-3B is a 3-billion-parameter protein language model, so full fine-tuning typically requires substantial compute and memory. LoRA makes adaptation practical by training only a small number of additional parameters, while still achieving strong performance on PSSP.

The team fine-tuned ESM2-3B for PSSP by adding a lightweight per-residue classification head (for Q3/Q8 labels) and training LoRA adapters through the PEFT library, while keeping the pretrained backbone weights frozen. For data, we used the curated splits released by the authors of the Porter 6 model and reported results on their provided test set. To maximize throughput, we enabled TE and sequence packing and ran the full training workflow on one NVIDIA RTX 6000 Blackwell Workstation Edition GPU in under one hour.

You can then plug this PEFT model into your training loop. The full recipe includes the dataloader, loss, and optimizer setup.

Table 1 summarizes Q3/Q8 test accuracy for the ESM2-3B plus LoRA model alongside strong published baselines reported in the Porter 6 paper. Table 1 reports the mean score over the top five validation checkpoints for ESM2-3B.

Table 1. Q3 and Q8 protein secondary structure prediction accuracy comparing ESM2-3B plus LoRA with published baseline models on the Porter 6 benchmark

Overall, LoRA fine-tuning reaches accuracy that is competitive with other state-of-the-art PSSP approaches. Figure 2 shows validation loss and accuracy versus fine-tuning steps.

Figure 2. Validation accuracy and loss curves for ESM2-3B plus LoRA during protein secondary structure prediction fine-tuning on Q3 and Q8 tasks

How can sequence packing yield higher utilization and throughput?

Protein datasets typically contain sequences of varying lengths. If they are batched naively (the padded BSHD format), they are padded to the maximum length in the batch and a large fraction of tokens become padding. This wastes compute and memory bandwidth inside attention and MLP layers.

Sequence packing (the packed/flattened THD format) reduces that waste by concatenating only the nonpadding tokens and tracking per-sequence boundaries with cumulative-length metadata. As a result, attention/MLP kernels operate on real tokens rather than padded tokens. For a deeper explanation of how packing works in practice (and how it interacts with TE packed formats), see Scale Biology Transformer Models with PyTorch and NVIDIA BioNeMo Recipes. Figure 3 shows throughput (tokens/sec) when fine-tuning with THD versus BSHD.

Figure 3. Throughput is 5.5x faster when fine-tuning with THD versus BSHD

In this setup, switching from BSHD to THD improved tokens/sec by ~5.5x, largely by removing padding overhead. The achieved speedup largely depends on the sequence length distribution, microbatch size, and GPU.

Beyond throughput, THD packing improves memory efficiency. It reduces the amount of activation or attention work spent on padding tokens, so a larger fraction of the GPU memory traffic and compute goes toward useful (non-padding) tokens.

For identical input sequences and batch size, THD typically uses less memory than BSHD because it avoids materializing padded tokens. In practice, that saved headroom is used to increase the number of real tokens processed per step.