marindna-exp472

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-v2
• R1: exp472-plantcad2-angiosperm-lr0p0005-wd0p1-train-s01-v1
• R2: exp472-plantcad2-angiosperm-lr0p0005-wd0p1-train-s02-v1

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

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.