Repository navigation
bioreason-demo: PRISM-BioReason KEGG demo (SFT, GRPO, eval, notebook) - #4
Merged
Merged
Conversation
…elected) Squashed diff of BaseMM_PRISM 00c641a..e992faf (bioreason_sft_no_projector_stage). Uncommitted notebook/narration/diagnostic files and the W&B launcher change are not included. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…ly.yaml Only dna_bioreason (wanglab/kegg) is enabled in zone_a (ts_qa now skip=true), so the per-run data override and the training presets' defaults lists are no longer needed. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Replaces the two hardcoded /lus/flare/... user paths with the site variable. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
batch_size 1, max_steps 100, warmup 10, grad-accum 1, save every 50. The projector launcher now takes these from the config instead of forcing its own defaults. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…ters, split PRISM and Jupyter setup - PROJ_DIR derived from the script location (no home-directory paths); BaseMM_PRISM -> PRISM - no hardcoded PBS allocation (qsub -A <project>; demo_helpers.qsub gains account=) - <project> placeholder in default site paths; notebook outputs sanitized - projector/SFT/GRPO launchers pass training.* overrides only when env vars are set - helpers/setup_prism_env.sh (PRISM venv) separated from setup_jupyter_kernel.sh (Jupyter + kernel) - notebook: rewritten setup cells, KEGG-only section 5 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…nd single data worker for GRPO Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…emo changes - ruff: notebook one-liners/magics, lambda, zip strict, ambiguous name - mypy: annotate build_collator's collator - regenerate tools/ci/site_path_allowlist.json - multimodal: tolerate tokenizers without eos_token (fallback <|endoftext|>) - tests: KEGG-only default, OLMo Tulu turn markers Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
With the KEGG-only default the dataset mix changed, so 4 draws could miss the requested modality by chance. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Edits needed to run the bioreason-demo (OLMo-1B + frozen NT-2.5B DNA encoder + projector + LoRA on KEGG variant→disease) end to end on Aurora/PBS: projector stage (optional), single-stage SFT, GRPO, KEGG eval, and a tutorial notebook. All changes were made for the bioreason-demo.
Defaults: single-stage SFT (random-init projector,
freeze_connector: false); a separate projector stage is opt-in viaSTAGE1_CKPT. GRPO needs then-getty/torchtune-aurorafork (providestorchtune.dev.rl.loss), installed byhelpers/setup_prism_env.sh.Files edited (all for the bioreason-demo)
Core (
src/)src/train.py– single-stage SFT with random-init projector, weights-only resume, tokenizer adoption.src/data/multimodal.py– DNA/KEGG sample handlers, OLMo turn format,<|endoftext|>EOS.src/training/trainer_zone_a.py– fixes false "no gradients" watchdog; inline KEGG eval off.src/training/distributed.py– weights-only checkpoint load onto CPU.src/config.py– new config fields used by the demo.src/data/datasets_config.json– KEGG-only (ts_qaskipped); variant-effectlocal_pathuses${PRISM_DATA_ROOT}.Configs (
src/conf/)model/prism_olmo_1b.yaml– OLMo-1B + NT DNA encoder model config.training/bioreason_projector.yaml– demo hyperparameters (batch 1, 100 steps).training/bioreason_sft.yaml– single-stage SFT defaults;defaults:block removed.training/bioreason_grpo.yaml–data_num_workers: 1(single-shard KEGG data made extra workers see an empty stream);defaults:block removed.data/bioreason_demo_mix.yaml– demo data mix (KEGG only).Launchers (
applications/bioreason/scripts/)training/submit_bioreason_projector.sh,training/submit_bioreason_sft.sh– moved/rewritten; site-neutral paths, no hardcoded PBS account, optionaltraining.*overrides.training/submit_bioreason_grpo_aurora.sh– new Aurora GRPO launcher.training/submit_bioreason_grpo.sh– moved from app root.eval/submit_eval_kegg_aurora.sh,eval/submit_eval_sft_sweep.sh– new eval launchers.eval/submit_eval_kegg.sh,test/verify_dna_dataset_mix.sh– moved from app root.submit_bioreason_projector.sh,submit_bioreason_sft.sh(replaced by the above).Eval / helpers (
applications/bioreason/)eval_kegg_aurora.py– KEGG eval with accuracy/macro-F1 and real/shuffled/blank DNA ablation.helpers/summarize_eval_sweep.py– summarizes eval sweeps.helpers/setup_prism_env.sh– new PRISM venv setup (torchtune-aurora fork, webdataset, torchao==0.5.0).helpers/setup_jupyter_kernel.sh– Jupyter-only kernel setup, kept separate from PRISM env.helpers/kernel_wrapper.sh– kernel launch wrapper.helpers/demo_helpers.py– notebook helpers (qsub, log tails, plots).helpers/download_bioreason_dna_datasets.py– moved intohelpers/.Notebook / docs
bioreason_demo.ipynb– tutorial notebook (setup Part A PRISM / Part B Jupyter, KEGG-only data, outputs sanitized).flow-diag.png– pipeline diagram.CONTRIBUTING.md,docs/applications/bioreason-eval.md,docs/applications/bioreason_grpo.md– updated paths/instructions.tools/ci/site_path_allowlist.json– allowlist update (stale; needscheck_site_paths.py --write).CI fixes (added after the first CI run; all checks now pass)
applications/bioreason/bioreason_demo.ipynb– ruff: magics/one-liners split onto separate lines, lambda →def, imports sorted.applications/bioreason/helpers/demo_helpers.py– ruff: renamed ambiguousl, explicitzip(strict=False).src/train.py– mypy: annotatecollatorinbuild_collator.src/data/multimodal.py– fall back to<|endoftext|>when the tokenizer has noeos_token.tools/ci/site_path_allowlist.json– regenerated withcheck_site_paths.py --write.tests/data/test_dna_bioreason_datasets_config.py– expect the KEGG-only demo default (variant-effect datasets skipped).tests/data/test_variant_effect_{coding,non_snv}_handler.py– expect OLMo<|assistant|>turn marker (was ChatML).tests/multimodal/test_modality_contracts.py– sample 32 items instead of 4 so a changed dataset mix can't miss the requested modality by chance.🤖 Generated with Claude Code