From 5bb87e7acda2d0f79924af42b464d4964f44ba40 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:15:38 -0400 Subject: [PATCH 1/2] Score several assays per Evo 2 checkpoint load The 40B checkpoint is 76.6 GiB and takes over a minute to load, so one process per assay spends most of its time loading. Add --row_ids, which takes the same range syntax as a slurm array (0-8,11-32) and scores each assay in turn after a single load. --row_id is unchanged, so score_evo2.sh and any other caller keep working exactly as before. Two guards come with it. --require_fp8 aborts unless the constructed model reports use_fp8_input_projections=True: Evo2 falls back to bf16 for 7B checkpoints when Transformer Engine is unavailable, and it does so silently, so without the check a run can be bf16 while its provenance records FP8. The bf16 fallback is only valid for the 7B in any case, since the package refuses to build the other checkpoints without FP8. Scored CSVs are now written to a temporary file and renamed into place. A run that dies mid-write would otherwise leave a truncated CSV that the resume logic reads as a finished assay. Token budgeting also had to learn what the model actually processes. --max_tokens_per_batch divided by the raw sequence length, but a BOS token adds one position and FP8 pads the sequence dimension to a multiple of 16 inside every input projection, so an assay of length 87 runs at 96. effective_length accounts for both. --- .../baselines/Evo/score_evo2_single_dms.py | 314 ++++++++++++------ fitness/merge_scoring_files.py | 10 + fitness/performance_fitness.py | 2 + 3 files changed, 233 insertions(+), 93 deletions(-) diff --git a/fitness/baselines/Evo/score_evo2_single_dms.py b/fitness/baselines/Evo/score_evo2_single_dms.py index a550541..2e31a2c 100644 --- a/fitness/baselines/Evo/score_evo2_single_dms.py +++ b/fitness/baselines/Evo/score_evo2_single_dms.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """ -Score a single RNAGym DMS assay with an Evo 2 model. +Score RNAGym DMS assays with an Evo 2 model. Evo 2 is an autoregressive genomic language model (StripedHyena 2). For each variant sequence we compute the mean per-token log-likelihood under the model @@ -9,17 +9,16 @@ This is the Evo 2 counterpart of ``score_evo_single_dms.py`` (Evo 1 / 1.5). It uses the official ``evo2`` package (https://github.com/ArcInstitute/evo2) rather -than the ``evo`` package, and supports the large ``evo2_40b`` checkpoint, which -requires FP8 via Transformer Engine on Hopper GPUs and is automatically sharded -across every visible GPU by the Vortex inference engine. +than the ``evo`` package, and supports every checkpoint in ``evo2.utils.MODEL_NAMES`` +(``evo2_1b_base``, ``evo2_7b``, ``evo2_20b``, ``evo2_40b``, ...). Notes on multi-GPU ------------------ Vortex places and (for large models) shards the model across all CUDA devices that are visible. Select the GPUs with ``CUDA_VISIBLE_DEVICES`` and do NOT move the model manually with ``.to(device)``. ``evo2_40b`` does not fit on a single -80 GB GPU and needs at least two (e.g. 2xH100-80GB — but note the -40B/20B/1B checkpoints require FP8 + Transformer Engine, i.e. a Hopper GPU). +80 GB GPU and needs at least two (e.g. 2xH100-80GB). The 40B/20B/7B/1B +checkpoints all request FP8 via Transformer Engine, i.e. a Hopper GPU. Offline weights --------------- @@ -28,6 +27,15 @@ pre-merge the checkpoint once (see ``download_weights.sh``) and pass the merged file via ``--local_path`` so no network access is needed at run time. +Batching +-------- +All variants of one assay share a single sequence length, so batches never need +padding between sequences and the batch size cannot change a score beyond +floating point roundoff. ``--max_tokens_per_batch`` sizes each batch by a token +budget rather than a sequence count. When FP8 input projections are enabled, +Vortex pads the sequence dimension up to a multiple of 16 inside every +projection, so the budget is applied to that padded length. + Usage ----- python score_evo2_single_dms.py \ @@ -38,11 +46,16 @@ --model_name evo2_40b \ --local_path /path/to/evo2_40b.pt \ --batch_size 1 + + # several assays in one process, so the checkpoint is loaded once + python score_evo2_single_dms.py --row_ids 0-8,11-32 ... """ import argparse +import math import os import sys +import tempfile from pathlib import Path import numpy as np @@ -63,17 +76,61 @@ def preprocess_sequence(sequence: str) -> str: return sequence.strip().upper().replace("U", "T") +def parse_row_ids(spec: str) -> list: + """Parse a row selection such as ``0-8,11-32,40`` into a sorted list. + + Empty components are rejected rather than skipped: ``0-8,,11`` is far more + likely to be a typo than an intention, and silently dropping it would score + a different set of assays than the caller asked for. + """ + rows = set() + for part in spec.split(","): + part = part.strip() + if not part: + raise ValueError(f"Empty component in --row_ids: {spec!r}") + if "-" in part.lstrip("-"): + start, end = part.split("-", 1) + start, end = int(start), int(end) + if end < start: + raise ValueError(f"Empty range in --row_ids: {part}") + rows.update(range(start, end + 1)) + else: + rows.add(int(part)) + if not rows: + raise ValueError(f"No rows selected by --row_ids {spec!r}") + return sorted(rows) + + +def effective_length(seq_len: int, prepend_bos: bool, fp8: bool) -> int: + """The sequence length the model actually processes. + + ``prepare_batch`` prepends one token when ``prepend_bos`` is set, and Vortex's + ``pad_to_multiple`` pads the sequence dimension to a multiple of 16 inside + every input projection when FP8 is enabled. + """ + length = seq_len + int(prepend_bos) + if fp8: + length = 16 * math.ceil(length / 16) + return length + + def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser( - description="Run Evo 2 inference on the sequences of a single DMS assay." + description="Run Evo 2 inference on the sequences of one or more DMS assays." ) - parser.add_argument( + rows = parser.add_mutually_exclusive_group(required=True) + rows.add_argument( "--row_id", type=int, - required=True, help="Row ID in the reference sheet to process", ) + rows.add_argument( + "--row_ids", + type=str, + help="Several reference sheet rows, e.g. '0-8,11-32'. They are scored in " + "one process so the checkpoint is loaded once.", + ) parser.add_argument( "--ref_sheet", type=str, @@ -120,9 +177,10 @@ def parse_args(): type=int, default=None, help="If set, the batch size is derived per assay as " - "max(1, max_tokens_per_batch // seq_len), overriding --batch_size. Keeps " - "GPU memory roughly constant across assays of very different lengths " - "while maximising throughput (e.g. 8192).", + "max(1, max_tokens_per_batch // effective_length), overriding " + "--batch_size. The effective length accounts for the BOS token and for " + "Vortex's multiple-of-16 padding under FP8. Keeps GPU memory roughly " + "constant across assays of very different lengths (e.g. 8192).", ) parser.add_argument( "--reduce_method", @@ -149,6 +207,14 @@ def parse_args(): "comparison. Pass --no-average_reverse_complement for forward strand " "only (~2x faster).", ) + parser.add_argument( + "--require_fp8", + action="store_true", + help="Abort unless the model was actually built with FP8 input " + "projections. Evo2.load_evo2_model silently falls back to bf16 for 7B " + "checkpoints when Transformer Engine is unavailable, so without this a " + "run can be bf16 while everything around it records FP8.", + ) parser.add_argument( "--overwrite", action="store_true", @@ -157,8 +223,8 @@ def parse_args(): return parser.parse_args() -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """Return the DMS_ID for ``row_id`` in the reference sheet.""" +def load_reference_data(ref_sheet_path: str, row_ids) -> list: + """Return the DMS_IDs for ``row_ids`` in the reference sheet.""" try: ref_df = pd.read_csv(ref_sheet_path) except FileNotFoundError: @@ -168,15 +234,18 @@ def load_reference_data(ref_sheet_path: str, row_id: int) -> str: ref_df.columns = [c.lstrip("") for c in ref_df.columns] if "DMS_ID" not in ref_df.columns: raise KeyError("Reference sheet must contain a 'DMS_ID' column") - if row_id < 0 or row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} out of range (reference sheet has {len(ref_df)} rows)" - ) - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - return str(dms_id) + dms_ids = [] + for row_id in row_ids: + if row_id < 0 or row_id >= len(ref_df): + raise ValueError( + f"Row ID {row_id} out of range (reference sheet has {len(ref_df)} rows)" + ) + dms_id = ref_df.loc[row_id, "DMS_ID"] + if pd.isna(dms_id): + raise ValueError(f"DMS_ID is missing for row {row_id}") + dms_ids.append(str(dms_id)) + return dms_ids def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: @@ -193,87 +262,146 @@ def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: return df +def write_csv_atomically(df, output_file): + """Write the scored assay, then rename it into place. + + These CSVs are the published scores, so a partial file must never appear + under the final name: an interrupted or out-of-quota write would otherwise + leave a truncated CSV that later looks like a completed assay to the + resume logic, and --overwrite would destroy a good file to produce it. + """ + output_file = Path(output_file) + handle, tmp_path = tempfile.mkstemp(dir=str(output_file.parent), + prefix=f".{output_file.name}.", suffix=".tmp") + os.close(handle) + try: + df.to_csv(tmp_path, index=False) + os.replace(tmp_path, output_file) + except BaseException: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + + +def score_one_assay(evo2_model, args, dms_id, fp8_enabled): + """Score one assay and write its CSV. Returns the Spearman correlation.""" + output_file = Path(args.output_dir_path) / f"{dms_id}.csv" + dms_df = load_dms_data(args.dms_dir_path, dms_id) + + # Preprocess sequences (RNA -> DNA), tracking any rows we cannot score. + print("Preprocessing sequences...") + raw = dms_df["sequence"] + valid_mask = raw.notna() & (raw.astype(str).str.strip() != "") + n_skipped = int((~valid_mask).sum()) + if n_skipped: + print(f"Skipping {n_skipped} rows with empty/NaN sequence") + sequences = [preprocess_sequence(s) for s in raw[valid_mask].astype(str)] + max_len = max((len(s) for s in sequences), default=0) + print(f"Scoring {len(sequences)} sequences (max length {max_len} nt)") + + # Choose the batch size (token-budget adaptive if requested). + batch_size = args.batch_size + if args.max_tokens_per_batch is not None and sequences: + eff_len = effective_length(max_len, args.prepend_bos, fp8_enabled) + batch_size = max(1, args.max_tokens_per_batch // eff_len) + print(f"Token budget {args.max_tokens_per_batch}: seq_len={max_len} " + f"prepend_bos={args.prepend_bos} fp8={fp8_enabled} " + f"effective_length={eff_len} -> batch_size={batch_size}") + + print(f"Running inference (batch_size={batch_size}, " + f"reduce_method={args.reduce_method}, prepend_bos={args.prepend_bos}, " + f"rc={args.average_reverse_complement})...") + scores = evo2_model.score_sequences( + sequences, + batch_size=batch_size, + prepend_bos=args.prepend_bos, + reduce_method=args.reduce_method, + average_reverse_complement=args.average_reverse_complement, + ) + scores = np.asarray(scores, dtype=float) + + # Write scores back onto the scored rows (NaN for skipped ones). + score_column = f"{args.model_name}_score" + dms_df[score_column] = np.nan + dms_df.loc[valid_mask, score_column] = scores + + # Spearman on the rows we actually scored. + scored = dms_df.loc[valid_mask, ["DMS_score", score_column]].dropna() + if len(scored) >= 2: + correlation, pvalue = spearmanr(scored["DMS_score"], scored[score_column]) + else: + correlation, pvalue = float("nan"), float("nan") + + write_csv_atomically(dms_df, output_file) + + print("\nSummary:") + print(f" DMS ID: {dms_id}") + print(f" Sequences scored: {len(sequences)}") + print(f" Score column: {score_column}") + print(f" Spearman vs DMS: {correlation:.3f} (p={pvalue:.2e})") + print(f" Saved to: {output_file}") + return correlation + + def main(): args = parse_args() output_dir = Path(args.output_dir_path) output_dir.mkdir(parents=True, exist_ok=True) - try: - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") + row_ids = [args.row_id] if args.row_id is not None else parse_row_ids(args.row_ids) + dms_ids = load_reference_data(args.ref_sheet, row_ids) + print(f"Rows {row_ids} -> DMS IDs: {dms_ids}") + + todo = [] + for row_id, dms_id in zip(row_ids, dms_ids): output_file = output_dir / f"{dms_id}.csv" if output_file.exists() and not args.overwrite: print(f"Output already exists (use --overwrite to redo): {output_file}") - return - - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (RNA -> DNA), tracking any rows we cannot score. - print("Preprocessing sequences...") - raw = dms_df["sequence"] - valid_mask = raw.notna() & (raw.astype(str).str.strip() != "") - n_skipped = int((~valid_mask).sum()) - if n_skipped: - print(f"Skipping {n_skipped} rows with empty/NaN sequence") - sequences = [preprocess_sequence(s) for s in raw[valid_mask].astype(str)] - print(f"Scoring {len(sequences)} sequences (max length " - f"{max((len(s) for s in sequences), default=0)} nt)") - - if not torch.cuda.is_available(): - print("WARNING: CUDA not available — Evo 2 requires a GPU.", - file=sys.stderr) - print(f"Visible GPUs: {torch.cuda.device_count()}") - - # Initialize model. Vortex handles device placement / multi-GPU sharding; - # do NOT call .to(device). - print(f"Loading Evo 2 model: {args.model_name} " - f"(local_path={args.local_path})...") - evo2_model = Evo2(args.model_name, local_path=args.local_path) - - # Choose the batch size (token-budget adaptive if requested). - batch_size = args.batch_size - if args.max_tokens_per_batch is not None and sequences: - max_len = max(len(s) for s in sequences) - batch_size = max(1, args.max_tokens_per_batch // max_len) - - print(f"Running inference (batch_size={batch_size}, " - f"reduce_method={args.reduce_method}, prepend_bos={args.prepend_bos}, " - f"rc={args.average_reverse_complement})...") - scores = evo2_model.score_sequences( - sequences, - batch_size=batch_size, - prepend_bos=args.prepend_bos, - reduce_method=args.reduce_method, - average_reverse_complement=args.average_reverse_complement, - ) - scores = np.asarray(scores, dtype=float) - - # Write scores back onto the scored rows (NaN for skipped ones). - score_column = f"{args.model_name}_score" - dms_df[score_column] = np.nan - dms_df.loc[valid_mask, score_column] = scores - - # Spearman on the rows we actually scored. - scored = dms_df.loc[valid_mask, ["DMS_score", score_column]].dropna() - if len(scored) >= 2: - correlation, pvalue = spearmanr(scored["DMS_score"], scored[score_column]) - else: - correlation, pvalue = float("nan"), float("nan") - - dms_df.to_csv(output_file, index=False) - - print("\nSummary:") - print(f" DMS ID: {dms_id}") - print(f" Sequences scored: {len(sequences)}") - print(f" Score column: {score_column}") - print(f" Spearman vs DMS: {correlation:.3f} (p={pvalue:.2e})") - print(f" Saved to: {output_file}") - - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - raise + continue + todo.append((row_id, dms_id)) + if not todo: + print("Nothing to score.") + return + + if not torch.cuda.is_available(): + print("WARNING: CUDA not available - Evo 2 requires a GPU.", file=sys.stderr) + print(f"Visible GPUs: {torch.cuda.device_count()}") + + # Initialize model. Vortex handles device placement / multi-GPU sharding; + # do NOT call .to(device). The checkpoint is loaded once for every assay. + print(f"Loading Evo 2 model: {args.model_name} " + f"(local_path={args.local_path})...") + evo2_model = Evo2(args.model_name, local_path=args.local_path) + # Always ask the built model, never the packaged YAML: load_evo2_model can + # turn FP8 off for 7B when Transformer Engine is missing, and the batch-size + # arithmetic below has to follow the config the model was actually built with. + config = evo2_model.model.config + fp8_enabled = bool(config.get("use_fp8_input_projections", False)) + print(f"use_fp8_input_projections={fp8_enabled}") + if args.require_fp8 and not fp8_enabled: + raise SystemExit( + "--require_fp8 was given but the model resolved to " + "use_fp8_input_projections=False. For a 7B checkpoint this happens " + "silently when Transformer Engine is unavailable; for the others it " + "means Transformer Engine is not providing FP8. Refusing to score, because " + "the surrounding provenance would claim FP8.") + + failures = [] + for row_id, dms_id in todo: + print(f"\n=== row {row_id}: {dms_id} ===") + try: + score_one_assay(evo2_model, args, dms_id, fp8_enabled) + except Exception as e: + print(f"Error scoring {dms_id}: {str(e)}", file=sys.stderr) + failures.append(dms_id) + if len(todo) == 1: + raise + + if failures: + print(f"\nFAILED assays ({len(failures)}): {failures}", file=sys.stderr) + sys.exit(1) if __name__ == "__main__": diff --git a/fitness/merge_scoring_files.py b/fitness/merge_scoring_files.py index d7dac22..479037d 100755 --- a/fitness/merge_scoring_files.py +++ b/fitness/merge_scoring_files.py @@ -164,6 +164,14 @@ def four_fill_entries(name, folder, column_stem): "aido_rna_25m": {"folder": "aido_rna_25m_4fill", "column": "aido_rna_score_wt_fill"}, "aido_rna_300m": {"folder": "aido_rna_300m_4fill", "column": "aido_rna_score_wt_fill"}, "aido_rna_650m": {"folder": "aido_rna_650m_4fill", "column": "aido_rna_score_wt_fill"}, + # Two more Evo 2 checkpoints, completing the family alongside the existing + # evo2 (7B) and evo2_40b entries. Each writes its own {model_name}_score + # column, so unlike the AIDO series these differ by column as well as by + # prediction folder. evo2_1b_base is the base-pretrained 8k-context + # checkpoint, not a mid-trained 1B, so it is not a size point comparable to + # the others. + "evo2_1b_base": "evo2_1b_base_score", + "evo2_20b": "evo2_20b_score", "EVmutation": "prediction_epistatic", } @@ -207,6 +215,8 @@ def four_fill_entries(name, folder, column_stem): "aido_rna_25m", "aido_rna_300m", "aido_rna_650m", + "evo2_1b_base", + "evo2_20b", ] diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index 5e256ae..8fa1c2c 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -528,6 +528,8 @@ def main(args): "evo1", "evo1.5", "evo2", + "evo2_1b_base", + "evo2_20b", "evo2_40b", "GenSLM", "NT", From 17c53c47d13b6e5974338df2063cd77a0aede4f5 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:15:38 -0400 Subject: [PATCH 2/2] Add the Evo 2 1B and 20B checkpoints to the leaderboard The leaderboard carried the 7B and the 40B. Adding the 1B and the 20B completes the family, so the four sizes can be read together on the same 31 ncRNA assays. evo2_20b scores 0.2129, second behind aido_rna_650m and ahead of the 40B. evo2_1b_base scores 0.1962, which needs a caveat: its Ribozyme mean is 0.0087 and the macro weights the three categories equally, so a strong tRNA mean over only three assays carries it. It is also the only base-pretrained 8k-context checkpoint here, where the 7B, 20B and 40B are mid-trained at 1M, so it is not a size point comparable to the rest. Both were scored with the packaged FP8 configuration, which is what the Evo 2 authors require for every checkpoint except the 7B. The 7B and the 40B were rescored the same way as a check and are not published: they reproduce the existing columns, the 40B bitwise on all 31 assays and the 7B to the resolution of its own float32 serialisation, so publishing them would add rows that duplicate what is already there. That reproduction is what allows the existing 7B and 40B columns to be read alongside these two as one comparable family. The row count in test_every_leaderboard_row_is_a_registered_model goes from 16 to 18 with them. --- leaderboard/fitness/leaderboard_signed_3ncRNA.csv | 2 ++ tests/test_masked_lm.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv index 0326f70..8913c4e 100644 --- a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv +++ b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv @@ -1,7 +1,9 @@ model,Ribozyme,tRNA,Aptamer,macro_3ncRNA aido_rna_650m,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 +evo2_20b,0.1104601596335257,0.4345270738070895,0.0937093504359844,0.2128988612921999 evo2_40b,0.1081020224381099,0.4309848333293184,0.09695482733614616,0.21201389436785814 aido_rna,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 +evo2_1b_base,0.0086596951000284,0.4466199808506075,0.1334226383526211,0.1962341047677523 RNAErnie,0.13427241327809655,0.416115138737617,0.030641733377440288,0.1936764284643846 evo2,0.06511905079221095,0.3867345488695757,0.11920502912796216,0.1903528762632496 rnagenesis,0.0750044159987944,0.4380555746849195,0.0343114577905016,0.1824571494914051 diff --git a/tests/test_masked_lm.py b/tests/test_masked_lm.py index 457173c..b0ab4c5 100644 --- a/tests/test_masked_lm.py +++ b/tests/test_masked_lm.py @@ -694,7 +694,7 @@ def test_every_leaderboard_row_is_a_registered_model(): board = Path(__file__).resolve().parent.parent / "leaderboard" / "fitness" with open(board / "leaderboard_signed_3ncRNA.csv") as handle: rows = list(csv.DictReader(handle)) - assert len(rows) == 16 + assert len(rows) == 18 unregistered = [r["model"] for r in rows if r["model"] not in SCORE_COLS] assert not unregistered, f"on the leaderboard but not in SCORE_COLS: {unregistered}" macros = [float(r["macro_3ncRNA"]) for r in rows]