Skip to content

Train MarinDNA-RAG 46M on five regions with order-deduplicated vertebrates #550

Description

@gonzalobenegas

TL;DR

All five RAG training datasets and the combined development VEP harness are complete; the datasets are publicly available at immutable Hugging Face revisions.
The eight-chip TPU run in europe-west4 has completed tokenization and begun training the 46M model for 100,000 updates with the agreed scaling heuristic and 200-document effective batch.
Production is advancing at approximately 1.6 seconds per update; at 900 updates on September 10, about 43–45 hours of training remained, plus validation, checkpoint, and interruption overhead.
Final-checkpoint development VEP and frozen probes take priority within the $30 data-construction/evaluation budget.

Status

All five datasets are built, audited, and publicly published at immutable Hugging Face revisions.
They contain 1,110,006 training documents and 2,000 chr18 validation documents; all chr18 anchors are excluded from training.
The combined development harness contains 51,623 variants and has SHA-256 6631d35f9ae0afc754c623a2e3960682e8a834a28001906279745882f99cb73b.
The verified input manifest pins all five public revisions and release checksums.
The assembly and publication evidence records exact row counts, payload hashes, species counts, and padding.
Expected training exposure is 126.34B unpadded positions under the fixed mixture, compared with 204.8B allocated positions; this is an expectation before training, not an observed token count.

The European eight-chip synthetic pilot completed all 20 updates on September 10.
All 20 locked training tests passed before dispatch, and independent review found no blocking issues.
Native and Hugging Face checkpoints at completed updates 5, 10, 15, and 20 are verified in the pilot receipt.
The last steady update took 1.55134 seconds, projecting 43.09 compute hours for 100,000 updates before tokenization, validation, checkpoint overhead, and interruptions.
A separate native resume from update 10 produced byte-identical final weights and exactly matching validation losses at updates 15 and 20; the resume receipt records the comparison.
The pilot W&B run records finite training loss and decreasing synthetic validation loss.

The active production job was accepted at 17:19:04 UTC from the same validated training snapshot.
It acquired eight v6e chips and began tokenization at approximately 17:21:42 UTC.
Tokenization completed and the production training loop began at approximately 19:24 UTC.
At 19:48:34 UTC, it had reached 900 of 100,000 updates (0.9%), generally around 1.6 seconds per update.
Reported training loss fell from approximately 1.10 at update 327 to 0.864 at update 900; biological performance has not yet been evaluated.
Rolling recovery checkpoints saved successfully at steps 360 and 737; the older temporary checkpoint was removed after the newer one committed.
The worker remained on its first attempt with no reported failure at the 19:48 live controller check.
Microbatch 5 with five accumulation steps preserves the 200-document batch, scaling heuristic, model, and 100,000-update schedule.
Checkpoints use gs://marin-eu-west4/MarinDNA/exp550_rag_five_regions/checkpoints/dna-exp550-rag46m-five-regions-v1/2026.09.10.9.
The coordinator has a seven-day timeout, with rolling native checkpoints and permanent 10,000-update exports configured for recovery.

The validated four-chip us-east5 fallback projected 75.16 compute hours.
It was canceled after the European worker completed four full production shards; its child was confirmed stopped at 17:43:07 UTC, before either production run reached optimizer updates.
Its partial version-8 caches remain in GCS.
The four-chip pilot receipt preserves the fallback validation.

The earlier eight-chip production attempts are canceled; none reached optimizer updates.
The final us-east1 attempt lost three workers with controller reason worker reconcile failure threshold exceeded; no out-of-memory failure was reported.
The recovery record preserves the failure and replacement evidence.
The paid recovery worker completed at 14:49:27 UTC and was terminated at 15:03:40 UTC after its data, receipts, and local logs were retained.
Its 59.5-minute runtime cost approximately $1.00 compute before disk costs.
No paid worker is active; the cumulative $30 CPU/GPU cap and final-checkpoint-first evaluation plan remain unchanged.
No biological VEP inference has run.

The dataset framework (PR #552), combined evaluation backend (PR #554), tokenizer compatibility fix (PR #559), and explicit inference precision fix (PR #562) pass CI and independent review and are ready for human review; none is merged.
PR #556 now adds execution-permissions preflight, persistent owner-controlled Full Access setup, and approval/budget/deadline handoffs for future sessions.
Local checks, CI, and independent review pass; PR #556 is ready and unmerged.
The guidance is also applied to the permanent experiment branch, with the current recovery record.

The synthetic TPU pilot completed 20 updates on eight v6e chips with microbatch 5 and the agreed 200-document effective batch.
Its last steady update took 1.55265 seconds, projecting 43.13 compute hours for 100,000 updates before setup, checkpoints, validation, and interruptions.
Native and HF milestones at completed updates 5, 10, 15, and 20 were verified.
A separate native resume from update 10 produced byte-identical final weights and exactly matching validation losses at updates 15 and 20.
The pilot evidence and validated training code include the tested workaround for the repeated-row-zero cache bug.
All 18 locked training tests and independent review pass.

The historical 46M baseline matches the three current development cohorts exactly: Mendelian 16,140, Complex Traits 11,630, and SGE 23,853 rows.
Existing zero-shot and frozen-probe results can be reused without baseline inference.
The completed synthetic A10G pilot selected compiled strict fp32 at batch 2 after full-bf16, bf16-autocast, and TF32 compiler candidates failed their predefined numerical gates.
Strict fp32 compilation matched eager fp32 with maximum LLR error 2.84e-5 and embedding error 9.54e-7.
Measured throughput was 0.898592 variants/second, projecting 15.96 hours for the final checkpoint's 51,623 development variants before setup and frozen probes.
Combining three synthetic batches saved 1.56 seconds across 96 rows with an already warm compiler cache, including model loads.
The GPU evidence and failed precision trials preserve the tolerances, runtime, memory, and timing definitions.
The Spot GPU was reclaimed after these results had been saved; final precision CI passed 426 tests, and a completed-output Snakemake fixture verified preservation of existing default results.

PR #565 registers the exact European version-9 final checkpoint and all three development probe cells.
All CI checks and independent review pass, including all five project test jobs and the credential-free full pipeline dry-run.
The PR is ready and remains unmerged.
The permanent experiment branch includes the complete evaluation consumer and final-checkpoint commands.
The final weights will undergo the same bounded synthetic numerical-parity gates before biological inference.
The reviewed lightweight checkpoint-staging helper passes ten mock tests and requires generation-pinned checksums plus a successful receipt before GPU launch; no final checkpoint has been transferred.
The latest logbook records the staging review, recovery state, and reconciled evaluation budget.
A separate pinned-source review reproduced a latent queue-full writer-closure hang and filed #566.
Production finished tokenization without evidence of that writer-closure failure.
The first permanent checkpoint and full chr18 LM validation are scheduled at 10,000 completed updates.
Hugging Face publication is complete; development VEP begins with the final checkpoint after training.

Description

Experiment #402 trained 46M and 104M causal models on seven fixed mammalian windows followed by human, using 255 bases per species and a 2,048-token document.
The accepted runs each completed 30,000 updates at 2,097,152 tokens per update, totaling 62,914,560,000 token presentations.
Validation loss was still falling at the final checkpoint.
Species identity, phylogenetic distance, and document position were confounded by the fixed order.

Experiment #517 selected region-specific anchor defaults and established a cohort with 39 non-human order representatives plus human.
Experiment #523 completed the mammalian chain release, and PR #549 merged a shared chain-based projection workflow for mammals and non-mammalian vertebrates.
These assets support constructing the follow-up without staging HAL or regenerating chains.

Hypothesis or Goal

Measure whether the combined changes in training footprint, species coverage, species order, and training exposure improve the 46M model's development VEP performance relative to the historical RAG baseline.
Track performance through training on Mendelian, Complex Traits, and SGE using the standard benchmark-specific metrics and eligible subsets.
Report gains, regressions, and uncertainty separately for each benchmark.
This first run tests the combined recipe; it cannot attribute a change to an individual intervention.

Implementation ownership and reuse

Implement the five RAG training datasets and a combined three-benchmark VEP context harness as additional rules in the existing snakemake/vertebrate_projection_dataset project.
Keep their Python implementation and tests in that project's src/marin_dna_vertebrate_projection/ package and tests/ directory.
Add RAG-specific rules, configuration, explicit targets, and output namespaces while preserving the existing rules and their outputs.

Reuse the pipeline's pinned anchor catalogs, order manifest, chain files, genome archives, and compatible projection and sequence-extraction outputs.
Read unsplit per-species sequence outputs, assemble complete locus documents, and then apply the chr18 RAG validation holdout and training-only RC augmentation.
The existing per-species-row train/validation splits are not the input to RAG assembly.
Keep shared intermediate artifacts under the existing pipeline's storage ownership and reference completed upstream outputs by their pinned identities.

Reuse a projected window only when its exact source coordinates, target species, chain and genome revisions, and projection contract match.
New variant-centered requests require chain queries when no matching result exists; a nearby training anchor is not a substitute.
Deduplicate identical projection requests across the training catalogs and combined VEP harness, and batch the outstanding requests per species to reuse staged chain and genome assets.
Do not treat historical HAL/MultiZ outputs as equivalent to a new chain result solely because they use the same species and window length.
Inspect the RAG dry-run to verify that cached inputs are consumed and only intended missing work is scheduled.

Keep model training in a self-contained Marin experiment project.
Keep VEP scoring, benchmark metrics, and frozen probes in the existing snakemake/analysis/evals_v2 project, consuming the combined harness.

Agreed data design

Human anchors and mixture

Build five datasets from the recommended anchor recipes:

  • CDS: uniform 255-bp windows at 128-bp stride with the established CDS assignment and exon-boundary context.
  • TSS / 5′ UTR: the same uniform grid with the combined protein-coding TSS-region and 5′-UTR assignment.
  • 3′ UTR: the same uniform grid with the established 3′-UTR assignment.
  • ncRNA: curated annotation-first ncRNA exons with the existing flanks, short-interval expansion, long-interval tiling, and ownership rules.
  • Enhancer: one window centered on each ENCODE V4 dELS/pELS element, preserving the exon-exclusion and source-arm ownership rules.

Carry forward the recipes' human conservation gate: at least 51 of 255 positions have phyloP447way ≥2.2162.
Pin the source anchor catalogs and preserve their membership when constructing the RAG datasets.
The existing recommended releases used different non-mammal projection backends; record the shared chain construction here as a new dataset recipe.
The previous recommendation did not establish a cross-region-exclusive hybrid catalog, so audit exact duplicates and overlap and document their treatment.

Use a uniform mixture across the five datasets.
Sample each dataset's train split with probability 20%.
Use the fixed 10,240-token document length, including padding, for all training and validation budget calculations.

Species and projection

Reuse the audited order cohort: 18 non-human mammalian representatives, 21 non-mammalian vertebrate representatives, and human.
Human is the sole Primates representative and appears once per locus.
Use every successfully projected representative at each locus without subsampling the available species.

Use the same pinned human-to-species chain files and assembly-matched genome archives for training and VEP.
Project the central human nucleotide, require an accepted unique target locus, and extract a 255-bp target window in human-anchor orientation.
The projected center must appear at index 127.
Keep internal genomic coordinates 0-based and half-open, with conversions at source-format boundaries.
Register the complete biological chain and genome manifests before execution; the merged workflow does not bundle that complete manifest.

Document layout and padding

Concatenate the available species windows with an atomic [SEQ] separator between adjacent windows.
Include human in the training permutation.
Shuffle species once when building each published row, using a reproducible random permutation that varies across rows.
Keep that order fixed on repeated training presentations.
For VEP, shuffle the available non-human windows and append human as the last biological sequence.
Choose the permutation without using labels, alleles, or phylogenetic distance.

Omit an unavailable species window entirely.
Preserve genuine N bases within available windows; do not insert an all-N placeholder for an unavailable species.
The tokenizer adds exactly one BOS token and right-pads to 10,240 token positions with a distinct padding token.
Padding targets contribute zero training loss, and padding keys cannot influence real-token attention.
Use ordinary fixed-shape batches without length bucketing or packing.

With k available species, the unpadded token count is 1 + 255k + (k - 1) = 256k.
The maximum is 40 × 256 = 10,240 tokens.
Retain loci with limited retrieval coverage, including human-only loci when all non-human projections fail.
Projection failure must not silently change the VEP cohort.

Minimal Hugging Face training datasets

Each of the five training datasets exposes only a string column named sequence in its training and validation splits.
Each HF dataset contains its own train and validation splits; the validation split contains up to 400 documents sampled from chr18 (all available if fewer).
Use these published validation splits for language-model monitoring throughout training.
The string contains the concatenated 255-bp windows and [SEQ] separators.
Store each window once in that concatenation.
Add BOS and padding during tokenization rather than storing them in the sequence strings.
The separators retain the boundaries needed to permute segments without separate species columns.

Keep coordinates, species identities and order, orientation, projection failures, QC, and publication row mappings in the Snakemake workflow's durable outputs.
Retain a mapping from published shard and row position to the producing records so the single-column dataset remains auditable.
The dataset card links the producing commit, configuration, and release manifest.
Consumers pin the immutable HF revision.

Language-model validation

Hold out all of chromosome 18 from training across all five datasets.
Sample validation documents randomly from chr18 and specify the validation budget using the fixed 10,240-token model input length.
Use up to 400 documents per region: at most 4,096,000 tokens per region and 20,480,000 tokens across all five datasets.
If a region has fewer than 400 chr18 anchors, use all available chr18 anchors and report the actual document and token counts; a smaller validation set is acceptable.
At 400 documents, this is 2.34% below the historical 16,384 × 256 = 4,194,304-token validation budget per region; regions with fewer chr18 anchors have a smaller validation budget.
Select complete locus documents and pad each to 10,240 tokens.
Exclude every chr18 window from training, including chr18 windows not selected for validation.
This also keeps all orientation copies and cross-region overlaps on chr18 out of training; no additional genome-wide overlap exclusion is required.
Choose the split once with a reproducible seed and keep the validation data fixed across checkpoints.
This chr18 language-model holdout is separate from the chromosome-defined labeled VEP development/test split.

Combined VEP harness and shared inference

Create one chain-derived VEP dataset containing the Mendelian, Complex Traits, and SGE records.
Include a benchmark identifier and stable source-row identity, preserving each benchmark's canonical variants, splits, labels, and metric-membership fields.
Namespace row identities and group metadata by benchmark so equal local identifiers cannot merge unrelated records.
The single-column publication decision applies to the training datasets; the combined evaluation harness retains the fields required for variant scoring and cohort matching.

Run one batched inference job per checkpoint over the combined development records.
Load the checkpoint once and reuse the model and supported compiled execution shapes across all three benchmarks.
The job performs the required REF/ALT and forward/RC computations and, at the final checkpoint, also collects the human-token embeddings needed for frozen probes.
Combining the jobs should reduce repeated setup, model loading, and compilation or warmup overhead; measure the end-to-end savings in the pilot.

Route the joint inference results back into the three canonical per-benchmark score outputs under the existing evals_v2 storage layout.
Assert exact source-cohort membership and row counts separately for each benchmark.
Reuse the existing benchmark-specific score transforms, metric functions, and probe protocols after routing the outputs.
Keep labels, matched groups, SGE accessions, uncertainty estimates, and probe fitting separate by benchmark.
Validate joint scoring and output routing against separate benchmark processing on a small development fixture before the full evaluation schedule.

Project each exact variant-centered human locus using the shared center-1 contract.
Assert the human REF allele against the pinned reference genome before scoring.
REF and ALT use the same retrieved windows, species permutation, and padding; only the designated human allele changes.
Keep human last in both forward and reverse-complement inputs.
For strand averaging, reverse-complement each species window consistently while preserving the chosen segment order.

Use only the official development cohorts for iteration: odd-numbered autosomes and chromosome X.
Filter the combined dataset to the authorized development split before inference or embedding extraction.
Preparing or caching held-out records is permitted, but evaluating held-out rows, inspecting their labels for analysis, or using their predictions requires separate explicit authorization.
Apply the maintained mature-miRNA exclusions, support gates, benchmark-specific score transforms, and metric definitions.
Register each exact model–dataset evaluation cell before running it through evals_v2.

Training and evaluation schedule

Start with one 46M model trained from scratch.
Use training seed 0 and data-construction seed 42 as the reproducibility defaults.
Carry forward training-only reverse-complement augmentation, applying the orientation consistently across all species windows in a document.
Publish one forward and one reverse-complement training row per retained training locus, with the species permutation fixed once for each published row.
Reverse-complement the DNA within each species window while preserving separator syntax.
Validation contains original-orientation documents only, and both training orientations of each locus remain in the same split.
Use the previous RAG architecture as the reference and validate the tokenizer, 10,240-token positional configuration, padding masks, and checkpoint exports together.
Resolve a current compatible Marin release and optimizer configuration before constructing the launch.

Use 200 documents per optimizer update: 2,048,000 allocated token positions.
Train for 100,000 optimizer updates: 204,800,000,000 allocated positions.
All stated token budgets include padding; padding targets remain excluded from the loss.

Use Iris TPUs for training; the user reports that these TPU resources are free.
Use separate CPU/GPU workers for dataset construction and VEP, with a combined spending ceiling of $30 across both activities, including pilots and retries.
Choose the least expensive suitable workers and reuse existing assets and cached outputs within that ceiling.
Run a bounded pilot to measure warmed throughput and memory use before the long training run.
Select per-device microbatch and accumulation from those measurements while preserving the agreed effective batch.
Record pilot throughput and costs; prioritize final-checkpoint VEP and decide intermediate evaluations from the remaining budget.

Save a resumable native checkpoint and HF model export every 10,000 completed optimizer updates, including update 100,000.
Run language-model validation at those same ten milestones.
Prioritize one combined development zero-shot VEP inference job and frozen probes on all three benchmarks at the final checkpoint.
After final-checkpoint evaluation, use the remaining $30 data-construction/VEP budget for earlier saved checkpoints; intermediate VEP milestones are conditional on remaining budget.
Keep the saved validation documents and VEP species permutations fixed across checkpoints.
Compare against the registered historical 46M RAG results only when variant cohorts and metric definitions match exactly.
Run frozen probes on all three development benchmarks at the final checkpoint, using the registered grouped protocols and pooling human-token representations only.

Expected training epochs

The run presents 20 million documents in total, with an expected 4 million from each region under the uniform mixture.
One RAG document groups all available species at a human anchor, so estimate dataset size from human-anchor counts rather than the old per-species row counts.
With forward and RC training rows, the epoch count for a region with A source anchors and H chr18 anchors is 4,000,000 / (2 × (A - H)).
The audited source catalogs give the following expected counts after excluding all chr18 anchors; these must be checked against the assembled train files before launch.

Region Source anchors Chr18 anchors Forward + RC train rows Expected epochs
CDS 295,561 4,933 581,256 6.88
TSS / 5′ UTR 57,418 1,048 112,740 35.48
3′ UTR 67,155 1,380 131,550 30.41
ncRNA 28,815 617 56,396 70.93
Enhancer 117,010 2,978 228,064 17.54

All five source catalogs support 400 chr18 validation documents.
The source counts and pins and biological producer configuration are immutable.
An epoch covers both published orientations; expected presentations per original training locus are twice the listed epoch values.
Uniform mixture sampling gives more repetitions to the smaller catalogs.

Execution budget

The agreed resource plan is free Iris TPU training and at most $30 combined for dataset construction and VEP.
Prioritize dataset construction, training, and final-checkpoint evaluation within that ceiling.
After evaluating the final checkpoint, assess the remaining budget and evaluate earlier checkpoints as resources allow.
The user explicitly made intermediate VEP evaluation conditional on remaining budget during September 9 execution setup.
Do not exceed the cap; a full ten-checkpoint VEP trajectory is no longer a prerequisite for proceeding.

Completion criteria

  • Five training datasets and one combined VEP harness covering all three benchmarks are produced with pinned anchors, order representatives, chains, genomes, and configuration.
  • RAG data construction extends the existing vertebrate pipeline with additive rules that reuse compatible precomputed artifacts and preserve existing outputs.
  • Training publications contain only sequence, with their row counts, sequence boundaries, and bytes verified against the pipeline-owned release manifests.
  • Coordinate bounds, center placement, strand orientation, missing-projection accounting, split isolation, token geometry, padding masks, and REF/ALT cohort parity pass their checks.
  • Combined VEP inference loads each checkpoint once, reproduces separate benchmark processing on the fixture, and emits the three standard score outputs for the existing metric and probe rules.
  • Changed Python projects pass their locked tests, and changed Snakemake workflows pass dry-runs before biological execution.
  • The approved training schedule completes with finite validation records, resumable checkpoints, and verified model exports.
  • Final development VEP comparisons report canonical cohorts, uncertainty, the agreed token budget, and per-region repetition; earlier-checkpoint trajectories are conditional on the remaining budget.
  • The experiment's outcome, including negative results, is recorded with reproducible artifact links and a reviewed knowledge-base disposition.

Links

Decision log

  • Use order-deduplicated representatives across the complete vertebrate cohort.
  • Keep 255-bp windows and include human at a shuffled training position.
  • Shuffle species once when materializing each published row and retain that permutation during training.
  • 2026-09-09: Hold out all chr18 anchors from training and sample up to 400 language-model validation documents per region from chr18 (all available if fewer), as requested by Gonzalo during execution setup.
  • Publish a separate train split and a validation split of up to 400 chr18 documents in each of the five HF datasets.
  • Sample the five train splits with equal probability.
  • Use 200 documents per optimizer update, totaling 2,048,000 allocated token positions.
  • Train for 100,000 updates, totaling 204.8 billion tokens under the agreed accounting.
  • Save checkpoints and run language-model validation every 10,000 updates, including the final checkpoint; VEP follows the final-first budget policy.
  • 2026-09-09: Evaluate the final checkpoint first (zero-shot VEP and frozen probes on all three development benchmarks), then use remaining budget for earlier checkpoints.
  • Use free Iris TPUs for training and a combined $30 cap for dataset construction and VEP.
  • Implement RAG dataset and VEP context construction in the existing vertebrate projection pipeline to share precomputed inputs and intermediate artifacts.
  • Combine the three VEP benchmarks into one harness and one inference job per checkpoint, then route results to the existing per-benchmark metrics and probes.
  • Count the fixed 10,240-token input length, including padding, in all training and validation budget calculations.
  • Force human last for VEP.
  • Omit missing species and use masked right-padding to a fixed 10,240-token model input.
  • Use ordinary batches without length-based packing or bucketing.
  • Publish each training dataset as concatenated sequence strings and retain detailed provenance with the Snakemake producer.
  • Publish the agreed experiment plan as a coordinating issue.

Conclusion

Dataset construction and publication are complete.
Production training and final-checkpoint development evaluation remain in progress.

Knowledge-base disposition

Pending results and interpretation review.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    agent-generatedCreated by an agentdataTopic: training-data construction, projection, labeling, and filteringevalsTopic: evaluation datasets, scoring protocols, and metricsexperimentKind: bounded unit of research with a hypothesis or goalmodelingTopic: architecture, objective, tokenization, loss weighting, and context

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions