diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e5b497..3b79aef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,10 @@ will make up v0.1.0. per-modality support-status table - `tools/ci/check_doc_links.py`, a link checker for relative Markdown links, run as a required `doc_links` CI job +- `docs/tutorials/`, an index of end-to-end tutorials grouped by application + area. The materials tutorial moved from `docs/materials_prism_tutorial.md` to + `docs/tutorials/material_science/materials_prism_tutorial.md`; notebooks stay + next to their scripts in `applications//` ### Changed diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6fcce00..312dabb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -274,6 +274,7 @@ descriptions and return types where they are not obvious from the signature. | `modalities/` | Encoders, projectors, output decoders | | `models/` | Backbone integrations | | `evaluation/` | Universal evaluator, vLLM inference | + | `tutorials/` | End-to-end walkthroughs, one folder per application area; notebooks stay in the top-level `applications//`, not `docs/applications/` | | `applications/` | Scientific applications (BioReason, KEGG, materials) | | `api/` | API reference for the public `src.*` surface | | `development/` | CI and contributor tooling | diff --git a/README.md b/README.md index 8d839a4..9463ec1 100644 --- a/README.md +++ b/README.md @@ -375,6 +375,7 @@ PRISM/ [docs/index.md](docs/index.md) indexes every page. Most-used entry points: - [docs/getting-started.md](docs/getting-started.md) — clone → environment → first run +- [docs/tutorials/README.md](docs/tutorials/README.md) — end-to-end tutorials by application area (BioReason, materials, …) - [docs/platforms/site_paths.md](docs/platforms/site_paths.md) — configuring the `PRISM_*` roots - [docs/platforms/aurora_operations.md](docs/platforms/aurora_operations.md) — Aurora ops, env, CCL tuning, troubleshooting - [docs/training/training.md](docs/training/training.md) — training stages diff --git a/docs/index.md b/docs/index.md index 4bb281a..cc0495d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -98,6 +98,19 @@ loadable as Claude agent skills or read as docs. See the | [getting-started.md](getting-started.md) | Clone, pick a platform, build an environment, run a first job | | [training/cli.md](training/cli.md) | Unified CLI (`prism`) reference: installation, commands, options, examples | +### Tutorials — `tutorials/` + +End-to-end walkthroughs, grouped by application area. Pages live in +`tutorials//`; notebooks live next to their scripts in the top-level +[`applications//`](../applications/) directory. + +| Document | Description | +|----------|-------------| +| [tutorials/README.md](tutorials/README.md) | Every tutorial by application area, and where a new one goes | +| [../applications/bioreason/bioreason_demo.ipynb](../applications/bioreason/bioreason_demo.ipynb) | BioReason notebook: train a DNA + language model through projector alignment, SFT, GRPO, and KEGG evaluation on Aurora | +| [tutorials/material_science/materials_prism_tutorial.md](tutorials/material_science/materials_prism_tutorial.md) | Materials: jointly train on crystal structures and text descriptions to predict a material property | +| [tutorials/vision_language/vlm_training_demo.md](tutorials/vision_language/vlm_training_demo.md) | Vision language: train an image projector between frozen SigLIP2 and OLMo-1B on PixMo-Cap, as a 2-node DDP or HSDP job on Aurora | + ### Platforms — `platforms/` Environment build and job submission, one page per machine, plus the diff --git a/docs/tutorials/README.md b/docs/tutorials/README.md new file mode 100644 index 0000000..370efad --- /dev/null +++ b/docs/tutorials/README.md @@ -0,0 +1,30 @@ +# Tutorials + +Hands-on walkthroughs of PRISM on a scientific application, grouped by +application area. Each one walks a single configuration from environment and +data through training. + +Tutorial pages live here, in `docs/tutorials//`. Notebooks live in +[`applications//`](../../applications/), next to the scripts they call. + +For reference material on the same applications, such as pipeline designs and +evaluation invocations, see [`docs/applications/`](../applications/). + +| Area | Tutorial | Modalities | Platform | +|---|---|---|---| +| Biology | [BioReason: train a DNA + language model](../../applications/bioreason/bioreason_demo.ipynb) (notebook): projector alignment, SFT, GRPO, and KEGG evaluation, submitted as chained PBS jobs | DNA + text | Aurora | +| Materials science | [Jointly train on material descriptions and crystal structures](material_science/materials_prism_tutorial.md): predict a material property from a CIF crystal graph and a text description | Crystal graph + text | CUDA, XPU, or CPU | +| Vision language | [Train a vision-language projector](vision_language/vlm_training_demo.md): project frozen SigLIP2 features into a frozen OLMo-1B on PixMo-Cap, launched as a 2-node DDP or HSDP job and logged to W&B | Image + text | Aurora | + +## Adding a tutorial + +- Put Markdown pages in `docs/tutorials//`, where `` is the + application's directory under `applications/`, or the name that directory + would take if the application has no scripts there yet. +- Put a notebook in `applications//`, next to the scripts it runs, and + link to it from this page. +- Add a row to the table above and to the Tutorials section of + [`docs/index.md`](../index.md). +- Run `make docs` before pushing. It checks relative links in Markdown files + the same way CI does. Links inside notebooks are not checked, so test those + by hand after moving a notebook. diff --git a/docs/materials_prism_tutorial.md b/docs/tutorials/material_science/materials_prism_tutorial.md similarity index 98% rename from docs/materials_prism_tutorial.md rename to docs/tutorials/material_science/materials_prism_tutorial.md index 9d4461e..e41a5eb 100644 --- a/docs/materials_prism_tutorial.md +++ b/docs/tutorials/material_science/materials_prism_tutorial.md @@ -89,7 +89,7 @@ python applications/material_science/stage_materials_text_model.py \ --cache-dir /path/to/shared/huggingface/hub ``` -Then point Hugging Face at that cache when launching training, according to your cluster's setup. The tokenizer and model name passed to `--text-model` must match the staged model. On Polaris, see [the materials guide](applications/materials_multimodal.md) for the PBS launcher and cache environment settings. +Then point Hugging Face at that cache when launching training, according to your cluster's setup. The tokenizer and model name passed to `--text-model` must match the staged model. On Polaris, see [the materials guide](../../applications/materials_multimodal.md) for the PBS launcher and cache environment settings. ## 5. What PRISM does with the two inputs diff --git a/docs/tutorials/vision_language/vlm_training_demo.md b/docs/tutorials/vision_language/vlm_training_demo.md new file mode 100644 index 0000000..b2ce510 --- /dev/null +++ b/docs/tutorials/vision_language/vlm_training_demo.md @@ -0,0 +1,340 @@ +# PRISM tutorial: train a vision-language projector + +A short Aurora walkthrough from environment setup to a W&B-logged +vision-language training run. The goal is to show the normal PRISM workflow: +install the repo, download models and data, convert data if needed, then launch +training through PRISM's existing launcher. + +No separate training script is used. + +## What you will build + +A small image-caption model that projects frozen SigLIP2 image features into a +frozen OLMo language backbone. Only the PRISM image projector trains. + +The demo uses: + +- language backbone: `allenai/OLMo-1B-0724-hf` +- image encoder: `google/siglip2-base-patch16-224` +- model config: `src/conf/model/prism_image_only.yaml` +- training recipe: `src/conf/training/projector_only.yaml` +- data format: image-caption WebDataset shards on Lustre +- logging: W&B online mode + +## Before you start + +You need an Aurora allocation, W&B credentials, and network access on the login +node for Hugging Face model and dataset downloads. + +The required data input is a WebDataset directory with this layout: + +```text +$WDS_DIR/ + manifest.json + shards/ + *.tar +``` + +`allenai/pixmo-cap` is a good public captioning source for this demo. PRISM +includes an application script to download PixMo-Cap images locally, then the +standard `prism data convert` command writes WebDataset shards. + +## 1. Clone PRISM + +Set site-specific paths first: + +```bash +export PROJECT_ALLOCATION=YOUR_ALLOCATION +export QUEUE_NAME=debug +export WANDB_ENTITY=YOUR_WANDB_ENTITY +export WORK_ROOT=/flare/${PROJECT_ALLOCATION}/${USER} + +export PRISM_DIR=$WORK_ROOT/PRISM +export DEMO_ROOT=$WORK_ROOT/prism-vlm-demo +export HF_HOME=$DEMO_ROOT/hf +export SHARED_HF_HOME=$HF_HOME/hub +export HF_HUB_CACHE=$SHARED_HF_HOME +export TRANSFORMERS_CACHE=$SHARED_HF_HOME +export HF_DATASETS_CACHE=$HF_HOME/datasets +export PIXMO_LOCAL_DIR=$DEMO_ROOT/pixmo-cap-local +export WDS_DIR=$DEMO_ROOT/pixmo-cap-webdataset + +mkdir -p "$WORK_ROOT" "$DEMO_ROOT" "$HF_HOME" "$SHARED_HF_HOME" "$HF_DATASETS_CACHE" +cd "$WORK_ROOT" +git clone --recursive https://github.com/AI-ModCon/PRISM.git +cd "$PRISM_DIR" +``` + +Keep the virtual environment, model cache, and data outside the git checkout. + +## 2. Install + +On Aurora: + +```bash +module use /opt/aurora/26.26.0/modulefiles +module load oneapi/release/2025.3.1 +module load mpich/prd/5.0.0.aurora_test.51a9474 +module load frameworks/2025.3.1 +curl -LsSf https://astral.sh/uv/install.sh | sh +export PATH="$HOME/.local/bin:$PATH" + +bash tools/setup_deepspeed_env.sh +source .venv-deepspeed/bin/activate +python -m pip install --no-deps -e . +tar -czf deepspeed_env.tar.gz .venv-deepspeed +``` + +If you open a new shell later, reload the Aurora modules before activating the +environment: + +```bash +module use /opt/aurora/26.26.0/modulefiles +module load oneapi/release/2025.3.1 +module load mpich/prd/5.0.0.aurora_test.51a9474 +module load frameworks/2025.3.1 +source .venv-deepspeed/bin/activate +``` + +Create `.env` for the launchers: + +```bash +cp .env.template .env +cat >> .env <> .env +``` + +The Aurora PBS launcher reads `WANDB_API_KEY` from `.env` on the compute nodes. + +## 3. Download Models + +```bash +huggingface-cli download allenai/OLMo-1B-0724-hf +huggingface-cli download google/siglip2-base-patch16-224 +``` + +## 4. Prepare Data + +Download a PixMo-Cap subset and write a local Arrow dataset with downloaded +image files: + +```bash +python applications/vision_language/prepare_pixmo_dataset.py \ + --hf-id allenai/pixmo-cap \ + --output-dir "$PIXMO_LOCAL_DIR" \ + --limit 250000 \ + --workers 64 +``` + +Convert it to WebDataset shards: + +```bash +prism data convert \ + --source-format arrow \ + --arrow-dir "$PIXMO_LOCAL_DIR" \ + --images-dir "$PIXMO_LOCAL_DIR/images" \ + --output-dir "$WDS_DIR" \ + --images-per-shard 100 \ + --workers 8 +``` + +If you already have compatible WebDataset shards instead, set: + +```bash +export WDS_DIR=/path/to/image_caption_webdataset +``` + +For other image-caption datasets, `prism data convert` also accepts Parquet +records with embedded images or local image paths: + +```bash +prism data convert \ + --source-format parquet \ + --input-dir /path/to/parquet_files \ + --dataset-name image-caption-demo \ + --output-dir "$WDS_DIR" \ + --images-per-shard 100 \ + --workers 8 +``` + +For a two-node run, use at least 24 shards so every Aurora rank has data to read. +For finite, no-repeat training, also make sure the number of usable samples in +`manifest.json` is at least: + +```text +required_samples = nodes * 12 * local_batch * max_steps +``` + +The launch commands below use two nodes, 400 steps, DDP local batch 10, and +HSDP local batch 16. That means the same WebDataset should contain at least +`2 * 12 * 16 * 400 = 153600` usable samples if you want both launches to run +without repeating data. Check this before submitting: + +```bash +python - <<'PY' +import json +import os + +manifest = os.path.join(os.environ["WDS_DIR"], "manifest.json") +with open(manifest) as f: + data = json.load(f) + +usable = int(data.get("total_written") or data.get("total_samples") or 0) +required = 2 * 12 * 16 * 400 +print(f"usable_samples={usable:,}") +print(f"required_samples_for_hsdp_demo={required:,}") +if usable < required: + raise SystemExit("Increase --limit and rebuild the WebDataset before training.") +PY +``` + +## 5. Launch Training + +The default launch path uses DDP. The second launch script switches to HSDP +and increases the local batch size. For this comparison to be meaningful, +first choose the largest local batch that runs under DDP, then use HSDP for +the first larger local batch that DDP cannot run. + +For example, if DDP runs at local batch 10 but fails at local batch 16, and +HSDP runs at local batch 16, set: + +```bash +export DEMO_MAX_STEPS=400 +export DDP_LOCAL_BATCH=10 +export HSDP_LOCAL_BATCH=16 +``` + +If DDP also runs at `HSDP_LOCAL_BATCH`, raise both values and repeat the short +smoke test. Otherwise, the comparison does not show a batch-size benefit from +HSDP. + +### DDP launch script + +```bash +python tools/launch_aurora_web.py \ + --file experiments/prism_designs.yaml \ + --id "PRISM-VLM-DEMO-OLMO1B-SIGLIP2-BASE-DDP-BS${DDP_LOCAL_BATCH}" \ + --design PRISM-IMAGE-ONLY-2N \ + --nodes 2 \ + --batch \ + --queue "$QUEUE_NAME" \ + --project "$PROJECT_ALLOCATION" \ + --walltime 01:00:00 \ + --webdataset-dir "$WDS_DIR" \ + --finite-webdataset \ + --packed-env "$PRISM_DIR/deepspeed_env.tar.gz" \ + --dist-strategy ddp \ + --wandb-project prism-vlm-demo \ + --max-steps "$DEMO_MAX_STEPS" \ + model=prism_image_only \ + training=projector_only \ + model.backbone_id=allenai/OLMo-1B-0724-hf \ + model.image_encoder_id=google/siglip2-base-patch16-224 \ + model.image_processor_id=google/siglip2-base-patch16-224 \ + training.batch_size="$DDP_LOCAL_BATCH" \ + training.gradient_accumulation_steps=1 \ + training.learning_rate=5e-4 \ + training.lr_connector=5e-4 \ + training.warmup_steps=40 \ + training.save_every_n_steps="$DEMO_MAX_STEPS" \ + training.viz_every_n_steps=100 \ + training.viz_num_examples=10 \ + training.eval_every_n_steps=1000 \ + training.device=xpu \ + wandb.mode=online \ + wandb.entity="$WANDB_ENTITY" +``` + +This is a two-node run. On Aurora, each node has 12 XPU ranks, so the global +batch is `2 * 12 * DDP_LOCAL_BATCH` samples/update. + +### HSDP launch script for a larger batch + +HSDP shards optimizer and gradient state within each node and replicates across +nodes. The script below keeps the same two nodes but uses +`HSDP_LOCAL_BATCH`, which should be larger than the validated DDP limit. The +global batch is `2 * 12 * HSDP_LOCAL_BATCH` samples/update. + +```bash +python tools/launch_aurora_web.py \ + --file experiments/prism_designs.yaml \ + --id "PRISM-VLM-DEMO-OLMO1B-SIGLIP2-BASE-HSDP-BS${HSDP_LOCAL_BATCH}" \ + --design PRISM-IMAGE-ONLY-2N \ + --nodes 2 \ + --batch \ + --queue "$QUEUE_NAME" \ + --project "$PROJECT_ALLOCATION" \ + --walltime 01:00:00 \ + --webdataset-dir "$WDS_DIR" \ + --finite-webdataset \ + --packed-env "$PRISM_DIR/deepspeed_env.tar.gz" \ + --dist-strategy hsdp \ + --fsdp-sharding shard_grad_op \ + --benchmark-mode \ + --wandb-project prism-vlm-demo \ + --max-steps "$DEMO_MAX_STEPS" \ + model=prism_image_only \ + training=projector_only \ + model.backbone_id=allenai/OLMo-1B-0724-hf \ + model.image_encoder_id=google/siglip2-base-patch16-224 \ + model.image_processor_id=google/siglip2-base-patch16-224 \ + training.batch_size="$HSDP_LOCAL_BATCH" \ + training.gradient_accumulation_steps=1 \ + training.learning_rate=5e-4 \ + training.lr_connector=5e-4 \ + training.warmup_steps=40 \ + training.save_every_n_steps="$DEMO_MAX_STEPS" \ + training.viz_every_n_steps=100 \ + training.viz_num_examples=10 \ + training.eval_every_n_steps=1000 \ + training.device=xpu \ + wandb.mode=online \ + wandb.entity="$WANDB_ENTITY" +``` + +If the HSDP run is unstable because the global batch changed, either lower +`training.batch_size` or scale the learning rate intentionally for the new +global batch. For an apples-to-apples comparison, keep the global batch fixed +by reducing the local batch or increasing/decreasing the node count together. + +## 6. Monitor + +Check PBS: + +```bash +qstat -u "$USER" +``` + +Check local logs: + +```bash +ls -lt logs/PRISM-IMAGE-ONLY-2N/ +``` + +The run should appear in W&B under `prism-vlm-demo`. The main fields to check are: + +- `loss` +- `lr` +- `loss/tokens_global` +- `grad_norm/projector_mean` +- `viz/predictions` + +`viz/predictions` logs image-caption examples every 100 training steps for this demo. diff --git a/src/conf/training/projector_only.yaml b/src/conf/training/projector_only.yaml index f61b5c7..ed87f81 100644 --- a/src/conf/training/projector_only.yaml +++ b/src/conf/training/projector_only.yaml @@ -23,6 +23,7 @@ gradient_accumulation_steps: 4 save_every_n_steps: 1000 eval_every_n_steps: 500 viz_every_n_steps: 500 +viz_num_examples: 4 resume_from_checkpoint: null resume_weights_only: null wandb_run_id: null diff --git a/src/config.py b/src/config.py index 0e3bb4b..d7cad05 100644 --- a/src/config.py +++ b/src/config.py @@ -67,6 +67,7 @@ class TrainingConfig: eval_enabled: bool = False validation_batches: int = 8 # Number of validation batches per eval event viz_every_n_steps: int = 500 # Run batch visualization every 500 steps + viz_num_examples: int = 4 # Number of rows to log in viz/predictions warmup_steps: int = 10 weight_decay: float = 0.01 gradient_accumulation_steps: int = 1 diff --git a/src/train.py b/src/train.py index 451175e..429b72a 100644 --- a/src/train.py +++ b/src/train.py @@ -402,6 +402,7 @@ def build_training_config( gradient_accumulation_steps=cfg.training.get("gradient_accumulation_steps", 1), save_every_n_steps=cfg.training.get("save_every_n_steps", 1000), viz_every_n_steps=cfg.training.get("viz_every_n_steps", 500), + viz_num_examples=cfg.training.get("viz_num_examples", 4), wandb_project=cfg.wandb.project, wandb_entity=cfg.wandb.get("entity", None), wandb_mode=cfg.wandb.get("mode", "online"), diff --git a/src/training/trainer_native.py b/src/training/trainer_native.py index 003cecd..371b314 100644 --- a/src/training/trainer_native.py +++ b/src/training/trainer_native.py @@ -2236,7 +2236,11 @@ def __iter__(self): # Prediction Monitoring try: - pred_table = visualizer.visualize_predictions(viz_batch, model) + pred_table = visualizer.visualize_predictions( + viz_batch, + model, + limit=getattr(config, "viz_num_examples", 4), + ) if pred_table: viz_results["viz/predictions"] = pred_table except Exception as e_gen: