Deploy DNABERT-2-117M: Full Install Guide

Why does every attempt to load DNABERT-2 throw a Triton compile error or a weird BertConfig complaint the moment you leave the toy notebook?

That friction is the real path for DNABERT-2 (117M multi-species genome foundation model). Weights: zhihan1996/DNABERT-2-117M on Hugging Face. Code and GUE tooling: MAGICS-LAB/DNABERT_2. Paper: arXiv:2306.15006 (ICLR 2024). Goal here is a local install of the public checkpoint that actually returns embeddings – without treating flash-attention as mandatory.

System requirements before you touch conda

Official setup wants Python 3.8. GPU helps a lot. Community reports (as of the 2025-2026 genomic FM threads) put typical human-genome-style inference in the 4-8 GB VRAM band. CPU-only runs short sequences; it crawls on anything longer.

The pins are old on purpose. “Latest CUDA + latest torch” is a conflict until you prove it isn’t. That mismatch – not the DNA math – is what stops most first installs.

Where to get DNABERT-2 (official sources only)

No release tarball. No official Docker image on the main repo. Two pieces only:

• Code + requirements: clone https://github.com/MAGICS-LAB/DNABERT_2
• Weights + tokenizer: zhihan1996/DNABERT-2-117M on Hugging Face (pulled by transformers)

GUE data is optional (Google Drive link in the README). Follow-ons named on the same repo – DNABERT-S (2024) and GenomeOcean (noted 2025) – are separate model IDs, not in-place upgrades of the 117M checkpoint.

Step-by-step install (recommended path)

Skip optional Triton on try one. Most “won’t load” tickets start there.

# 1. Fresh environment
conda create -n dna python=3.8 -y
conda activate dna

# 2. Clone code
git clone https://github.com/MAGICS-LAB/DNABERT_2.git
cd DNABERT_2

# 3. Install pinned deps (do NOT freestyle newer versions yet)
python3 -m pip install -r requirements.txt
# as of the repo requirements.txt: torch==1.13.1 transformers==4.29.2
# peft==0.3.0 einops==0.6.1 omegaconf==2.3.0 accelerate==0.20.3 ...

Want flash later? README still shows Triton-from-source (clone openai/triton, cmake, editable install). On many current drivers that build hangs or blows up at kernel compile time. Optional. Painful.

The catch is you don’t need it. Community no-flash weights – e.g. quietflamingo/dnabert2-no-flashattention – drop Triton/FlashAttention integration. Public demos match; dependency surface shrinks.

First-time load configuration that actually works

Load fails after a transformers bump? Inject BertConfig first. Official Quick Start documents two paths: plain AutoModel around 4.28, and config-first once you are past that. The repo pins 4.29.2, so use the safer recipe by default.

import torch
from transformers import AutoTokenizer, AutoModel
from transformers.models.bert.configuration_bert import BertConfig

model_id = "zhihan1996/DNABERT-2-117M"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

config = BertConfig.from_pretrained(model_id)
model = AutoModel.from_pretrained(
 model_id,
 trust_remote_code=True,
 config=config
)
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

trust_remote_code=True is mandatory – custom Bert layers and the BPE tokenizer ship inside the model repo, not vanilla transformers.

Pro tip: After the first successful download, point HF_HOME or TRANSFORMERS_CACHE at a fast disk. Re-pulling the 117M checkpoint after every env wipe is a waste of time.

Finetune flag worth getting right once: --model_max_length ≈ 0.25 × raw nucleotide length. BPE shrinks sequences about 5× (README finetune notes + paper). 1000 bases → start near 250 tokens. Full base count as max length mostly pads.

Verify the install in under a minute

dna = "ATGCGATCGATCGATCGATCGAAATTTGGGCCC" # short synthetic stretch
inputs = tokenizer(dna, return_tensors="pt")["input_ids"].to(device)
with torch.no_grad():
 hidden = model(inputs)[0] # [1, seq_len_tokens, 768]
emb = torch.mean(hidden[0], dim=0)
print(emb.shape) # torch.Size([768])
print(float(emb.mean())) # finite float, not NaN

Shape prints and the mean is finite? Stack is live. Then try a tiny finetune run against whatever sample CSV the repo ships, low epoch count, before you touch production labels.

Common install errors and fixes

Same five failures show up in GitHub threads again and again.

• Triton / flash-attn compiler error or hang – Don’t install optional Triton. Use no-flash community weights, or stay on the standard attention path. Issues such as #57 and #140 track this on the README inference snippet.
• AutoModel fails or “unexpected” config after upgrading transformers – BertConfig.from_pretrained before AutoModel when version > 4.28. Dual recipe is in official Quick Start.
• CUDA / torch mismatch after pip install – torch 1.13.1 wheel doesn’t match the driver. Recreate the env; install a matching CUDA build of torch 1.13.1 before the rest of requirements.txt, or debug on CPU first.
• Download fails for DNABERT-2-117M – HF connectivity, auth if rate-limited, or manual snapshot into the cache.
• LoRA / PEFT breakage – Pinned transformers 4.29.2 + peft 0.3.0 is awkward for some LoRA stacks (see open issues around #152-style reports). Full finetune path first; upgrade PEFT/transformers only after smoke test is green.

Still haunted by half-broken site-packages? New conda env name. Not dna again.

Upgrade, migration, and uninstall

No semantic version tags on the core 117M weights – you pull whatever HF serves for that model id. Successor models get new ids.

cd DNABERT_2
git pull
pip install -r requirements.txt --force-reinstall

conda deactivate
conda env remove -n dna
rm -rf ~/.cache/huggingface/hub/models--zhihan1996--DNABERT-2-117M
# same pattern for any no-flash fork cache entry

FAQ

Do I need a GPU to try DNABERT-2?

No for a short embedding smoke test. Yes for real finetune batches or anything you care about finishing today.

Why does the README say model_max_length = 0.25 × my DNA length?

BPE merges frequent segments. A 1000-base string often lands near ~200 tokens, not 1000. Set max length around a quarter of the raw bases so padding stays sane and you don’t truncate the tokenized form by accident. Using the full nucleotide count mostly buys you pad tokens and VRAM burn.

Is the flash-attention / Triton path required for correctness?

No – efficiency only. People treat the README’s optional Triton build like a required dependency, then burn an afternoon on compile errors. It isn’t required. Start on standard attention or a no-flash fork; confirm the 768-d mean-pool check; only then try Triton on a machine where that build is known to work. Matching embeddings on reference sequences matter more than winning a kernel compile.

Clone the repo, python=3.8 env, pinned requirements, BertConfig + trust_remote_code, 768-d mean-pool check. When that prints, the deploy is yours – CSV finetune or embedding pipeline next.