feat(evo2-sae): streaming smoke test (producer/consumer) + Lepton convergence config by polinabinder1 · Pull Request #1668 · NVIDIA-BioNeMo/bionemo-recipes

polinabinder1 · 2026-06-30T03:06:47Z

Evo2 analogue of the ESM2 SAE smoke (#1663). Adds a producer/consumer streaming path — Megatron predict forward passes feed activations into a bounded queue (sae.streaming); the SAE Trainer consumes them in-process, no activation store on disk. Self-contained like the ESM2 smoke: the 1B checkpoint is fetched by identifier and converted in-job (nothing pre-staged).

2 GPUs (producer/consumer, not data-parallel): when a second GPU is visible, the Evo2 producer runs its forwards on cuda:0 and the SAE consumer trains on cuda:1, so the two stages run concurrently instead of contending for one device (num_devices: 2, resource_shape → gpu.2xh100-sxm). It stays one predict rank — --nproc_per_node 1, --dp-size 1; activations cross the boundary as CPU tensors via the queue, so no NCCL/DDP.

Files

scripts/train_streaming.py — Evo2ActivationProducer bridges predict's per-batch callback to a pull-generator (daemon thread + queue); SAE trains off the stream. Auto-places the SAE consumer on a 2nd GPU when available; wraps the predict forward in bf16 autocast (see Robustness).
scripts/prepare_1b_checkpoint.py — bionemo_load("evo2/1b-8k-bf16") → run_nemo2_to_mbridge (idempotent). Honors EVO2_CKPT_DIR.
ci/lepton/model_convergence/configs/recipes/evo2_sae_smoke.yaml + convergence-tests.yml — 2×H100 convergence config, plus the evo2_sae_smoke dispatch option.

Review feedback addressed

Deduped against /base — dropped container (identical 26.02 image), node_group, mount_from.
Dropped TE-recipe markers (te_enabled/fp8_enabled/extras) — the run_script never references them (verified the launcher doesn't either).
Bumped to 2 GPUs — via the producer/consumer split above (gpu.2xh100-sxm).
No Docker — checkout_script runs evo2_megatron/.ci_build.sh and loads an MBridge checkpoint; nothing uses Docker.

Robustness

--input-dim is validated against the layer's true residual width on the first streamed chunk — a mismatch fails with a clear message instead of an opaque matmul error inside the SAE encoder.
bf16 autocast around the predict forward — Evo2's params are bf16 and predict's minimal-arg path sets no autocast region, so a fp32_residual_connection=true checkpoint (e.g. a 7B EVO2_CKPT_DIR override) would otherwise trip a TransformerEngine param/input dtype assertion. Harmless no-op on the default 1b-8k-bf16 (bf16-residual) path.
--init-pre-bias is scoped out of the streaming path and fails fast pointing at --no-init-pre-bias (it would require a second Evo2 predict pass, which doesn't compose with Megatron's process-global init). A single-pass implementation is a natural follow-up.

Validation

Streaming train (Evo2-1B / layer 12): NGC pull + MBridge convert, then 20 steps → loss ↓ to ~0.94, dead_latents 0%, checkpoint_final.pt written.
2-GPU split validated: on 2×H100 the run places producer→cuda:0 / consumer→cuda:1 and trains to a loss identical to the single-GPU run to 4 decimals — the split changes device placement, not the math.
CPU unit tests (tests/test_streaming.py, fake predict module — no GPU/Megatron): queue semantics (backpressure, sentinel, exception propagation), the --input-dim mismatch, and the --init-pre-bias rejection.
Lepton (end-to-end, 2×H100): dispatched to vfco61g2 / yo-bom-lepton-001, gpu.2xh100-sxm; built via .ci_build.sh, ran Evo2 predict on GPU, streamed-trained 20 steps → SMOKE OK (checkpoint_final.pt written).

Launching on Lepton

Runs on the existing convergence system (convergence-tests.yml → launch_job.py), not the #1667 L4 unit-test lane. This PR adds evo2_sae_smoke to the workflow_dispatch model_config options.

Dispatch: Actions → BioNeMo Model Convergence Tests → Run workflow → model_config = evo2_sae_smoke, gpu_type = h100-sxm. Success = checkpoint_final.pt written. Dispatch-only (a 20-step liveness smoke, telemetry off) — deliberately not on the biweekly cron.

Ordering

Stacked on #1622 (pbinder/evo2-sae-serve), which adds the recipe's .ci_build.sh. Also uses chunk_fasta.py, already on main. Review/merge after #1622. Do not fold into #1667 (the presence-guarded unit-test lane, which must merge before the recipe).

🤖 Generated with Claude Code

…vergence config Evo2 analogue of the ESM2 SAE smoke (#1663): Megatron predict forward passes feed activations into a bounded queue (sae.streaming); the SAE Trainer consumes in-process, no activation store on disk. Self-contained like the ESM2 smoke - the 1B checkpoint is fetched by identifier and converted in-job. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

copy-pr-bot · 2026-06-30T03:06:51Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-06-30T03:06:53Z

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: eeff36b3-55a4-48c0-a71a-5201bfe3a631

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch pbinder/evo2-sae-streaming-smoke

_{Comment @coderabbitai help to get the list of available commands.}

…ut-dim GPU-validated on Evo2-1B (1-GPU H100); baseline smoke unchanged (loss 0.9407). #1 --init-pre-bias hard-crashed: it ran a *second* Evo2 `predict` pass to sample activations, but `predict` leaves Megatron global state initialized with no teardown, so the second pass died with "num microbatches calculator is already initialized". Remove the broken second-pass init and fail fast with a clear NotImplementedError pointing at --no-init-pre-bias. Single-pass init (sample the first rows off the one training stream) tracked as a follow-up. #4 Guard --input-dim against the layer's true residual width on the first streamed chunk, so a mismatch fails clearly instead of as an opaque matmul error in the encoder. Add tests/test_streaming.py: CPU guards (fake predict module, no GPU/Megatron) for the --input-dim mismatch, happy-path streaming, predict-failure propagation, and the --init-pre-bias rejection. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

copy-pr-bot · 2026-06-30T17:39:14Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Adds `evo2_sae_smoke` to the workflow_dispatch model_config options so the Lepton 1B streaming smoke (config already in this PR at ci/lepton/.../recipes/evo2_sae_smoke.yaml) can be launched manually. NOT added to the biweekly schedule matrix: it's a liveness smoke (20 steps, log_to_wandb/log_to_kratos=false), so it collects no convergence metrics and doesn't belong on the convergence cron. (Because kratos is off, the convergence_tests.sh wrapper isn't invoked, so the wandb_dir wrapper bug is moot here — it'd only matter for a real telemetry-on convergence run.) Config composition verified locally (hydra compose over base.yaml resolves). The 1B streaming run itself is validated on H100 (see PR description); the actual Lepton submission is exercised by a first workflow_dispatch post-merge. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: polinabinder1 <pbinder@nvidia.com>

Signed-off-by: polinabinder1 <pbinder@nvidia.com>

…12.4-driver nodes) The workspace's convergence nodes (e.g. yo-bom-lepton-001) run an NVIDIA driver on CUDA 12.4. The evo2 GPU-lane image (svcbionemo…pytorch26.04-py3-squashed) needs driver >= 12.8, so torch.cuda.is_available() is False there and Evo2 predict aborts with "Inference requires CUDA" (predict.py:236) — the first dispatch failed exactly this way after building + converting the checkpoint. Point the smoke at the base NGC image (nvcr.io/nvidia/pytorch:26.02-py3, what the other convergence recipes use), whose cuda-compat libs run on the 12.4 driver. Megatron/TE are added in-job by .ci_build.sh, so the base image suffices. This is also the only stack the recipe is validated on (system torch is nv26.02). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: polinabinder1 <pbinder@nvidia.com>

polinabinder1 · 2026-07-02T22:27:41Z

@@ -0,0 +1,208 @@
+# @package _global_
+defaults:


Make sure not redundant to base

Done — removed container (identical 26.02 image), node_group, and mount_from; all inherited from /base.

jwilber · 2026-07-02T22:30:35Z

+# on those nodes with "Inference requires CUDA". Megatron/TransformerEngine are
+# added in-job by .ci_build.sh (checkout_script), so this base image is enough.
+############################################################
+container:


these are in the base config so can remove

Done — removed; inherited from /base.

jwilber · 2026-07-02T22:31:27Z

+device_type: gpu
+num_devices: 1
+gpu_type: h100-sxm
+resource_shape: gpu.h100-sxm


bump to 2 gpus: (resource shape needs to change for the)

Done — num_devices: 2 + templated resource_shape (→ gpu.2xh100-sxm). Implemented as a producer/consumer split rather than data-parallel: the Evo2 producer runs on cuda:0 and the SAE consumer trains on cuda:1 (activations cross as CPU tensors via the queue), so it stays one predict rank / --dp-size 1. Validated on Lepton — the job ran on gpu.2xh100-sxm and hit SMOKE OK, loss identical to the 1-GPU run.

jwilber · 2026-07-02T22:32:16Z

+framework: sae
+precision: bf16
+te_enabled: true
+fp8_enabled: false


we don't need these markers for these tests (these are for TE recupes)

Done — removed te_enabled/fp8_enabled/extras; the run_script never references them and the launcher does not either, so it is launcher-safe.

jwilber · 2026-07-02T22:33:44Z

+  fi
+  # Builds evo2_megatron's venv (bionemo.evo2 / predict) + installs sae + evo2_sae on top.
+  cd interpretability/sparse_autoencoders/recipes/evo2
+  bash .ci_build.sh


Docker in Lepton must run on a privileged node (make sure we have that set)

It looks like instead we actually are loading a checkpoint, so we don't need docker:

Successfully installed evo2-sae-0.1.0 fastapi-0.139.0 starlette-1.3.1 uvicorn-0.49.0 WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning. Reusing existing MBridge checkpoint: /data/interpretability/sae/evo2_smoke/outputs/evo2_1b_mbridge/iter_0000001 Chunked 4 sequences -> 4 chunks (16,384 bp) at window=8192 Using device: cuda SAE: topk, input_dim=1920, hidden_dim=15360

(from lepton logs: https://dashboard.dgxc-lepton.nvidia.com/workspace/vfco61g2/compute/jobs/detail/evo2-sae-1b-smoke-59w4/replicas/list#/job/evo2-sae-1b-smoke-59w4/logs
)

It looks like the ci build actually comes from this: https://github.com/NVIDIA-BioNeMo/bionemo-recipes/blob/main/recipes/evo2_megatron/.ci_build.sh

Thus, if you're able to load a mbridge checkpoint and run evo2 inference (like what seems to be done here), maybe your other scripts don't need Docker to run? We don't use Docker here, can you double-check.

Confirmed — no Docker. checkout_script runs recipes/evo2_megatron/.ci_build.sh and the run loads an MBridge checkpoint (the log you pasted is that path). Separately hardened the forward with a bf16 autocast wrap so an fp32_residual_connection=true ckpt override does not trip the TE dtype assertion.

jstjohn · 2026-07-02T23:38:21Z

/ok to test aabc352

jwilber · 2026-07-03T00:20:50Z

Run logs:
This pr edits this gh action script: https://github.com/NVIDIA-BioNeMo/bionemo-recipes/blob/aabc352953793ea7470ae385cb04e083f6f1ea54/.github/workflows/convergence-tests.yml

Which triggers runs on Lepton. See example run: https://dashboard.dgxc-lepton.nvidia.com/workspace/vfco61g2/compute/jobs/detail/evo2-sae-1b-smoke-59w4/replicas/list#/job/evo2-sae-1b-smoke-59w4/logs

…eaming Lepton config (evo2_sae_smoke.yaml): - Dedupe against /base: drop container (identical 26.02 image), node_group, mount_from. - Drop te_enabled/fp8_enabled/extras (TE-recipe markers; the run_script never uses them). - num_devices: 2 + templated resource_shape (-> gpu.2xh100-sxm), matching sibling recipes. train_streaming.py: - 2-GPU as a producer/consumer pipeline (NOT data-parallel): the Evo2 producer runs on cuda:0 and the SAE consumer trains on cuda:1 when a 2nd GPU is visible, so the stages stop contending for one device. --dp-size stays 1; --device overrides. - Wrap predict.main() in bf16 autocast: robustness for fp32_residual_connection=true checkpoints (e.g. a 7B EVO2_CKPT_DIR override) that otherwise trip a TransformerEngine param/input dtype assertion. No-op for the default 1b-8k-bf16 (bf16-residual) path. - Add two missing docstrings (ruff D107/D102). Validated on 2xH100 with the real 1b-8k-bf16 checkpoint: producer cuda:0 / consumer cuda:1, trains, writes checkpoint_final.pt (exit 0). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

polinabinder1 and others added 3 commits July 1, 2026 21:08

Merge remote-tracking branch 'origin/pbinder/evo2-sae-serve' into HEAD

6754ea1

Signed-off-by: polinabinder1 <pbinder@nvidia.com>

polinabinder1 marked this pull request as ready for review July 2, 2026 19:42

polinabinder1 requested review from jstjohn, jwilber, pstjohn, savitha-eng and trvachov as code owners July 2, 2026 19:42

polinabinder1 commented Jul 2, 2026

View reviewed changes

jwilber reviewed Jul 2, 2026

View reviewed changes

Uh oh!

Conversation

polinabinder1 commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Files

Review feedback addressed

Robustness

Validation

Launching on Lepton

Ordering

Uh oh!

copy-pr-bot Bot commented Jun 30, 2026

Uh oh!

coderabbitai Bot commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

copy-pr-bot Bot commented Jun 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jstjohn commented Jul 2, 2026

Uh oh!

jwilber commented Jul 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

polinabinder1 commented Jun 30, 2026 •

edited

Loading

coderabbitai Bot commented Jun 30, 2026 •

edited

Loading