feat(evo2-sae): streaming smoke test (producer/consumer) + Lepton convergence config #1668
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,149 @@ | ||
| 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). | ||
| # Use `if` (not `[ -f ] && echo`): the short-circuit leaves exit 1 when the last dir | ||
| # lacks a .ci_build.sh, which `set -e` would turn into a job failure instead of an | ||
| # empty (green no-op) matrix. | ||
| ALL=$(for d in "$RROOT"/*/; do if [ -f "${d}.ci_build.sh" ]; then echo "${d%/}"; fi; 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/ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,194 @@ | ||
| # @package _global_ | ||
| defaults: | ||
| - /base | ||
| - _self_ | ||
|
|
||
| ############################################################ | ||
| # lepton job info (node_group + mount_from are inherited from /base) | ||
| ############################################################ | ||
| num_nodes: 1 | ||
| device_type: gpu | ||
| # Two GPUs as a producer/consumer pipeline (NOT data parallel): the Evo2 producer runs | ||
| # forwards on cuda:0 and the SAE consumer trains on cuda:1 (train_streaming.py auto-places | ||
| # the consumer on the 2nd GPU). Still one predict rank: --nproc_per_node 1, --dp-size 1. | ||
| num_devices: 2 | ||
| gpu_type: h100-sxm | ||
| resource_shape: "${device_type}.${num_devices}x${gpu_type}" | ||
|
|
||
| job_name: "evo2-sae-1b-smoke" | ||
|
|
||
| ############################################################ | ||
| # Telemetry | ||
| ############################################################ | ||
| # Disabled by default because this is an example smoke run. Set both | ||
| # log_to_wandb=true and train_wandb_enabled=true to collect W&B metrics. | ||
| log_to_wandb: false | ||
| log_to_kratos: false | ||
| kratos_subject: "interpretability_sae_examples_v0.0.1" | ||
| wandb_dir: /workspace/bionemo-framework/interpretability/sparse_autoencoders/recipes/evo2/wandb | ||
|
|
||
| ############################################################ | ||
| # recipe identifiers | ||
| ############################################################ | ||
| recipe_subdir: interpretability/sparse_autoencoders/recipes/evo2 | ||
| model_type: evo2_sae | ||
| variant: smoke | ||
| framework: sae | ||
| precision: bf16 | ||
|
|
||
| total_gpus: ${multiply:${num_devices},${num_nodes}} | ||
|
|
||
| wandb_init_args: | ||
| project: "interpretability_sae_examples__${sanitize:${branch}}" | ||
| group: "${model_type}__${variant}__${total_gpus}gpus__${sanitize:${gpu_type}}" | ||
| job_type: "${recipe_subdir}" | ||
| name: null | ||
| entity: clara-discovery | ||
| mode: online | ||
|
|
||
| ############################################################ | ||
| # Evo2 SAE streaming smoke settings | ||
| ############################################################ | ||
| run_name: lepton_evo2_1b_sae_smoke | ||
|
|
||
| # Fetched by identifier and converted to MBridge in-job; nothing is pre-staged. | ||
| model_tag: evo2/1b-8k-bf16:1.0 | ||
| model_size: evo2_1b_base | ||
| seq_length: 8192 | ||
| bionemo_data_source: "" # set to "pbss" if NGC is unreachable from the job | ||
| ckpt_dir: ${output_base}/evo2_1b_mbridge | ||
| layer: 12 | ||
| input_dim: 1920 # 1B residual-stream width (Hyena1bModelProvider.hidden_size) | ||
| micro_batch_size: 4 | ||
| max_tokens: 100000 | ||
| chunk_bp: 8192 | ||
| extract_dtype: fp32 | ||
|
|
||
| dp_size: 1 | ||
| train_n_epochs: 1 | ||
| max_steps: 20 | ||
| train_batch_size: 256 | ||
| expansion_factor: 8 | ||
| top_k: 32 | ||
| auxk: 256 | ||
| learning_rate: 1e-4 | ||
| log_interval: 10 | ||
| checkpoint_steps: 999999 | ||
| init_pre_bias: false | ||
| train_wandb_enabled: false | ||
| queue_size: 4 | ||
| shuffle_buffer_size: 8192 | ||
|
|
||
| data_dir: /data/interpretability/sae/evo2_smoke/data | ||
| output_base: /data/interpretability/sae/evo2_smoke/outputs | ||
|
|
||
| products: | ||
| - config: evo2_1b_smoke | ||
|
|
||
| ############################################################ | ||
| # Checkout Script | ||
| ############################################################ | ||
| checkout_script: | | ||
| set -euo pipefail | ||
| cd /workspace | ||
| git clone https://github.com/NVIDIA/bionemo-framework.git | ||
| cd bionemo-framework | ||
| if [ -n "${commit_sha}" ]; then | ||
| echo "Checking out commit: ${commit_sha}" | ||
| git checkout "${commit_sha}" | ||
| elif [ "${branch}" != "main" ]; then | ||
| echo "Checking out branch: ${branch}" | ||
| git checkout "${branch}" | ||
| 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 | ||
|
Collaborator
There was a problem hiding this comment. Docker in Lepton must run on a privileged node (make sure we have that set)
Collaborator
There was a problem hiding this comment. It looks like instead we actually are loading a checkpoint, so we don't need docker:
Collaborator
There was a problem hiding this comment.
Collaborator
There was a problem hiding this comment. 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.
Collaborator
Author
There was a problem hiding this comment. Confirmed — no Docker. |
||
|
|
||
| ############################################################ | ||
| # run script | ||
| ############################################################ | ||
| run_script: | | ||
| cd /workspace/bionemo-framework/interpretability/sparse_autoencoders/recipes/evo2 | ||
| # The recipe shares evo2_megatron's venv (see .ci_build.sh / .ci_test_env.sh). | ||
| source ../../../../recipes/evo2_megatron/.venv/bin/activate | ||
|
|
||
| WANDB_FLAG="--no-wandb" | ||
| if [ "${train_wandb_enabled}" = "true" ]; then | ||
| WANDB_FLAG="--wandb" | ||
| fi | ||
|
|
||
| INIT_PRE_BIAS_FLAG="--no-init-pre-bias" | ||
| if [ "${init_pre_bias}" = "true" ]; then | ||
| INIT_PRE_BIAS_FLAG="--init-pre-bias" | ||
| fi | ||
|
|
||
| # Read EVO2_CKPT_DIR via printenv so Hydra does not treat it as an interpolation. | ||
| CKPT_DIR="$(printenv EVO2_CKPT_DIR || true)" | ||
| if [ -z "$CKPT_DIR" ]; then | ||
| if [ -n "${bionemo_data_source}" ]; then | ||
| export BIONEMO_DATA_SOURCE="${bionemo_data_source}" | ||
| fi | ||
| python scripts/prepare_1b_checkpoint.py \ | ||
| --mbridge-ckpt-dir ${ckpt_dir} \ | ||
| --model-tag ${model_tag} \ | ||
| --model-size ${model_size} \ | ||
| --seq-length ${seq_length} | ||
| CKPT_DIR="${ckpt_dir}" | ||
| fi | ||
|
|
||
| # Synthesize a tiny FASTA and chunk it to <= chunk_bp bp. | ||
| mkdir -p "${data_dir}" | ||
| RAW_FASTA="${data_dir}/smoke_raw.fasta" | ||
| CHUNKED_FASTA="${data_dir}/smoke_chunked${chunk_bp}.fasta" | ||
| python - "$RAW_FASTA" <<'PY' | ||
| import random, sys | ||
| random.seed(42) | ||
| with open(sys.argv[1], "w") as f: | ||
| for i in range(4): | ||
| seq = "".join(random.choice("ACGT") for _ in range(4096)) | ||
| f.write(f">smoke_{i}\n{seq}\n") | ||
| PY | ||
| python scripts/chunk_fasta.py --input "$RAW_FASTA" --output "$CHUNKED_FASTA" --window ${chunk_bp} | ||
|
|
||
| # Producer/consumer streaming: predict feeds activations into a queue; the SAE | ||
| # Trainer consumes them. Under torchrun (predict needs the distributed env). | ||
| torchrun --nproc_per_node 1 scripts/train_streaming.py \ | ||
| --ckpt-dir "$CKPT_DIR" \ | ||
| --fasta "$CHUNKED_FASTA" \ | ||
| --embedding-layer ${layer} \ | ||
| --input-dim ${input_dim} \ | ||
| --micro-batch-size ${micro_batch_size} \ | ||
| --max-tokens ${max_tokens} \ | ||
| --dtype ${extract_dtype} \ | ||
| --model-type topk \ | ||
| --expansion-factor ${expansion_factor} \ | ||
| --top-k ${top_k} \ | ||
| --normalize-input \ | ||
| --auxk ${auxk} \ | ||
| --auxk-coef 0.03125 \ | ||
| $INIT_PRE_BIAS_FLAG \ | ||
| --n-epochs ${train_n_epochs} \ | ||
| --max-steps ${max_steps} \ | ||
| --batch-size ${train_batch_size} \ | ||
| --lr ${learning_rate} \ | ||
| --log-interval ${log_interval} \ | ||
| --queue-size ${queue_size} \ | ||
| --shuffle-buffer-size ${shuffle_buffer_size} \ | ||
| $WANDB_FLAG \ | ||
| --wandb-project ${wandb_init_args.project} \ | ||
| --wandb-run-name ${run_name} \ | ||
| --wandb-group ${wandb_init_args.group} \ | ||
| --wandb-job-type ${wandb_init_args.job_type} \ | ||
| --dp-size ${dp_size} \ | ||
| --seed 42 \ | ||
| --output-dir ${output_base}/${run_name} \ | ||
| --checkpoint-dir ${output_base}/${run_name}/checkpoints \ | ||
| --checkpoint-steps ${checkpoint_steps} | ||
|
|
||
| # Fail loudly on a silent no-op (streaming produced nothing / training skipped). | ||
| CKPT="${output_base}/${run_name}/checkpoints/checkpoint_final.pt" | ||
| if [ ! -f "$CKPT" ]; then | ||
| echo "ERROR: expected SAE checkpoint not written: $CKPT" >&2 | ||
| exit 1 | ||
| fi | ||
| echo "SMOKE OK: $CKPT" | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Make sure not redundant to base
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done — removed
container(identical 26.02 image),node_group, andmount_from; all inherited from/base.