evo2 SAE recipe: streaming extract + train (7B layer-26)

Add the Evo2 SAE recipe under recipes/evo2: a streaming activation extractor
(extract.py -- reuses predict_evo2 but writes a parquet ActivationStore directly,
no intermediate .pt) and the SAE trainer (train.py), chained by 7b.sh which
reproduces the layer-26 7B normalize_input run.

train.py wires the merged sae training-quality fixes as opt-in flags
(--aggregate-loss / --dead-count-global / --mix-shards / --presample-shards),
all defaulting to the previous behavior so a baseline run is reproduced exactly.
Assumes a local Evo2 MBridge checkpoint (NGC pull / nemo2 convert is a documented
prerequisite, not recipe code). CPU tests cover arg defaults + SAE construction.

Supersedes #1579 and #1583.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Polina Binder <pbinder@nvidia.com>

Author: polinabinder1

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/README.md

@@ -0,0 +1,61 @@
+# Evo2 SAE recipe
+
+Train a Sparse Autoencoder (SAE) on [Evo2](https://github.com/ArcInstitute/evo2)
+residual-stream activations. The pipeline streams layer-L activations out of an Evo2
+model into a parquet `ActivationStore`, then trains a TopK SAE on them:
+
+```
+chunk FASTA  ->  stream-extract (extract.py)  ->  train (train.py)
+```
+
+`scripts/7b.sh` runs all three and reproduces our layer-26 7B (`normalize_input`) run.
+
+## Prerequisites
+
+This recipe assumes you already have the model weights and the runtime; it does **not**
+download or convert anything.
+
+- An Evo2 **MBridge checkpoint directory**. Obtain it from NGC, e.g.
+  `ngc registry model download-version "nvidia/clara/evo2:7b_<ver>" --dest <dir>`,
+  or convert a nemo2 checkpoint with the `evo2_megatron` converter.
+- `bionemo-recipes/recipes/evo2_megatron` built (`.ci_build.sh`) with its `.venv` active —
+  this provides `predict_evo2`, which `extract.py` drives.
+- The `sae` workspace package importable in that same venv.
+
+## Quickstart (7B, layer 26)
+
+```bash
+CKPT_DIR=<evo2-7b-mbridge-ckpt> FASTA=<prok_euk_sequences.fasta> bash scripts/7b.sh
+```
+
+## Steps (manual)
+
+1. **Chunk** sequences to the extraction context:
+   `python scripts/chunk_fasta.py --input seqs.fasta --output chunked.fasta --window 8192`
+2. **Extract** (streaming, no intermediate `.pt`):
+   `torchrun --nproc_per_node 8 scripts/extract.py --ckpt-dir <ckpt> --embedding-layer 26 \`
+   `--fasta chunked.fasta --activation-store-dir OUT --max-tokens 1000000000 --dtype fp32`
+3. **Train**:
+   `torchrun --nproc_per_node 8 scripts/train.py --cache-dir OUT --model-path <ckpt> --layer 26 \`
+   `--expansion-factor 16 --top-k 128 --normalize-input --auxk 2048 --init-pre-bias --dp-size 8 ...`
+
+`train.py` never loads the model — it reads only the activation cache; `--model-path`/`--layer`
+are used solely to validate the cache metadata.
+
+## Opt-in training-quality fixes
+
+These come from the `sae` package and **all default to the previous behavior**, so omitting
+them reproduces a baseline run exactly. Opt in when you want them (the 7B recipe does):
+
+| flag                   | effect                                                                      |
+| ---------------------- | --------------------------------------------------------------------------- |
+| `--aggregate-loss`     | batch-level FVU + AuxK loss instead of the per-token ratio                  |
+| `--dead-count-global`  | count dead-latent inactivity in total tokens (× world_size) under DDP       |
+| `--mix-shards N`       | shuffle + blend N shards per batch (replaces the old `--shards-per-buffer`) |
+| `--presample-shards N` | spread the pre-bias-init sample across N shards (needs `--init-pre-bias`)   |
+
+## Tests
+
+```bash
+pytest tests/        # CPU only: arg parsing + SAE construction (no GPU, no model)
+```

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "evo2-sae"
+version = "0.1.0"
+description = "Sparse Autoencoders for the Evo2 DNA language model"
+readme = "README.md"
+requires-python = ">=3.10"
+
+dependencies = [
+    "sae",
+    "torch>=2.0",
+    "numpy>=1.20",
+    "tqdm>=4.60",
+    "pyarrow>=10.0",
+]
+
+# No package code lives here yet — the recipe is just an entry-point for
+# scripts/ that depends on the shared `sae` workspace package. Declare no
+# packages so setuptools doesn't try to discover anything.
+[tool.setuptools]
+packages = []
+
+[tool.uv.sources]
+sae = { workspace = true }

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh

@@ -0,0 +1,95 @@
+#!/bin/bash
+# Evo2 7B layer-26 SAE recipe: chunk FASTA -> stream-extract activations -> train SAE.
+# This reproduces the layer26_7B (normalize_input) run.
+#
+# Prerequisites (this recipe does NOT download or convert the model):
+#   - An Evo2 7B *MBridge* checkpoint directory (CKPT_DIR). Obtain it from NGC, e.g.:
+#         ngc registry model download-version "nvidia/clara/evo2:7b_<ver>" --dest "${WORK_ROOT}/checkpoints"
+#     (or convert a nemo2 checkpoint to MBridge with the evo2_megatron converter).
+#   - bionemo-recipes/recipes/evo2_megatron built (.ci_build.sh) with its .venv active,
+#     providing `predict_evo2`.
+#   - The `sae` workspace package importable in that same venv.
+#
+# Override any of these by exporting before invocation.
+
+set -euo pipefail
+
+EVO2_MEGATRON_DIR="${EVO2_MEGATRON_DIR:-/workspace/bionemo-framework/bionemo-recipes/recipes/evo2_megatron}"
+RECIPE_DIR="$(cd "$(dirname "$0")/.." && pwd)"
+
+LAYER="${LAYER:-26}"
+# Context length the activations were extracted at (the model is context-extended; we
+# trained the SAE on 8192-bp chunks).
+CHUNK_BP="${CHUNK_BP:-8192}"
+
+# An Evo2 7B MBridge checkpoint directory (see prerequisites above).
+CKPT_DIR="${CKPT_DIR:?Set CKPT_DIR to an Evo2 7B MBridge checkpoint directory (see header)}"
+FASTA="${FASTA:?Set FASTA to the (prok+euk) input sequences}"
+WORK_ROOT="${WORK_ROOT:-/data/interp/evo2}"
+
+NPROC="${NPROC:-8}"            # GPUs / DP ranks
+MAX_TOKENS="${MAX_TOKENS:-1000000000}"
+
+PARQUET_DIR="${WORK_ROOT}/activations/evo2_7b_layer${LAYER}_parquet"
+OUTPUT_DIR="${WORK_ROOT}/sae/evo2_7b_layer${LAYER}"
+
+source "${EVO2_MEGATRON_DIR}/.venv/bin/activate"
+
+echo "============================================================"
+echo "STEP 0: Chunk FASTA to <=${CHUNK_BP} bp"
+echo "============================================================"
+INPUT_STEM="$(basename "$FASTA")"; INPUT_STEM="${INPUT_STEM%.gz}"; INPUT_STEM="${INPUT_STEM%.fasta}"
+CHUNKED_FASTA="${WORK_ROOT}/scratch/${INPUT_STEM}_chunked${CHUNK_BP}.fasta"
+if [[ -f "$CHUNKED_FASTA" ]]; then
+    echo "Reusing existing chunked FASTA: $CHUNKED_FASTA"
+else
+    python "${RECIPE_DIR}/scripts/chunk_fasta.py" --input "$FASTA" --output "$CHUNKED_FASTA" --window "$CHUNK_BP"
+fi
+
+echo "============================================================"
+echo "STEP 1: Stream-extract layer-${LAYER} activations -> parquet ActivationStore (no .pt)"
+echo "============================================================"
+if [[ -f "${PARQUET_DIR}/metadata.json" ]]; then
+    echo "Reusing existing parquet shards at $PARQUET_DIR"
+else
+    torchrun --nproc_per_node="$NPROC" "${RECIPE_DIR}/scripts/extract.py" \
+        --ckpt-dir "$CKPT_DIR" \
+        --embedding-layer "$LAYER" \
+        --fasta "$CHUNKED_FASTA" \
+        --activation-store-dir "$PARQUET_DIR" \
+        --max-tokens "$MAX_TOKENS" \
+        --micro-batch-size 4 \
+        --dtype fp32
+fi
+
+echo "============================================================"
+echo "STEP 2: Train TopK SAE (layer26_7B normalize_input config)"
+echo "============================================================"
+# unset a leaked key so ~/.netrc wins; clara-discovery is the wandb entity.
+unset WANDB_API_KEY || true
+export WANDB_ENTITY="${WANDB_ENTITY:-clara-discovery}"
+torchrun --nproc_per_node="$NPROC" "${RECIPE_DIR}/scripts/train.py" \
+    --cache-dir "$PARQUET_DIR" \
+    --model-path "$CKPT_DIR" \
+    --layer "$LAYER" \
+    --model-type topk \
+    --expansion-factor 16 --top-k 128 \
+    --normalize-input \
+    --auxk 2048 --auxk-coef 0.03125 \
+    --dead-tokens-threshold 10000000 \
+    --init-pre-bias \
+    --n-epochs 1 \
+    --batch-size 1024 \
+    --lr 1e-4 --lr-schedule cosine --lr-min 1e-5 --warmup-steps 1000 \
+    --max-grad-norm 1.0 \
+    --mix-shards 10 \
+    --dp-size "$NPROC" \
+    --log-interval 100 \
+    --wandb --wandb-project evo2-sae-v2-diverse --wandb-run-name "layer${LAYER}_7B_normalize_input" \
+    --output-dir "$OUTPUT_DIR" \
+    --checkpoint-dir "${OUTPUT_DIR}/checkpoints" \
+    --checkpoint-steps 2000
+
+echo "============================================================"
+echo "DONE: SAE checkpoint at ${OUTPUT_DIR}/checkpoints/checkpoint_final.pt"
+echo "============================================================"

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/scripts/chunk_fasta.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Chunk a FASTA into <=N-bp windows so predict_evo2 stays inside the model's trained context.
+
+Evo2 1B was trained with seq_length=8192; longer inputs OOM in the Hyena
+fftconv path (intermediates scale super-linearly with L). For 7B/40B raise
+--window to whatever those checkpoints were context-extended to.
+
+Non-overlapping windows by default. Each chunk gets a header of the form
+">{orig_id}:{start}-{end}" so downstream parquet can be back-mapped.
+"""
+
+import argparse
+import gzip
+from pathlib import Path
+
+
+def parse_fasta(path: Path):
+    """Yield (seq_id, sequence) tuples from a FASTA file (transparently handles .gz)."""
+    opener = gzip.open if path.suffix == ".gz" else open
+    seq_id, parts = None, []
+    with opener(path, "rt") as f:
+        for line in f:
+            line = line.rstrip()
+            if line.startswith(">"):
+                if seq_id is not None:
+                    yield seq_id, "".join(parts)
+                seq_id = line[1:].split()[0]
+                parts = []
+            else:
+                parts.append(line)
+        if seq_id is not None:
+            yield seq_id, "".join(parts)
+
+
+def main():
+    """Read input FASTA, write non-overlapping <=window-bp chunks to output FASTA."""
+    p = argparse.ArgumentParser()
+    p.add_argument("--input", type=Path, required=True)
+    p.add_argument("--output", type=Path, required=True)
+    p.add_argument("--window", type=int, default=8192)
+    args = p.parse_args()
+
+    n_in = n_out = bp_out = 0
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with open(args.output, "w") as out:
+        for seq_id, seq in parse_fasta(args.input):
+            n_in += 1
+            for start in range(0, len(seq), args.window):
+                end = min(start + args.window, len(seq))
+                chunk = seq[start:end]
+                out.write(f">{seq_id}:{start}-{end}\n{chunk}\n")
+                n_out += 1
+                bp_out += len(chunk)
+
+    print(f"Chunked {n_in} sequences -> {n_out} chunks ({bp_out:,} bp) at window={args.window}")
+
+
+if __name__ == "__main__":
+    main()