diff --git a/.github/workflows/unit-tests-interpretability-recipes.yaml b/.github/workflows/unit-tests-interpretability-recipes.yaml new file mode 100644 index 0000000000..a171934c1e --- /dev/null +++ b/.github/workflows/unit-tests-interpretability-recipes.yaml @@ -0,0 +1,146 @@ +name: "BioNeMo Interpretability Recipes CI" + +# CI for the SAE interpretability recipes under interpretability/sparse_autoencoders/recipes/*. +# Generic matrix, modeled on the repo-wide unit-tests-recipes.yml but scoped to the interp subtree: +# * `changed-dirs` (cheap ubuntu) discovers which interp recipes to test and emits a matrix. +# * `unit-tests` runs each selected recipe on the L4: its own `.ci_build.sh` + `pytest tests/`. +# +# Eligible recipes = any interp recipe with its own `.ci_build.sh`. Today that's `evo2`; codonfm/esm2 +# have no `.ci_build.sh`/tests yet, so they are skipped until they add them. This also makes the lane +# a green no-op before the evo2 SAE recipe lands (#1622) — the presence guard is "has a .ci_build.sh". +# +# What runs when: +# * change under a recipe's own dir -> that recipe. +# * change to the shared `sae` lib, or to this workflow file, or the nightly schedule +# -> ALL eligible recipes (they all depend on sae). +# * change to an unrelated recipe / elsewhere -> nothing (empty matrix, green no-op). +# Each recipe's `.ci_build.sh` owns its own build (evo2 -> mbridge bionemo.evo2; esm2 -> HF; etc.), +# so a codonfm/esm2 change never triggers the Evo2 megatron build, and vice-versa. + +on: + push: + branches: + - "pull-request/[0-9]+" + - "dependabot/**" + merge_group: + types: [checks_requested] + schedule: + - cron: "0 9 * * *" # Runs at 9 AM UTC daily (2 AM MST) + +defaults: + run: + shell: bash -x -e -u -o pipefail {0} + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + changed-dirs: + # Cheap ubuntu pre-job: decide which interp recipes to test and emit the matrix. + runs-on: ubuntu-latest + permissions: + contents: read + outputs: + dirs: ${{ steps.set-dirs.outputs.dirs }} + steps: + - name: Checkout repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Determine which interp recipes to run + id: set-dirs + env: + EVENT_NAME: ${{ github.event_name }} + run: | + RROOT="interpretability/sparse_autoencoders/recipes" + SAE="interpretability/sparse_autoencoders/sae" + WF=".github/workflows/unit-tests-interpretability-recipes.yaml" + + # Eligible recipes = those with their own .ci_build.sh (others are green no-ops until + # they add one; empty before #1622 lands -> whole lane is a no-op). + ALL=$(for d in "$RROOT"/*/; do [ -f "${d}.ci_build.sh" ] && echo "${d%/}"; done \ + | jq -R -s -c 'split("\n") | map(select(length > 0))') + echo "Eligible recipes (have .ci_build.sh): $ALL" + + MERGE_BASE=$(git merge-base HEAD origin/main) + CHANGED=$(git diff --name-only "$MERGE_BASE" HEAD) + echo "Changed files:"; echo "$CHANGED" | sed 's/^/ - /' + + # sae lib / this workflow / nightly -> run EVERY eligible recipe (all depend on sae). + if [ "$EVENT_NAME" = "schedule" ] || echo "$CHANGED" | grep -qE "^(${SAE}/|${WF}$)"; then + DIRS="$ALL" + else + # otherwise -> only eligible recipes that have a changed file under them. + DIRS=$(echo "$ALL" | jq -c --arg changed "$CHANGED" ' + ($changed | split("\n")) as $cf + | map(select(. as $d | $cf | any(startswith($d + "/")))) + ') + fi + echo "Recipes to run: $DIRS" + + # Attach the runner image (the dockerhub-cached squashed pytorch, like the repo-wide lane). + IMG="svcbionemo023/bionemo-framework:pytorch26.04-py3-squashed" + MATRIX=$(echo "$DIRS" | jq -c --arg img "$IMG" \ + 'map({dir: ., name: (. | split("/") | last), image: $img})') + echo "dirs=$MATRIX" >> "$GITHUB_OUTPUT" + echo "matrix: $MATRIX" + + unit-tests: + needs: changed-dirs + if: ${{ needs.changed-dirs.outputs.dirs != '' && needs.changed-dirs.outputs.dirs != '[]' }} + runs-on: linux-amd64-gpu-l4-latest-1 + name: "interp-unit-tests (${{ matrix.recipe.name }})" + permissions: + contents: read + container: + image: ${{ matrix.recipe.image }} + options: --shm-size=16G + env: + CI: true + HF_TOKEN: ${{ secrets.HF_TOKEN }} + HF_HOME: /cache/huggingface + strategy: + fail-fast: false + matrix: + recipe: ${{ fromJson(needs.changed-dirs.outputs.dirs) }} + + steps: + - name: Show GPU info + run: nvidia-smi + + - name: Setup proxy cache + uses: nv-gha-runners/setup-proxy-cache@14229018fe157c83e03c008f27d183d8e99bc67c + + - name: Checkout repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 + with: + # Interp recipes live under interpretability/sparse_autoencoders (recipe + shared sae); + # evo2 additionally builds on recipes/evo2_megatron. Check out both so any recipe's + # .ci_build.sh has what it needs. + sparse-checkout: | + interpretability/sparse_autoencoders + recipes/evo2_megatron + sparse-checkout-cone-mode: false + persist-credentials: false + + - name: Cache Hugging Face models + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 + with: + path: /cache/huggingface + key: ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}-${{ github.sha }} + restore-keys: | + ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}- + ${{ runner.os }}-huggingface- + + - name: Install dependencies + working-directory: ${{ matrix.recipe.dir }} + run: bash .ci_build.sh + + - name: Run tests + working-directory: ${{ matrix.recipe.dir }} + run: | + [ -f .ci_test_env.sh ] && source .ci_test_env.sh + pytest -v tests/ diff --git a/interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh b/interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh new file mode 100755 index 0000000000..c00387a739 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Build the Evo2-SAE recipe environment for CI. +# +# The inference engine imports `bionemo.evo2` — which lives ONLY in the mbridge +# `recipes/evo2_megatron` recipe (there is no prebaked image that ships it). So we build +# that environment with evo2_megatron's OWN CI build script, inheriting its pinned +# megatron-bridge / causal-conv1d / TransformerEngine handling verbatim (no fork that +# could drift), then install the SAE library + this recipe on top of the same venv. +# +# The CI checkout must also provide the sibling recipe (recipes/evo2_megatron); it pulls its +# own bionemo-core / bionemo-recipeutils deps from git, so nothing else is required locally. +# +# Phases (arg, default "all") — let the Dockerfile cache the expensive megatron build separately +# from the cheap, frequently-changing SAE installs (so a code edit doesn't rebuild megatron): +# env build the bionemo.evo2 (mbridge) env only (depends on recipes/evo2_megatron) +# install install the SAE library + this recipe only (depends on the SAE source) +# all both (CI / default) +set -euo pipefail + +phase="${1:-all}" + +HERE="$(cd "$(dirname "$0")" && pwd)" +SAE_LIB="$HERE/../../sae" # sparse_autoencoders/sae +MEGATRON="$HERE/../../../../recipes/evo2_megatron" + +# Fail loudly if the sibling recipe isn't checked out, rather than emitting a cryptic path +# error deep inside the build. In CI this means the workflow's sparse-checkout MUST include +# recipes/evo2_megatron. +if [[ ! -f "$MEGATRON/.ci_build.sh" ]]; then + echo "ERROR: evo2_megatron not found at $MEGATRON" >&2 + echo " evo2-sae builds on the mbridge bionemo.evo2 recipe; ensure the checkout" >&2 + echo " includes recipes/evo2_megatron." >&2 + exit 1 +fi + +# 1. Build the bionemo.evo2 (mbridge) environment. Creates $MEGATRON/.venv with the +# system-site torch/TE plus the full megatron stack. +if [[ "$phase" == "env" || "$phase" == "all" ]]; then + ( cd "$MEGATRON" && bash .ci_build.sh ) +fi + +# 2. Add the generic SAE library + this recipe into that venv. PIP_CONSTRAINT= clears the +# TE pin constraint evo2_megatron sets (it must not block our pure-Python installs). +if [[ "$phase" == "install" || "$phase" == "all" ]]; then + source "$MEGATRON/.venv/bin/activate" + PIP_CONSTRAINT= pip install -e "$SAE_LIB" + PIP_CONSTRAINT= pip install -e "$HERE" +fi diff --git a/interpretability/sparse_autoencoders/recipes/evo2/.ci_test_env.sh b/interpretability/sparse_autoencoders/recipes/evo2/.ci_test_env.sh new file mode 100755 index 0000000000..ec7f92e487 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/.ci_test_env.sh @@ -0,0 +1,4 @@ +# Sourced by CI before running pytest (mirrors evo2_megatron/.ci_test_env.sh). +# The recipe shares evo2_megatron's venv (see .ci_build.sh), which lives in the sibling +# recipe dir rather than here. +source ../../../../recipes/evo2_megatron/.venv/bin/activate diff --git a/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile b/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile new file mode 100644 index 0000000000..111e42588f --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile @@ -0,0 +1,36 @@ +# syntax=docker/dockerfile:1.4 +# +# Runnable image for the Evo2 SAE recipe. Thin and non-forking: the engine imports bionemo.evo2 +# (the mbridge evo2_megatron recipe), so .ci_build.sh delegates to evo2_megatron's OWN build +# (inheriting its pinned megatron-bridge / causal-conv1d / TransformerEngine versions), then +# installs the sae library + this recipe on top. We don't reimplement the megatron build here. +# +# Build from the REPO ROOT (the context must include the recipes/evo2_megatron sibling): +# docker build -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae . +# Run (needs a GPU + checkpoints; see the recipe README): +# docker run --gpus all -it evo2-sae bash -lc "source .ci_test_env.sh && pytest tests/" +ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3 +FROM ${BASE_IMAGE} + +WORKDIR /workspace +# Two layers so editing SAE code doesn't rebuild the ~30-min megatron stack: +# +# 1. Expensive, rarely-changing: build the bionemo.evo2 (mbridge) env. Depends ONLY on +# recipes/evo2_megatron (which pulls its bionemo-core/recipeutils deps from git), so this layer +# stays cached across SAE code edits. Copy just the recipe's .ci_build.sh to drive the build. +COPY recipes/evo2_megatron recipes/evo2_megatron +COPY interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh \ + interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh +WORKDIR /workspace/interpretability/sparse_autoencoders/recipes/evo2 +RUN bash .ci_build.sh env + +# 2. Cheap, frequently-changing: add the SAE library + this recipe, install editable. A code +# change invalidates only this layer (the two pip installs, seconds), not the build above. +WORKDIR /workspace +COPY interpretability/sparse_autoencoders interpretability/sparse_autoencoders +WORKDIR /workspace/interpretability/sparse_autoencoders/recipes/evo2 +RUN bash .ci_build.sh install + +# Default to the built venv so `docker run … pytest tests/` (or importing Evo2SAE) just works. +ENV VIRTUAL_ENV=/workspace/recipes/evo2_megatron/.venv +ENV PATH="$VIRTUAL_ENV/bin:$PATH" diff --git a/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile.dockerignore b/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile.dockerignore new file mode 100644 index 0000000000..f6b7eb2da4 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/Dockerfile.dockerignore @@ -0,0 +1,11 @@ +# Per-Dockerfile ignore (BuildKit uses .dockerignore when present). Paths are relative +# to the build context (the repo root). Keep generated/heavy trees out of the context so they +# neither bloat the upload nor bust the COPY cache when they change locally. +**/node_modules/ +**/dist/ +**/.vite/ +**/__pycache__/ +**/*.pyc +**/.venv/ +**/.pytest_cache/ +.git/ diff --git a/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml b/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml index f7b23d8eff..2862d8a59d 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml +++ b/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml @@ -12,14 +12,17 @@ dependencies = [ "sae", "torch>=2.0", "numpy>=1.20", + "pandas>=2.0", # cli.py `batch` -> parquet "pyarrow>=23.0.0", + "fastapi>=0.110", # server.py + "uvicorn>=0.29", # serve CLI + "anyio>=4.0", # server.py: threadpool concurrency cap (ships with fastapi, declared explicitly) ] -# 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 = [] +# The `evo2_sae` package (src/) holds the live inference engine + steering hook + the FastAPI +# server and CLI; scripts/ (extract, train) are standalone entry points alongside it. +[tool.setuptools.packages.find] +where = ["src"] [tool.uv.sources] sae = { workspace = true } diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh b/interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh index 88c0fc6451..436bac7a55 100755 --- a/interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/7b.sh @@ -25,7 +25,7 @@ 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}" +WORK_ROOT="${WORK_ROOT:?Set WORK_ROOT to a (large, multi-TB) output dir for activations + SAE checkpoints}" NPROC="${NPROC:-8}" # GPUs / DP ranks MAX_TOKENS="${MAX_TOKENS:-1000000000}" @@ -72,7 +72,6 @@ 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 \ diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh b/interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh new file mode 100755 index 0000000000..dc9476b41b --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# Launch the Evo2 SAE inference engine. One engine, four modes: +# +# ./launch_inference.sh serve # live HTTP server on :8001 (viz backend) +# ./launch_inference.sh encode --sequence ATGC... # annotate ONE sequence -> top features +# ./launch_inference.sh batch --fasta in.fa --out out.parquet # MANY sequences -> parquet +# ./launch_inference.sh generate --prompt ATGC... --clamp 29244:300 # steer + generate DNA +# +# Steering loop: `encode` a sequence to find an active feature id, then +# `generate --clamp ID:STRENGTH` (strength ~2-3x the feature's max_activation; repeat --clamp). +# +# Config via env. Required: EVO2_CKPT_DIR, SAE_CKPT_PATH. Optional (have defaults): +# FEATURE_ANNOTATIONS, EMBEDDING_LAYER (26), DEVICE, PORT, CUDA_VISIBLE_DEVICES. +# +# Run INSIDE the evo2_megatron venv (it provides bionemo.evo2 + megatron) — same venv the tests +# use. In the Docker image it's already active (on PATH), so just run this. On bare metal, activate +# it first: source /.venv/bin/activate +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +RECIPE_DIR="$(cd "$HERE/.." && pwd)" # recipes/evo2 — so the evo2_sae package imports + +# Required (no hardcoded defaults — supply your own paths via env): +export EVO2_CKPT_DIR="${EVO2_CKPT_DIR:?Set EVO2_CKPT_DIR to an Evo2 MBridge checkpoint directory}" +export SAE_CKPT_PATH="${SAE_CKPT_PATH:?Set SAE_CKPT_PATH to a trained SAE checkpoint (.pt)}" +# Optional: feature-label parquet (empty = features are unlabeled). Layer defaults to 26. +export FEATURE_ANNOTATIONS="${FEATURE_ANNOTATIONS:-}" +export EMBEDDING_LAYER="${EMBEDDING_LAYER:-26}" + +export PYTHONPATH="$RECIPE_DIR/src${PYTHONPATH:+:$PYTHONPATH}" # find evo2_sae without a pip install + +# We assume the evo2_megatron venv is already active (Docker: on PATH; bare metal: you sourced it). +# Fail with a clear message rather than a deep ImportError if it isn't. +if ! python -c "import bionemo.evo2" 2>/dev/null; then + echo "ERROR: bionemo.evo2 not importable — activate the evo2_megatron venv first" >&2 + echo " (source /.venv/bin/activate) or run inside the Docker image." >&2 + exit 1 +fi + +# One-shot modes (encode/batch/generate) just run once -> exec, signals pass straight through. +# For `serve`, a bad request can trip a CUDA device-side assert that poisons the process context +# (unrecoverable in-process — restart is the only cure). So tell the engine to exit the worker on +# that fault (EXIT_ON_CUDA_WEDGE) and respawn it here, making recovery host-independent. +# +# The worker runs in the BACKGROUND and we `wait` on it so this supervisor stays responsive to +# signals: a trap forwards SIGTERM/SIGINT to the worker (triggering uvicorn's graceful shutdown) +# and stops the respawn loop — important when this script is PID 1 under `docker stop`/k8s, where +# bash would otherwise not forward the signal and orphan the worker. Crash/wedge -> respawn with +# backoff, capped so a persistent failure (e.g. port already bound) doesn't loop forever. +if [[ "${1:-}" == "serve" ]]; then + export EXIT_ON_CUDA_WEDGE=1 + child="" + stop=0 + trap 'stop=1; [[ -n "$child" ]] && kill -TERM "$child" 2>/dev/null || true' TERM INT + rc=0 + fails=0 + while [[ "$stop" -eq 0 ]]; do + python -m evo2_sae.cli "$@" & + child=$! + wait "$child" && rc=0 || rc=$? + child="" + # stop on a clean exit or when a signal asked us to (143 SIGTERM / 130 SIGINT belt-and-suspenders). + [[ "$stop" -eq 1 || "$rc" -eq 0 || "$rc" -eq 143 || "$rc" -eq 130 ]] && break + fails=$((fails + 1)) + if [[ "$fails" -ge 10 ]]; then + echo "[launch] serve exited ($rc) $fails times — giving up; fix the cause." >&2 + break + fi + backoff=$([[ "$fails" -lt 5 ]] && echo 2 || echo 10) + echo "[launch] serve exited ($rc); restart $fails/10 in ${backoff}s…" >&2 + sleep "$backoff" & + wait $! 2>/dev/null || true # interruptible: a signal during backoff stops the loop promptly + done + exit "$rc" +fi +exec python -m evo2_sae.cli "$@" diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/train.py b/interpretability/sparse_autoencoders/recipes/evo2/scripts/train.py index 7edcd2ef6e..c17f539468 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/scripts/train.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/train.py @@ -50,7 +50,7 @@ import torch import torch.distributed as dist from sae.activation_store import load_activations -from sae.architectures import ReLUSAE, TopKSAE +from sae.architectures import TopKSAE from sae.perf_logger import PerfLogger from sae.training import ParallelConfig, Trainer, TrainingConfig, WandbConfig from sae.utils import get_device, set_seed @@ -69,7 +69,6 @@ def parse_args(): # noqa: D103 # SAE architecture sae_group = p.add_argument_group("SAE model") - sae_group.add_argument("--model-type", type=str, default="topk", choices=["topk", "relu"]) sae_group.add_argument("--expansion-factor", type=int, default=8) sae_group.add_argument("--top-k", type=int, default=32) sae_group.add_argument("--normalize-input", action=argparse.BooleanOptionalAction, default=False) @@ -77,7 +76,6 @@ def parse_args(): # noqa: D103 sae_group.add_argument("--auxk-coef", type=float, default=1 / 32) sae_group.add_argument("--dead-tokens-threshold", type=int, default=10_000_000) sae_group.add_argument("--init-pre-bias", action=argparse.BooleanOptionalAction, default=False) - sae_group.add_argument("--l1-coeff", type=float, default=1e-2, help="L1 coefficient (relu only)") # Opt-in training-quality fixes (sae package). Defaults reproduce previous behavior. sae_group.add_argument( "--aggregate-loss", @@ -164,27 +162,17 @@ def parse_args(): # noqa: D103 def build_sae(args, input_dim: int) -> torch.nn.Module: # noqa: D103 hidden_dim = input_dim * args.expansion_factor - - if args.model_type == "topk": - return TopKSAE( - input_dim=input_dim, - hidden_dim=hidden_dim, - top_k=args.top_k, - normalize_input=args.normalize_input, - auxk=args.auxk, - auxk_coef=args.auxk_coef, - dead_tokens_threshold=args.dead_tokens_threshold, - aggregate_loss=args.aggregate_loss, - dead_count_global=args.dead_count_global, - ) - elif args.model_type == "relu": - return ReLUSAE( - input_dim=input_dim, - hidden_dim=hidden_dim, - l1_coeff=args.l1_coeff, - ) - else: - raise ValueError(f"Unknown model type: {args.model_type}") + return TopKSAE( + input_dim=input_dim, + hidden_dim=hidden_dim, + top_k=args.top_k, + normalize_input=args.normalize_input, + auxk=args.auxk, + auxk_coef=args.auxk_coef, + dead_tokens_threshold=args.dead_tokens_threshold, + aggregate_loss=args.aggregate_loss, + dead_count_global=args.dead_count_global, + ) def build_training_config(args, device: str) -> TrainingConfig: # noqa: D103 @@ -267,7 +255,7 @@ def main(): # noqa: D103 input_dim = meta["hidden_dim"] sae = build_sae(args, input_dim) - print(f"SAE: {args.model_type}, input_dim={input_dim}, hidden_dim={sae.hidden_dim}") + print(f"SAE: TopK, input_dim={input_dim}, hidden_dim={sae.hidden_dim}") # Initialize pre_bias from the geometric median of a sample of activations. With # --presample-shards N>1, draw the sample across N shards spanning the store (avoids diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/__init__.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/__init__.py new file mode 100644 index 0000000000..7664577534 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/__init__.py @@ -0,0 +1,39 @@ +# 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. + +"""Evo2 + SAE inference engine — reused by the live server, the batch CLI, and the viz backend.""" + +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: # for type checkers / ruff F822 — runtime access goes through __getattr__ below + from .core import DEFAULT_ORGANISM_TAGS, Evo2SAE, clean_dna + + +__all__ = ["DEFAULT_ORGANISM_TAGS", "Evo2SAE", "clean_dna"] + + +def __getattr__(name: str): + """Lazily pull the heavy engine symbols from ``.core`` (importing ``.core`` loads torch). + + Keeps ``import evo2_sae`` (and lightweight submodules like ``evo2_sae.fasta``) cheap so + callers that only need the helpers don't drag in torch, while ``from evo2_sae import + Evo2SAE`` still works. + """ + if name in __all__: + from . import core + + return getattr(core, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py new file mode 100644 index 0000000000..a164ec33ea --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py @@ -0,0 +1,192 @@ +# 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. + +"""Evo2 SAE inference CLI — one engine, four modes. + + serve : start the FastAPI server (one sequence at a time, interactive) + encode : annotate ONE sequence -> top features (stdout JSON) + batch : run a FASTA of MANY sequences -> parquet of per-sequence top features + generate: generate DNA, optionally steering SAE features (stdout JSON) + +They all build the same `Evo2SAE` engine; config comes from flags or env +(EVO2_CKPT_DIR / SAE_CKPT_PATH / FEATURE_ANNOTATIONS / EMBEDDING_LAYER). +""" + +from __future__ import annotations + +import argparse +import json +import os + + +def _add_common(p: argparse.ArgumentParser) -> None: + """Register the shared inference arguments (checkpoints, layer, device) on a parser. + + Defaults come from env vars (``EVO2_CKPT_DIR``, ``SAE_CKPT_PATH``, ``FEATURE_ANNOTATIONS``, + ``EMBEDDING_LAYER``, ``DEVICE``, ``MAX_SEQ_LEN``); pass the flags to override. No hardcoded + paths — the checkpoints must be supplied via flag or env. + + Args: + p: The argparse parser (or subparser) to add the shared arguments to. + + Returns: + None. Mutates ``p`` in place. + """ + p.add_argument("--evo2-ckpt-dir", default=os.environ.get("EVO2_CKPT_DIR")) + p.add_argument("--sae-ckpt-path", default=os.environ.get("SAE_CKPT_PATH")) + p.add_argument("--feature-annotations", default=os.environ.get("FEATURE_ANNOTATIONS")) + # int() the env defaults explicitly: argparse's type= only coerces values passed on the command + # line, never the default — so an env-sourced (or absent) value would otherwise stay a str. + p.add_argument("--layer", type=int, default=int(os.environ.get("EMBEDDING_LAYER", "26"))) + p.add_argument("--device", default=os.environ.get("DEVICE", "cuda")) + p.add_argument("--max-seq-len", type=int, default=int(os.environ.get("MAX_SEQ_LEN", "8192"))) + + +def _engine(args): + """Construct an Evo2SAE engine from parsed CLI args. + + Args: + args: Parsed argparse namespace with ``evo2_ckpt_dir``, ``sae_ckpt_path``, ``layer``, + ``device``, ``max_seq_len``, ``feature_annotations``. + + Returns: + An (unloaded) ``Evo2SAE`` instance — call ``.load()`` before use. + """ + from .core import Evo2SAE + + return Evo2SAE( + evo2_ckpt_dir=args.evo2_ckpt_dir, + sae_ckpt_path=args.sae_ckpt_path, + layer=args.layer, + device=args.device, + max_seq_len=args.max_seq_len, + feature_annotations=args.feature_annotations, + ) + + +def main(): + """Parse args and dispatch to the serve / encode / batch subcommand.""" + ap = argparse.ArgumentParser(description="Evo2 SAE inference (serve | encode | batch | generate)") + sub = ap.add_subparsers(dest="cmd", required=True) + + ps = sub.add_parser("serve", help="start the FastAPI inference server") + _add_common(ps) + ps.add_argument("--host", default="0.0.0.0") + ps.add_argument( + "--port", type=int, default=int(os.environ.get("PORT", "8001")) + ) # int: uvicorn.run needs an int port + + pe = sub.add_parser("encode", help="annotate ONE sequence -> top features (JSON)") + _add_common(pe) + pe.add_argument("--sequence", required=True) + pe.add_argument("--organism", default="None (raw DNA)") + pe.add_argument("--top-k", type=int, default=8) + + pb = sub.add_parser("batch", help="MANY sequences (FASTA) -> parquet of per-sequence top features") + _add_common(pb) + pb.add_argument("--fasta", required=True) + pb.add_argument("--out", required=True) + pb.add_argument("--top-k", type=int, default=16) + pb.add_argument("--batch-size", type=int, default=8) + + pg = sub.add_parser("generate", help="generate DNA, optionally steering SAE features") + _add_common(pg) + pg.add_argument("--prompt", default="", help="DNA to seed; steering applies to the continuation") + pg.add_argument("--organism", default="None (raw DNA)") + pg.add_argument( + "--clamp", + action="append", + default=[], + metavar="FEATURE_ID[:STRENGTH]", + help="clamp a feature on the continuation; repeatable (e.g. --clamp 29244:300). " + "Find feature ids with `encode`.", + ) + pg.add_argument("--n-tokens", type=int, default=120) + pg.add_argument("--temperature", type=float, default=1.0) + pg.add_argument("--top-k", type=int, default=0) + pg.add_argument("--compare-baseline", action="store_true", help="also generate unsteered, for comparison") + + args = ap.parse_args() + + if args.cmd == "serve": + import uvicorn + + from .server import build_app + + uvicorn.run(build_app(_engine(args)), host=args.host, port=args.port, log_level="info") + return + + from . import core + + eng = _engine(args).load() + + if args.cmd == "encode": + try: + dna, _tag, codes, tag_len = core.annotate(eng, args.sequence, args.organism) + except ValueError as e: + raise SystemExit(str(e)) + feats = eng.top_features(codes, tag_len=tag_len, k=args.top_k) + print( + json.dumps( + {"sequence": dna, "organism": args.organism, "bases": len(dna), "top_features": feats}, indent=2 + ) + ) + + elif args.cmd == "batch": + import pandas as pd + + from .fasta import read_fasta + + ids, seqs = [], [] + for sid, seq in read_fasta(args.fasta): + ids.append(sid) + seqs.append(seq) + print(f"[batch] {len(seqs)} sequences from {args.fasta}; encoding (batch_size={args.batch_size})…") + codes_list = eng.encode_batch(seqs, batch_size=args.batch_size) + rows = [] + for sid, codes in zip(ids, codes_list): + for rank, ft in enumerate(eng.top_features(codes, k=args.top_k)): + rows.append({"sequence_id": sid, "bp": int(codes.shape[0]), "rank": rank, **ft}) + df = pd.DataFrame(rows) + df.to_parquet(args.out, index=False) + print(f"[batch] wrote {len(df)} rows for {len(seqs)} sequences -> {args.out}") + + elif args.cmd == "generate": + try: + out = eng.generate( + prompt=args.prompt, + organism=args.organism, + features=args.clamp, # raw "ID[:STRENGTH]" strings; core.parse_clamp_spec normalizes + n_tokens=args.n_tokens, + temperature=args.temperature, + top_k=args.top_k, + compare_baseline=args.compare_baseline, + ) + except ValueError as e: + raise SystemExit(str(e)) + result = { + "prompt": out["prompt"], + "organism": out["organism"], + "steered": out["steered"], + "features": out["features"], + "sequence": out["generation"]["sequence"], + } + if out.get("baseline"): + result["baseline_sequence"] = out["baseline"]["sequence"] + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/core.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/core.py new file mode 100644 index 0000000000..df5a16d5b4 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/core.py @@ -0,0 +1,557 @@ +# 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. + +"""Evo2 + SAE inference core — one importable engine for live and batch use. + +`Evo2SAE` loads a base Evo2 model and a trained SAE once, then exposes: + + encode(dna) -> codes [S, n_features] # ONE sequence (interactive) + encode_batch(seqs) -> list of codes [S_i, n_features] # MANY sequences (batched on GPU) + feature_tracks(dna, f) -> {feature_id: [per-base activation]} + generate(...) -> autoregressive DNA generation with optional additive + SAE-feature clamping on the generated continuation + +It has NO web dependency: the FastAPI server (`server.py`) and the batch CLI +(`cli.py`) are thin wrappers over this class, and the viz backend imports it too. + +The heavy Evo2 machinery is reused from the recipe: a SINGLE inference engine +(`infer.setup_inference_engine`, run eager with `cuda_graph_impl="none"` so the +residual-stream steering hook applies) serves both paths. Encode/highlight runs a normal +full-sequence forward on the engine's model and reads layer `layer` off a forward hook; +generation uses the same model via `infer.generate`. This module only adds the SAE layer: +encode, feature labels, and the decode-only feature-clamp hook. +""" + +from __future__ import annotations + +import logging +import math +import os +import re +import sys +import threading +from pathlib import Path +from typing import Optional + +import torch + + +logger = logging.getLogger("evo2_sae_infer") + +# Disable Inductor CUDA graphs before torch initializes inductor — graph capture +# conflicts with the residual-stream forward hook (which replaces the layer output) +# and with re-feeding a growing sequence each decode step. +os.environ.setdefault("TORCHINDUCTOR_CUDAGRAPHS", "0") + +# Make the local `sae` package importable (sparse_autoencoders/sae/src). +_SAE_SRC = os.environ.get("SAE_SRC", str(Path(__file__).resolve().parents[4] / "sae" / "src")) +if _SAE_SRC not in sys.path: + sys.path.insert(0, _SAE_SRC) + +_VALID_BASES = re.compile(r"[^ACGTN]") + +# Steering clamp targets are absolute SAE-code values; this SAE's features peak ~100-300. +# Cap the magnitude so an extreme target can't blow the logits to NaN (which device-asserts +# and wedges the process). Generous headroom for amplification, bounded against runaway. +MAX_CLAMP_STRENGTH = float(os.environ.get("MAX_CLAMP_STRENGTH", "300")) + +# Phylogenetic-tag prefixes per organism (Evo2 was trained with lineage tags). +DEFAULT_ORGANISM_TAGS = { + "None (raw DNA)": "", + "Human": "|d__Eukaryota;p__Chordata;c__Mammalia;o__Primates;f__Hominidae;g__Homo;s__Homo sapiens|", + "E. coli": "|d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Escherichia;s__Escherichia coli|", + "S. cerevisiae": "|d__Eukaryota;p__Ascomycota;c__Saccharomycetes;o__Saccharomycetales;f__Saccharomycetaceae;g__Saccharomyces;s__Saccharomyces cerevisiae|", +} + + +def clean_dna(seq: str) -> str: + """Uppercase and strip everything that isn't a nucleotide.""" + return _VALID_BASES.sub("", (seq or "").upper()) + + +def annotate(engine, sequence: str, organism: str = "None (raw DNA)", tag: Optional[str] = None): + """Shared encode path for the CLI ``encode`` and the server ``/annotate`` (topk). + + Cleans the sequence, resolves the phylo tag, encodes once, and computes the tag length + (the leading tag tokens to skip, ignored if it would drop the whole sequence). Returns + ``(dna, resolved_tag, codes, tag_len)``; callers add their own top-k / per-base presentation. + + Raises ``ValueError`` on an empty/non-DNA sequence or an unknown organism — the server maps + these to HTTP 400, the CLI to a clean exit. Takes the engine as an argument (rather than + living on ``Evo2SAE``) so it composes with any object exposing ``resolve_tag``/``encode``. + """ + dna = clean_dna(sequence) + if not dna: + raise ValueError("No valid nucleotides in sequence") + resolved_tag = engine.resolve_tag(organism, tag) + if resolved_tag is None: + raise ValueError(f"Unknown organism '{organism}' and no custom tag") + full = resolved_tag + dna + # Reject over-length input rather than letting encode() silently truncate to max_seq_len — + # the per-base `activations` would then be shorter than `bases` and the viz would misalign. + if len(full) > engine.max_seq_len: + raise ValueError(f"sequence too long: {len(full)} > max_seq_len ({engine.max_seq_len})") + codes = engine.encode(full) + if codes.shape[0] != len(full): # belt-and-suspenders: encode must be 1:1 with the input here + raise ValueError("encoded length != input length (tokenizer truncated)") + tag_len = len(resolved_tag) if codes.shape[0] >= len(resolved_tag) else 0 + return dna, resolved_tag, codes, tag_len + + +def parse_clamp_spec(spec) -> dict: + """Normalize one feature-clamp spec into ``{"feature_id": int, "strength": float}``. + + Accepts a ``"FEATURE_ID[:STRENGTH]"`` string (CLI ``--clamp``; strength defaults to 1.0, + e.g. ``29244:300`` or ``29244``), or a mapping with ``feature_id``/``strength`` (server + ``FeatureClamp``). Single source of truth so the CLI and the API parse clamps identically. + Magnitude/finiteness/range are enforced later by ``_sanitize_steering``. + + Raises ``ValueError`` on a malformed string (caller decides whether that's a 400 or an exit). + """ + if isinstance(spec, str): + fid, sep, strength = spec.partition(":") + try: + return {"feature_id": int(fid), "strength": float(strength) if (sep and strength) else 1.0} + except ValueError: + raise ValueError(f"invalid clamp {spec!r}: expected FEATURE_ID[:STRENGTH] with numeric values") + return {"feature_id": int(spec["feature_id"]), "strength": float(spec.get("strength", 1.0))} + + +def _is_unrecoverable_cuda(e: Exception) -> bool: + """True for CUDA errors that poison the process context (device-side assert / sticky CUDA error). + + These don't clear on the next request — the worker must be recycled — so the caller marks the + engine not-ready rather than retrying into a stream of 500s. + """ + s = str(e).lower() + return isinstance(e, RuntimeError) and ("device-side assert" in s or "cuda error" in s) + + +def _sanitize_steering(features, n_features, temperature, top_k): + """Validate/normalize steering inputs (pure, no GPU); raise on bad input. + + Each guard prevents a CUDA device-side assert that would corrupt the context and wedge + the server until restart: + + * feature id outside [0, n_features) indexes off the SAE codes -> ValueError (server -> 400); + * |strength| beyond MAX_CLAMP_STRENGTH blows the logits to inf/NaN -> capped; + * temperature <= 0 makes the recipe's sampler divide logits by temperature (NaN under + multinomial) -> coerce to greedy top-1, which is deterministic and skips that path; + * negative top_k is an invalid sampler argument -> coerce to 0 (no top-k filtering). + + Returns ``(clamps: dict[int, float], fids: list[int], temperature: float, top_k: int)``. + """ + bad = sorted({int(f["feature_id"]) for f in features if not (0 <= int(f["feature_id"]) < n_features)}) + if bad: + raise ValueError(f"feature_id(s) {bad} out of range [0, {n_features})") + clamps = {} + for f in features: + s = float(f.get("strength", 1.0)) + if not math.isfinite(s): # NaN/±inf would blow the logits to NaN -> neutralize explicitly + s = 0.0 + clamps[int(f["feature_id"])] = max(-MAX_CLAMP_STRENGTH, min(MAX_CLAMP_STRENGTH, s)) + temperature = float(temperature) + top_k = max(0, int(top_k)) # negative top_k is an invalid sampler arg -> 0 (no filtering) + if temperature <= 0: # greedy — avoid the sampler's logits/temperature division (NaN) + top_k = max(top_k, 1) + return clamps, list(clamps), temperature, top_k + + +def _init_single_process_distributed() -> None: + """Set the env vars Megatron's distributed init expects for a 1-GPU process.""" + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29577") + + +class Evo2SAE: + """Persistent Evo2 + SAE inference engine (single-sequence and batched).""" + + def __init__( + self, + evo2_ckpt_dir: str, + sae_ckpt_path: str, + layer: int, + device: str = "cuda", + max_seq_len: int = 8192, + feature_annotations: Optional[str] = None, + organism_tags: Optional[dict] = None, + ): + """Record config; call .load() to actually load the model + SAE onto the GPU.""" + self.evo2_ckpt_dir = evo2_ckpt_dir + self.sae_ckpt_path = sae_ckpt_path + self.layer = int(layer) + self.device = device + self.max_seq_len = int(max_seq_len) + self.feature_annotations = feature_annotations + self.organism_tags = dict(organism_tags) if organism_tags else dict(DEFAULT_ORGANISM_TAGS) + + self.model = None # the engine's mcore model (full): encode forward + steering hook target + self.gen_components = None # recipe inference engine (the single model), built in load() + self.tokenizer = None + self.sae = None + self.n_features = None + self.labels: dict[int, str] = {} + self.peaks: dict[int, float] = {} + self._lock = threading.Lock() # serialize GPU access (Megatron isn't thread-safe) + self.ready = False + + # ------------------------------------------------------------------ loading + def load(self) -> "Evo2SAE": + """Load the single Evo2 inference engine + SAE + feature labels (one-time, ~1 min). + + One model serves both paths: encode/highlight runs a full-sequence forward and reads + layer ``layer`` off a hook (see ``_forward_hidden``); generate drives the same model's + dynamic-decode engine with a steering hook on the same layer. + """ + from megatron.core.utils import unwrap_model + + _init_single_process_distributed() + comp = self._ensure_engine() # the one engine; reused by generate() + self.model = unwrap_model(comp.model) # mcore model: encode forward + steering-hook target + self.tokenizer = comp.tokenizer + self.sae, self.n_features = self._load_sae() + # Fail loudly now if the SAE doesn't fit the model, instead of a cryptic matmul error on the + # first encode. We can only catch a dim mismatch (wrong SAE/model); a wrong layer with the + # same hidden size is undetectable here (the SAE checkpoint records no training layer). + self._check_dim(self._sae_input_dim, self._model_hidden_size(), self.layer) + self.labels, self.peaks = self._load_feature_meta() + self.ready = True + logger.info("Evo2SAE ready: layer=%d n_features=%d n_labels=%d", self.layer, self.n_features, len(self.labels)) + return self + + def _ensure_engine(self): + """Lazily build the recipe's inference engine (eager/hookable) for generation. + + cuda_graph_impl="none" keeps decode eager so the residual-stream steering hook + takes effect (a CUDA-graph-captured decode would replay frozen ops and ignore it). + """ + if self.gen_components is None: + from bionemo.evo2.run import infer as INF + + # Defensive: if anything earlier initialized the global num-microbatches calculator, + # setup_inference_engine re-inits it and asserts unless we tear the singleton down first. + try: + from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator + + destroy_num_microbatches_calculator() + except Exception: + pass + + self.gen_components = INF.setup_inference_engine( + Path(self.evo2_ckpt_dir), max_seq_length=self.max_seq_len, cuda_graph_impl="none" + ) + return self.gen_components + + def _model_hidden_size(self): + """The model's hidden size at this layer (== encode's feature dim). + + Config first (cheap); else a 1-token forward (ground truth); ``None`` if neither works + (then the dim check is skipped, never blocking load). + """ + try: + from megatron.core.utils import unwrap_model + + cfg = getattr(unwrap_model(self.model), "config", None) + if cfg is not None and getattr(cfg, "hidden_size", None): + return int(cfg.hidden_size) + except Exception: + pass + try: + probe = self._forward_hidden([self.tokenize("A")]) + if probe and probe[0].numel(): + return int(probe[0].shape[-1]) + except Exception: + pass + return None + + @staticmethod + def _check_dim(sae_input_dim, hidden, layer): + """Raise if the SAE input_dim != the model hidden size (skip if hidden is unknown).""" + if hidden is not None and hidden != sae_input_dim: + raise ValueError( + f"SAE input_dim={sae_input_dim} does not match the Evo2 hidden size={hidden} at " + f"layer {layer} — wrong SAE/model pairing (check --sae-ckpt-path / --layer)." + ) + + def _load_sae(self): + ckpt = torch.load(self.sae_ckpt_path, map_location="cpu", weights_only=False) + cfg = dict(ckpt["model_config"]) + state = ckpt["model_state_dict"] + if any(k.startswith("module.") for k in state): + state = {k.removeprefix("module."): v for k, v in state.items()} + from sae.architectures import TopKSAE + + sae = TopKSAE(**cfg) + sae.load_state_dict(state) + sae.eval().to(self.device) + self._sae_input_dim = int(cfg["input_dim"]) # must equal the model hidden size (checked in load) + logger.info("SAE loaded: TopKSAE input_dim=%d n_features=%d", cfg["input_dim"], cfg["hidden_dim"]) + return sae, int(cfg["hidden_dim"]) + + def _load_feature_meta(self): + """feature_id -> (label, natural peak) from the annotation **parquet** (produced by probe.py annotate).""" + labels: dict[int, str] = {} + peaks: dict[int, float] = {} + if not self.feature_annotations: + return labels, peaks + path = Path(self.feature_annotations) + if not path.exists(): + logger.warning("Feature annotations %s not found — features unlabeled", path) + return labels, peaks + if path.suffix.lower() != ".parquet": + logger.warning("Feature annotations %s: only .parquet is supported — features unlabeled", path) + return labels, peaks + import pyarrow.parquet as pq + + tbl = pq.read_table(path).to_pydict() + ids = tbl.get("feature_id", []) + names = tbl.get("label", tbl.get("annotation", [None] * len(ids))) + pk = tbl.get("max_activation", [None] * len(ids)) + for i, n, p in zip(ids, names, pk): + if n is not None: + labels[int(i)] = str(n) + if p is not None: + peaks[int(i)] = float(p) + logger.info("Loaded %d labels from %s", len(labels), path) + return labels, peaks + + # ------------------------------------------------------------------ tokenize + def tokenize(self, text: str) -> list[int]: + """Tokenize text to token ids, truncated to max_seq_len.""" + tok = self.tokenizer + ids = tok.tokenize(text) if hasattr(tok, "tokenize") else tok.text_to_ids(text) + return ids[: self.max_seq_len] + + def resolve_tag(self, organism: str, tag: Optional[str]) -> Optional[str]: + """Explicit custom `tag` wins; else look up the organism preset.""" + if tag is not None: + return tag + return self.organism_tags.get(organism) + + # ------------------------------------------------------------------ encode + @torch.no_grad() + def encode(self, dna: str) -> torch.Tensor: + """ONE sequence -> SAE codes [seq_len, n_features] on CPU. No phylo tag.""" + return self.encode_batch([dna])[0] + + @torch.no_grad() + def encode_batch(self, seqs: list[str], batch_size: int = 8) -> list[torch.Tensor]: + """MANY sequences -> list of SAE codes [S_i, n_features], batched on the GPU. + + Sequences are padded to the longest in each micro-batch; padding is masked + out before SAE-encoding so each result has the true per-base length. + + Work is length-bucketed (processed in token-length order) so each micro-batch holds + similar-length sequences and wastes little padding on mixed-length inputs; results are + written back by original index, so the returned order matches the input order. + """ + out: list[torch.Tensor] = [None] * len(seqs) # type: ignore + order = [(i, self.tokenize(s)) for i, s in enumerate(seqs)] + order.sort(key=lambda it: len(it[1])) # length-bucket to minimize padding (out[orig_i] un-sorts) + with self._lock: + for start in range(0, len(order), batch_size): + chunk = order[start : start + batch_size] + id_lists = [ids for _, ids in chunk] + hiddens = self._forward_hidden(id_lists) # list of [S_i, H] + for (orig_i, ids), h in zip(chunk, hiddens): + out[orig_i] = ( + self.sae.encode(h.to(self.device)).detach().cpu() + if h.shape[0] > 0 + else torch.empty(0, self.n_features) + ) + return out + + @torch.no_grad() + def _forward_hidden(self, id_lists: list[list[int]]) -> list[torch.Tensor]: + """Run the engine's model on a (padded) batch of token-id lists. + + Returns the unpadded layer-`layer` hidden states [S_i, H] per sequence. + + The single engine model is ``post_process=True`` (it produces logits for generation), so + we can't ask ``_predict_step`` for embeddings. Instead we drive a normal full-sequence + forward and read layer ``layer`` off a forward hook — the same module + ``[S, B, H]`` + layout the steering ``clamp_hook`` uses, so encode and steer see identical activations. + """ + from bionemo.evo2.run import predict as P + + lens = [len(ids) for ids in id_lists] + maxlen = max(lens) if lens else 0 + if maxlen == 0: + return [torch.empty(0, 0) for _ in id_lists] + b = len(id_lists) + tokens = torch.zeros(b, maxlen, dtype=torch.long, device=self.device) + loss_mask = torch.zeros(b, maxlen, dtype=torch.long, device=self.device) + for i, ids in enumerate(id_lists): + if ids: + tokens[i, : len(ids)] = torch.tensor(ids, dtype=torch.long, device=self.device) + loss_mask[i, : len(ids)] = 1 + batch = { + "tokens": tokens, + "position_ids": torch.arange(maxlen, dtype=torch.long, device=self.device).unsqueeze(0).expand(b, -1), + "loss_mask": loss_mask, + "seq_idx": torch.arange(b, dtype=torch.long, device=self.device), + } + # Capture layer-`layer` output via a forward hook (engine model has no embedding-only mode). + captured: dict[str, torch.Tensor] = {} + + def _capture(_module, _inputs, output): + captured["h"] = output[0] if isinstance(output, tuple) else output # [S, B, H] + + handle = self.model.decoder.layers[self.layer].register_forward_hook(_capture) + # Evo2 runs bf16; TransformerEngine asserts param/input dtypes match unless inside an + # autocast region. predict()'s loop relies on this too, so wrap the direct _predict_step. + device_type = "cuda" if self.device.startswith("cuda") else "cpu" + try: + with torch.autocast(device_type=device_type, dtype=torch.bfloat16): + P._predict_step(model=self.model, batch=batch, output_embeddings=False) + finally: + handle.remove() + hidden = captured["h"].transpose(0, 1).contiguous() # [S, B, H] -> [B, S, H] + return [hidden[i, : lens[i]].float() for i in range(b)] + + def feature_tracks(self, dna: str, fids: list[int]) -> dict: + """Per-base activation of several features on `dna`. {fid: [..]} (encoded once).""" + if not dna: + return {int(f): [] for f in fids} + codes = self.encode(dna) + return {int(f): [round(float(v), 4) for v in codes[:, int(f)].tolist()] for f in fids} + + def top_features(self, codes: torch.Tensor, tag_len: int = 0, k: int = 8) -> list[dict]: + """Top-k features by per-base max activation over the DNA region (excluding the tag). + + `codes` is [S, n_features] from `encode`/`encode_batch`; `tag_len` skips the leading + phylo-tag tokens (ignored if it would drop the whole sequence). Returns the strictly + positive features as [{feature_id, label, max_activation}], used by the CLI and server. + """ + if codes.shape[0] == 0: + return [] + region = codes[tag_len:] if codes.shape[0] > tag_len else codes + per = region.max(dim=0).values + idx = per.topk(min(int(k), per.numel())).indices.tolist() + return [ + {"feature_id": int(i), "label": self.labels.get(int(i)), "max_activation": round(float(per[i]), 4)} + for i in idx + if per[i].item() > 0 + ] + + # ------------------------------------------------------------------ generate + def generate( + self, + prompt="", + organism="None (raw DNA)", + tag=None, + features=None, + n_tokens=120, + temperature=1.0, + top_k=0, + compare_baseline=False, + ) -> dict: + """Autoregressively generate DNA, optionally clamping features on the continuation. + + `features` = list of {"feature_id": int, "strength": float} (or []). Generation runs + through the recipe's inference engine (`infer.generate`, eager so the hook applies); + steering is a decode-only forward hook on layer `layer`. Returns + {generation:{sequence,activations}, baseline:..|None, features, steered}. + """ + from bionemo.evo2.run import infer as INF + + # Accept CLI "ID[:STRENGTH]" strings and server FeatureClamp dicts identically, then + # validate/cap below via _sanitize_steering. + features = [parse_clamp_spec(f) for f in (features or [])] + resolved_tag = self.resolve_tag(organism, tag) + if resolved_tag is None: + raise ValueError(f"Unknown organism '{organism}' and no custom tag") + dna = clean_dna(prompt) + full_prompt = resolved_tag + dna + if not full_prompt: + raise ValueError("Provide a prompt or pick an organism (need >=1 token to seed)") + # Reject an over-context prompt rather than silently truncating it in tokenize() (parity + # with annotate; the server maps "too long" -> 413). Raise MAX_SEQ_LEN to allow longer. + if len(full_prompt) > self.max_seq_len: + raise ValueError(f"prompt too long: {len(full_prompt)} > max_seq_len ({self.max_seq_len})") + # Cap to the engine's configured context budget (prompt + generation must fit max_seq_len), + # not an arbitrary constant. Raise MAX_SEQ_LEN at launch to generate longer — the 7B is the + # long-context (1M) model, so it's memory-bound, not architecture-bound (OOD past training len). + budget = max(1, self.max_seq_len - len(self.tokenize(full_prompt))) + n_tokens = max(1, min(int(n_tokens), budget)) + # Validate/normalize steering inputs — out-of-range ids, extreme clamps, and temperature + # 0 each trigger CUDA device-side asserts that wedge the server (see _sanitize_steering). + clamps, fids, temperature, top_k = _sanitize_steering(features, self.n_features, temperature, top_k) + + try: + with self._lock: + comp = self._ensure_engine() + hook_layer = self.model.decoder.layers[self.layer] # same module encode reads; steer here + from evo2_sae.steering import clamp_hook + + feat_meta = [{"id": fid, "label": self.labels.get(fid), "strength": s} for fid, s in clamps.items()] + + def _run(steer: bool) -> str: + handle = ( + hook_layer.register_forward_hook(clamp_hook(self.sae, clamps, decode_only=True)) + if (steer and clamps) + else None + ) + try: + out = INF.generate( + comp, [full_prompt], max_new_tokens=n_tokens, temperature=temperature, top_k=top_k + ) + return clean_dna(INF._unwrap_result(out[0]).generated_text) + finally: + if handle is not None: + handle.remove() + + main_dna = _run(steer=True) + base_dna = _run(steer=False) if (compare_baseline and clamps) else None + except Exception as e: + # PURELY DEFENSIVE: the known client-reachable CUDA-assert triggers are all neutralized + # earlier by _sanitize_steering (out-of-range id, magnitude cap, non-finite strength, + # temperature<=0, negative top_k), so a wedge here implies a genuine hardware/driver + # fault, NOT a crafted request — it is not client-inducible (so exit+restart is not a + # remote DoS vector). If that ever stops holding, a new trigger must be added to + # _sanitize_steering, not handled by leaning on this path. + # + # When it does happen, the device-side assert poisons the CUDA context for the whole + # process — unrecoverable in-process (re-init won't clear it), so restart is the only + # cure. Mark not-ready so /health flips to 503; and if EXIT_ON_CUDA_WEDGE is set + # (launch_inference.sh sets it for `serve`, which runs under a restart loop), exit the + # worker so ANY restart-on-exit supervisor (the launch loop / docker --restart / + # systemd / k8s) respawns it — host-independent recovery, no readiness probe required. + # Default (unset): fail-closed at 503 (safe for library/CLI/test use — must not exit). + if _is_unrecoverable_cuda(e): + self.ready = False + logger.exception("unrecoverable CUDA error in generate() — marking engine not-ready") + if os.environ.get("EXIT_ON_CUDA_WEDGE") == "1": + logger.critical("EXIT_ON_CUDA_WEDGE=1 — exiting the worker for the supervisor to restart") + os._exit(1) + raise + + resp = { + "prompt": dna, + "organism": organism, + "tag": resolved_tag, + "tag_len": len(resolved_tag), + "n_tokens": n_tokens, + "features": feat_meta, + "steered": bool(clamps), + "generation": {"sequence": main_dna, "activations": self.feature_tracks(main_dna, fids)}, + "baseline": None, + } + if base_dna is not None: + resp["baseline"] = {"sequence": base_dna, "activations": self.feature_tracks(base_dna, fids)} + return resp diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/fasta.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/fasta.py new file mode 100644 index 0000000000..46db93066c --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/fasta.py @@ -0,0 +1,56 @@ +# 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. + +"""Shared FASTA reader for the evo2 SAE recipe (stdlib-only; no torch import). + +One streaming parser reused by the batch CLI (`cli.py`) and the FASTA chunker +(`scripts/chunk_fasta.py`) so the header/`.gz`/concat logic lives in one place. +""" + +from __future__ import annotations + +import gzip +from collections.abc import Iterator +from pathlib import Path + + +def read_fasta(path: str | Path) -> Iterator[tuple[str, str]]: + """Yield ``(seq_id, sequence)`` for each record in a FASTA file. + + Args: + path: Path to a FASTA file; a ``.gz`` suffix is decompressed transparently. + + Yields: + ``(seq_id, sequence)``: the first whitespace-delimited token of the header, + or a generated ``seq_`` when the header carries no token (e.g. ``">"`` or + ``"> "``), paired with the record's concatenated sequence lines. + """ + opener = gzip.open if str(path).endswith(".gz") else open + seq_id: str | None = None + parts: list[str] = [] + n = 0 # records yielded so far — used to name token-less headers + 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) + n += 1 + header = line[1:].strip().split() + seq_id, parts = (header[0] if header else f"seq_{n}"), [] + else: + parts.append(line) + if seq_id is not None: + yield seq_id, "".join(parts) diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py new file mode 100644 index 0000000000..c25dc5fdd5 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py @@ -0,0 +1,241 @@ +# 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. + +"""FastAPI server over the Evo2SAE engine — the live backend the viz talks to. + +Endpoints (under /api): /api/health, /api/features, /api/annotate (per-base activations for a +pasted sequence), /api/generate (autoregressive generation + optional SAE-feature clamp). The +/api prefix lets a prebuilt frontend be served from "/" on the same origin (single-container +deploy); when no frontend is configured the server is API-only. This is a thin layer; all model +work lives in `core.Evo2SAE`. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Optional + +import anyio +from fastapi import APIRouter, FastAPI, HTTPException, Response +from fastapi.responses import JSONResponse +from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel + +from . import core +from .core import Evo2SAE + + +logger = logging.getLogger("evo2_sae_infer.server") + + +def _resolve_static_dir(static_dir: Optional[str]) -> Optional[str]: + """Directory of a prebuilt frontend to serve at ``/``, or None for an API-only server. + + Explicit ``static_dir`` arg wins, else the ``DASHBOARD_DIST`` env var. Returns None unless the + path is an existing directory — so a server with no built frontend (dev, or an image built + without one) just serves the API and ``/`` 404s, instead of crashing. This layer is generic: + it serves whatever static dir it's pointed at and knows nothing about the dashboard (the + dashboard recipe supplies the dir via DASHBOARD_DIST / the Docker build). + """ + cand = static_dir or os.getenv("DASHBOARD_DIST") + return cand if (cand and Path(cand).is_dir()) else None + + +class AnnotateRequest(BaseModel): + """Request body for /annotate (top-k feature scan or an explicit feature pick).""" + + sequence: str + organism: str = "None (raw DNA)" + tag: Optional[str] = None + mode: str = "topk" # "topk" | "pick" + k: int = 8 + feature_ids: Optional[list[int]] = None + + +class FeatureClamp(BaseModel): + """A single SAE-feature steering clamp (feature id + target strength).""" + + feature_id: int + strength: float = 1.0 + + +class GenerateRequest(BaseModel): + """Request body for /generate (autoregressive generation + optional SAE-feature clamps).""" + + prompt: str = "" + organism: str = "None (raw DNA)" + tag: Optional[str] = None + features: list[FeatureClamp] = [] + n_tokens: int = 120 + temperature: float = 1.0 + top_k: int = 0 + compare_baseline: bool = False + + +def build_app(engine: Evo2SAE, static_dir: Optional[str] = None) -> FastAPI: + """Build the FastAPI app; the engine is loaded once in the lifespan handler. + + API routes live under ``/api`` (so the dashboard and the API can be served from one origin: + the frontend always calls ``/api/*``, in dev via the Vite proxy and in production from the + same server). If a built frontend is found (``static_dir`` / ``DASHBOARD_DIST``), it is mounted + at ``/`` so a single container serves both the UI and the API; otherwise the server is API-only. + """ + # One GPU (the engine serializes model calls with a lock), so cap how many sync requests run + # at once: excess requests wait for a worker instead of piling up dozens of parked threads. + # NOTE: generation is bounded only by the context window now, so a single /generate can run + # long — under concurrent load requests queue behind it. Sync endpoints run in Starlette's + # AnyIO threadpool (default 40); shrink it. Tune MAX_CONCURRENCY. + max_concurrency = int(os.getenv("MAX_CONCURRENCY", "8")) + + @asynccontextmanager + async def lifespan(app: FastAPI): + anyio.to_thread.current_default_thread_limiter().total_tokens = max_concurrency + try: + engine.load() + logger.info("engine ready") + except Exception: + logger.exception("engine startup failed — /health stays not-ready") + yield + + app = FastAPI(title="Evo2 SAE inference", lifespan=lifespan) + + # No CORS middleware: the dashboard always reaches the backend same-origin (Vite proxies + # /api -> :8001), so cross-origin is never used. CORS is browser-only and not an access + # control anyway — scripts ignore it; SSO + the limits below are what gate this endpoint. + + # Reject oversized request bodies up front (a multi-MB sequence would be read into memory + # before per-field validation could reject it). Default 16 MiB; override with MAX_BODY_BYTES. + # NOTE: advisory — this trusts the Content-Length header, so a chunked request (no length) or a + # lying header bypasses it; it guards well-behaved clients, not a hard cap. Fine behind SSO; + # real enforcement would count streamed bytes. + max_body = int(os.getenv("MAX_BODY_BYTES", str(16 * 1024 * 1024))) + + @app.middleware("http") + async def _limit_body(request, call_next): + cl = request.headers.get("content-length") + if cl is not None and cl.isdigit() and int(cl) > max_body: + return JSONResponse({"detail": f"request body too large (> {max_body} bytes)"}, status_code=413) + return await call_next(request) + + def _require_ready(): + if not engine.ready: + raise HTTPException(503, "Backend not ready") + + # All endpoints under /api (mounted below) so the SPA at "/" never collides with them. + api = APIRouter() + + @api.get("/health") + def health(response: Response): + if not engine.ready: + response.status_code = 503 # readiness probes shed this pod until load finishes (body still informative) + return { + "ready": bool(engine.ready), + "layer": engine.layer, + "n_features": engine.n_features, + "n_labels": len(engine.labels), + "organisms": list(engine.organism_tags.keys()), + "organism_tags": engine.organism_tags, + "device": engine.device, + "max_seq_len": engine.max_seq_len, # context budget — UI caps generation length to this + } + + @api.get("/features") + def features(): + _require_ready() + rows = [ + {"id": int(f), "label": lab, "natural_peak": engine.peaks.get(int(f))} for f, lab in engine.labels.items() + ] + rows.sort(key=lambda r: r["id"]) + return rows + + @api.post("/annotate") + def annotate(req: AnnotateRequest): + _require_ready() + try: + dna, tag, codes, tag_len = core.annotate(engine, req.sequence, req.organism, req.tag) + except ValueError as e: + raise HTTPException(413 if "too long" in str(e) else 400, str(e)) + full = tag + dna + if req.mode not in ("pick", "topk"): + raise HTTPException(400, f"Invalid mode {req.mode!r}: must be 'pick' or 'topk'") + if req.mode == "pick": + if not req.feature_ids: + raise HTTPException(400, "mode='pick' requires feature_ids") + chosen = [int(i) for i in req.feature_ids] + # Pick ids are user-supplied; an out-of-range id would IndexError (500) and a negative + # one would silently index the wrong feature via torch negative-indexing. Reject -> 400. + bad = sorted({i for i in chosen if not (0 <= i < engine.n_features)}) + if bad: + raise HTTPException(400, f"feature_id(s) {bad} out of range [0, {engine.n_features})") + else: + k = max(1, min(int(req.k), 64)) + chosen = [ft["feature_id"] for ft in engine.top_features(codes, tag_len=tag_len, k=k)] + feats = [] + for fid in chosen: + col = codes[:, fid] + feats.append( + { + "feature_id": fid, + "label": engine.labels.get(fid), + "max_activation": float(col[tag_len:].max().item()) + if codes.shape[0] > tag_len + else float(col.max().item()), + "activations": [round(float(v), 4) for v in col.tolist()], + } + ) + return { + "sequence": dna, + "organism": req.organism, + "tag": tag, + "tag_len": tag_len, + "bases": list(full), + "n_tokens": codes.shape[0], + "layer": engine.layer, + "features": feats, + } + + @api.post("/generate") + def generate(req: GenerateRequest): + _require_ready() + try: + return engine.generate( + prompt=req.prompt, + organism=req.organism, + tag=req.tag, + features=[f.model_dump() for f in req.features], + n_tokens=req.n_tokens, + temperature=req.temperature, + top_k=req.top_k, + compare_baseline=req.compare_baseline, + ) + except ValueError as e: + raise HTTPException(413 if "too long" in str(e) else 400, str(e)) + + app.include_router(api, prefix="/api") + + # Serve a prebuilt frontend at "/" when present, so one container serves UI + API. Mounted + # AFTER the API router, so /api/* always resolves to the API; unknown /api/* paths 404 here + # (StaticFiles only serves index.html for "/", not as a SPA catch-all — the UI is tabs at "/"). + resolved = _resolve_static_dir(static_dir) + if resolved: + app.mount("/", StaticFiles(directory=resolved, html=True), name="dashboard") + logger.info("serving dashboard from %s", resolved) + else: + logger.info("no frontend mounted (API-only)") + + return app diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steering.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steering.py new file mode 100644 index 0000000000..edec423bf5 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steering.py @@ -0,0 +1,88 @@ +# 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. + +"""Causal feature steering for SAEs — clamp features in code-space, inject only the delta. + +A forward hook on the layer the SAE was trained on: it re-encodes the layer output through +the SAE, overrides chosen features in code-space, decodes, and adds the **delta** back to the +activation. Because we add ``decode(clamped) - decode(original)`` (not the recon itself), the +SAE's reconstruction error cancels and only the clamped feature's decoder contribution moves +the activation. Model-agnostic: needs only the SAE (``encode`` / ``decode``) and the module to +hook. Measure the effect (e.g. ΔP of a target token) by running the model with vs. without the +hook. +""" + +from contextlib import contextmanager +from typing import Any, Callable, Dict, Iterator + +import torch + + +def clamp_hook(sae: Any, clamps: Dict[int, float], decode_only: bool = False) -> Callable: + """Build a forward hook that clamps ``{feature_idx: value}`` via the delta method. + + The hook adds ``decode(clamped_codes) - decode(original_codes)`` to the hooked module's + output, so the SAE reconstruction error cancels. ``value=0`` ablates a feature; a negative + value reverses its decoder direction. Works whether the module returns a tensor or a tuple + whose first element is the hidden state. + + Args: + sae: A trained TopK SAE exposing ``encode(x) -> codes``, ``encode_pre_act(x) -> (_, info)``, + and ``decode(codes, info)``. + clamps: Map of feature index -> absolute code value to force at every position. + decode_only: If True, steer only autoregressive *decode* steps and leave the prompt + prefill untouched (continuation-only steering). Assumes a ``(sequence, batch, hidden)`` + layout — the convention for Evo2/megatron decoder layers — and applies the clamp only + when the sequence dimension is 1 (a single new token). + + Returns: + A ``register_forward_hook``-compatible ``hook(module, inputs, output)``. + """ + items = [(int(f), float(v)) for f, v in clamps.items()] + + def hook(module, inputs, output): + h, rest = (output[0], output[1:]) if isinstance(output, tuple) else (output, None) + if decode_only and h.shape[0] != 1: # prefill (seq dim > 1) — leave untouched + return output + dtype, shape = h.dtype, h.shape + h_flat = h.reshape(-1, h.shape[-1]).float() + with torch.no_grad(): + # Encode through the SAE itself (canonical encode) instead of re-deriving relu+topk on + # sae.top_k: no hardcoded sparsity, and it can never drift from the model's true + # encoding. encode_pre_act supplies the normalization info decode needs to map back to + # the original activation scale. + _, info = sae.encode_pre_act(h_flat) + codes_orig = sae.encode(h_flat) + codes_clamped = codes_orig.clone() + n_feat = codes_orig.shape[-1] + for f, v in items: + if not 0 <= f < n_feat: + raise ValueError(f"clamp feature {f} out of range [0, {n_feat})") + codes_clamped[:, f] = v + delta = sae.decode(codes_clamped, info) - sae.decode(codes_orig, info) + h_out = (h_flat + delta).to(dtype).reshape(shape) + return (h_out, *rest) if rest is not None else h_out + + return hook + + +@contextmanager +def steer(module: "torch.nn.Module", sae: Any, clamps: Dict[int, float], decode_only: bool = False) -> Iterator[None]: + """Register the clamp hook on ``module`` for the duration of the ``with`` block, then remove it.""" + handle = module.register_forward_hook(clamp_hook(sae, clamps, decode_only=decode_only)) + try: + yield + finally: + handle.remove() diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/conftest.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/conftest.py new file mode 100644 index 0000000000..da2cfdaad3 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/conftest.py @@ -0,0 +1,201 @@ +# 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. + +"""Shared checkpoint fixtures for the Evo2-SAE GPU tests. + +The GPU steering tests need an mbridge Evo2 checkpoint + an SAE checkpoint. These fixtures +resolve them two ways: + + * **Manual / 7B:** if ``EVO2_CKPT_DIR`` / ``SAE_CKPT_PATH`` are set, use those as-is. + * **CI:** otherwise build the **1B-8k-bf16** the way ``evo2_megatron``'s own tests do + (``bionemo_load`` -> ``run_nemo2_to_mbridge``) and synthesize a tiny random TopK SAE + matching the 1B layer dim. Gated by a GPU-memory check (not ``skipif(CI)``), so it runs + on the L4 the same way evo2's ``test_batch_generate`` 1B case does, and skips on smaller + GPUs / when the checkpoint can't be fetched. + +The fixtures are lazy (heavy imports + the build happen only when a GPU test requests the +engine), so the CPU-only sanitize tests collect and run without CUDA or bionemo.evo2. +""" + +import os + +import pytest +import torch +from sae.architectures import TopKSAE + + +# The 1B (Hyena1bModelProvider) needs ~18 GB to load+generate; the L4 (24 GB) fits, smaller +# GPUs skip. The 7B path is reached only via EVO2_CKPT_DIR (manual), never auto-built here. +_MIN_GPU_GB = 20 +_EVO2_1B_HIDDEN = 1920 # Hyena1bModelProvider.hidden_size -> SAE input_dim + + +@pytest.fixture(scope="session") +def evo2_ckpt_dir(tmp_path_factory) -> str: + """MBridge Evo2 checkpoint dir for the engine. + + Honors ``EVO2_CKPT_DIR``; otherwise builds the 1B-8k-bf16 from NGC (memory-gated). + """ + env = os.environ.get("EVO2_CKPT_DIR") + if env: + return env + + if not torch.cuda.is_available(): + pytest.skip("building/loading the 1B requires CUDA (or set EVO2_CKPT_DIR)") + total_gb = torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory / 1024**3 + if total_gb < _MIN_GPU_GB: + pytest.skip(f"need >= {_MIN_GPU_GB} GB GPU for the 1B; have {total_gb:.1f} GB") + + try: + # The NGC loader lives in bionemo.common on the migrated recipes layout; older framework + # builds expose it as bionemo.core. Support both so the fixture works either way. + try: + from bionemo.common.data.load import load as bionemo_load + except ImportError: + from bionemo.core.data.load import load as bionemo_load + from bionemo.evo2.data.dataset_tokenizer import DEFAULT_HF_TOKENIZER_MODEL_PATH_512 + from bionemo.evo2.utils.checkpoint.nemo2_to_mbridge import run_nemo2_to_mbridge + except ImportError as e: # pragma: no cover - environment guard + pytest.skip(f"bionemo.evo2 not importable (build via .ci_build.sh): {e}") + + try: + nemo2_ckpt_path = bionemo_load("evo2/1b-8k-bf16:1.0") + except ValueError as e: + pytest.skip(f"could not fetch evo2/1b-8k-bf16 (try BIONEMO_DATA_SOURCE=pbss): {e}") + + out_dir = tmp_path_factory.mktemp("evo2_1b_mbridge_session") + mbridge_ckpt_dir = run_nemo2_to_mbridge( + nemo2_ckpt_dir=nemo2_ckpt_path, + tokenizer_path=DEFAULT_HF_TOKENIZER_MODEL_PATH_512, + mbridge_ckpt_dir=out_dir / "evo2_1b_mbridge", + model_size="evo2_1b_base", + seq_length=8192, + mixed_precision_recipe="bf16_mixed", + vortex_style_fp8=False, + ) + return str(mbridge_ckpt_dir / "iter_0000001") + + +@pytest.fixture(scope="session") +def sae_ckpt_path(tmp_path_factory) -> str: + """SAE checkpoint for the engine. + + Honors ``SAE_CKPT_PATH``; otherwise synthesizes a tiny random TopK SAE whose ``input_dim`` + matches the 1B layer (1920). Random weights — this exercises the load/encode/steer code + path, not SAE fidelity, which is the point of the CI smoke. + """ + env = os.environ.get("SAE_CKPT_PATH") + if env: + return env + + torch.manual_seed(0) + sae = TopKSAE(input_dim=_EVO2_1B_HIDDEN, hidden_dim=256, top_k=16) + with torch.no_grad(): + # Real (random) signal so encode produces nonzero codes and steering has something to clamp. + sae.decoder.weight.normal_(0, 0.02) + sae.encoder.weight.copy_(sae.decoder.weight.t()) + + out = tmp_path_factory.mktemp("tiny_sae") / "tiny_sae_1b.pt" + torch.save({"model_config": sae._get_config(), "model_state_dict": sae.state_dict()}, out) + return str(out) + + +@pytest.fixture(scope="session") +def embedding_layer() -> int: + """Layer whose residual stream the SAE reads/steers (1B has 25 layers; default 19).""" + return int(os.environ.get("EMBEDDING_LAYER", "19")) + + +# --------------------------------------------------------------------------------------------- +# CPU fixtures for the serve layer (#1637). `FakeEngine` is the one mock the CLI tests +# (test_cli.py) and the server contract tests (test_server.py) both drive — no model, CPU-only — +# so the two suites stay in lockstep on the engine surface the server/CLI touch. +# --------------------------------------------------------------------------------------------- +from evo2_sae import core # noqa: E402 + + +class FakeEngine: + """Minimal stand-in for Evo2SAE exposing only what the CLI + server touch.""" + + def __init__(self): + self.ready = True + self.layer = 19 + self.n_features = 4 + self.labels = {0: "feat0", 1: "feat1"} + self.peaks = {0: 0.5} + self.organism_tags = {"None (raw DNA)": "", "Human": "|tag|"} + self.device = "cpu" + self.sae_ckpt_path = "fake.pt" + self.max_seq_len = 8192 + self.gen_kwargs = None # records the last generate() call for CLI assertions + self.last_k = None # records the k top_features() was last called with + + def load(self): + self.ready = True + return self # the CLI uses `_engine(args).load()` + + def resolve_tag(self, organism, tag): + return tag if tag is not None else self.organism_tags.get(organism) + + def encode(self, full): + codes = torch.zeros(len(full), self.n_features) + codes[:, 0] = 1.0 # feature 0 fires everywhere + return codes + + def encode_batch(self, seqs, batch_size=8): + return [self.encode(s) for s in seqs] + + def top_features(self, codes, tag_len=0, k=8): + self.last_k = k # so tests can assert the server clamped k into [1, 64] + feats = [{"feature_id": i, "label": self.labels.get(i), "max_activation": 1.0} for i in range(self.n_features)] + return feats[:k] + + def generate(self, **kw): + self.gen_kwargs = kw + # Mirror the real engine: clamps normalize through the shared parser (a malformed --clamp + # raises), out-of-range ids are rejected (the wedge guard), and a seedless request is too. + specs = [core.parse_clamp_spec(f) for f in (kw.get("features") or [])] + bad = [s["feature_id"] for s in specs if not (0 <= s["feature_id"] < self.n_features)] + if bad: + raise ValueError(f"feature_id(s) {bad} out of range [0, {self.n_features})") + prompt = core.clean_dna(kw.get("prompt", "")) + if not prompt and kw.get("organism") == "None (raw DNA)" and not kw.get("tag"): + raise ValueError("need a seed") + if len(prompt) > self.max_seq_len: # mirror the real over-context reject (server -> 413) + raise ValueError(f"prompt too long: {len(prompt)} > max_seq_len ({self.max_seq_len})") + # Match the real engine's response: features are feat_meta dicts keyed {id, label, strength} + # (NOT feature_id) — this is the shape the dashboard consumes through /generate. + feat_meta = [ + {"id": s["feature_id"], "label": self.labels.get(s["feature_id"]), "strength": s["strength"]} + for s in specs + ] + resp = { + "prompt": prompt, + "organism": kw.get("organism"), + "tag": kw.get("tag"), + "steered": bool(specs), + "features": feat_meta, + "generation": {"sequence": "ACGT", "activations": {0: [1.0, 1.0, 1.0, 1.0]}}, + "baseline": None, + } + if kw.get("compare_baseline") and specs: + resp["baseline"] = {"sequence": "TTTT", "activations": {}} + return resp + + +@pytest.fixture +def fake_engine(): + """A fresh CPU FakeEngine instance.""" + return FakeEngine() diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_clamp.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_clamp.py new file mode 100644 index 0000000000..a19a226387 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_clamp.py @@ -0,0 +1,76 @@ +# 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. + +"""CPU tests for evo2_sae.steering: the delta-clamp adds exactly decode(clamped) - decode(orig).""" + +import torch +from evo2_sae.steering import clamp_hook, steer +from sae.architectures import TopKSAE +from torch import nn + + +def _sae(): + torch.manual_seed(0) + return TopKSAE(input_dim=8, hidden_dim=16, top_k=4, normalize_input=False) + + +def test_delta_clamp_is_exact_and_cancels_recon(): + """No-op clamp leaves the activation unchanged (recon error cancels); a real clamp shifts + it by exactly decode(clamped) - decode(orig) — the two halves of the delta-clamp contract.""" + sae, m, x = _sae(), nn.Identity(), torch.randn(5, 8) + + # No-op: decode(orig) != x, but the added delta is 0, so the output is unchanged. + with steer(m, sae, {}): + assert torch.allclose(m(x), x, atol=1e-5) + + # Real clamp: output == x + (decode(clamped) - decode(orig)), recon error cancelled. + with torch.no_grad(): + pre, info = sae.encode_pre_act(x.float()) + codes = torch.relu(pre) + kv, ki = torch.topk(codes, sae.top_k, dim=-1) + co = torch.zeros_like(codes).scatter(-1, ki, kv) + cc = co.clone() + cc[:, 3] = 5.0 + expected = x + (sae.decode(cc, info) - sae.decode(co, info)) + with steer(m, sae, {3: 5.0}): + assert torch.allclose(m(x), expected, atol=1e-4) + + +def test_tuple_output_steers_only_hidden_state(): + """When the hooked module returns a tuple, only element 0 is steered; the rest passes through.""" + + class M(nn.Module): + def forward(self, x): + return (x, "meta") + + sae, x = _sae(), torch.randn(3, 8) + m = M() + handle = m.register_forward_hook(clamp_hook(sae, {0: 2.0})) + out = m(x) + handle.remove() + assert isinstance(out, tuple) and out[1] == "meta" + assert out[0].shape == x.shape and not torch.allclose(out[0], x) # clamp moved it + + +def test_decode_only_skips_prefill(): + """decode_only steers single-token decode steps ([1,B,H]) but leaves multi-token prefill alone.""" + sae, m = _sae(), nn.Identity() + prefill = torch.randn(5, 2, 8) # [S=5, B, H] — prompt prefill, must pass through + decode = torch.randn(1, 2, 8) # [S=1, B, H] — a single new token, must be steered + handle = m.register_forward_hook(clamp_hook(sae, {3: 5.0}, decode_only=True)) + out_prefill, out_decode = m(prefill), m(decode) + handle.remove() + assert torch.allclose(out_prefill, prefill, atol=1e-5) # prefill untouched + assert not torch.allclose(out_decode, decode) # decode step steered diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_cli.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_cli.py new file mode 100644 index 0000000000..d541d3eeed --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_cli.py @@ -0,0 +1,121 @@ +# 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. + +"""CLI contract tests — the `encode | batch | generate` subcommands, CPU-only. + +A mocked engine (no model) drives `cli.main()` so these run in CI alongside test_server.py. +They lock the stdout JSON shapes, the clean-exit error handling (no tracebacks), and the +shared clamp/annotate paths in `core` that the CLI and the server now both go through. +Real model inference is covered by test_steering.py. +""" + +import json + +import pandas as pd +import pytest +from evo2_sae import cli, core + + +# FakeEngine lives in conftest.py — shared with test_server.py so both suites mock the same surface. + + +@pytest.fixture +def fake(monkeypatch, fake_engine): + """Patch the CLI's engine factory so every subcommand runs against the shared CPU fake.""" + monkeypatch.setattr(cli, "_engine", lambda args: fake_engine) + return fake_engine + + +def run(monkeypatch, *argv): + """Invoke `cli.main()` with the given argv (no leading prog name).""" + monkeypatch.setattr("sys.argv", ["evo2-sae", *argv]) + cli.main() + + +# ----------------------------------------------------------------------------- encode +def test_encode_outputs_top_features(fake, monkeypatch, capsys): + run(monkeypatch, "encode", "--sequence", "ACGTACGT") + out = json.loads(capsys.readouterr().out) + assert out["bases"] == 8 + assert out["top_features"][0]["feature_id"] == 0 + + +def test_encode_rejects_non_dna(fake, monkeypatch): + with pytest.raises(SystemExit): # clean exit, not a traceback + run(monkeypatch, "encode", "--sequence", "ZZZZ") + + +def test_encode_rejects_unknown_organism(fake, monkeypatch): + with pytest.raises(SystemExit): + run(monkeypatch, "encode", "--sequence", "ACGT", "--organism", "Martian") + + +# --------------------------------------------------------------------------- generate +def test_generate_outputs_sequence(fake, monkeypatch, capsys): + run(monkeypatch, "generate", "--prompt", "ACGT") + out = json.loads(capsys.readouterr().out) + assert out["sequence"] == "ACGT" + assert out["steered"] is False + + +def test_generate_passes_raw_clamp_strings(fake, monkeypatch, capsys): + run(monkeypatch, "generate", "--prompt", "ACGT", "--clamp", "0:5", "--clamp", "1") + # The CLI hands raw "ID[:STRENGTH]" strings straight to the engine; core normalizes them. + assert fake.gen_kwargs["features"] == ["0:5", "1"] + out = json.loads(capsys.readouterr().out) + # generate() echoes feat_meta keyed {id, label, strength} — the shape the viz consumes. + assert out["features"] == [ + {"id": 0, "label": "feat0", "strength": 5.0}, + {"id": 1, "label": "feat1", "strength": 1.0}, + ] + + +def test_generate_rejects_malformed_clamp(fake, monkeypatch): + with pytest.raises(SystemExit): + run(monkeypatch, "generate", "--prompt", "ACGT", "--clamp", "notanumber") + + +# ------------------------------------------------------------------------------ batch +def test_batch_writes_parquet(fake, monkeypatch, capsys, tmp_path): + fasta = tmp_path / "in.fa" + fasta.write_text(">a\nACGTACGT\n>b\nTTTT\n") + out = tmp_path / "out.parquet" + run(monkeypatch, "batch", "--fasta", str(fasta), "--out", str(out)) + df = pd.read_parquet(out) + assert set(df["sequence_id"]) == {"a", "b"} + assert {"sequence_id", "bp", "rank", "feature_id"} <= set(df.columns) + + +# ------------------------------------------------------- shared core helpers (CLI + API) +def test_parse_clamp_spec_string_and_dict(): + assert core.parse_clamp_spec("29244:300") == {"feature_id": 29244, "strength": 300.0} + assert core.parse_clamp_spec("7") == {"feature_id": 7, "strength": 1.0} # default strength + assert core.parse_clamp_spec({"feature_id": 7, "strength": 2.5}) == {"feature_id": 7, "strength": 2.5} + assert core.parse_clamp_spec({"feature_id": 7}) == {"feature_id": 7, "strength": 1.0} + + +def test_parse_clamp_spec_rejects_garbage(): + with pytest.raises(ValueError): + core.parse_clamp_spec("x:y") + + +def test_annotate_skips_tag_and_rejects_bad_input(fake_engine): + dna, tag, codes, tag_len = core.annotate(fake_engine, "ACGT", organism="Human") + assert dna == "ACGT" and tag == "|tag|" and tag_len == len("|tag|") + assert codes.shape == (len(tag) + len(dna), fake_engine.n_features) + with pytest.raises(ValueError): + core.annotate(fake_engine, "ZZZZ") # non-DNA + with pytest.raises(ValueError): + core.annotate(fake_engine, "ACGT", organism="Martian") # unknown organism diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_core.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_core.py new file mode 100644 index 0000000000..bf25681435 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_core.py @@ -0,0 +1,120 @@ +# 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. + +"""CPU unit tests for the engine logic that needs no model load. + +These exercise the pure / orchestration parts of ``evo2_sae.core`` — DNA cleaning, organism-tag +resolution, top-feature selection, SAE checkpoint loading, and ``generate`` input guards — on an +*unloaded* engine (``__init__`` only records config; it touches no GPU and imports no Evo2). So +they run in CI even when the GPU/1B path skips, and cover the code paths the model smoke can't. +""" + +import pytest +import torch +from evo2_sae import Evo2SAE +from sae.architectures import TopKSAE + + +def _engine(sae_ckpt_path="unused.pt", **kw): + """An unloaded engine — __init__ records config only (no GPU, no Evo2/SAE load).""" + return Evo2SAE(evo2_ckpt_dir="unused", sae_ckpt_path=str(sae_ckpt_path), layer=0, device="cpu", **kw) + + +# ------------------------------------------------------------------------------- top_features +def test_top_features_ranks_positive_only_and_skips_tag(): + eng = _engine() # self.labels defaults to {} + codes = torch.zeros(4, 5) + codes[:, 2] = 3.0 # strongest over the DNA region + codes[:, 4] = 1.0 # weaker, still positive + codes[0, 1] = 9.0 # lives only in row 0, which tag_len=1 skips + + feats = eng.top_features(codes, tag_len=1, k=8) + ids = [f["feature_id"] for f in feats] + assert ids == [2, 4] # ranked by per-base max, positive-only + assert 1 not in ids # tag region (row 0) excluded by tag_len + assert 0 not in ids and 3 not in ids # zero-activation features dropped + assert feats[0]["max_activation"] == 3.0 and feats[0]["label"] is None + + +def test_top_features_empty_codes_returns_empty(): + assert _engine().top_features(torch.zeros(0, 5)) == [] + + +# ------------------------------------------------------------------------------- _load_sae +def _save_topk(path, prefix=""): + """Write a tiny TopK SAE checkpoint in the format _load_sae expects (optionally DDP-prefixed).""" + torch.manual_seed(0) + sae = TopKSAE(input_dim=8, hidden_dim=16, top_k=4) + state = {prefix + k: v for k, v in sae.state_dict().items()} + torch.save({"model_config": sae._get_config(), "model_state_dict": state}, path) + + +def test_load_sae_topk_returns_sae_and_n_features(tmp_path): + ckpt = tmp_path / "tiny_topk.pt" + _save_topk(ckpt) + sae, n_features = _engine(ckpt)._load_sae() + assert isinstance(sae, TopKSAE) and n_features == 16 + + +def test_load_sae_strips_module_prefix(tmp_path): + ckpt = tmp_path / "ddp_topk.pt" + _save_topk(ckpt, prefix="module.") # as saved under DDP + sae, n_features = _engine(ckpt)._load_sae() + assert isinstance(sae, TopKSAE) and n_features == 16 + + +def test_check_dim_rejects_sae_model_mismatch(): + # an SAE whose input_dim != the model's hidden size would matmul-fail on the first encode; + # load() catches it with a clear error instead. (hidden=None -> unknown -> skip, never blocks.) + with pytest.raises(ValueError, match="does not match"): + Evo2SAE._check_dim(sae_input_dim=4096, hidden=1920, layer=26) + Evo2SAE._check_dim(sae_input_dim=1920, hidden=1920, layer=26) # match -> no raise + Evo2SAE._check_dim(sae_input_dim=1920, hidden=None, layer=26) # unknown -> skip + + +def test_encode_batch_length_bucketing_preserves_order(): + # encode_batch sorts work by length (bucketing) but must return results in INPUT order. + # Stub the model so this runs on CPU: each sequence's first token carries its (distinct) length + # as a marker, _forward_hidden echoes it, and the SAE is identity — so out[i] should carry the + # marker of seqs[i] regardless of the internal length-sort. + import types + + eng = _engine() + seqs = ["AC", "ACGTACGT", "A" * 20, "ACG", "A" * 15] # lengths 2,8,20,3,15 (distinct markers) + eng.tokenize = lambda s: [len(s)] + [0] * (len(s) - 1) + eng._forward_hidden = lambda id_lists: [torch.tensor([[float(ids[0])]]) for ids in id_lists] + eng.sae = types.SimpleNamespace(encode=lambda h: h) + eng.n_features = 1 + out = eng.encode_batch(seqs, batch_size=2) + assert len(out) == len(seqs) + assert [int(o[0, 0].item()) for o in out] == [len(s) for s in seqs] # input order preserved + + +# ------------------------------------------------------------------------------- generate input guards +def test_generate_rejects_unknown_organism(): + with pytest.raises(ValueError, match="organism"): + _engine().generate(prompt="ACGT", organism="Klingon") + + +def test_generate_rejects_empty_prompt(): + # "None (raw DNA)" has an empty tag, so an empty prompt leaves nothing to seed generation. + with pytest.raises(ValueError): + _engine().generate(prompt="", organism="None (raw DNA)") + + +def test_generate_rejects_overlong_prompt(): + # An over-context prompt is rejected (server -> 413), not silently truncated by tokenize(). + with pytest.raises(ValueError, match="too long"): + _engine(max_seq_len=16).generate(prompt="A" * 32, organism="None (raw DNA)") diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py new file mode 100644 index 0000000000..b9f826bd41 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py @@ -0,0 +1,190 @@ +# 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. + +"""Server contract tests — the API the feature-explorer viz consumes. + +A mocked engine (no model, CPU-only) drives the FastAPI app so these run in CI and lock the +response shapes + error codes the dashboard depends on: /health, /features, /annotate (per-base +activations), /generate. Real model inference is covered by test_steering.py. +""" + +import pytest +from evo2_sae.server import build_app +from fastapi.testclient import TestClient + + +# FakeEngine lives in conftest.py — shared with test_cli.py so both suites mock the same surface. + + +@pytest.fixture +def client(fake_engine): + with TestClient(build_app(fake_engine)) as c: + yield c + + +@pytest.fixture +def dist_dir(tmp_path): + """A minimal stand-in for a built frontend (feature_explorer/dist) — index + one asset.""" + d = tmp_path / "dist" + (d / "assets").mkdir(parents=True) + (d / "index.html").write_text('
') + (d / "assets" / "app.js").write_text("console.log('dashboard')") + return d + + +@pytest.fixture +def static_client(fake_engine, dist_dir): + """A client for the single-container shape: API under /api + the frontend served at /.""" + with TestClient(build_app(fake_engine, static_dir=str(dist_dir))) as c: + yield c + + +def test_health(client): + r = client.get("/api/health") + assert r.status_code == 200 # 200 only when ready + b = r.json() + assert b["ready"] is True and b["layer"] == 19 + assert "None (raw DNA)" in b["organisms"] + + +def test_annotate_rejects_too_long(client, fake_engine): + seq = "A" * (fake_engine.max_seq_len + 1) # exceeds the context budget + assert client.post("/api/annotate", json={"sequence": seq}).status_code == 413 + + +def test_features(client): + rows = client.get("/api/features").json() + assert {"id", "label", "natural_peak"} <= set(rows[0]) + + +def test_annotate_returns_per_base_activations(client): + b = client.post("/api/annotate", json={"sequence": "ACGTACGT", "organism": "None (raw DNA)"}).json() + assert {"sequence", "features", "bases", "tag_len", "layer", "n_tokens"} <= set(b) + assert b["features"][0]["activations"] # the per-base track the viz plots + + +def test_annotate_rejects_non_dna(client): + assert client.post("/api/annotate", json={"sequence": "ZZZZ"}).status_code == 400 + + +def test_annotate_pick_mode(client): + b = client.post("/api/annotate", json={"sequence": "ACGTACGT", "mode": "pick", "feature_ids": [1]}).json() + assert [f["feature_id"] for f in b["features"]] == [1] + assert b["features"][0]["activations"] # per-base track returned for the picked feature + + +def test_annotate_pick_requires_ids(client): + assert client.post("/api/annotate", json={"sequence": "ACGT", "mode": "pick"}).status_code == 400 + + +def test_annotate_pick_rejects_out_of_range_id(client, fake_engine): + # user-supplied pick ids: an over-range id must 400 (not 500/IndexError), a negative one must + # 400 (not silently index the wrong feature via torch negative-indexing). + over = client.post( + "/api/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [fake_engine.n_features]} + ) + assert over.status_code == 400 + neg = client.post("/api/annotate", json={"sequence": "ACGT", "mode": "pick", "feature_ids": [-1]}) + assert neg.status_code == 400 + + +def test_annotate_rejects_invalid_mode(client): + assert client.post("/api/annotate", json={"sequence": "ACGT", "mode": "bogus"}).status_code == 400 + + +def test_annotate_clamps_k_into_range(client, fake_engine): + client.post("/api/annotate", json={"sequence": "ACGT", "k": 999}) + assert fake_engine.last_k == 64 # upper bound + client.post("/api/annotate", json={"sequence": "ACGT", "k": 0}) + assert fake_engine.last_k == 1 # lower bound + + +def test_generate_returns_sequence(client): + b = client.post("/api/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).json() + assert b["generation"]["sequence"] + + +def test_generate_rejects_out_of_range_feature(client): + r = client.post("/api/generate", json={"prompt": "ACGT", "features": [{"feature_id": 999}]}) + assert r.status_code == 400 # the wedge guard, surfaced to the client + + +def test_generate_rejects_too_long(client, fake_engine): + seq = "A" * (fake_engine.max_seq_len + 1) # exceeds the context budget -> 413 (parity w/ annotate) + assert client.post("/api/generate", json={"prompt": seq}).status_code == 413 + + +def test_generate_compare_baseline(client): + b = client.post( + "/api/generate", + json={"prompt": "ACGT", "features": [{"feature_id": 0, "strength": 5.0}], "compare_baseline": True}, + ).json() + assert b["baseline"]["sequence"] # unsteered comparison returned alongside the steered one + + +def test_rejects_oversized_body(monkeypatch, fake_engine): + monkeypatch.setenv("MAX_BODY_BYTES", "100") + with TestClient(build_app(fake_engine)) as c: + assert c.post("/api/annotate", json={"sequence": "ACGT" * 100}).status_code == 413 + + +def test_endpoints_503_until_ready(fake_engine): + fake_engine.ready = False + fake_engine.load = lambda: None # startup leaves it not-ready + with TestClient(build_app(fake_engine)) as c: + assert c.get("/api/health").status_code == 503 # readiness probe sheds the pod + assert c.get("/api/features").status_code == 503 + assert c.post("/api/annotate", json={"sequence": "ACGT"}).status_code == 503 + assert c.post("/api/generate", json={"prompt": "ACGT", "organism": "None (raw DNA)"}).status_code == 503 + + +# ---------------------------------------------------------- single-container: SPA at / + API at /api +def test_serves_spa_index(static_client): + """With a built frontend mounted, GET / returns the SPA index (HTML), not an API response.""" + r = static_client.get("/") + assert r.status_code == 200 + assert "text/html" in r.headers["content-type"] + assert 'id="root"' in r.text + + +def test_serves_static_asset(static_client): + """Frontend assets are served from the mount (so the SPA's JS/CSS load same-origin).""" + r = static_client.get("/assets/app.js") + assert r.status_code == 200 + assert "dashboard" in r.text + + +def test_api_reachable_under_prefix_with_frontend(static_client): + """The API stays fully reachable under /api even with the SPA mounted at / (no shadowing).""" + b = static_client.get("/api/health").json() + assert b["ready"] is True and b["layer"] == 19 + + +def test_unknown_api_path_is_404_not_spa(static_client): + """An unknown /api/* path 404s — the /api namespace never falls through to the SPA index.""" + r = static_client.get("/api/does-not-exist") + assert r.status_code == 404 + assert 'id="root"' not in r.text + + +def test_api_only_when_no_frontend(fake_engine, monkeypatch): + """No static_dir and no DASHBOARD_DIST -> API-only: / 404s (nothing mounted), /api still works. + + Guards dev / a frontend-less image: a missing build degrades to API-only, never a crash. + """ + monkeypatch.delenv("DASHBOARD_DIST", raising=False) + with TestClient(build_app(fake_engine)) as c: + assert c.get("/").status_code == 404 + assert c.get("/api/health").status_code == 200 diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steering.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steering.py new file mode 100644 index 0000000000..c934b9890a --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steering.py @@ -0,0 +1,237 @@ +# 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. + +"""GPU inference tests: SAE encode + feature steering. + +`test_clamp_math` (CPU) checks the steering arithmetic the forward hook applies; it needs no +model and runs in CI. The GPU tests (slow, checkpoint-gated) drive the real engine: +`test_encode_smoke` guards the bf16 encode forward, and `test_unsteered_is_dna` / +`test_steering_changes_continuation` drive `generate()` through the recipe inference engine +and assert steering (on a discovered active feature) changes the continuation only. Gated by +EVO2_CKPT_DIR + SAE_CKPT_PATH. +""" + +import math + +import pytest +import torch +from evo2_sae import Evo2SAE +from evo2_sae.core import MAX_CLAMP_STRENGTH, _is_unrecoverable_cuda, _sanitize_steering + + +# The delta-clamp math + decode-only/prefill behavior is covered against the production +# hook in tests/test_clamp.py (evo2_sae.steering.clamp_hook). This file covers the +# engine-level input guards (CPU) and the live steering smoke tests (GPU). + + +# --------------------------------------------------- CPU: steering input guards (no model) +# Each guards against a CUDA device-side assert that would corrupt the context and wedge the +# server: out-of-range id (indexes off the SAE codes), extreme clamp (NaN logits), and +# temperature 0 (sampler divides logits by temperature -> NaN under multinomial). +def test_sanitize_rejects_out_of_range_id(): + """feature_id outside [0, n_features) -> ValueError (server maps to 400, not a wedge).""" + with pytest.raises(ValueError): + _sanitize_steering([{"feature_id": 70000, "strength": 1.0}], n_features=65536, temperature=1.0, top_k=0) + with pytest.raises(ValueError): + _sanitize_steering([{"feature_id": -1, "strength": 1.0}], n_features=65536, temperature=1.0, top_k=0) + + +def test_sanitize_caps_clamp_magnitude(): + """An extreme target is capped to ±MAX_CLAMP_STRENGTH (prevents NaN logits).""" + clamps, _, _, _ = _sanitize_steering([{"feature_id": 5, "strength": 99999.0}], 65536, 1.0, 0) + assert clamps[5] == MAX_CLAMP_STRENGTH + clamps, _, _, _ = _sanitize_steering([{"feature_id": 5, "strength": -99999.0}], 65536, 1.0, 0) + assert clamps[5] == -MAX_CLAMP_STRENGTH + + +def test_sanitize_temperature_zero_forces_greedy_topk(): + """temperature<=0 (divide-by-zero in the sampler) is coerced to greedy top-1.""" + _, _, _, top_k = _sanitize_steering([{"feature_id": 5, "strength": 1.0}], 65536, 0.0, 0) + assert top_k == 1 # bumped from 0 so the logits/temperature path is skipped + + +def test_sanitize_passthrough(): + """Valid inputs pass through unchanged — no spurious capping or top_k change.""" + clamps, fids, temp, top_k = _sanitize_steering([{"feature_id": 5, "strength": 2.0}], 65536, 0.8, 0) + assert clamps == {5: 2.0} and fids == [5] and temp == 0.8 and top_k == 0 + + +def test_sanitize_neutralizes_nonfinite_strength(): + """NaN / ±inf are capped to a finite magnitude <= MAX_CLAMP_STRENGTH, never propagated. + + A non-finite clamp would blow the logits to NaN (a CUDA device-assert that wedges the + server) — exactly what this guard exists to prevent. + """ + for bad in (math.nan, math.inf, -math.inf): + clamps, _, _, _ = _sanitize_steering([{"feature_id": 3, "strength": bad}], 65536, 1.0, 0) + assert math.isfinite(clamps[3]) and abs(clamps[3]) <= MAX_CLAMP_STRENGTH + + +def test_sanitize_feature_contract(): + """Duplicate ids collapse (last wins), missing strength defaults to 1.0, strength 0 is kept.""" + clamps, fids, _, _ = _sanitize_steering( + [ + {"feature_id": 5, "strength": 1.0}, + {"feature_id": 5, "strength": 2.0}, # duplicate id + {"feature_id": 7}, # missing strength + {"feature_id": 9, "strength": 0.0}, # suppress-to-zero + ], + 65536, + 1.0, + 0, + ) + assert clamps[5] == 2.0 # duplicate id -> last value wins + assert clamps[7] == 1.0 # missing strength -> default 1.0 + assert clamps[9] == 0.0 # strength 0 retained, not dropped + assert fids == [5, 7, 9] # one entry per distinct feature, in insertion order + + +def test_sanitize_clamps_negative_top_k(): + """A negative top_k is an invalid sampler arg -> coerced to 0 (no top-k filtering).""" + _, _, _, top_k = _sanitize_steering([{"feature_id": 5, "strength": 1.0}], 65536, 1.0, -5) + assert top_k == 0 + + +def test_is_unrecoverable_cuda(): + """Only sticky CUDA faults flip the engine not-ready; ordinary errors propagate normally.""" + assert _is_unrecoverable_cuda(RuntimeError("CUDA error: device-side assert triggered")) + assert not _is_unrecoverable_cuda(RuntimeError("shape mismatch")) + assert not _is_unrecoverable_cuda(ValueError("bad input")) + + +# --------------------------------------------------------------------- GPU: real generation +# These load the real Evo2 model: skip unless a GPU is present. (The conftest fixtures further skip +# on too-little GPU memory or an unfetchable/unimportable checkpoint.) +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU (loads Evo2 + SAE)") + +_PROMPT = "ATGGCCGAATTCGGCACGAGGACGTGCTGAAAGCTAGCTAGGCTAACCGGTTACGTGCAT" +_ORG = "Human" + + +@pytest.fixture(scope="module") +def engine(evo2_ckpt_dir, sae_ckpt_path, embedding_layer): + """Load the Evo2 + SAE engine once (skips unless CUDA is available). + + ``evo2_ckpt_dir`` / ``sae_ckpt_path`` come from conftest: they honor EVO2_CKPT_DIR / + SAE_CKPT_PATH for manual or 7B runs, and otherwise build the 1B-8k-bf16 + a synthesized + tiny SAE in CI (memory-gated). Those fixtures skip if no checkpoint can be produced. + """ + if not torch.cuda.is_available(): + pytest.skip("steering tests require CUDA") + return Evo2SAE(evo2_ckpt_dir=evo2_ckpt_dir, sae_ckpt_path=sae_ckpt_path, layer=embedding_layer).load() + + +def _gen(engine, features): + torch.manual_seed(0) + return engine.generate(prompt=_PROMPT, organism=_ORG, features=features, n_tokens=48, temperature=0.0, top_k=1) + + +def _tag(engine): + return engine.resolve_tag(_ORG, None) or "" + + +@requires_gpu +def test_encode_smoke(engine): + """encode runs the engine's bf16 forward and returns finite per-feature codes (>=1 firing). + + Guards the TransformerEngine bf16/fp32 autocast path: a dtype mismatch would crash here. + """ + codes = engine.encode(_tag(engine) + _PROMPT) + assert codes.ndim == 2 and codes.shape[1] == engine.n_features + assert torch.isfinite(codes).all() + assert (codes > 0).any() + + +@requires_gpu +def test_unsteered_is_dna(engine): + """Unsteered generation yields a non-empty ACGT string (Evo2 stays in-distribution).""" + seq = _gen(engine, [])["generation"]["sequence"] + assert seq and set(seq) <= set("ACGTN") + + +@requires_gpu +def test_steering_changes_continuation(engine): + """Clamping a KNOWN-ACTIVE feature hard changes the continuation; empty clamp is a no-op. + + Discovers the most-active feature on the prompt (SAE-agnostic) so the clamp has real signal — + an arbitrary/dead feature would leave greedy decoding unchanged and make this test useless. + Also exercises ``compare_baseline=True`` (returns the steered generation + the unsteered baseline + in one call). + """ + per = engine.encode(_tag(engine) + _PROMPT).max(dim=0).values + fid, peak = int(per.argmax()), float(per.max()) + feature = {"feature_id": fid, "strength": max(peak * 3.0, 50.0)} + base = _gen(engine, [])["generation"]["sequence"] + steered = _gen(engine, [feature])["generation"]["sequence"] + assert steered != base # the clamp on an active feature changed the continuation + assert _gen(engine, [])["generation"]["sequence"] == base # determinism + empty-clamp no-op + + # compare_baseline=True returns both halves in one call. + both = engine.generate( + prompt=_PROMPT, organism=_ORG, features=[feature], n_tokens=48, temperature=0.0, top_k=1, compare_baseline=True + ) + assert both["generation"]["sequence"] == steered # steered half matches + assert both["baseline"]["sequence"] == base # baseline half == unsteered + + +@requires_gpu +def test_highlight_steer_interleaving_no_bleed(engine): + """Single engine: encode (highlight) and generate (steer) drive ONE shared model, so + interleaving them must not corrupt either path. + + (1) encode is bit-identical before and after a steered generate runs in between — the + full-sequence highlight forward is unaffected by the decode dynamic-context / KV state + the generate path builds and tears down on the same model. + (2) a baseline (unsteered) generate is identical whether or not an encode + a steered + generate ran first — no steering hook leaked and no decode state carried over. + """ + full = _tag(engine) + _PROMPT + codes_before = engine.encode(full) + per = codes_before.max(dim=0).values + fid, peak = int(per.argmax()), float(per.max()) + feature = {"feature_id": fid, "strength": max(peak * 3.0, 50.0)} + + base_clean = _gen(engine, [])["generation"]["sequence"] # reference baseline + _ = _gen(engine, [feature]) # steered generate between the two encodes + codes_after = engine.encode(full) + assert torch.equal(codes_before, codes_after) # (1) highlight forward unchanged by the steer + + _ = engine.encode(full) # more interleaving: encode then steered-gen before a baseline + _ = _gen(engine, [feature]) + base_after = _gen(engine, [])["generation"]["sequence"] + assert base_after == base_clean # (2) baseline unaffected by prior encode/steer history + + +@requires_gpu +def test_encode_batch_preserves_order_and_lengths(engine): + """Batched encode keeps input order + per-sequence length (incl. an empty sequence), with + finite codes (real model, no mock).""" + seqs = ["", _tag(engine) + _PROMPT[:20], _tag(engine) + _PROMPT, _tag(engine) + _PROMPT[:40]] + batch = engine.encode_batch(seqs) + assert [c.shape[0] for c in batch] == [len(engine.tokenize(s)) for s in seqs] # order + length (empty -> 0) + assert batch[0].shape == (0, engine.n_features) # empty sequence -> empty codes, correct width + assert all(c.shape[1] == engine.n_features and torch.isfinite(c).all() for c in batch) + + +@requires_gpu +def test_max_clamp_strength_stays_in_distribution(engine): + """Clamping at MAX_CLAMP_STRENGTH still produces valid DNA — proves the cap is a *safe* level. + + The CPU sanitize test only checks the value is capped; this proves the capped magnitude + doesn't blow the logits to NaN (a CUDA device-assert that would wedge the server). + """ + fid = int(engine.encode(_tag(engine) + _PROMPT).max(dim=0).values.argmax()) + seq = _gen(engine, [{"feature_id": fid, "strength": MAX_CLAMP_STRENGTH}])["generation"]["sequence"] + assert seq and set(seq) <= set("ACGTN")