evo2 SAE recipe: streaming extract.py + env-overridable sweep runner

[polinabinder1]

Read this before reviewing

This PR is stacked on top of #1579 (evo2-sae-recipe). Both branches live on the same fork, so GitHub can't show a base-on-fork comparison — the diff below includes #1579's commits at the bottom. The new work here is the top three commits only:

Commit	Scope
a5b1db7 evo2 sae recipe: streaming-extract pipeline, drop pt_to_parquet shim	1b.sh rewrite, pt_to_parquet.py deletion, README update
9b7856a evo2 sae: streaming extractor + prok+euk FASTA composer	extract.py (+206) and compose_prokeuk_fasta.py (+142), both new
89ff40c evo2_megatron predict: skip empty batches on DP shard boundary	8-line bug fix in predict.py

Net diff for the new work: 6 files, +439 / −107. Once #1579 merges into main, this PR's diff auto-cleans to just these three commits.

What changes

The recipe's extraction pipeline collapses from two steps to one. Before:

predict_evo2  --embedding-layer N  (FASTA in, .pt out)
        |
pt_to_parquet shim  (.pt -> ActivationStore parquet shards)
        |
train.py

After:

extract.py  (FASTA in, ActivationStore parquet shards out)
        |
train.py

extract.py monkey-patches predict_evo2's writer so activations stream straight into the SAE ActivationStore parquet format during inference. No .pt intermediate, no shim. This is the pipeline used for the current 100M / 500M layer-22 prok+euk training runs — the old .pt path was never used by any real run, so pt_to_parquet.py is deleted.

Files in the new work

• scripts/extract.py (new, +206) — streaming activation extractor. Re-uses bionemo.evo2.run.predict for all the heavy machinery (Megatron model load, DP/CP/TP/PP, FASTA dataloader, inference loop) but swaps the per-batch .pt writer for an in-process ActivationStore. Handles multi-rank merge with a file-based wait (dist.barrier() no-ops after predict.main() tears down the process group).
• scripts/compose_prokeuk_fasta.py (new, +142) — composes a prokaryotic + eukaryotic FASTA mix from OpenGenome2 shards with a hard token budget. Excludes metagenome sources; matches the training mix used in 100M / 500M runs.
• scripts/pt_to_parquet.py (deleted) — obsoleted by extract.py.
• scripts/1b.sh (rewritten) — 3-step pipeline (convert -> extract -> train) instead of 4. Env-overridable hyperparams (RUN_TAG, LAYER, MAX_TOKENS, MICRO_BATCH, DEVICES, EXPANSION_FACTOR, TOP_K, AUXK, AUXK_COEF, DEAD_TOKENS_THRESHOLD, N_EPOCHS, LR) so the same script drives a multi-config sweep. TRAIN_ONLY=1 skips extraction against a cached parquet. WANDB_API_KEY gates wandb.
• evo2_megatron/.../predict.py (+8) — skip empty batches on DP shard boundaries. Fixes a hang where the last DP rank could receive 0 sequences and stall predict_evo2.
• README.md — updated pipeline diagram and quick-start examples for the new flow.

How to use

# Default 1B prok+euk run on layer 12, 4 GPU
bash scripts/1b.sh

# Layer 22, custom FASTA, tagged paths, more epochs
LAYER=22 \
  FASTA=/data/interp/evo2/OpenGenome2/fasta/prokeuk_25M.fasta \
  RUN_TAG=25M_prokeuk \
  N_EPOCHS=4 \
  bash scripts/1b.sh

# Re-train against an existing parquet dir without re-extracting
TRAIN_ONLY=1 \
  PARQUET_DIR=/data/interp/evo2/activations/evo2_1b_base_layer22_parquet_25M_prokeuk \
  AUXK=2048 RUN_TAG=auxk2048 \
  bash scripts/1b.sh

Compose a custom mix first if needed:

python scripts/compose_prokeuk_fasta.py \
  --out /data/interp/evo2/OpenGenome2/fasta/prokeuk_25M.fasta \
  --target-tokens 25_000_000

Test plan

• bash scripts/1b.sh end-to-end on a small FASTA: convert step skips if MBridge ckpt exists, extract writes parquet shards + metadata.json, train completes and writes checkpoint_final.pt.
• compose_prokeuk_fasta.py --target-tokens 1_000_000 writes a FASTA close to the target (±genome-record granularity) and excludes metagenome sources.
• TRAIN_ONLY=1 against an existing parquet skips chunk/convert/extract and runs train directly.
• predict_evo2 on a sequence count not divisible by DP world size completes instead of hanging.
• All env overrides (LAYER, RUN_TAG, AUXK, N_EPOCHS, etc.) surface in the printed step banners and the underlying CLI flags.

🤖 Generated with Claude Code

[copy-pr-bot]

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.

[coderabbitai]

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: Pro Plus

Run ID: c96d3048-8b9b-479a-81fb-4b1cb4731915

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

---

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