evo2 SAE: inference engine + steering, CI lane, and tests (on migrated layout)

Re-lands #1622 onto the post-#1633 top-level layout (interpretability/sparse_autoencoders/).
A prior `merge main` had orphaned all of #1622's added files at the old bionemo-recipes/
path while the rest of the tree migrated; this replays the net delta at the correct paths.

- Evo2SAE inference engine + steering hook (src/evo2_sae/, sae/src/sae/steering.py)
- dedicated path-gated GPU CI lane + .ci_build.sh (reuses evo2_megatron's build, no fork)
  + .ci_test_env.sh
- conftest 1B-load fixture (bionemo_load -> run_nemo2_to_mbridge + synth tiny SAE, memory-gated)
- engine unit tests (test_core.py, CPU) + steering/encode GPU tests (test_steering.py)

Validated on the 1B: test_core 7 passed (CPU), test_steering 9 passed (GPU).

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

Author: polinabinder1

.github/workflows/unit-tests-interpretability-recipes.yaml

@@ -0,0 +1,77 @@
+name: "BioNeMo Interpretability Recipes CI"
+
+# CI for the SAE interpretability recipes (interpretability/sparse_autoencoders).
+# Uses the same squashed megatron image + L4 GPU as the mbridge lane, but the interpretability
+# tree is a single uv workspace (not N *megatron recipes), so there is no per-recipe matrix —
+# just a path-gated GPU job. Currently scoped to the evo2 recipe, whose engine builds on the
+# mbridge bionemo.evo2 recipe, so the job also checks out recipes/evo2_megatron (which pulls its
+# bionemo-core / bionemo-recipeutils deps from git).
+
+on:
+  push:
+    branches:
+      - "pull-request/[0-9]+"
+      - "dependabot/**"
+    paths:
+      - "interpretability/sparse_autoencoders/**"
+      - "recipes/evo2_megatron/**" # engine builds on bionemo.evo2
+      - ".github/workflows/unit-tests-interpretability-recipes.yaml"
+  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:
+  unit-tests:
+    runs-on: linux-amd64-gpu-l4-latest-1
+    name: "interpretability-unit-tests (evo2)"
+    container:
+      image: svcbionemo023/bionemo-framework:pytorch26.04-py3-squashed
+      options: --shm-size=16G
+      env:
+        CI: true
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        HF_HOME: /cache/huggingface
+
+    steps:
+      - name: Show GPU info
+        run: nvidia-smi
+      - name: Setup proxy cache
+        uses: nv-gha-runners/setup-proxy-cache@main
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          # The evo2 recipe builds on bionemo.evo2 (recipes/evo2_megatron); evo2_megatron pulls
+          # its bionemo-core / bionemo-recipeutils deps from git, so only these two are needed.
+          sparse-checkout: |
+            interpretability/sparse_autoencoders
+            recipes/evo2_megatron
+          sparse-checkout-cone-mode: false
+
+      - name: Cache Hugging Face models
+        uses: actions/cache@v4
+        with:
+          path: /cache/huggingface
+          key: ${{ runner.os }}-huggingface-interp-evo2-${{ github.sha }}
+          restore-keys: |
+            ${{ runner.os }}-huggingface-interp-evo2-
+            ${{ runner.os }}-huggingface-
+
+      - name: Install dependencies
+        working-directory: interpretability/sparse_autoencoders/recipes/evo2
+        run: bash .ci_build.sh
+
+      - name: Run tests
+        working-directory: interpretability/sparse_autoencoders/recipes/evo2
+        run: |
+          source .ci_test_env.sh
+          pytest -v tests/

interpretability/sparse_autoencoders/recipes/evo2/.ci_build.sh

@@ -0,0 +1,36 @@
+#!/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.
+set -euo pipefail
+
+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.
+( cd "$MEGATRON" && bash .ci_build.sh )
+
+# 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).
+source "$MEGATRON/.venv/bin/activate"
+PIP_CONSTRAINT= pip install -e "$SAE_LIB"
+PIP_CONSTRAINT= pip install -e "$HERE"

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

interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml

@@ -15,11 +15,11 @@ dependencies = [
     "pyarrow>=23.0.0",
 ]
 
-# No package code lives here yet — the recipe is just an entry-point for
-# scripts/ that depends on the shared `sae` workspace package. Declare no
-# packages so setuptools doesn't try to discover anything.
-[tool.setuptools]
-packages = []
+# The `evo2_sae` package (src/) holds the live inference engine + steering hook;
+# scripts/ (extract, train) are standalone entry points alongside it. The FastAPI
+# server + CLI (and their fastapi/uvicorn deps) are added by the serve PR (#1637).
+[tool.setuptools.packages.find]
+where = ["src"]
 
 [tool.uv.sources]
 sae = { workspace = true }

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}")