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", 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]