MarinDNA exp472
This repository archives experimental MarinDNA checkpoints from exp472: approximately 973M-parameter Qwen3 causal language models pretrained from scratch on plant DNA. The models use an 8,192-nucleotide context, a seven-token nucleotide vocabulary, and the plantcad/Angiosperm_65_genomes_8192bp corpus.
The primary series uses learning rate 5e-4 and weight decay 0.1. For inference, start with the latest primary Hugging Face checkpoint at step 535985. PlantCAD2 evaluation results and methodology are tracked in eric-czech/plantcad2#1.
These are research artifacts rather than official PlantCAD releases.
Primary model series: LR 5e-4 / WD 0.1
Training was continued across three run IDs:
R0:exp472-plantcad2-angiosperm-lr0p0005-wd0p1-v2R1:exp472-plantcad2-angiosperm-lr0p0005-wd0p1-train-s01-v1R2:exp472-plantcad2-angiosperm-lr0p0005-wd0p1-train-s02-v1
| Run | Step | Exact tokens seen | Approx. tokens | Approx. epoch | Phase / significance | Available format and path |
|---|---|---|---|---|---|---|
| R0 | 20,615 | 21,617,442,816 | 0.02162T | 1.00 | Retained main-training snapshot | Levanter-only: checkpoint |
| R0 | 41,230 | 43,233,837,056 | 0.04323T | 2.00 | Retained main-training snapshot | Levanter-only: checkpoint |
| R0 | 82,460 | 86,466,625,536 | 0.08647T | 4.00 | Retained main-training snapshot | Levanter-only: checkpoint |
| R0 | 164,920 | 172,932,202,496 | 0.17293T | 8.00 | Main-training endpoint; cooldown and R1 source | Levanter-only: checkpoint |
| R0 | 206,144 | 216,158,699,520 | 0.21616T | 10.00 | Cooldown endpoint; evaluated | HF + Levanter: HF · Levanter |
| R1 | 329,840 | 345,863,356,416 | 0.34586T | 16.00 | Main-training endpoint; cooldown and R2 source | Levanter-only: checkpoint |
| R1 | 371,065 | 389,090,902,016 | 0.38909T | 18.00 | Cooldown endpoint; evaluated | HF + Levanter: HF · Levanter |
| R2 | 494,760 | 518,794,510,336 | 0.51879T | 24.00 | Main-training endpoint; cooldown source | Levanter-only: checkpoint |
| R2 | 535,985 | 562,022,055,936 | 0.56202T | 26.00 | Cooldown endpoint; evaluated | HF + Levanter: HF · Levanter |
Token counts are cumulative training tokens, including repeated corpus passes and cooldown. They are computed as (step + 1) × 128 global sequences/step × 8,192 tokens/sequence. Approximate epoch is the cumulative token count divided by the 21,615,869,952 tokens in the training split; it is a corpus-relative count, including cooldown, rather than a run-name label.
Other checkpoints
| Configuration / run | Step | Exact tokens seen | Approx. epoch | Phase / significance | Available format and path |
|---|---|---|---|---|---|
LR 1e-4, WD 0.2 · exp472-plantcad2-angiosperm-lr0p0001-wd0p2-v2 |
75,046 | 78,692,483,072 (0.07869T) | 3.64 | Main-training snapshot; evaluated in the initial pilot | HF + Levanter: HF · Levanter |
LR 2e-4, WD 0.1 · exp472-plantcad2-angiosperm-lr0p0002-wd0p1-v2 |
206,144 | 216,158,699,520 (0.21616T) | 10.00 | Cooldown endpoint; evaluated | HF-only: checkpoint |
LR 1e-4, WD 0.2 · exp472-plantcad2-angiosperm-lr0p0001-wd0p2-train-s02-v1 |
535,985 | 562,022,055,936 (0.56202T) | 26.00 | Cooldown endpoint; evaluated | HF-only: checkpoint |
Formats
- HF directories contain a Transformers config, tokenizer, and
model.safetensors; they are intended for inference and evaluation. - Levanter directories under
checkpoints/preserve native training state for resuming or converting with Marin/Levanter. - Format labels describe what is present in this repository. A missing format is not implied to exist elsewhere.
Recommended inference path for this ~1B model on an H100: BF16 model compute, explicit external FlashAttention-2, use_cache=False, batches up to about 32 equal-length sequences, and FP32 softmax over only the A/C/G/T logits.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
REPO = "eczech/marindna-exp472"
SUBDIR = "exp472-plantcad2-angiosperm-lr0p0005-wd0p1-train-s02-v1/hf/step-535985"
hf = dict(subfolder=SUBDIR, trust_remote_code=True)
tok = AutoTokenizer.from_pretrained(REPO, **hf)
model = AutoModelForCausalLM.from_pretrained(
REPO,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2", # requires flash-attn; no SDPA fallback
**hf,
).eval().to("cuda")
model.config.use_cache = False
seqs = [...] # ≤32 equal-length DNA sequences
ids = tok(seqs, add_special_tokens=False, padding=False, return_attention_mask=False, return_tensors="pt").input_ids.cuda()
vocab = tok.get_vocab()
acgt = torch.tensor([vocab[b] if b in vocab else vocab[b.lower()] for b in "ACGT"], device="cuda")
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
logits = model(input_ids=ids, use_cache=False).logits
probs = logits[:, :-1].index_select(-1, acgt).float().softmax(-1)
# probs[:, t-1] = P(A,C,G,T at position t | sequence positions < t)
This preserves the PlantCAD2 causal scoring definition: normalize only the A/C/G/T logits, with softmax computed in FP32. Explicitly requesting flash_attention_2 prevents silent use of the slower SDPA path; disabling the unused generation KV cache avoids a substantial throughput penalty at larger batch sizes. These settings follow the exp472 execution-sensitivity study.
Training details
- Architecture: Qwen3 causal decoder, ~973M parameters, 16 layers, hidden size 2,048, 16 attention heads, 4 KV heads, and MLP size 8,192.
- Sequence length: 8,192 nucleotides.
- Tokenizer:
kuleshov-group/PlantCAD2-Small-l24-d0768, vocabulary[PAD],[MASK],[UNK],a,c,g,t. - Training corpus:
plantcad/Angiosperm_65_genomes_8192bp, 2,638,656 training sequences. - Augmentation: deterministic reverse complement with probability 0.5.
- Training code:
Open-Athena/marin-dna, exp472 experiment.