GenomeOcean MoE 8x2.4B — upcycled from GenomeOcean-4B

A sparse mixture-of-experts (MoE) causal language model for microbial genomic sequences, initialized from DOEJGI/GenomeOcean-4B and continued through short-context and 10,240-token training phases. This release contains the final checkpoint at long-context update 8,061.

Model details

The 8x2.4B name uses the shared backbone plus one expert per layer as the approximately 2.4B component size. Shared attention and embedding weights are counted once in the 15.13B total. Each expert's feed-forward weights across all 24 layers contain approximately 1.81B parameters.

The source checkpoint is GenomeOcean-4B revision cd63253071beaebe0a60910986f64fb59a578275. Each dense 16,384-wide feed-forward network was split into two complementary 8,192-wide experts, replicated four times, with scaled down projections and paired router initialization. Attention, embeddings, and normalization weights were inherited. All model weights were subsequently trained.

Quick start

The export uses the native Transformers MixtralForCausalLM implementation. The release validation uses Transformers 4.53.3. No remote modeling code is required.

pip install "transformers==4.53.3" accelerate safetensors torch

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "DOEJGI/genomeocean-moe-8x2.4b-upcycled-from-4b"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",
).eval()

sequence = "ATGCGATCGATCGATCGATCG"
inputs = tokenizer(sequence, return_tensors="pt").to(model.device)
with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
    )
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

BF16 weights occupy approximately 30.25 GB (28.17 GiB) before runtime memory and the attention cache. Keep input plus generated output within 10,240 tokens. This is a genomic sequence model; prompts should contain genomic sequence tokens.

Token and numerical compatibility

Original GenomeOcean token IDs are preserved: [UNK] = 0, [CLS] = 1, [SEP] = 2, [PAD] = 3, and [MASK] = 4. Packed contigs use [SEP] boundaries. The Megatron checkpoint has 4,097 embedding/output rows because its NullTokenizer reserves an additional end-of-document ID. The export retains that row as <|endoftext|> = 4,096. Generation stops at [SEP] or this reserved token. This preserves all checkpoint weights and leaves ordinary DNA tokenization unchanged.

IUPAC / nucleotide handling

genomeocean-moe-8x2.4b-upcycled-from-4b uses the GenomeOcean-4B tokenizer (4096-token BPE, transformers PreTrainedTokenizerFast, normalizer: null, byte_fallback: false). Inputs must be uppercase and are expected to match the training preprocessing before tokenization:

def normalize_sequence(s: str) -> str:
    s = s.upper().replace("U", "T")
    return "".join(b if b in "ACGTNRYWSKMBDHV" else "N" for b in s)

# training-faithful encode: no [CLS], append [SEP]
ids = tokenizer(normalize_sequence(seq), add_special_tokens=False)["input_ids"] + [tokenizer.sep_token_id]

Ambiguity codes are not represented.The vocabulary contains single-letter tokens only for A, C, G, T, N (ids 5–9). There is no token for any IUPAC ambiguity code (R/S/W/Y/K/M/B/D/H/V), no U, and byte_fallback=False. The consequence is that ambiguity positions carry no residue-specific signal:

• A/C/G/T/N tokenize normally. U is mapped to T by normalization.
• Any IUPAC ambiguity code that reaches the tokenizer is encoded as [UNK] (id 0)​ — the normalizer preserves R/S/W/Y/K/M/B/D/H/V, but they are absent from the vocabulary.
• Characters outside ACGTNRYWSKMBDHV (e.g. ., -, *, non-ACGT letters) are mapped to N by normalization.
• Lowercase is not tokenized (no lowercase tokens).

Net behavior:​ IUPAC positions are effectively collapsed to the unknown token — the model treats them as a neutral, information-free gap rather than a specific base. This means:

• Fidelity is high on canonical A/C/G/T/N sequence (the common case for complete genomes).
• On sequence rich in ambiguity codes (draft assemblies, organelle/mitochondrial references, or long N/IUPAC spans), those positions don't bias the output (they read as [UNK]/N), but they also carry no base identity.
• For consistent scoring, either keep the ambiguity codes (they become [UNK], matching the packed-fragment training recipe) or map them to N yourself (N is id 8). Both are valid; [UNK] is the data-faithful choice.

Training computed router logits in FP32. Native Transformers 4.53.3 computes the router linear projection in the loaded model dtype and softmax in FP32, so BF16 inference can produce small routing and logit differences. The export preserves the trained BF16 router weights. See validation_report.json for a measured comparison on a fixed 128-token sample; it is not a full Megatron/HF equivalence test.

Training

Training used NERSC Perlmutter and Megatron-LM with Transformer Engine. Both production phases used 256 nodes / 1,024 NVIDIA A100 GPUs.

These counts exclude the original dense model's pretraining. They count token positions processed, including repeated and reverse-complement examples, and do not represent unique biological sequence exposure.

The long-context phase began from short-context update 34,286 with a fresh optimizer and schedule. It used Adam (betas 0.9/0.95, epsilon 1e-8), weight decay 0.1, gradient clipping at 1.0, a 100-update warmup to 5e-5, and cosine decay to 5e-6 at update 8,000. Router auxiliary loss coefficient was 0.01. Training ended at update 8,061 after reaching the allocation's duration limit and saving the final checkpoint. No skipped or NaN updates were recorded in the logged intervals.

Long-context data

The long-context corpus combines microbial metagenomic contigs in two groups:

• Contigs at least 51,200 bp long, emitted as 51,200-bp chunks.
• Contigs from 10,000 bp to less than 51,200 bp, packed with [SEP] boundaries.

Both include forward and reverse-complement examples. The combined corpus has 2,055,904 documents and 17,156,333,064 stored tokens. The split uses 2,035,344 training documents and 20,560 validation documents, preserving adjacent forward/reverse-complement pairs at the split boundary. The training portion contains 16,985,188,896 tokens; the long-context run processed approximately 2.49 training-split pass equivalents. Attention and positions were not reset at packed contig boundaries. Training data is not included in this repository.

Evaluation

The last scheduled validation was at update 8,000, immediately before this release's final update 8,061:

There was no separate scheduled validation at update 8,061. The 8,000 metrics describe the latest evaluated checkpoint, not an exact measurement of the released weights. Evaluations used changing batches. The exported model's loading and inference checks are recorded separately in validation_report.json; single-sample losses there are technical checks, not benchmark results.

Limitations

• The validation split contains packed-contig documents only. It does not separately establish performance on continuous long contigs.
• Independence from earlier training data, reverse complements, and homologous sequences has not been established. Pair-safe splitting alone is insufficient to establish biological independence.
• Downstream biological benchmarks, controlled comparisons with the dense model, expert specialization, and short-context retention have not been evaluated for this release. Loss differences across training phases use different datasets and context lengths and do not isolate an upcycling or context-length benefit.
• Context beyond 10,240 tokens has not been validated. Generated sequences are model predictions and do not establish biological function.

Intended use

Research on microbial genomic sequence modeling, sequence likelihoods, generation, and adaptation to downstream genomics tasks. Suitability for a particular downstream application requires task-specific evaluation.

Release files and provenance

The release contains BF16 safetensors, a native Mixtral configuration, tokenizer, generation settings, model card, license, export manifest, and validation report. export_manifest.json records checkpoint provenance, parameter count, per-shard SHA-256 checksums, and verification that every exported tensor exactly matches its corresponding source tensor after layout conversion. Optimizer states are not part of the inference release.

License

This release retains the source GenomeOcean LBNL BSD-style license, including its additional enhancements clause, in LICENSE. It uses the custom Hub label genomeocean-lbnl-bsd because the full license text is more specific than a generic BSD-3-Clause label. Redistribution and modification are permitted subject to the conditions in that file; preserve its notices and disclaimer.

The original copyright belongs to The Regents of the University of California, through Lawrence Berkeley National Laboratory, and Northwestern University. The source license was retrieved from the exact source-model revision above.

Attribution

This model builds on GenomeOcean. Please cite the original GenomeOcean work and identify this repository and its revision when reporting results with these MoE weights.