From 24d7058f63a9d934ea0b1c09af38b97dcde9c2f9 Mon Sep 17 00:00:00 2001 From: Sezgin Er Date: Thu, 30 Jul 2026 15:01:39 +0200 Subject: [PATCH 1/4] Add single-writer MR report training strategy Adds source-grounded complete-report training with exact online/cached paths, strict MIL/encoder provenance checks, Slurm launchers, and CPU/GPU integration gates. --- report-generation-training/.gitignore | 9 + report-generation-training/README.md | 96 ++++ report-generation-training/configs/base.yaml | 51 ++ report-generation-training/pyproject.toml | 23 + .../scripts/build_exact_cache.sh | 7 + .../scripts/slurm_train.sh | 34 ++ .../scripts/synthetic_gpu_e2e.sbatch | 36 ++ .../scripts/train_cached.sh | 7 + .../scripts/train_node.sh | 64 +++ .../scripts/train_online.sh | 7 + .../src/mrrate_report_training/__init__.py | 3 + .../src/mrrate_report_training/build_cache.py | 165 +++++++ .../src/mrrate_report_training/cache.py | 120 +++++ .../src/mrrate_report_training/config.py | 36 ++ .../src/mrrate_report_training/mil.py | 84 ++++ .../src/mrrate_report_training/model.py | 346 ++++++++++++++ .../src/mrrate_report_training/online.py | 202 ++++++++ .../src/mrrate_report_training/preflight.py | 96 ++++ .../src/mrrate_report_training/provenance.py | 118 +++++ .../src/mrrate_report_training/targets.py | 69 +++ .../src/mrrate_report_training/train.py | 445 ++++++++++++++++++ report-generation-training/tests/gpu_e2e.py | 301 ++++++++++++ .../tests/smoke_checks.py | 258 ++++++++++ .../tests/test_cache.py | 62 +++ .../tests/test_report_writer.py | 67 +++ .../tests/test_targets.py | 17 + .../tests/test_training_policy.py | 16 + 27 files changed, 2739 insertions(+) create mode 100644 report-generation-training/.gitignore create mode 100644 report-generation-training/README.md create mode 100644 report-generation-training/configs/base.yaml create mode 100644 report-generation-training/pyproject.toml create mode 100755 report-generation-training/scripts/build_exact_cache.sh create mode 100644 report-generation-training/scripts/slurm_train.sh create mode 100644 report-generation-training/scripts/synthetic_gpu_e2e.sbatch create mode 100755 report-generation-training/scripts/train_cached.sh create mode 100644 report-generation-training/scripts/train_node.sh create mode 100755 report-generation-training/scripts/train_online.sh create mode 100644 report-generation-training/src/mrrate_report_training/__init__.py create mode 100644 report-generation-training/src/mrrate_report_training/build_cache.py create mode 100644 report-generation-training/src/mrrate_report_training/cache.py create mode 100644 report-generation-training/src/mrrate_report_training/config.py create mode 100644 report-generation-training/src/mrrate_report_training/mil.py create mode 100644 report-generation-training/src/mrrate_report_training/model.py create mode 100644 report-generation-training/src/mrrate_report_training/online.py create mode 100644 report-generation-training/src/mrrate_report_training/preflight.py create mode 100644 report-generation-training/src/mrrate_report_training/provenance.py create mode 100644 report-generation-training/src/mrrate_report_training/targets.py create mode 100644 report-generation-training/src/mrrate_report_training/train.py create mode 100644 report-generation-training/tests/gpu_e2e.py create mode 100644 report-generation-training/tests/smoke_checks.py create mode 100644 report-generation-training/tests/test_cache.py create mode 100644 report-generation-training/tests/test_report_writer.py create mode 100644 report-generation-training/tests/test_targets.py create mode 100644 report-generation-training/tests/test_training_policy.py diff --git a/report-generation-training/.gitignore b/report-generation-training/.gitignore new file mode 100644 index 0000000..4d19af9 --- /dev/null +++ b/report-generation-training/.gitignore @@ -0,0 +1,9 @@ +artifacts/ +logs/ +runs/ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ + diff --git a/report-generation-training/README.md b/report-generation-training/README.md new file mode 100644 index 0000000..d821ae8 --- /dev/null +++ b/report-generation-training/README.md @@ -0,0 +1,96 @@ +# MR-RATE single-writer report training + +This is a standalone report-generation training layer for MR-RATE. It does not +modify FORA or an existing MR-RATE checkout. + +The target is the complete ordered list of `extracted_sentences` supplied by +MR-RATE. The trainer does not infer abnormal/healthy labels, split reports into +subtasks, or fabricate supervision that is absent from the dataset. + +The two execution modes are deliberately equivalent: + +- `online`: a frozen MR-RATE encoder produces the complete projected visual + token bag during training. +- `cached`: the same complete projected token bag is read from the upstream + ragged memmap cache. + +In both modes the same frozen 74-label `ClassifyThenAggregate` MIL head consumes +the full bag, a trainable query resampler makes 512 language-prefix tokens, and +one Gemma LoRA writer learns the complete findings text. The MIL probabilities +are soft conditioning context; they are not report targets. + +## Intentional training policy + +- One source-grounded report target per study, preserving statement order. +- Natural one-pass coverage: every train study occurs exactly once per epoch. +- No replacement sampler and no pathology oversampling. +- No MIL proposal dropout. +- No localization target, localization token, or localization loss. +- No MR-specific disease loss. +- Empty source findings use the explicit target ``. +- The encoder and MIL head are frozen. +- Exact cached training rejects `max_tokens_per_study != 0`. +- Startup requires the encoder SHA-256 and configuration recorded by MIL + training to match. Cached MIL additionally requires the training-cache + fingerprint to match the report token cache. + +## Required artifacts + +Set paths in `configs/base.yaml`: + +1. MR-RATE pretraining checkpoint used by the trained MIL model. +2. The corresponding `mil_head.pt`. +3. A local Gemma 3 model. +4. The MR-RATE extracted-findings JSONL, labels CSV, and splits CSV. +5. For cached mode, `token_features_{split}.json` plus its ragged memmap files. + +Preflight rejects pooled features, token-capped caches, mismatched +dimensions/classes, missing reports, or mismatched MIL/encoder provenance: + +```bash +python -m mrrate_report_training.preflight --config configs/base.yaml --mode cached +``` + +## Build exact caches + +```bash +bash scripts/build_exact_cache.sh train +bash scripts/build_exact_cache.sh val +``` + +The cache builder always uses full projected token bags with +`max_tokens_per_study=0`. + +## Train + +```bash +bash scripts/train_online.sh configs/base.yaml +bash scripts/train_cached.sh configs/base.yaml +``` + +For a quick real-data test, add `--max-studies 2 --max-updates 1` to the Python +command in either launcher. Checkpoints contain the query +resampler/connector, report LoRA, optimizer/scheduler state, configuration, +data position, and per-rank RNG state. + +For multi-node Slurm training, use `scripts/slurm_train.sh`. It launches one +`torchrun` agent per node, stages Gemma plus the encoder/MIL checkpoints once +per node, and places CUDA, Triton, and TorchInductor caches under node-local +`/tmp`. + +The online reader supports the released `batchXX/.zip` layout. It +extracts only the current study to node-local `/tmp`, applies upstream MR-RATE +preprocessing, encodes it, and removes the extraction. + +## Tests + +```bash +python tests/smoke_checks.py +``` + +The tests cover complete ordered targets, strict cache validation, +online/cached numerical equivalence, frozen MIL behavior, single-prefix +construction, optimizer updates, strict weight provenance, and checkpoint +resume. `scripts/synthetic_gpu_e2e.sbatch` runs the integration gate on a +Slurm GPU with a deterministic dummy dataset. + diff --git a/report-generation-training/configs/base.yaml b/report-generation-training/configs/base.yaml new file mode 100644 index 0000000..a60baf0 --- /dev/null +++ b/report-generation-training/configs/base.yaml @@ -0,0 +1,51 @@ +seed: 17 +mode: cached +output_dir: runs/mrrate_single_writer_v1 + +upstream_root: /hnvme/workspace/b180dc51-sezgin/MR-RATE-linearprobe/contrastive-pretraining +encoder_checkpoint: /path/to/mrrate_pretraining_checkpoint.pt +mil_checkpoint: /path/to/mil_head.pt +llm_path: /hnvme/data/LLM-TZ/hub/models--google--gemma-3-12b-it/snapshots/96b6f1eccf38110c56df3a15bffe176da04bfd80 + +data: + data_folder: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/mri + jsonl_file: /hnvme/workspace/b180dc51-sezgin/MR-RATE/contrastive-pretraining/data/findings_sentences.jsonl + labels_file: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/mrrate_neurovfm74_labels.csv + splits_csv: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/splits.csv + cached_tokens_dir: artifacts/exact_tokens + use_preprocessed: false + preprocessed_dir: null + normalizer: zscore + space: native_space + +encoder: + name: vjepa2 + vjepa21_checkpoint: null + chunk_size: 64 + fusion_mode: late + pooling_strategy: simple_attn + extra_latent_projection: false + dim_latent: 512 + +writer: + num_visual_queries: 512 + resampler_depth: 2 + resampler_heads: 8 + lora_r: 16 + lora_alpha: 32 + max_target_tokens: 384 + mil_conditioning: all_classes + mil_proposal_dropout: 0.0 + localization: false + +training: + epochs: 1 + batch_size: 1 + gradient_accumulation: 1 + learning_rate: 0.0001 + weight_decay: 0.01 + warmup_ratio: 0.03 + num_workers: 2 + checkpoint_every: 500 + shuffle: true + replacement_sampling: false diff --git a/report-generation-training/pyproject.toml b/report-generation-training/pyproject.toml new file mode 100644 index 0000000..26813a9 --- /dev/null +++ b/report-generation-training/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mrrate-report-training" +version = "0.1.0" +description = "Exact online and cached single-writer report training for MR-RATE" +requires-python = ">=3.10" +dependencies = [ + "numpy", + "pyyaml", + "torch", + "transformers", + "peft", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +pythonpath = ["src"] +testpaths = ["tests"] diff --git a/report-generation-training/scripts/build_exact_cache.sh b/report-generation-training/scripts/build_exact_cache.sh new file mode 100755 index 0000000..627f5af --- /dev/null +++ b/report-generation-training/scripts/build_exact_cache.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +set -euo pipefail +project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$project_dir" +export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" +exec python -m mrrate_report_training.build_cache \ + --config "${MRRATE_REPORT_CONFIG:-configs/base.yaml}" --split "${1:?split required}" diff --git a/report-generation-training/scripts/slurm_train.sh b/report-generation-training/scripts/slurm_train.sh new file mode 100644 index 0000000..d142a27 --- /dev/null +++ b/report-generation-training/scripts/slurm_train.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mrrate_report +#SBATCH --partition=h200 +#SBATCH --nodes=2 +#SBATCH --gres=gpu:4 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=32 +#SBATCH --time=24:00:00 +#SBATCH --signal=B:USR1@300 +#SBATCH --output=logs/mrrate_report_%j.out +#SBATCH --error=logs/mrrate_report_%j.err +set -euo pipefail + +: "${MODE:?set MODE=online or cached}" +: "${CONFIG:?set CONFIG to an absolute yaml path}" +: "${LLM_PATH:?set LLM_PATH}" +: "${ENCODER_CHECKPOINT:?set ENCODER_CHECKPOINT}" +: "${MIL_CHECKPOINT:?set MIL_CHECKPOINT}" +case "$MODE" in online|cached) ;; *) exit 2 ;; esac + +project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +sif="${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate-ib3.sif}" +mkdir -p "$project_dir/logs" +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1)" +export MASTER_PORT="${MASTER_PORT:-29541}" +export MODE CONFIG LLM_PATH ENCODER_CHECKPOINT MIL_CHECKPOINT +export GPUS_PER_NODE="${GPUS_PER_NODE:-4}" + +exec srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ + --kill-on-bad-exit=1 \ + singularity exec --nv -B /hnvme:/hnvme,/tmp:/tmp "$sif" \ + bash "$project_dir/scripts/train_node.sh" + diff --git a/report-generation-training/scripts/synthetic_gpu_e2e.sbatch b/report-generation-training/scripts/synthetic_gpu_e2e.sbatch new file mode 100644 index 0000000..54b249a --- /dev/null +++ b/report-generation-training/scripts/synthetic_gpu_e2e.sbatch @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_syn_e2e +#SBATCH --partition=h200 +#SBATCH --nodes=1 +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=00:30:00 +#SBATCH --output=logs/synthetic_gpu_e2e_%j.out +#SBATCH --error=logs/synthetic_gpu_e2e_%j.err +set -euo pipefail + +project=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training +sif=/hnvme/workspace/b180dc51-sezgin/mrrate-ib3.sif +mkdir -p "$project/logs" +cache="/tmp/mrrate_synthetic_${SLURM_JOB_ID}" +mkdir -p "$cache"/{cuda,triton,inductor,xdg} +export CUDA_CACHE_PATH="$cache/cuda" +export TRITON_CACHE_DIR="$cache/triton" +export TORCHINDUCTOR_CACHE_DIR="$cache/inductor" +export XDG_CACHE_HOME="$cache/xdg" +export MRRATE_REPORT_STREAM_THRESHOLD=128 +export MRRATE_REPORT_STREAM_CHUNK=64 + +exec singularity exec --nv -B /hnvme:/hnvme,/tmp:/tmp "$sif" bash -lc " + set -euo pipefail + export PYTHONPATH=/hnvme/workspace/b180dc51-sezgin/extra-pip:$project/src + export CUDA_CACHE_PATH=$CUDA_CACHE_PATH + export TRITON_CACHE_DIR=$TRITON_CACHE_DIR + export TORCHINDUCTOR_CACHE_DIR=$TORCHINDUCTOR_CACHE_DIR + export XDG_CACHE_HOME=$XDG_CACHE_HOME + export MRRATE_REPORT_STREAM_THRESHOLD=$MRRATE_REPORT_STREAM_THRESHOLD + export MRRATE_REPORT_STREAM_CHUNK=$MRRATE_REPORT_STREAM_CHUNK + cd $project + python tests/gpu_e2e.py +" diff --git a/report-generation-training/scripts/train_cached.sh b/report-generation-training/scripts/train_cached.sh new file mode 100755 index 0000000..aea2a51 --- /dev/null +++ b/report-generation-training/scripts/train_cached.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +set -euo pipefail +project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$project_dir" +export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" +exec python -m torch.distributed.run --standalone --nproc-per-node="${GPUS_PER_NODE:-1}" \ + -m mrrate_report_training.train --mode cached --config "${1:-configs/base.yaml}" diff --git a/report-generation-training/scripts/train_node.sh b/report-generation-training/scripts/train_node.sh new file mode 100644 index 0000000..b021e08 --- /dev/null +++ b/report-generation-training/scripts/train_node.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +set -euo pipefail +: "${MODE:?MODE must be online or cached}" +: "${CONFIG:?CONFIG is required}" +: "${MASTER_ADDR:?MASTER_ADDR is required}" +: "${MASTER_PORT:?MASTER_PORT is required}" + +project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +extra_pip="${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip}" +export PYTHONPATH="$extra_pip:$project_dir/src" +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 TOKENIZERS_PARALLELISM=false + +cache_root="/tmp/mrrate_report_${SLURM_JOB_ID:-local}" +mkdir -p "$cache_root"/{cuda,triton,inductor,xdg,model} +export CUDA_CACHE_PATH="$cache_root/cuda" +export TRITON_CACHE_DIR="$cache_root/triton" +export TORCHINDUCTOR_CACHE_DIR="$cache_root/inductor" +export XDG_CACHE_HOME="$cache_root/xdg" + +stage_path() { + local source="$1" name="$2" destination="$cache_root/model/$name" + ( + flock 9 + if [[ -d "$source" ]]; then + if [[ ! -f "$destination/.complete" ]]; then + mkdir -p "$destination" + cp -aL "$source"/. "$destination"/ + touch "$destination/.complete" + fi + else + mkdir -p "$destination" + local target="$destination/$(basename "$source")" + if [[ ! -f "$target" || $(stat -Lc %s "$target") -ne $(stat -Lc %s "$source") ]]; then + cp -aL "$source" "$target" + fi + printf '%s\n' "$target" + fi + ) 9>"$cache_root/model/$name.lock" +} + +llm_local="$cache_root/model/llm" +stage_path "$LLM_PATH" llm >/dev/null +encoder_local="$(stage_path "$ENCODER_CHECKPOINT" encoder)" +mil_local="$(stage_path "$MIL_CHECKPOINT" mil)" + +nnodes="${SLURM_NNODES:-1}" +node_rank="${SLURM_NODEID:-0}" +gpus="${GPUS_PER_NODE:-4}" +command=(python -m torch.distributed.run \ + --nnodes="$nnodes" \ + --node-rank="$node_rank" \ + --nproc-per-node="$gpus" \ + --master-addr="$MASTER_ADDR" \ + --master-port="$MASTER_PORT" \ + -m mrrate_report_training.train \ + --mode "$MODE" \ + --config "$CONFIG" \ + --llm-path "$llm_local" \ + --encoder-checkpoint "$encoder_local" \ + --mil-checkpoint "$mil_local") +if [[ -n "${RESUME:-}" ]]; then + command+=(--resume "$RESUME") +fi +exec "${command[@]}" diff --git a/report-generation-training/scripts/train_online.sh b/report-generation-training/scripts/train_online.sh new file mode 100755 index 0000000..fc04779 --- /dev/null +++ b/report-generation-training/scripts/train_online.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +set -euo pipefail +project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$project_dir" +export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" +exec python -m torch.distributed.run --standalone --nproc-per-node="${GPUS_PER_NODE:-1}" \ + -m mrrate_report_training.train --mode online --config "${1:-configs/base.yaml}" diff --git a/report-generation-training/src/mrrate_report_training/__init__.py b/report-generation-training/src/mrrate_report_training/__init__.py new file mode 100644 index 0000000..fe4f81c --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/__init__.py @@ -0,0 +1,3 @@ +"""MR-RATE exact online/cached single-writer report training.""" + +__version__ = "0.1.0" diff --git a/report-generation-training/src/mrrate_report_training/build_cache.py b/report-generation-training/src/mrrate_report_training/build_cache.py new file mode 100644 index 0000000..366397d --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/build_cache.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import uuid +from pathlib import Path + +import numpy as np +import torch + +from .config import load_config +from .online import OnlineSource + + +def atomic_text(path: Path, value: str) -> None: + temporary = path.with_name(path.name + f".tmp.{os.getpid()}") + temporary.write_text(value) + os.replace(temporary, path) + + +def atomic_npy(path: Path, value: np.ndarray) -> None: + temporary = path.with_name(path.name + f".tmp.{os.getpid()}") + with temporary.open("wb") as handle: + np.save(handle, value) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + + +def cache_provenance(config: dict) -> tuple[dict, str]: + checkpoint = Path(config["encoder_checkpoint"]).resolve() + stat = checkpoint.stat() + encoder = config["encoder"] + data = config["data"] + provenance = { + "version": 1, + "checkpoint": { + "path": str(checkpoint), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + }, + "encoder": encoder, + "jsonl_file": str(Path(data["jsonl_file"]).resolve()), + "labels_file": str(Path(data["labels_file"]).resolve()), + "splits_csv": str(Path(data["splits_csv"]).resolve()), + "data_folder": ( + str(Path(data["data_folder"]).resolve()) if data.get("data_folder") else None + ), + "preprocessed_dir": ( + str(Path(data["preprocessed_dir"]).resolve()) + if data.get("preprocessed_dir") + else None + ), + "use_preprocessed": bool(data.get("use_preprocessed", False)), + "normalizer": data.get("normalizer", "zscore"), + "space": data.get("space", "native_space"), + "cache_dtype": "float16", + "max_tokens_per_study": 0, + } + serialized = json.dumps(provenance, sort_keys=True, separators=(",", ":")) + return provenance, hashlib.sha256(serialized.encode()).hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--split", choices=("train", "val", "test"), required=True) + args = parser.parse_args() + config = load_config(args.config) + if not torch.cuda.is_available(): + raise RuntimeError("Exact cache generation requires a CUDA GPU") + torch.cuda.set_device(0) + source = OnlineSource(config, torch.device("cuda", 0), split=args.split) + out = Path(config["data"]["cached_tokens_dir"]).resolve() + out.mkdir(parents=True, exist_ok=True) + provenance, fingerprint = cache_provenance(config) + for manifest_path in out.glob("token_features_*.json"): + existing = json.loads(manifest_path.read_text()) + if existing.get("cache_fingerprint") != fingerprint: + raise ValueError( + f"{manifest_path} was built with a different encoder/preprocessing" + ) + + cache_id = uuid.uuid4().hex[:12] + token_path = out / f"tokens_{args.split}_{cache_id}.bin" + token_temporary = token_path.with_name(token_path.name + f".tmp.{os.getpid()}") + offsets = [0] + full_counts, series_counts, labels, subject_ids = [], [], [], [] + try: + with token_temporary.open("wb") as handle: + for index in range(len(source)): + item = source.get(index) + array = item["tokens"].float().cpu().numpy().astype( + np.float16, copy=False + ) + if not len(array): + raise ValueError(f"{item['subject_id']} produced an empty token bag") + array.tofile(handle) + offsets.append(offsets[-1] + len(array)) + full_counts.append(len(array)) + series_counts.append(item["series_count"]) + labels.append(item["mil_labels"].numpy().astype(np.float32)) + subject_ids.append(item["subject_id"]) + if (index + 1) % 100 == 0: + print( + f"encoded={index + 1}/{len(source)} tokens={offsets[-1]}", + flush=True, + ) + handle.flush() + os.fsync(handle.fileno()) + os.replace(token_temporary, token_path) + except BaseException: + token_temporary.unlink(missing_ok=True) + raise + + offsets_path = out / f"token_offsets_{args.split}_{cache_id}.npy" + labels_path = out / f"labels_{args.split}_{cache_id}.npy" + ids_path = out / f"subject_ids_{args.split}_{cache_id}.txt" + full_path = out / f"full_token_counts_{args.split}_{cache_id}.npy" + series_path = out / f"series_counts_{args.split}_{cache_id}.npy" + atomic_npy(offsets_path, np.asarray(offsets, dtype=np.int64)) + atomic_npy(labels_path, np.stack(labels).astype(np.float32)) + atomic_npy(full_path, np.asarray(full_counts, dtype=np.int64)) + atomic_npy(series_path, np.asarray(series_counts, dtype=np.int32)) + atomic_text(ids_path, "\n".join(subject_ids) + "\n") + label_names = [str(value) for value in source.dataset.label_columns] + names_path = out / "label_names.json" + if names_path.exists() and json.loads(names_path.read_text()) != label_names: + raise ValueError("Existing cache label_names.json differs") + atomic_text(names_path, json.dumps(label_names, indent=2) + "\n") + manifest = { + "format": "raw_numpy_memmap", + "format_version": 2, + "feature_level": "projected_per_series_visual_tokens", + "split": args.split, + "tokens_file": token_path.name, + "offsets_file": offsets_path.name, + "labels_file": labels_path.name, + "subject_ids_file": ids_path.name, + "full_token_counts_file": full_path.name, + "series_counts_file": series_path.name, + "dtype": "float16", + "dim": int(config["encoder"]["dim_latent"]), + "num_studies": len(subject_ids), + "num_tokens": offsets[-1], + "max_tokens_per_study": 0, + "cache_fingerprint": fingerprint, + "provenance": provenance, + } + # Publish the canonical manifest last. + atomic_text( + out / f"token_features_{args.split}.json", + json.dumps(manifest, indent=2) + "\n", + ) + print( + f"exact cache complete: studies={len(subject_ids)} tokens={offsets[-1]}", + flush=True, + ) + + +if __name__ == "__main__": + main() + diff --git a/report-generation-training/src/mrrate_report_training/cache.py b/report-generation-training/src/mrrate_report_training/cache.py new file mode 100644 index 0000000..e9802f2 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/cache.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from .targets import ReportTarget + + +class ExactRaggedTokenDataset(Dataset): + """Exact projected visual-token bags aligned with complete report targets.""" + + def __init__( + self, + cache_dir: str | Path, + split: str, + targets: dict[str, ReportTarget], + *, + expected_dim: int = 512, + expected_label_names: list[str] | None = None, + ) -> None: + self.cache_dir = Path(cache_dir) + self.split = split + manifest_path = self.cache_dir / f"token_features_{split}.json" + if not manifest_path.exists(): + raise FileNotFoundError(f"Missing exact token manifest: {manifest_path}") + self.metadata = json.loads(manifest_path.read_text()) + if self.metadata.get("format") != "raw_numpy_memmap": + raise ValueError("Cached training needs the raw_numpy_memmap token cache") + if self.metadata.get("feature_level") != "projected_per_series_visual_tokens": + raise ValueError("Pooled features cannot be used for report attention") + if int(self.metadata.get("max_tokens_per_study", -1)) != 0: + raise ValueError("Exact training requires max_tokens_per_study=0") + self.dim = int(self.metadata["dim"]) + if self.dim != int(expected_dim): + raise ValueError(f"Token dim {self.dim} != expected {expected_dim}") + self.dtype = np.dtype(self.metadata["dtype"]) + if self.dtype not in (np.dtype("float16"), np.dtype("float32")): + raise ValueError(f"Unsupported cache dtype: {self.dtype}") + + def load_array(key: str) -> np.ndarray: + return np.load(self.cache_dir / self.metadata[key]) + + self.offsets = load_array("offsets_file").astype(np.int64, copy=False) + self.labels = load_array("labels_file").astype(np.float32, copy=False) + ids_path = self.cache_dir / self.metadata["subject_ids_file"] + self.subject_ids = ids_path.read_text().strip().splitlines() + self.full_counts = load_array("full_token_counts_file").astype( + np.int64, copy=False + ) + self.series_counts = load_array("series_counts_file").astype( + np.int64, copy=False + ) + if len(self.subject_ids) != len(set(self.subject_ids)): + raise ValueError(f"Duplicate subject IDs in {ids_path}") + if self.offsets.shape != (len(self.subject_ids) + 1,): + raise ValueError("Offsets and subject IDs are misaligned") + if self.labels.shape[0] != len(self.subject_ids): + raise ValueError("Labels and subject IDs are misaligned") + if self.full_counts.shape != (len(self.subject_ids),): + raise ValueError("Full token counts are misaligned") + if self.series_counts.shape != (len(self.subject_ids),): + raise ValueError("Series counts are misaligned") + cached_counts = np.diff(self.offsets) + if np.any(cached_counts <= 0) or not np.array_equal( + cached_counts, self.full_counts + ): + raise ValueError("Cache is empty, capped, or has inconsistent token counts") + missing = [subject_id for subject_id in self.subject_ids if subject_id not in targets] + if missing: + raise ValueError( + f"{len(missing)} cached studies lack report targets; first={missing[:5]}" + ) + self.targets = targets + if expected_label_names is not None: + names_path = self.cache_dir / "label_names.json" + if not names_path.exists(): + raise FileNotFoundError(f"Missing label schema: {names_path}") + names = json.loads(names_path.read_text()) + if names != list(expected_label_names): + raise ValueError("Token cache label schema differs from MIL checkpoint") + self.num_tokens = int(self.offsets[-1]) + token_path = self.cache_dir / self.metadata["tokens_file"] + expected_bytes = self.num_tokens * self.dim * self.dtype.itemsize + if token_path.stat().st_size != expected_bytes: + raise ValueError( + f"Token memmap has {token_path.stat().st_size} bytes, " + f"expected {expected_bytes}" + ) + self.tokens = np.memmap( + token_path, + mode="r", + dtype=self.dtype, + shape=(self.num_tokens, self.dim), + ) + + def __len__(self) -> int: + return len(self.subject_ids) + + def __getitem__(self, index: int) -> dict: + start, end = int(self.offsets[index]), int(self.offsets[index + 1]) + subject_id = self.subject_ids[index] + return { + "subject_id": subject_id, + "tokens": torch.from_numpy(np.array(self.tokens[start:end], copy=True)), + "mil_labels": torch.from_numpy(np.array(self.labels[index], copy=True)), + "target": self.targets[subject_id], + "series_count": int(self.series_counts[index]), + } + + +def collate_single_study(batch: list[dict]) -> dict: + if len(batch) != 1: + raise ValueError( + "MR token bags are intentionally trained one study/GPU at a time" + ) + return batch[0] diff --git a/report-generation-training/src/mrrate_report_training/config.py b/report-generation-training/src/mrrate_report_training/config.py new file mode 100644 index 0000000..6e8cd31 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/config.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from pathlib import Path + +import yaml + + +def load_config(path: str | Path) -> dict: + config_path = Path(path).resolve() + config = yaml.safe_load(config_path.read_text()) + if not isinstance(config, dict): + raise ValueError(f"Config must be a mapping: {config_path}") + project_root = config_path.parent.parent + for container, key in ( + (config, "output_dir"), + (config.get("data", {}), "cached_tokens_dir"), + ): + if key in container and not Path(container[key]).is_absolute(): + container[key] = str((project_root / container[key]).resolve()) + config["_config_path"] = str(config_path) + return config + + +def require_training_policy(config: dict) -> None: + writer = config["writer"] + training = config["training"] + if float(writer.get("mil_proposal_dropout", -1)) != 0.0: + raise ValueError("MR-RATE strategy fixes MIL proposal dropout at zero") + if bool(writer.get("localization", True)): + raise ValueError("MR-RATE strategy has localization completely disabled") + if bool(training.get("replacement_sampling", True)): + raise ValueError("Replacement sampling violates exact epoch coverage") + if int(training.get("epochs", 0)) <= 0: + raise ValueError("epochs must be positive") + if int(training.get("batch_size", 0)) != 1: + raise ValueError("Ragged MR studies require batch_size=1 per GPU") diff --git a/report-generation-training/src/mrrate_report_training/mil.py b/report-generation-training/src/mrrate_report_training/mil.py new file mode 100644 index 0000000..2a64480 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/mil.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import torch +from torch import nn + + +def _import_mil_class(upstream_root: str | Path): + scripts = Path(upstream_root) / "scripts" + if not (scripts / "mil_probe.py").exists(): + raise FileNotFoundError(f"Missing upstream MIL code: {scripts / 'mil_probe.py'}") + sys.path.insert(0, str(scripts)) + try: + from mil_probe import ClassifyThenAggregate + finally: + sys.path.pop(0) + return ClassifyThenAggregate + + +def load_frozen_mil( + checkpoint: str | Path, + upstream_root: str | Path, + *, + expected_dim: int = 512, +) -> tuple[nn.Module, list[str], torch.Tensor]: + package = torch.load(checkpoint, map_location="cpu", weights_only=False) + architecture = package.get("architecture") + state = package.get("state_dict") or package.get("head_state_dict") + label_names = [str(value) for value in package.get("label_names", [])] + if not isinstance(architecture, dict) or not isinstance(state, dict): + raise ValueError("MIL checkpoint lacks architecture/state_dict") + dim = int(architecture["dim"]) + classes = int(architecture["n_classes"]) + if dim != int(expected_dim): + raise ValueError(f"MIL input dim {dim} != expected {expected_dim}") + if len(label_names) != classes or len(set(label_names)) != classes: + raise ValueError("MIL checkpoint has an invalid label schema") + cls = _import_mil_class(upstream_root) + head = cls( + dim=dim, + n_classes=classes, + hidden_dim=int(architecture.get("hidden_dim", 512)), + mlp_hidden_dims=tuple(architecture.get("mlp_hidden_dims", [384])), + drop_rate=float(architecture.get("drop_rate", 0.0)), + init_std=float(architecture.get("init_std", 0.02)), + use_gating=bool(architecture.get("use_gating", True)), + use_norm=bool(architecture.get("use_norm", False)), + use_output_bias_scale=bool( + architecture.get("use_output_bias_scale", True) + ), + ) + head.load_state_dict(state, strict=True) + loaded = head.state_dict() + for name, expected in state.items(): + actual = loaded[name].detach().cpu() + expected = expected.detach().cpu() + if not torch.equal(actual, expected): + raise RuntimeError(f"MIL tensor changed while loading: {name}") + if actual.is_floating_point() and not torch.isfinite(actual).all(): + raise ValueError(f"MIL checkpoint contains non-finite tensor: {name}") + head.requires_grad_(False) + head.eval() + thresholds = torch.as_tensor( + package.get("validation_thresholds", [0.0] * classes), + dtype=torch.float32, + ) + if thresholds.shape != (classes,) or not torch.isfinite(thresholds).all(): + raise ValueError("MIL validation thresholds are invalid") + # The upstream selector stores thresholds in logit space. Conditioning + # compares sigmoid probabilities, so convert the boundary exactly once. + return head, label_names, thresholds.sigmoid() + + +@torch.no_grad() +def infer_mil(head: nn.Module, tokens: torch.Tensor) -> torch.Tensor: + cu_seqlens = torch.tensor( + [0, tokens.shape[0]], dtype=torch.long, device=tokens.device + ) + logits = head(tokens, cu_seqlens) + if logits.ndim != 2 or logits.shape[0] != 1: + raise RuntimeError(f"Unexpected MIL output: {tuple(logits.shape)}") + return logits.float() diff --git a/report-generation-training/src/mrrate_report_training/model.py b/report-generation-training/src/mrrate_report_training/model.py new file mode 100644 index 0000000..393cc4b --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/model.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +import math +import os +from collections.abc import Sequence + +import torch +import torch.nn.functional as F +from peft import LoraConfig, TaskType, get_peft_model +from torch import nn +from torch.utils.checkpoint import checkpoint +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from .targets import ReportTarget + + +class CrossAttentionBlock(nn.Module): + """Exact all-token attention with a bounded-memory streaming fallback.""" + + def __init__(self, dim: int, heads: int) -> None: + super().__init__() + if dim % heads: + raise ValueError("resampler dimension must be divisible by heads") + self.dim, self.heads, self.head_dim = dim, heads, dim // heads + self.query_norm = nn.LayerNorm(dim) + self.memory_norm = nn.LayerNorm(dim) + self.self_attention = nn.MultiheadAttention(dim, heads, batch_first=True) + self.cross_norm = nn.LayerNorm(dim) + self.cross_q = nn.Linear(dim, dim, bias=False) + self.cross_k = nn.Linear(dim, dim, bias=False) + self.cross_v = nn.Linear(dim, dim, bias=False) + self.cross_out = nn.Linear(dim, dim, bias=False) + self.mlp_norm = nn.LayerNorm(dim) + self.mlp = nn.Sequential( + nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) + ) + + def _stream_step( + self, + running_max: torch.Tensor, + denominator: torch.Tensor, + numerator: torch.Tensor, + query: torch.Tensor, + memory: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + normalized = self.memory_norm(memory) + key = self.cross_k(normalized).view( + 1, -1, self.heads, self.head_dim + ).transpose(1, 2) + value = self.cross_v(normalized).view( + 1, -1, self.heads, self.head_dim + ).transpose(1, 2) + scores = torch.matmul(query, key.transpose(-2, -1)) + scores = scores.mul(self.head_dim**-0.5).float() + chunk_max = scores.amax(dim=-1).detach() + new_max = torch.maximum(running_max, chunk_max) + old_scale = torch.exp(running_max - new_max) + weights = torch.exp(scores - new_max.unsqueeze(-1)) + denominator = denominator * old_scale + weights.sum(dim=-1) + numerator = ( + numerator * old_scale.unsqueeze(-1) + + torch.matmul(weights, value.float()) + ) + return new_max, denominator, numerator + + def _stream(self, query: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + shape = (1, self.heads, query.shape[2]) + running_max = torch.full( + shape, -torch.inf, dtype=torch.float32, device=query.device + ) + denominator = torch.zeros_like(running_max) + numerator = torch.zeros( + *shape, self.head_dim, dtype=torch.float32, device=query.device + ) + chunk_size = int(os.environ.get("MRRATE_REPORT_STREAM_CHUNK", "8192")) + if chunk_size <= 0: + raise ValueError("MRRATE_REPORT_STREAM_CHUNK must be positive") + for start in range(0, memory.shape[1], chunk_size): + chunk = memory[:, start : start + chunk_size] + if self.training and torch.is_grad_enabled(): + running_max, denominator, numerator = checkpoint( + self._stream_step, + running_max, + denominator, + numerator, + query, + chunk, + use_reentrant=False, + ) + else: + running_max, denominator, numerator = self._stream_step( + running_max, denominator, numerator, query, chunk + ) + return (numerator / denominator.clamp_min(1e-20).unsqueeze(-1)).to( + query.dtype + ) + + def forward(self, latents: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + normalized = self.query_norm(latents) + latents = latents + self.self_attention( + normalized, normalized, normalized, need_weights=False + )[0] + query = self.cross_q(self.cross_norm(latents)).view( + 1, -1, self.heads, self.head_dim + ).transpose(1, 2) + threshold = int(os.environ.get("MRRATE_REPORT_STREAM_THRESHOLD", "131072")) + if threshold > 0 and memory.shape[1] > threshold: + attended = self._stream(query, memory) + else: + normalized_memory = self.memory_norm(memory) + key = self.cross_k(normalized_memory).view( + 1, -1, self.heads, self.head_dim + ).transpose(1, 2) + value = self.cross_v(normalized_memory).view( + 1, -1, self.heads, self.head_dim + ).transpose(1, 2) + attended = F.scaled_dot_product_attention(query, key, value) + attended = attended.transpose(1, 2).reshape(1, -1, self.dim) + latents = latents + self.cross_out(attended) + return latents + self.mlp(self.mlp_norm(latents)) + + +class QueryResampler(nn.Module): + def __init__(self, dim: int, num_queries: int, depth: int, heads: int) -> None: + super().__init__() + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) * 0.02) + self.blocks = nn.ModuleList( + CrossAttentionBlock(dim, heads) for _ in range(depth) + ) + self.output_norm = nn.LayerNorm(dim) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + if tokens.ndim != 2 or not tokens.shape[0]: + raise ValueError("tokens must be a non-empty [N,D] tensor") + memory = tokens.unsqueeze(0) + latents = self.latents.to(memory.dtype) + for block in self.blocks: + latents = ( + checkpoint(block, latents, memory, use_reentrant=False) + if self.training and torch.is_grad_enabled() + else block(latents, memory) + ) + return self.output_norm(latents) + + +def activate_adapter(model: nn.Module, name: str) -> None: + model.set_adapter(name) + # PEFT freezes the inactive adapter during set_adapter. Both adapters are + # used before one backward pass, so keep both optimizer-visible. + for parameter_name, parameter in model.named_parameters(): + if "lora_" in parameter_name: + parameter.requires_grad_(True) + + +def build_gemma_writer( + model_path: str, + device: torch.device, + *, + lora_r: int, + lora_alpha: int, +) -> tuple[nn.Module, AutoTokenizer, int]: + tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) + llm = Gemma3ForConditionalGeneration.from_pretrained( + model_path, + local_files_only=True, + torch_dtype=torch.bfloat16, + attn_implementation="sdpa", + low_cpu_mem_usage=True, + ) + llm.model.vision_tower = None + llm.model.multi_modal_projector = None + llm.requires_grad_(False) + llm.config.use_cache = False + llm.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + lora = LoraConfig( + r=int(lora_r), + lora_alpha=int(lora_alpha), + lora_dropout=0.05, + target_modules=[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + bias="none", + task_type=TaskType.CAUSAL_LM, + ) + llm = get_peft_model(llm, lora, adapter_name="report") + activate_adapter(llm, "report") + llm.to(device) + return llm, tokenizer, int(llm.config.text_config.hidden_size) + + +@torch.no_grad() +def label_semantic_embeddings( + llm: nn.Module, tokenizer: AutoTokenizer, names: Sequence[str] +) -> torch.Tensor: + embedding = llm.get_input_embeddings() + vectors = [] + device = embedding.weight.device + for name in names: + ids = tokenizer( + f"MRI finding: {name}", add_special_tokens=False, return_tensors="pt" + ).input_ids.to(device) + vectors.append(embedding(ids).float().mean(dim=1).squeeze(0)) + return torch.stack(vectors) + + +class ReportWriter(nn.Module): + REPORT_PROMPT = ( + "Write the complete MRI findings supported by the visual evidence. " + "Preserve both positive findings and explicit negative findings. " + "Do not invent location details. If the source report has no findings, " + "write .\nFindings:" + ) + + def __init__( + self, + llm: nn.Module, + tokenizer, + label_embeddings: torch.Tensor, + *, + visual_dim: int = 512, + num_visual_queries: int = 512, + resampler_depth: int = 2, + resampler_heads: int = 8, + max_target_tokens: int = 384, + ) -> None: + super().__init__() + self.llm = llm + self.tokenizer = tokenizer + self.max_target_tokens = int(max_target_tokens) + llm_dim = int(label_embeddings.shape[1]) + self.resampler = QueryResampler( + visual_dim, num_visual_queries, resampler_depth, resampler_heads + ) + self.visual_projection = nn.Sequential( + nn.LayerNorm(visual_dim), nn.Linear(visual_dim, llm_dim) + ) + self.image_start = nn.Parameter(torch.randn(1, 1, llm_dim) * 0.02) + self.image_end = nn.Parameter(torch.randn(1, 1, llm_dim) * 0.02) + self.register_buffer( + "label_embeddings", label_embeddings.float(), persistent=True + ) + self.mil_value_projection = nn.Sequential( + nn.Linear(2, 64), nn.GELU(), nn.Linear(64, llm_dim) + ) + self.mil_norm = nn.LayerNorm(llm_dim) + + def shared_prefix( + self, + tokens: torch.Tensor, + mil_logits: torch.Tensor, + thresholds: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + tokens = tokens.to(self.image_start.dtype) + visual = self.visual_projection(self.resampler(tokens)) + visual_prefix = torch.cat((self.image_start, visual, self.image_end), dim=1) + probability = mil_logits.sigmoid().reshape(-1) + thresholds = thresholds.to(probability).reshape(-1) + if probability.numel() != self.label_embeddings.shape[0]: + raise ValueError("MIL logits and label semantics differ") + values = torch.stack((probability, probability - thresholds), dim=-1) + mil_tokens = self.mil_norm( + self.label_embeddings.to(values) + self.mil_value_projection(values) + ).unsqueeze(0).to(visual_prefix.dtype) + return visual_prefix, mil_tokens + + def _token_ids(self, text: str, *, append_eos: bool) -> torch.Tensor: + ids = self.tokenizer( + text, + add_special_tokens=False, + truncation=True, + max_length=self.max_target_tokens, + return_tensors="pt", + ).input_ids[0] + if append_eos: + eos = torch.tensor([self.tokenizer.eos_token_id], dtype=ids.dtype) + ids = torch.cat((ids, eos)) + return ids.to(self.image_start.device) + + def _report_loss( + self, prefix: torch.Tensor, prompt: str, target: str + ) -> torch.Tensor: + activate_adapter(self.llm, "report") + prompt_ids = self._token_ids(prompt, append_eos=False) + target_ids = self._token_ids(target, append_eos=True) + embedding = self.llm.get_input_embeddings() + inputs = torch.cat( + ( + prefix, + embedding(prompt_ids).unsqueeze(0), + embedding(target_ids).unsqueeze(0), + ), + dim=1, + ) + attention_mask = torch.ones( + 1, inputs.shape[1], dtype=torch.long, device=inputs.device + ) + token_type_ids = torch.zeros_like(attention_mask) + token_type_ids[:, : prefix.shape[1]] = 1 + outputs = self.llm( + inputs_embeds=inputs, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + use_cache=False, + logits_to_keep=target_ids.numel() + 1, + ) + logits = outputs.logits[:, :-1].float() + return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target_ids) + + def forward( + self, + tokens: torch.Tensor, + mil_logits: torch.Tensor, + thresholds: torch.Tensor, + target: ReportTarget, + *, + loss_scale: float = 1.0, + ) -> dict[str, torch.Tensor]: + visual_prefix, mil_tokens = self.shared_prefix( + tokens, mil_logits, thresholds + ) + report_loss = self._report_loss( + torch.cat((visual_prefix, mil_tokens), dim=1), + self.REPORT_PROMPT, + target.text, + ) + return {"loss": report_loss * float(loss_scale), "report_loss": report_loss} + + +def trainable_state_dict(module: nn.Module) -> dict[str, torch.Tensor]: + trainable = {name for name, value in module.named_parameters() if value.requires_grad} + # label_embeddings are deterministic but retaining them makes class-schema + # drift immediately visible during resume. + trainable.add("label_embeddings") + return { + name: value.detach().cpu() + for name, value in module.state_dict().items() + if name in trainable + } diff --git a/report-generation-training/src/mrrate_report_training/online.py b/report-generation-training/src/mrrate_report_training/online.py new file mode 100644 index 0000000..f578618 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/online.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import argparse +import os +import shutil +import sys +import tempfile +import zipfile +from pathlib import Path + +import numpy as np +import torch + + +class OnlineSource: + """Frozen upstream MR-RATE encoder source with exact full token bags.""" + + def __init__( + self, config: dict, device: torch.device, *, split: str = "train" + ) -> None: + upstream = Path(config["upstream_root"]).resolve() + search_paths = [ + upstream, + upstream / "scripts", + upstream / "mr_rate", + upstream / "vision_encoder", + ] + for path in reversed(search_paths): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + # Keep these paths for lazy imports inside build_encoder (notably + # vision_encoder); this is a dedicated training process. + from extract_features import _load_and_verify, build_encoder + from mil_probe_online import build_dataset, encode_study + from data_inference import collate_fn_infer + from data import SPACE_TO_IMG_SUBDIR + + data = config["data"] + encoder_config = config["encoder"] + args = argparse.Namespace( + weights_path=config["encoder_checkpoint"], + encoder=encoder_config["name"], + vjepa21_checkpoint=encoder_config.get("vjepa21_checkpoint"), + chunk_size=int(encoder_config.get("chunk_size", 64)), + fusion_mode=encoder_config["fusion_mode"], + pooling_strategy=encoder_config["pooling_strategy"], + extra_latent_projection=bool( + encoder_config.get("extra_latent_projection", False) + ), + dim_latent=int(encoder_config["dim_latent"]), + data_folder=data.get("data_folder"), + jsonl_file=data["jsonl_file"], + labels_file=data["labels_file"], + splits_csv=data["splits_csv"], + space=data.get("space", "native_space"), + normalizer=data.get("normalizer", "zscore"), + preprocessed_dir=data.get("preprocessed_dir"), + use_preprocessed=bool(data.get("use_preprocessed", False)), + cache_allow_mismatch=False, + ) + if args.fusion_mode != "late": + raise ValueError("Exact MR MIL/report tokens require fusion_mode=late") + self.dataset = build_dataset(args, split) + if len(self.dataset) == 0 and args.data_folder: + zip_paths = sorted(Path(args.data_folder).glob("batch*/*.zip")) + if zip_paths: + self.dataset.samples = self._zip_samples( + zip_paths, self.dataset, SPACE_TO_IMG_SUBDIR + ) + self.dataset._mrrate_zip_mode = True + self.dataset._mrrate_original_getitem = self.dataset.__class__.__getitem__ + self.dataset.__class__ = self._zip_dataset_class( + self.dataset.__class__, SPACE_TO_IMG_SUBDIR + ) + print( + f"[online] ZIP streaming enabled for {len(self.dataset)} " + f"{split} studies", + flush=True, + ) + if len(self.dataset) == 0: + raise ValueError(f"Online {split} source contains no eligible studies") + self.subject_ids = [ + str(sample["subject_id"]) for sample in self.dataset.samples + ] + self.encoder, dim = build_encoder(args) + if int(dim) != int(encoder_config["dim_latent"]): + raise ValueError("Constructed encoder dimension differs from config") + _load_and_verify(self.encoder, args.weights_path, strict_missing=True) + try: + visual = self.encoder.visual_transformer + if hasattr(visual, "model") and hasattr(visual.model, "merge_and_unload"): + visual.model.merge_and_unload() + except Exception as error: + print(f"[online] LoRA merge skipped: {error}", flush=True) + self.encoder.to(device=device, dtype=torch.bfloat16) + self.encoder.requires_grad_(False) + self.encoder.eval() + self._encode_study = encode_study + self._collate = collate_fn_infer + self.device = device + + @staticmethod + def _zip_samples(zip_paths, dataset, _space_mapping) -> list[dict]: + samples = [] + for zip_path in zip_paths: + subject_id = zip_path.stem + if subject_id not in dataset.subject_to_sentences: + continue + sample = { + "subject_id": subject_id, + "zip_path": str(zip_path), + "sentences": dataset.subject_to_sentences[subject_id], + } + if subject_id in dataset.subject_to_labels: + sample["labels"] = dataset.subject_to_labels[subject_id] + samples.append(sample) + return samples + + @staticmethod + def _zip_dataset_class(base_class, space_mapping): + class ZipStreamingDataset(base_class): + def __getitem__(self, index): + sample = self.samples[index] + image_subdir = space_mapping.get(self.space, "img") + temp_parent = Path( + os.environ.get( + "MRRATE_ZIP_TEMP", + f"/tmp/mrrate_zip_{os.environ.get('SLURM_JOB_ID', 'local')}", + ) + ) + temp_parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=temp_parent) as directory: + image_paths = [] + with zipfile.ZipFile(sample["zip_path"]) as archive: + members = sorted( + value + for value in archive.namelist() + if value.endswith(".nii.gz") + and f"/{image_subdir}/" in f"/{value}" + ) + if not members: + raise ValueError( + f"{sample['subject_id']} has no {image_subdir} NIfTIs" + ) + for member_index, member in enumerate(members): + filename = f"{member_index:03d}_{Path(member).name}" + destination = Path(directory) / filename + with archive.open(member) as source, destination.open( + "wb" + ) as target: + shutil.copyfileobj(source, target) + image_paths.append(str(destination)) + volumes = [] + for image_path in image_paths: + volume = self.load_and_resample_nii(image_path) + volume = self.normalize_volume(volume) + volumes.append(self.crop_or_pad(volume)) + stack = torch.stack(volumes, dim=0) + mask = torch.ones(stack.shape[0], dtype=torch.bool) + labels = sample.get("labels", np.array([], dtype=np.float32)) + return ( + stack, + sample["sentences"], + sample["subject_id"], + mask, + labels, + ) + + ZipStreamingDataset.__name__ = "ZipStreamingMRReportDatasetInfer" + return ZipStreamingDataset + + def __len__(self) -> int: + return len(self.dataset) + + @torch.no_grad() + def get(self, index: int) -> dict: + batch = self._collate([self.dataset[index]]) + encoded = self._encode_study( + self.encoder, + batch, + self.device, + torch.bfloat16, + 0, + keep_mapping=False, + ) + if encoded.full_token_count != encoded.tokens.shape[0]: + raise RuntimeError("Online token path unexpectedly capped a study") + return { + "subject_id": encoded.subject_id, + "tokens": encoded.tokens, + "mil_labels": encoded.target.squeeze(0).detach().cpu(), + "series_count": encoded.series_count, + } + + +def verify_frozen_encoder(source: OnlineSource) -> None: + if source.encoder.training: + raise RuntimeError("Online encoder must remain in eval mode") + if any(parameter.requires_grad for parameter in source.encoder.parameters()): + raise RuntimeError("Online encoder must remain frozen") + if any(parameter.grad is not None for parameter in source.encoder.parameters()): + raise RuntimeError("Frozen online encoder accumulated gradients") diff --git a/report-generation-training/src/mrrate_report_training/preflight.py b/report-generation-training/src/mrrate_report_training/preflight.py new file mode 100644 index 0000000..be35d74 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/preflight.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from .config import load_config, require_training_policy +from .mil import load_frozen_mil +from .provenance import verify_mil_encoder_provenance +from .targets import load_target_index + + +def check(config: dict, mode: str) -> dict: + require_training_policy(config) + required = { + "upstream_root": Path(config["upstream_root"]), + "encoder_checkpoint": Path(config["encoder_checkpoint"]), + "mil_checkpoint": Path(config["mil_checkpoint"]), + "llm_path": Path(config["llm_path"]), + "jsonl_file": Path(config["data"]["jsonl_file"]), + "labels_file": Path(config["data"]["labels_file"]), + "splits_csv": Path(config["data"]["splits_csv"]), + } + missing = {name: str(path) for name, path in required.items() if not path.exists()} + if missing: + raise FileNotFoundError(f"Missing configured artifacts: {missing}") + targets = load_target_index(required["jsonl_file"]) + _, labels, thresholds = load_frozen_mil( + required["mil_checkpoint"], + required["upstream_root"], + expected_dim=int(config["encoder"]["dim_latent"]), + ) + result = { + "mode": mode, + "report_targets": len(targets), + "report_statements": sum( + len(value.statements) for value in targets.values() + ), + "empty_report_targets": sum( + not value.statements for value in targets.values() + ), + "mil_classes": len(labels), + "mil_thresholds": int(thresholds.numel()), + } + if mode == "cached": + from .cache import ExactRaggedTokenDataset + + dataset = ExactRaggedTokenDataset( + config["data"]["cached_tokens_dir"], + "train", + targets, + expected_dim=int(config["encoder"]["dim_latent"]), + expected_label_names=labels, + ) + result.update( + train_studies=len(dataset), + train_tokens=dataset.num_tokens, + cache_fingerprint=dataset.metadata.get("cache_fingerprint"), + ) + result["provenance"] = verify_mil_encoder_provenance( + required["mil_checkpoint"], + required["encoder_checkpoint"], + config["encoder"], + cache_metadata=dataset.metadata, + ) + elif mode == "online": + if config["encoder"]["fusion_mode"] != "late": + raise ValueError("Online exact token training requires late fusion") + online_data = ( + config["data"].get("preprocessed_dir") + if config["data"].get("use_preprocessed") + else config["data"].get("data_folder") + ) + if not online_data or not Path(online_data).exists(): + raise FileNotFoundError(f"Missing online MR data source: {online_data}") + result["train_source"] = "frozen encoder" + result["provenance"] = verify_mil_encoder_provenance( + required["mil_checkpoint"], + required["encoder_checkpoint"], + config["encoder"], + ) + else: + raise ValueError("mode must be online or cached") + return result + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--mode", choices=("online", "cached"), required=True) + args = parser.parse_args() + print(json.dumps(check(load_config(args.config), args.mode), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/report-generation-training/src/mrrate_report_training/provenance.py b/report-generation-training/src/mrrate_report_training/provenance.py new file mode 100644 index 0000000..194ea06 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/provenance.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any + +import torch + + +def sha256_file(path: str | Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _mil_encoder_contract(package: dict) -> tuple[dict | None, dict | None]: + """Return recorded encoder file identity and architecture configuration.""" + + online = package.get("data_provenance") + if isinstance(online, dict): + return online.get("encoder_checkpoint"), online.get("encoder") + caches = package.get("cache_provenance") + if isinstance(caches, dict): + train = caches.get("train") + if isinstance(train, dict): + provenance = train.get("provenance") + if isinstance(provenance, dict): + encoder = { + "name": provenance.get("encoder"), + "chunk_size": provenance.get("chunk_size"), + "fusion_mode": provenance.get("fusion_mode"), + "pooling_strategy": provenance.get("pooling_strategy"), + "dim_latent": provenance.get("dim_latent"), + "extra_latent_projection": provenance.get( + "extra_latent_projection" + ), + } + return provenance.get("checkpoint"), encoder + return None, None + + +def _compare_config(recorded: dict, configured: dict) -> None: + configured_contract = { + "name": configured.get("name"), + "chunk_size": configured.get("chunk_size"), + "fusion_mode": configured.get("fusion_mode"), + "pooling_strategy": configured.get("pooling_strategy"), + "dim_latent": configured.get("dim_latent"), + "extra_latent_projection": bool( + configured.get("extra_latent_projection", False) + ), + } + for key, current in configured_contract.items(): + previous = recorded.get(key) + if previous is None: + continue + if key == "extra_latent_projection": + previous, current = bool(previous), bool(current) + if previous != current: + raise ValueError( + f"MIL encoder configuration mismatch for {key}: " + f"checkpoint={previous!r}, configured={current!r}" + ) + + +def verify_mil_encoder_provenance( + mil_checkpoint: str | Path, + encoder_checkpoint: str | Path, + encoder_config: dict, + *, + cache_metadata: dict | None = None, +) -> dict[str, Any]: + """Prove that MIL, encoder, and optional token cache share one origin.""" + + package = torch.load(mil_checkpoint, map_location="cpu", weights_only=False) + recorded_file, recorded_config = _mil_encoder_contract(package) + if not isinstance(recorded_file, dict) or not recorded_file.get("sha256"): + raise ValueError( + "MIL checkpoint has no verifiable encoder SHA-256 provenance" + ) + actual_path = Path(encoder_checkpoint).resolve() + actual_sha = sha256_file(actual_path) + if actual_sha != recorded_file["sha256"]: + raise ValueError( + "MIL checkpoint was trained with a different encoder checkpoint: " + f"recorded={recorded_file['sha256']}, actual={actual_sha}" + ) + if isinstance(recorded_config, dict): + _compare_config(recorded_config, encoder_config) + + cache_fingerprint = None + if cache_metadata is not None: + cache_fingerprint = cache_metadata.get("cache_fingerprint") + if not cache_fingerprint: + raise ValueError("Current token cache has no verified fingerprint") + cached_provenance = package.get("cache_provenance") + if isinstance(cached_provenance, dict): + train = cached_provenance.get("train") + recorded_fingerprint = ( + train.get("cache_fingerprint") if isinstance(train, dict) else None + ) + if not recorded_fingerprint: + raise ValueError( + "Cached MIL checkpoint lacks its training-cache fingerprint" + ) + if recorded_fingerprint != cache_fingerprint: + raise ValueError( + "Report token cache differs from the cache used to train MIL" + ) + return { + "encoder_sha256": actual_sha, + "encoder_config_verified": isinstance(recorded_config, dict), + "cache_fingerprint": cache_fingerprint, + "cache_verified": cache_metadata is not None, + } + diff --git a/report-generation-training/src/mrrate_report_training/targets.py b/report-generation-training/src/mrrate_report_training/targets.py new file mode 100644 index 0000000..e9d7e71 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/targets.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + + +_SPACE = re.compile(r"\s+") + + +def clean_statement(value: object) -> str: + return _SPACE.sub(" ", str(value or "")).strip() + + +@dataclass(frozen=True) +class ReportTarget: + """All source findings in their original order, with no inferred labels.""" + + subject_id: str + statements: tuple[str, ...] + + @property + def text(self) -> str: + return " ".join(self.statements) if self.statements else "" + + def validate(self) -> None: + if not self.subject_id: + raise ValueError("subject_id cannot be empty") + if any(not value for value in self.statements): + raise ValueError(f"{self.subject_id}: report contains an empty statement") + + +def make_report_target( + subject_id: str, statements: Iterable[object] +) -> ReportTarget: + target = ReportTarget( + str(subject_id), + tuple(text for value in statements if (text := clean_statement(value))), + ) + target.validate() + return target + + +def load_target_index(path: str | Path) -> dict[str, ReportTarget]: + targets: dict[str, ReportTarget] = {} + with Path(path).open() as handle: + for line_number, line in enumerate(handle, 1): + if not line.strip(): + continue + row = json.loads(line) + subject_id = clean_statement( + row.get("volume_name") or row.get("study_uid") or row.get("subject_id") + ) + if not subject_id: + raise ValueError(f"{path}:{line_number}: missing study identifier") + if subject_id in targets: + raise ValueError(f"{path}:{line_number}: duplicate {subject_id}") + statements = row.get("extracted_sentences") + if not isinstance(statements, list): + raise ValueError( + f"{path}:{line_number}: extracted_sentences must be a list" + ) + targets[subject_id] = make_report_target(subject_id, statements) + if not targets: + raise ValueError(f"No report targets found in {path}") + return targets + diff --git a/report-generation-training/src/mrrate_report_training/train.py b/report-generation-training/src/mrrate_report_training/train.py new file mode 100644 index 0000000..d03a656 --- /dev/null +++ b/report-generation-training/src/mrrate_report_training/train.py @@ -0,0 +1,445 @@ +from __future__ import annotations + +import argparse +import contextlib +import json +import math +import os +import random +import signal +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel + +from .cache import ExactRaggedTokenDataset +from .config import load_config, require_training_policy +from .mil import infer_mil, load_frozen_mil +from .model import ( + ReportWriter, + build_gemma_writer, + label_semantic_embeddings, + trainable_state_dict, +) +from .online import OnlineSource, verify_frozen_encoder +from .provenance import verify_mil_encoder_provenance +from .targets import ReportTarget, load_target_index + + +DUMMY_TARGET = ReportTarget("__padding__", ()) +_CHECKPOINT_AND_STOP = False + + +def _request_checkpoint(_signal_number, _frame) -> None: + global _CHECKPOINT_AND_STOP + _CHECKPOINT_AND_STOP = True + + +def distributed_setup() -> tuple[int, int, int, torch.device]: + world = int(os.environ.get("WORLD_SIZE", "1")) + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if world > 1: + dist.init_process_group("nccl") + if not torch.cuda.is_available(): + raise RuntimeError("MR-RATE report training requires CUDA") + torch.cuda.set_device(local_rank) + return rank, world, local_rank, torch.device("cuda", local_rank) + + +def exact_rank_indices( + length: int, epoch: int, seed: int, world: int, rank: int, shuffle: bool +) -> list[int]: + """Every real index exactly once; only explicit -1 no-op slots are padded.""" + + if length <= 0 or world <= 0 or not 0 <= rank < world: + raise ValueError("invalid exact-shard arguments") + if shuffle: + generator = torch.Generator().manual_seed(int(seed) + int(epoch)) + indices = torch.randperm(length, generator=generator).tolist() + else: + indices = list(range(length)) + padded_length = math.ceil(length / world) * world + indices.extend([-1] * (padded_length - length)) + return indices[rank::world] + + +def cosine_schedule( + optimizer: torch.optim.Optimizer, total_updates: int, warmup_ratio: float +): + warmup = int(total_updates * float(warmup_ratio)) + + def multiplier(step: int) -> float: + if warmup and step < warmup: + return float(step + 1) / warmup + progress = (step - warmup) / max(1, total_updates - warmup) + return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, multiplier) + + +def rng_state() -> dict: + return { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "cuda": torch.cuda.get_rng_state(), + } + + +def restore_rng(state: dict) -> None: + random.setstate(state["python"]) + np.random.set_state(state["numpy"]) + torch.set_rng_state(state["torch"]) + torch.cuda.set_rng_state(state["cuda"]) + + +def save_checkpoint( + path: Path, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + config: dict, + label_names: list[str], + *, + epoch: int, + next_slot: int, + update: int, + rank: int, + world: int, +) -> None: + local_rng = rng_state() + if world > 1: + all_rng: list[dict | None] = [None] * world + dist.all_gather_object(all_rng, local_rng) + else: + all_rng = [local_rng] + if rank: + return + bare = model.module if isinstance(model, DistributedDataParallel) else model + package = { + "format_version": 1, + "trainable_state_dict": trainable_state_dict(bare), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "config": {key: value for key, value in config.items() if key != "_config_path"}, + "label_names": label_names, + "epoch": epoch, + "next_slot": next_slot, + "update": update, + "rng_by_rank": all_rng, + } + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + torch.save(package, temporary) + os.replace(temporary, path) + + +def load_checkpoint( + path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + labels: list[str], + rank: int, +) -> tuple[int, int, int]: + package = torch.load(path, map_location="cpu", weights_only=False) + if package.get("label_names") != labels: + raise ValueError("Resume checkpoint MIL label schema differs") + bare = model.module if isinstance(model, DistributedDataParallel) else model + incompatible = bare.load_state_dict(package["trainable_state_dict"], strict=False) + unexpected = list(incompatible.unexpected_keys) + if unexpected: + raise ValueError(f"Unexpected resume tensors: {unexpected[:5]}") + optimizer.load_state_dict(package["optimizer"]) + scheduler.load_state_dict(package["scheduler"]) + states = package["rng_by_rank"] + restore_rng(states[rank] if rank < len(states) else states[0]) + return int(package["epoch"]), int(package["next_slot"]), int(package["update"]) + + +def main() -> None: + signal.signal(signal.SIGUSR1, _request_checkpoint) + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--mode", choices=("online", "cached"), required=True) + parser.add_argument("--resume") + parser.add_argument("--max-studies", type=int, default=0) + parser.add_argument("--max-updates", type=int, default=0) + parser.add_argument("--llm-path") + parser.add_argument("--encoder-checkpoint") + parser.add_argument("--mil-checkpoint") + args = parser.parse_args() + config = load_config(args.config) + config["mode"] = args.mode + if args.llm_path: + config["llm_path"] = args.llm_path + if args.encoder_checkpoint: + config["encoder_checkpoint"] = args.encoder_checkpoint + if args.mil_checkpoint: + config["mil_checkpoint"] = args.mil_checkpoint + require_training_policy(config) + rank, world, _, device = distributed_setup() + seed = int(config["seed"]) + random.seed(seed + rank) + np.random.seed(seed + rank) + torch.manual_seed(seed + rank) + torch.cuda.manual_seed(seed + rank) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + targets = load_target_index(config["data"]["jsonl_file"]) + mil_head, label_names, thresholds = load_frozen_mil( + config["mil_checkpoint"], + config["upstream_root"], + expected_dim=int(config["encoder"]["dim_latent"]), + ) + mil_head.to(device).eval() + thresholds = thresholds.to(device) + + online_source = None + if args.mode == "cached": + source = ExactRaggedTokenDataset( + config["data"]["cached_tokens_dir"], + "train", + targets, + expected_dim=int(config["encoder"]["dim_latent"]), + expected_label_names=label_names, + ) + subject_ids = source.subject_ids + else: + online_source = OnlineSource(config, device) + source = online_source + subject_ids = online_source.subject_ids + missing = [value for value in subject_ids if value not in targets] + if missing: + raise ValueError( + f"{len(missing)} online studies lack report targets; first={missing[:5]}" + ) + cache_metadata = source.metadata if args.mode == "cached" else None + if rank == 0: + try: + provenance_result = verify_mil_encoder_provenance( + config["mil_checkpoint"], + config["encoder_checkpoint"], + config["encoder"], + cache_metadata=cache_metadata, + ) + provenance_message = {"ok": True, "result": provenance_result} + except Exception as error: + provenance_message = { + "ok": False, + "error": f"{type(error).__name__}: {error}", + } + else: + provenance_message = None + if world > 1: + messages = [provenance_message] + dist.broadcast_object_list(messages, src=0) + provenance_message = messages[0] + if not provenance_message["ok"]: + raise RuntimeError( + f"MIL/encoder provenance verification failed: " + f"{provenance_message['error']}" + ) + if rank == 0: + print( + "[provenance] " + json.dumps(provenance_message["result"]), + flush=True, + ) + source_length = len(source) + if args.max_studies: + source_length = min(source_length, int(args.max_studies)) + + writer_config = config["writer"] + llm, tokenizer, _ = build_gemma_writer( + config["llm_path"], + device, + lora_r=int(writer_config["lora_r"]), + lora_alpha=int(writer_config["lora_alpha"]), + ) + semantics = label_semantic_embeddings(llm, tokenizer, label_names) + model = ReportWriter( + llm, + tokenizer, + semantics, + visual_dim=int(config["encoder"]["dim_latent"]), + num_visual_queries=int(writer_config["num_visual_queries"]), + resampler_depth=int(writer_config["resampler_depth"]), + resampler_heads=int(writer_config["resampler_heads"]), + max_target_tokens=int(writer_config["max_target_tokens"]), + ).to(device) + trainable = [value for value in model.parameters() if value.requires_grad] + optimizer = torch.optim.AdamW( + trainable, + lr=float(config["training"]["learning_rate"]), + weight_decay=float(config["training"]["weight_decay"]), + ) + epochs = int(config["training"]["epochs"]) + accumulation = int(config["training"]["gradient_accumulation"]) + slots_per_epoch = math.ceil(source_length / world) + updates_per_epoch = math.ceil(slots_per_epoch / accumulation) + scheduler = cosine_schedule( + optimizer, + epochs * updates_per_epoch, + float(config["training"]["warmup_ratio"]), + ) + if world > 1: + model = DistributedDataParallel( + model, device_ids=[device.index], find_unused_parameters=False + ) + + start_epoch = start_slot = update = 0 + if args.resume: + start_epoch, start_slot, update = load_checkpoint( + args.resume, model, optimizer, scheduler, label_names, rank + ) + output_dir = Path(config["output_dir"]).resolve() + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + print( + json.dumps( + { + "mode": args.mode, + "studies": source_length, + "world_size": world, + "slots_per_rank": slots_per_epoch, + "epochs": epochs, + "visual_queries": writer_config["num_visual_queries"], + "mil_classes": len(label_names), + "localization": False, + "replacement_sampling": False, + } + ), + flush=True, + ) + + model.train() + optimizer.zero_grad(set_to_none=True) + stop = False + for epoch in range(start_epoch, epochs): + local_indices = exact_rank_indices( + source_length, + epoch, + seed, + world, + rank, + bool(config["training"]["shuffle"]), + ) + first_slot = start_slot if epoch == start_epoch else 0 + for slot in range(first_slot, len(local_indices)): + index = local_indices[slot] + is_real = index >= 0 + if is_real: + if args.mode == "cached": + item = source[index] + tokens = item["tokens"].to( + device=device, dtype=torch.bfloat16, non_blocking=True + ) + subject_id = item["subject_id"] + else: + item = online_source.get(index) + tokens = item["tokens"] + subject_id = item["subject_id"] + target = targets[subject_id] + else: + tokens = torch.zeros( + 1, + int(config["encoder"]["dim_latent"]), + dtype=torch.bfloat16, + device=device, + ) + target = DUMMY_TARGET + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + mil_logits = infer_mil(mil_head, tokens) + with torch.autocast("cuda", dtype=torch.bfloat16): + losses = model( + tokens, + mil_logits, + thresholds, + target, + loss_scale=1.0 if is_real else 0.0, + ) + loss = losses["loss"] / accumulation + loss.backward() + flush = (slot + 1) % accumulation == 0 or slot + 1 == len(local_indices) + if flush: + torch.nn.utils.clip_grad_norm_(trainable, 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + update += 1 + reduced = torch.stack( + ( + losses["loss"].detach(), + losses["report_loss"].detach(), + torch.tensor(float(is_real), device=device), + ) + ) + if world > 1: + dist.all_reduce(reduced) + reduced /= world + if rank == 0 and (update == 1 or update % 10 == 0): + print( + f"epoch={epoch + 1} update={update} " + f"loss={reduced[0]:.4f} report={reduced[1]:.4f} " + f"real_fraction={reduced[2]:.3f}", + flush=True, + ) + checkpoint_every = int(config["training"]["checkpoint_every"]) + if checkpoint_every and update % checkpoint_every == 0: + save_checkpoint( + output_dir / f"checkpoint-{update:08d}.pt", + model, + optimizer, + scheduler, + config, + label_names, + epoch=epoch, + next_slot=slot + 1, + update=update, + rank=rank, + world=world, + ) + stop_requested = bool(_CHECKPOINT_AND_STOP) or bool( + args.max_updates and update >= args.max_updates + ) + stop_tensor = torch.tensor( + int(stop_requested), dtype=torch.int32, device=device + ) + if world > 1: + dist.all_reduce(stop_tensor, op=dist.ReduceOp.MAX) + if stop_tensor.item(): + stop = True + break + if online_source is not None: + verify_frozen_encoder(online_source) + if any(parameter.grad is not None for parameter in mil_head.parameters()): + raise RuntimeError("Frozen MIL head accumulated gradients") + start_slot = 0 + if stop: + break + + save_checkpoint( + output_dir / "last.pt", + model, + optimizer, + scheduler, + config, + label_names, + epoch=epoch if "epoch" in locals() else 0, + next_slot=(slot + 1) if "slot" in locals() else 0, + update=update, + rank=rank, + world=world, + ) + if rank == 0: + print(f"Saved {output_dir / 'last.pt'}", flush=True) + if world > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/report-generation-training/tests/gpu_e2e.py b/report-generation-training/tests/gpu_e2e.py new file mode 100644 index 0000000..81e8d86 --- /dev/null +++ b/report-generation-training/tests/gpu_e2e.py @@ -0,0 +1,301 @@ +"""One-GPU end-to-end integration gate using a deterministic dummy MR dataset.""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +from mrrate_report_training.cache import ExactRaggedTokenDataset +from mrrate_report_training.mil import infer_mil +from mrrate_report_training.model import ReportWriter +from mrrate_report_training.targets import make_report_target +from mrrate_report_training.train import ( + cosine_schedule, + load_checkpoint, + save_checkpoint, +) + + +class TinyTokenizer: + eos_token_id = 1 + + def __call__(self, text, **kwargs): + maximum = int(kwargs.get("max_length", 128)) + values = [2 + (ord(char) % 61) for char in text][:maximum] or [2] + return SimpleNamespace(input_ids=torch.tensor([values], dtype=torch.long)) + + +class TinyReportAdapterLLM(nn.Module): + """Small causal decoder surface with one actual low-rank report adapter.""" + + def __init__(self, hidden: int = 32, vocabulary: int = 64, rank: int = 4): + super().__init__() + self.embedding = nn.Embedding(vocabulary, hidden) + self.output = nn.Linear(hidden, vocabulary) + self.lora_A_report = nn.Parameter(torch.randn(hidden, rank) * 0.02) + self.lora_B_report = nn.Parameter(torch.zeros(rank, hidden)) + self.active_adapter = "report" + self.embedding.requires_grad_(False) + self.output.requires_grad_(False) + + def set_adapter(self, name: str) -> None: + if str(name) != "report": + raise ValueError(f"unexpected adapter: {name}") + self.active_adapter = str(name) + + def get_input_embeddings(self): + return self.embedding + + def forward(self, inputs_embeds, logits_to_keep, **_): + hidden = inputs_embeds[:, -int(logits_to_keep) :] + hidden = hidden + hidden @ self.lora_A_report @ self.lora_B_report + return SimpleNamespace(logits=self.output(hidden)) + + +def write_dummy_cache(root: Path) -> tuple[dict, list[torch.Tensor], list[str]]: + generator = np.random.default_rng(41) + subject_ids = [f"dummy_{index:02d}" for index in range(4)] + token_counts = [257, 129, 301, 193] + bags = [ + generator.normal(size=(count, 512)).astype(np.float16) + for count in token_counts + ] + labels = generator.integers(0, 2, size=(len(bags), 74)).astype(np.float32) + label_names = [f"synthetic_mr_finding_{index:02d}" for index in range(74)] + root.mkdir(parents=True, exist_ok=True) + tokens_path = root / "tokens_train_dummy.bin" + with tokens_path.open("wb") as handle: + for bag in bags: + bag.tofile(handle) + offsets = np.concatenate(([0], np.cumsum(token_counts))).astype(np.int64) + np.save(root / "offsets.npy", offsets) + np.save(root / "labels.npy", labels) + np.save(root / "full_counts.npy", np.asarray(token_counts, dtype=np.int64)) + np.save(root / "series_counts.npy", np.ones(len(bags), dtype=np.int32)) + (root / "subject_ids.txt").write_text("\n".join(subject_ids) + "\n") + (root / "label_names.json").write_text(json.dumps(label_names)) + (root / "token_features_train.json").write_text( + json.dumps( + { + "format": "raw_numpy_memmap", + "format_version": 2, + "feature_level": "projected_per_series_visual_tokens", + "split": "train", + "tokens_file": tokens_path.name, + "offsets_file": "offsets.npy", + "labels_file": "labels.npy", + "subject_ids_file": "subject_ids.txt", + "full_token_counts_file": "full_counts.npy", + "series_counts_file": "series_counts.npy", + "dtype": "float16", + "dim": 512, + "num_studies": len(bags), + "num_tokens": int(offsets[-1]), + "max_tokens_per_study": 0, + "cache_fingerprint": "synthetic_gpu_e2e_v1", + }, + indent=2, + ) + ) + targets = { + subject_ids[0]: make_report_target( + subject_ids[0], + ["There is no hemorrhage.", "A small chronic infarct is present."], + ), + subject_ids[1]: make_report_target( + subject_ids[1], ["There is no acute intracranial abnormality."] + ), + subject_ids[2]: make_report_target( + subject_ids[2], + ["Possible demyelinating lesion is present.", "The ventricles are normal."], + ), + subject_ids[3]: make_report_target( + subject_ids[3], + ["Postoperative change is present.", "There is no hydrocephalus."], + ), + } + return targets, [torch.from_numpy(value.copy()) for value in bags], label_names + + +def build_mil(upstream_root: Path, device: torch.device) -> nn.Module: + scripts = upstream_root / "scripts" + sys.path.insert(0, str(scripts)) + from mil_probe import ClassifyThenAggregate + + torch.manual_seed(9) + head = ClassifyThenAggregate( + dim=512, + n_classes=74, + hidden_dim=64, + mlp_hidden_dims=(48,), + drop_rate=0.0, + use_gating=True, + use_norm=False, + use_output_bias_scale=True, + ).to(device) + head.requires_grad_(False) + head.eval() + return head + + +def build_writer(device: torch.device) -> ReportWriter: + torch.manual_seed(13) + llm = TinyReportAdapterLLM() + writer = ReportWriter( + llm, + TinyTokenizer(), + torch.randn(74, 32), + visual_dim=512, + num_visual_queries=512, + resampler_depth=2, + resampler_heads=8, + max_target_tokens=96, + ) + return writer.to(device=device, dtype=torch.bfloat16) + + +def main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("This integration gate must run on a Slurm GPU") + device = torch.device("cuda", 0) + torch.cuda.set_device(device) + torch.manual_seed(3) + torch.cuda.manual_seed(3) + job_id = os.environ.get("SLURM_JOB_ID", "local") + project = Path(__file__).resolve().parents[1] + output = project / "runs" / f"synthetic_gpu_e2e_{job_id}" + cache_root = output / "dummy_cache" + targets, online_bags, label_names = write_dummy_cache(cache_root) + cached = ExactRaggedTokenDataset( + cache_root, + "train", + targets, + expected_dim=512, + expected_label_names=label_names, + ) + mil = build_mil( + Path( + "/hnvme/workspace/b180dc51-sezgin/" + "MR-RATE-linearprobe/contrastive-pretraining" + ), + device, + ) + thresholds = torch.full((74,), 0.5, device=device) + writer = build_writer(device) + + # Exact online/cached equivalence through MIL, resampler, and both writers. + writer.eval() + online_tokens = online_bags[0].to(device) + cached_tokens = cached[0]["tokens"].to(device) + assert torch.equal(online_tokens, cached_tokens) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + online_logits = infer_mil(mil, online_tokens) + cached_logits = infer_mil(mil, cached_tokens) + assert torch.equal(online_logits, cached_logits) + online_loss = writer( + online_tokens, online_logits, thresholds, targets["dummy_00"] + ) + cached_loss = writer( + cached_tokens, cached_logits, thresholds, targets["dummy_00"] + ) + for name in ("loss", "report_loss"): + torch.testing.assert_close( + online_loss[name], cached_loss[name], rtol=0.0, atol=0.0 + ) + + # Real CUDA optimizer steps using the cached path. + writer.train() + trainable = [value for value in writer.parameters() if value.requires_grad] + optimizer = torch.optim.AdamW(trainable, lr=1e-2, weight_decay=0.0) + scheduler = cosine_schedule(optimizer, total_updates=8, warmup_ratio=0.0) + adapter_before = writer.llm.lora_B_report.detach().clone() + losses_seen = [] + for index in range(len(cached)): + item = cached[index] + tokens = item["tokens"].to(device) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + logits = infer_mil(mil, tokens) + with torch.autocast("cuda", dtype=torch.bfloat16): + losses = writer(tokens, logits, thresholds, item["target"]) + optimizer.zero_grad(set_to_none=True) + losses["loss"].backward() + optimizer.step() + scheduler.step() + losses_seen.append(float(losses["loss"].detach())) + assert not torch.equal(adapter_before, writer.llm.lora_B_report) + assert all(parameter.grad is None for parameter in mil.parameters()) + + checkpoint = output / "checkpoint.pt" + save_checkpoint( + checkpoint, + writer, + optimizer, + scheduler, + {"synthetic": True}, + label_names, + epoch=0, + next_slot=len(cached), + update=len(cached), + rank=0, + world=1, + ) + + # Fresh model/optimizer, exact resume, then one more update. + resumed = build_writer(device) + resumed_trainable = [ + value for value in resumed.parameters() if value.requires_grad + ] + resumed_optimizer = torch.optim.AdamW( + resumed_trainable, lr=1e-2, weight_decay=0.0 + ) + resumed_scheduler = cosine_schedule( + resumed_optimizer, total_updates=8, warmup_ratio=0.0 + ) + epoch, slot, update = load_checkpoint( + str(checkpoint), + resumed, + resumed_optimizer, + resumed_scheduler, + label_names, + rank=0, + ) + assert (epoch, slot, update) == (0, len(cached), len(cached)) + item = cached[0] + tokens = item["tokens"].to(device) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + logits = infer_mil(mil, tokens) + with torch.autocast("cuda", dtype=torch.bfloat16): + resumed_losses = resumed(tokens, logits, thresholds, item["target"]) + resumed_optimizer.zero_grad(set_to_none=True) + resumed_losses["loss"].backward() + resumed_optimizer.step() + result = { + "status": "PASS", + "gpu": torch.cuda.get_device_name(0), + "cuda": torch.version.cuda, + "studies": len(cached), + "tokens": cached.num_tokens, + "visual_queries": 512, + "mil_classes": 74, + "online_cached_exact": True, + "optimizer_updates_before_resume": len(cached), + "post_resume_update": True, + "report_adapter_updated": True, + "losses": losses_seen, + "resumed_loss": float(resumed_losses["loss"].detach()), + "peak_memory_gib": torch.cuda.max_memory_allocated() / 2**30, + "checkpoint": str(checkpoint), + } + (output / "result.json").write_text(json.dumps(result, indent=2) + "\n") + print("MRRATE_SYNTHETIC_GPU_E2E_PASS " + json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/report-generation-training/tests/smoke_checks.py b/report-generation-training/tests/smoke_checks.py new file mode 100644 index 0000000..10bf62f --- /dev/null +++ b/report-generation-training/tests/smoke_checks.py @@ -0,0 +1,258 @@ +"""Dependency-light single-writer checks runnable without pytest.""" + +import hashlib +import json +import sys +import tempfile +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +from mrrate_report_training.cache import ExactRaggedTokenDataset +from mrrate_report_training.mil import load_frozen_mil +from mrrate_report_training.model import ReportWriter +from mrrate_report_training.provenance import verify_mil_encoder_provenance +from mrrate_report_training.targets import make_report_target +from mrrate_report_training.train import exact_rank_indices + + +class TinyTokenizer: + eos_token_id = 1 + + def __call__(self, text, **_): + values = [2 + (ord(char) % 13) for char in text][:20] or [2] + return SimpleNamespace(input_ids=torch.tensor([values], dtype=torch.long)) + + +class TinyLLM(nn.Module): + def __init__(self): + super().__init__() + self.embedding = nn.Embedding(16, 8) + self.output = nn.Linear(8, 16) + self.lora_report = nn.Parameter(torch.tensor(0.1)) + self.active = "report" + for value in (*self.embedding.parameters(), *self.output.parameters()): + value.requires_grad_(False) + + def set_adapter(self, value): + self.active = value + + def get_input_embeddings(self): + return self.embedding + + def forward(self, inputs_embeds, logits_to_keep, **_): + return SimpleNamespace( + logits=self.output( + inputs_embeds[:, -logits_to_keep:] + self.lora_report + ) + ) + + +def check_targets() -> None: + statements = [ + "There is no acute intracranial abnormality.", + "A chronic infarct is present.", + "Cannot exclude hemorrhage.", + ] + target = make_report_target("s", statements) + assert target.statements == tuple(statements) + assert target.text == " ".join(statements) + assert not hasattr(target, "abnormal") + assert not hasattr(target, "normal") + + +def check_coverage() -> None: + shards = [ + exact_rank_indices(11, 0, 7, 4, rank, True) for rank in range(4) + ] + real = [value for shard in shards for value in shard if value >= 0] + assert sorted(real) == list(range(11)) + assert len(real) == len(set(real)) + + +def check_cache_and_path_equivalence() -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tokens = np.arange(30, dtype=np.float16).reshape(6, 5) + tokens.tofile(root / "tokens.bin") + np.save(root / "offsets.npy", [0, 2, 6]) + np.save(root / "labels.npy", [[1, 0], [0, 1]]) + np.save(root / "full.npy", [2, 4]) + np.save(root / "series.npy", [1, 2]) + (root / "ids.txt").write_text("a\nb\n") + (root / "label_names.json").write_text(json.dumps(["x", "y"])) + (root / "token_features_train.json").write_text( + json.dumps( + { + "format": "raw_numpy_memmap", + "feature_level": "projected_per_series_visual_tokens", + "max_tokens_per_study": 0, + "dim": 5, + "dtype": "float16", + "tokens_file": "tokens.bin", + "offsets_file": "offsets.npy", + "labels_file": "labels.npy", + "subject_ids_file": "ids.txt", + "full_token_counts_file": "full.npy", + "series_counts_file": "series.npy", + } + ) + ) + targets = { + value: make_report_target( + value, + ["There is no abnormality.", "A chronic finding is present."], + ) + for value in ("a", "b") + } + cache = ExactRaggedTokenDataset( + root, + "train", + targets, + expected_dim=5, + expected_label_names=["x", "y"], + ) + online = torch.from_numpy(tokens[:2].copy()) + cached = cache[0]["tokens"] + assert torch.equal(cached, online) + model = ReportWriter( + TinyLLM(), + TinyTokenizer(), + torch.randn(2, 8), + visual_dim=5, + num_visual_queries=2, + resampler_depth=1, + resampler_heads=1, + max_target_tokens=20, + ).eval() + logits = torch.tensor([[0.2, -0.7]]) + thresholds = torch.tensor([0.5, 0.5]) + with torch.no_grad(): + online_losses = model(online, logits, thresholds, targets["a"]) + cached_losses = model(cached, logits, thresholds, targets["a"]) + for name in ("loss", "report_loss"): + assert torch.equal(online_losses[name], cached_losses[name]), name + + +def check_single_writer_update() -> None: + llm = TinyLLM() + model = ReportWriter( + llm, + TinyTokenizer(), + torch.randn(3, 8), + visual_dim=8, + num_visual_queries=2, + resampler_depth=1, + resampler_heads=2, + max_target_tokens=20, + ) + calls = [] + hook = model.resampler.register_forward_hook(lambda *_: calls.append(1)) + optimizer = torch.optim.AdamW( + [value for value in model.parameters() if value.requires_grad], lr=1e-3 + ) + before = llm.lora_report.detach().clone() + losses = model( + torch.randn(5, 8), + torch.randn(1, 3), + torch.full((3,), 0.5), + make_report_target( + "s", ["There is no hemorrhage.", "A small infarct is present."] + ), + ) + losses["loss"].backward() + hook.remove() + assert len(calls) == 1 + assert llm.lora_report.grad is not None + optimizer.step() + assert not torch.equal(before, llm.lora_report) + lora_names = [name for name, _ in llm.named_parameters() if "lora_" in name] + assert lora_names == ["lora_report"] + + +def check_strict_weight_and_provenance_loading() -> None: + upstream = Path( + "/hnvme/workspace/b180dc51-sezgin/" + "MR-RATE-linearprobe/contrastive-pretraining" + ) + sys.path.insert(0, str(upstream / "scripts")) + from mil_probe import ClassifyThenAggregate + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + encoder = root / "encoder.pt" + encoder.write_bytes(b"synthetic encoder identity") + encoder_sha = hashlib.sha256(encoder.read_bytes()).hexdigest() + head = ClassifyThenAggregate( + dim=8, + n_classes=2, + hidden_dim=4, + mlp_hidden_dims=(3,), + drop_rate=0.0, + use_gating=True, + use_norm=False, + use_output_bias_scale=True, + ) + checkpoint = root / "mil_head.pt" + architecture = { + "dim": 8, + "n_classes": 2, + "hidden_dim": 4, + "mlp_hidden_dims": [3], + "drop_rate": 0.0, + "use_gating": True, + "use_norm": False, + "use_output_bias_scale": True, + } + encoder_config = { + "name": "vjepa2", + "chunk_size": 64, + "fusion_mode": "late", + "pooling_strategy": "simple_attn", + "dim_latent": 8, + "extra_latent_projection": False, + } + torch.save( + { + "state_dict": head.state_dict(), + "architecture": architecture, + "label_names": ["a", "b"], + "validation_thresholds": [0.0, 0.0], + "data_provenance": { + "encoder_checkpoint": {"sha256": encoder_sha}, + "encoder": encoder_config, + }, + }, + checkpoint, + ) + loaded, labels, thresholds = load_frozen_mil( + checkpoint, upstream, expected_dim=8 + ) + assert labels == ["a", "b"] + assert torch.equal(thresholds, torch.full((2,), 0.5)) + for name, value in head.state_dict().items(): + assert torch.equal(value, loaded.state_dict()[name]) + verified = verify_mil_encoder_provenance( + checkpoint, encoder, encoder_config + ) + assert verified["encoder_sha256"] == encoder_sha + encoder.write_bytes(b"wrong encoder") + try: + verify_mil_encoder_provenance(checkpoint, encoder, encoder_config) + except ValueError as error: + assert "different encoder checkpoint" in str(error) + else: + raise AssertionError("Mismatched encoder provenance was accepted") + + +if __name__ == "__main__": + check_targets() + check_coverage() + check_cache_and_path_equivalence() + check_single_writer_update() + check_strict_weight_and_provenance_loading() + print("single-writer smoke checks: PASS") + diff --git a/report-generation-training/tests/test_cache.py b/report-generation-training/tests/test_cache.py new file mode 100644 index 0000000..0b31658 --- /dev/null +++ b/report-generation-training/tests/test_cache.py @@ -0,0 +1,62 @@ +import json + +import numpy as np +import pytest +import torch + +from mrrate_report_training.cache import ExactRaggedTokenDataset +from mrrate_report_training.targets import make_report_target + + +def make_cache(root, max_tokens=0): + tokens = np.arange(30, dtype=np.float16).reshape(6, 5) + token_file = root / "tokens_train_test.bin" + tokens.tofile(token_file) + np.save(root / "offsets.npy", np.array([0, 2, 6], dtype=np.int64)) + np.save(root / "labels.npy", np.array([[1, 0], [0, 1]], dtype=np.float32)) + np.save(root / "full.npy", np.array([2, 4], dtype=np.int64)) + np.save(root / "series.npy", np.array([1, 2], dtype=np.int32)) + (root / "ids.txt").write_text("a\nb\n") + (root / "label_names.json").write_text(json.dumps(["x", "y"])) + manifest = { + "format": "raw_numpy_memmap", + "feature_level": "projected_per_series_visual_tokens", + "max_tokens_per_study": max_tokens, + "dim": 5, + "dtype": "float16", + "tokens_file": token_file.name, + "offsets_file": "offsets.npy", + "labels_file": "labels.npy", + "subject_ids_file": "ids.txt", + "full_token_counts_file": "full.npy", + "series_counts_file": "series.npy", + } + (root / "token_features_train.json").write_text(json.dumps(manifest)) + return torch.from_numpy(tokens.copy()) + + +def test_exact_cache_preserves_online_tokens(tmp_path): + online = make_cache(tmp_path) + targets = { + value: make_report_target(value, ["There is no abnormality."]) + for value in ("a", "b") + } + cached = ExactRaggedTokenDataset( + tmp_path, + "train", + targets, + expected_dim=5, + expected_label_names=["x", "y"], + ) + assert torch.equal(cached[0]["tokens"], online[:2]) + assert torch.equal(cached[1]["tokens"], online[2:]) + + +def test_capped_cache_is_rejected(tmp_path): + make_cache(tmp_path, max_tokens=100) + targets = { + value: make_report_target(value, ["There is no abnormality."]) + for value in ("a", "b") + } + with pytest.raises(ValueError, match="max_tokens_per_study=0"): + ExactRaggedTokenDataset(tmp_path, "train", targets, expected_dim=5) diff --git a/report-generation-training/tests/test_report_writer.py b/report-generation-training/tests/test_report_writer.py new file mode 100644 index 0000000..1f6cf59 --- /dev/null +++ b/report-generation-training/tests/test_report_writer.py @@ -0,0 +1,67 @@ +from types import SimpleNamespace + +import torch +from torch import nn + +from mrrate_report_training.model import ReportWriter +from mrrate_report_training.targets import make_report_target + + +class TinyTokenizer: + eos_token_id = 1 + + def __call__(self, text, **_): + values = [2 + (ord(char) % 13) for char in text][:20] or [2] + return SimpleNamespace(input_ids=torch.tensor([values], dtype=torch.long)) + + +class TinyLLM(nn.Module): + def __init__(self, hidden=8, vocabulary=16): + super().__init__() + self.embedding = nn.Embedding(vocabulary, hidden) + self.output = nn.Linear(hidden, vocabulary) + self.lora_report = nn.Parameter(torch.tensor(0.1)) + self.active = "report" + + def set_adapter(self, value): + self.active = value + + def get_input_embeddings(self): + return self.embedding + + def forward(self, inputs_embeds, logits_to_keep, **_): + hidden = inputs_embeds[:, -logits_to_keep:] + self.lora_report + return SimpleNamespace(logits=self.output(hidden)) + + +def test_prefix_runs_once_and_report_adapter_receives_gradient(): + llm = TinyLLM() + llm.embedding.weight.requires_grad_(False) + llm.output.weight.requires_grad_(False) + llm.output.bias.requires_grad_(False) + model = ReportWriter( + llm, + TinyTokenizer(), + torch.randn(3, 8), + visual_dim=8, + num_visual_queries=2, + resampler_depth=1, + resampler_heads=2, + max_target_tokens=20, + ) + calls = [] + handle = model.resampler.register_forward_hook(lambda *_: calls.append(1)) + output = model( + torch.randn(5, 8), + torch.tensor([[0.2, -0.4, 1.1]]), + torch.tensor([0.5, 0.5, 0.5]), + make_report_target( + "x", ["There is no hemorrhage.", "A small infarct is present."] + ), + ) + output["loss"].backward() + handle.remove() + assert len(calls) == 1 + assert llm.lora_report.grad is not None + assert torch.isfinite(llm.lora_report.grad) + diff --git a/report-generation-training/tests/test_targets.py b/report-generation-training/tests/test_targets.py new file mode 100644 index 0000000..7ccf880 --- /dev/null +++ b/report-generation-training/tests/test_targets.py @@ -0,0 +1,17 @@ +from mrrate_report_training.targets import make_report_target + + +def test_complete_report_preserves_statement_order(): + statements = [ + "There is no acute intracranial abnormality.", + "A small chronic infarct is present.", + "Cannot exclude a tiny focus of hemorrhage.", + ] + target = make_report_target("study", statements) + assert target.statements == tuple(statements) + assert target.text == " ".join(statements) + + +def test_empty_report_has_explicit_none(): + assert make_report_target("study", []).text == "" + diff --git a/report-generation-training/tests/test_training_policy.py b/report-generation-training/tests/test_training_policy.py new file mode 100644 index 0000000..2dd8b77 --- /dev/null +++ b/report-generation-training/tests/test_training_policy.py @@ -0,0 +1,16 @@ +import torch + +from mrrate_report_training.train import exact_rank_indices + + +def test_distributed_epoch_has_full_coverage_without_replacement(): + shards = [ + exact_rank_indices(11, epoch=2, seed=4, world=4, rank=rank, shuffle=True) + for rank in range(4) + ] + real = [value for shard in shards for value in shard if value >= 0] + assert sorted(real) == list(range(11)) + assert len(real) == len(set(real)) + assert {len(shard) for shard in shards} == {3} + assert sum(value == -1 for shard in shards for value in shard) == 1 + From 4eea05a1a669ec2526756652c19fedf50931538f Mon Sep 17 00:00:00 2001 From: Sezgin Er Date: Thu, 30 Jul 2026 15:10:11 +0200 Subject: [PATCH 2/4] Use natural report findings as writer targets Reads all_reports.csv findings by study_uid, retains findings_sentences.jsonl only for encoder/MIL cohort compatibility, and bounds one corrupted long-report outlier. --- report-generation-training/README.md | 22 +++--- report-generation-training/configs/base.yaml | 5 +- .../src/mrrate_report_training/preflight.py | 10 ++- .../src/mrrate_report_training/targets.py | 68 +++++++++---------- .../src/mrrate_report_training/train.py | 4 +- report-generation-training/tests/gpu_e2e.py | 8 +-- .../tests/smoke_checks.py | 21 +++--- .../tests/test_cache.py | 4 +- .../tests/test_report_writer.py | 3 +- .../tests/test_targets.py | 22 +++--- 10 files changed, 86 insertions(+), 81 deletions(-) diff --git a/report-generation-training/README.md b/report-generation-training/README.md index d821ae8..d7c8c7a 100644 --- a/report-generation-training/README.md +++ b/report-generation-training/README.md @@ -3,9 +3,11 @@ This is a standalone report-generation training layer for MR-RATE. It does not modify FORA or an existing MR-RATE checkout. -The target is the complete ordered list of `extracted_sentences` supplied by -MR-RATE. The trainer does not infer abnormal/healthy labels, split reports into -subtasks, or fabricate supervision that is absent from the dataset. +The writer target is the natural `findings` field from +`MR-RATE-validation/reports/all_reports.csv`. The standardized +`findings_sentences.jsonl` remains an encoder/MIL data dependency only; it is +never used as report-generation supervision. The trainer does not infer +abnormal/healthy labels or split reports into artificial subtasks. The two execution modes are deliberately equivalent: @@ -21,13 +23,17 @@ are soft conditioning context; they are not report targets. ## Intentional training policy -- One source-grounded report target per study, preserving statement order. +- One source-grounded raw findings target per study, preserving line order. +- Raw findings cover all 97,896 split IDs. The encoder/MIL-compatible cohort is + the same 97,887 studies retained by upstream MR-RATE's JSONL loader. - Natural one-pass coverage: every train study occurs exactly once per epoch. - No replacement sampler and no pathology oversampling. - No MIL proposal dropout. - No localization target, localization token, or localization loss. - No MR-specific disease loss. -- Empty source findings use the explicit target ``. +- The 1,536-token target ceiling preserves every normal findings report. It + bounds one corrupted 32k-token outlier (`5NIUCVXWHA`) containing embedded + editing dialogue. - The encoder and MIL head are frozen. - Exact cached training rejects `max_tokens_per_study != 0`. - Startup requires the encoder SHA-256 and configuration recorded by MIL @@ -41,7 +47,8 @@ Set paths in `configs/base.yaml`: 1. MR-RATE pretraining checkpoint used by the trained MIL model. 2. The corresponding `mil_head.pt`. 3. A local Gemma 3 model. -4. The MR-RATE extracted-findings JSONL, labels CSV, and splits CSV. +4. `all_reports.csv`, plus the MR-RATE JSONL, labels CSV, and splits CSV used + to reconstruct the frozen encoder/MIL cohort. 5. For cached mode, `token_features_{split}.json` plus its ragged memmap files. Preflight rejects pooled features, token-capped caches, mismatched @@ -88,9 +95,8 @@ preprocessing, encodes it, and removes the extraction. python tests/smoke_checks.py ``` -The tests cover complete ordered targets, strict cache validation, +The tests cover natural findings targets, strict cache validation, online/cached numerical equivalence, frozen MIL behavior, single-prefix construction, optimizer updates, strict weight provenance, and checkpoint resume. `scripts/synthetic_gpu_e2e.sbatch` runs the integration gate on a Slurm GPU with a deterministic dummy dataset. - diff --git a/report-generation-training/configs/base.yaml b/report-generation-training/configs/base.yaml index a60baf0..d66e6d2 100644 --- a/report-generation-training/configs/base.yaml +++ b/report-generation-training/configs/base.yaml @@ -10,6 +10,7 @@ llm_path: /hnvme/data/LLM-TZ/hub/models--google--gemma-3-12b-it/snapshots/96b6f1 data: data_folder: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/mri jsonl_file: /hnvme/workspace/b180dc51-sezgin/MR-RATE/contrastive-pretraining/data/findings_sentences.jsonl + reports_csv: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/reports/all_reports.csv labels_file: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/mrrate_neurovfm74_labels.csv splits_csv: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/splits.csv cached_tokens_dir: artifacts/exact_tokens @@ -33,7 +34,9 @@ writer: resampler_heads: 8 lora_r: 16 lora_alpha: 32 - max_target_tokens: 384 + # Covers every normal raw findings report. One corrupted 32k-token outlier + # (5NIUCVXWHA) is safely bounded by this ceiling. + max_target_tokens: 1536 mil_conditioning: all_classes mil_proposal_dropout: 0.0 localization: false diff --git a/report-generation-training/src/mrrate_report_training/preflight.py b/report-generation-training/src/mrrate_report_training/preflight.py index be35d74..3367d15 100644 --- a/report-generation-training/src/mrrate_report_training/preflight.py +++ b/report-generation-training/src/mrrate_report_training/preflight.py @@ -18,13 +18,14 @@ def check(config: dict, mode: str) -> dict: "mil_checkpoint": Path(config["mil_checkpoint"]), "llm_path": Path(config["llm_path"]), "jsonl_file": Path(config["data"]["jsonl_file"]), + "reports_csv": Path(config["data"]["reports_csv"]), "labels_file": Path(config["data"]["labels_file"]), "splits_csv": Path(config["data"]["splits_csv"]), } missing = {name: str(path) for name, path in required.items() if not path.exists()} if missing: raise FileNotFoundError(f"Missing configured artifacts: {missing}") - targets = load_target_index(required["jsonl_file"]) + targets = load_target_index(required["reports_csv"]) _, labels, thresholds = load_frozen_mil( required["mil_checkpoint"], required["upstream_root"], @@ -33,11 +34,8 @@ def check(config: dict, mode: str) -> dict: result = { "mode": mode, "report_targets": len(targets), - "report_statements": sum( - len(value.statements) for value in targets.values() - ), - "empty_report_targets": sum( - not value.statements for value in targets.values() + "findings_characters": sum( + len(value.findings) for value in targets.values() ), "mil_classes": len(labels), "mil_thresholds": int(thresholds.numel()), diff --git a/report-generation-training/src/mrrate_report_training/targets.py b/report-generation-training/src/mrrate_report_training/targets.py index e9d7e71..2b85914 100644 --- a/report-generation-training/src/mrrate_report_training/targets.py +++ b/report-generation-training/src/mrrate_report_training/targets.py @@ -1,69 +1,69 @@ from __future__ import annotations -import json +import csv import re from dataclasses import dataclass from pathlib import Path -from typing import Iterable -_SPACE = re.compile(r"\s+") +_INLINE_SPACE = re.compile(r"[^\S\n]+") -def clean_statement(value: object) -> str: - return _SPACE.sub(" ", str(value or "")).strip() +def clean_findings(value: object) -> str: + """Normalize whitespace while retaining report line boundaries.""" + + lines = [] + for line in str(value or "").splitlines(): + cleaned = _INLINE_SPACE.sub(" ", line).strip() + if cleaned: + lines.append(cleaned) + return "\n".join(lines) @dataclass(frozen=True) class ReportTarget: - """All source findings in their original order, with no inferred labels.""" + """Natural findings text from all_reports.csv; no inferred labels.""" subject_id: str - statements: tuple[str, ...] + findings: str @property def text(self) -> str: - return " ".join(self.statements) if self.statements else "" + return self.findings def validate(self) -> None: if not self.subject_id: raise ValueError("subject_id cannot be empty") - if any(not value for value in self.statements): - raise ValueError(f"{self.subject_id}: report contains an empty statement") + if not self.findings: + raise ValueError(f"{self.subject_id}: findings cannot be empty") -def make_report_target( - subject_id: str, statements: Iterable[object] -) -> ReportTarget: - target = ReportTarget( - str(subject_id), - tuple(text for value in statements if (text := clean_statement(value))), - ) +def make_report_target(subject_id: str, findings: object) -> ReportTarget: + target = ReportTarget(str(subject_id), clean_findings(findings)) target.validate() return target def load_target_index(path: str | Path) -> dict[str, ReportTarget]: + """Load natural report findings keyed by study_uid.""" + targets: dict[str, ReportTarget] = {} - with Path(path).open() as handle: - for line_number, line in enumerate(handle, 1): - if not line.strip(): - continue - row = json.loads(line) - subject_id = clean_statement( - row.get("volume_name") or row.get("study_uid") or row.get("subject_id") - ) + with Path(path).open(newline="") as handle: + reader = csv.DictReader(handle) + required = {"study_uid", "findings"} + missing = required.difference(reader.fieldnames or ()) + if missing: + raise ValueError(f"{path}: missing columns {sorted(missing)}") + for line_number, row in enumerate(reader, 2): + subject_id = str(row["study_uid"]).strip() if not subject_id: - raise ValueError(f"{path}:{line_number}: missing study identifier") + raise ValueError(f"{path}:{line_number}: missing study_uid") if subject_id in targets: raise ValueError(f"{path}:{line_number}: duplicate {subject_id}") - statements = row.get("extracted_sentences") - if not isinstance(statements, list): - raise ValueError( - f"{path}:{line_number}: extracted_sentences must be a list" - ) - targets[subject_id] = make_report_target(subject_id, statements) + findings = clean_findings(row.get("findings")) + if not findings: + continue + targets[subject_id] = ReportTarget(subject_id, findings) if not targets: - raise ValueError(f"No report targets found in {path}") + raise ValueError(f"No report findings found in {path}") return targets - diff --git a/report-generation-training/src/mrrate_report_training/train.py b/report-generation-training/src/mrrate_report_training/train.py index d03a656..9b72bc0 100644 --- a/report-generation-training/src/mrrate_report_training/train.py +++ b/report-generation-training/src/mrrate_report_training/train.py @@ -28,7 +28,7 @@ from .targets import ReportTarget, load_target_index -DUMMY_TARGET = ReportTarget("__padding__", ()) +DUMMY_TARGET = ReportTarget("__padding__", "") _CHECKPOINT_AND_STOP = False @@ -190,7 +190,7 @@ def main() -> None: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True - targets = load_target_index(config["data"]["jsonl_file"]) + targets = load_target_index(config["data"]["reports_csv"]) mil_head, label_names, thresholds = load_frozen_mil( config["mil_checkpoint"], config["upstream_root"], diff --git a/report-generation-training/tests/gpu_e2e.py b/report-generation-training/tests/gpu_e2e.py index 81e8d86..ac01330 100644 --- a/report-generation-training/tests/gpu_e2e.py +++ b/report-generation-training/tests/gpu_e2e.py @@ -107,18 +107,18 @@ def write_dummy_cache(root: Path) -> tuple[dict, list[torch.Tensor], list[str]]: targets = { subject_ids[0]: make_report_target( subject_ids[0], - ["There is no hemorrhage.", "A small chronic infarct is present."], + "There is no hemorrhage.\nA small chronic infarct is present.", ), subject_ids[1]: make_report_target( - subject_ids[1], ["There is no acute intracranial abnormality."] + subject_ids[1], "There is no acute intracranial abnormality." ), subject_ids[2]: make_report_target( subject_ids[2], - ["Possible demyelinating lesion is present.", "The ventricles are normal."], + "Possible demyelinating lesion is present.\nThe ventricles are normal.", ), subject_ids[3]: make_report_target( subject_ids[3], - ["Postoperative change is present.", "There is no hydrocephalus."], + "Postoperative change is present.\nThere is no hydrocephalus.", ), } return targets, [torch.from_numpy(value.copy()) for value in bags], label_names diff --git a/report-generation-training/tests/smoke_checks.py b/report-generation-training/tests/smoke_checks.py index 10bf62f..511fb04 100644 --- a/report-generation-training/tests/smoke_checks.py +++ b/report-generation-training/tests/smoke_checks.py @@ -52,14 +52,14 @@ def forward(self, inputs_embeds, logits_to_keep, **_): def check_targets() -> None: - statements = [ - "There is no acute intracranial abnormality.", - "A chronic infarct is present.", - "Cannot exclude hemorrhage.", - ] - target = make_report_target("s", statements) - assert target.statements == tuple(statements) - assert target.text == " ".join(statements) + findings = ( + "There is no acute intracranial abnormality.\n" + "A chronic infarct is present.\n" + "Cannot exclude hemorrhage." + ) + target = make_report_target("s", findings) + assert target.findings == findings + assert target.text == findings assert not hasattr(target, "abnormal") assert not hasattr(target, "normal") @@ -104,7 +104,7 @@ def check_cache_and_path_equivalence() -> None: targets = { value: make_report_target( value, - ["There is no abnormality.", "A chronic finding is present."], + "There is no abnormality.\nA chronic finding is present.", ) for value in ("a", "b") } @@ -160,7 +160,7 @@ def check_single_writer_update() -> None: torch.randn(1, 3), torch.full((3,), 0.5), make_report_target( - "s", ["There is no hemorrhage.", "A small infarct is present."] + "s", "There is no hemorrhage.\nA small infarct is present." ), ) losses["loss"].backward() @@ -255,4 +255,3 @@ def check_strict_weight_and_provenance_loading() -> None: check_single_writer_update() check_strict_weight_and_provenance_loading() print("single-writer smoke checks: PASS") - diff --git a/report-generation-training/tests/test_cache.py b/report-generation-training/tests/test_cache.py index 0b31658..8380d95 100644 --- a/report-generation-training/tests/test_cache.py +++ b/report-generation-training/tests/test_cache.py @@ -38,7 +38,7 @@ def make_cache(root, max_tokens=0): def test_exact_cache_preserves_online_tokens(tmp_path): online = make_cache(tmp_path) targets = { - value: make_report_target(value, ["There is no abnormality."]) + value: make_report_target(value, "There is no abnormality.") for value in ("a", "b") } cached = ExactRaggedTokenDataset( @@ -55,7 +55,7 @@ def test_exact_cache_preserves_online_tokens(tmp_path): def test_capped_cache_is_rejected(tmp_path): make_cache(tmp_path, max_tokens=100) targets = { - value: make_report_target(value, ["There is no abnormality."]) + value: make_report_target(value, "There is no abnormality.") for value in ("a", "b") } with pytest.raises(ValueError, match="max_tokens_per_study=0"): diff --git a/report-generation-training/tests/test_report_writer.py b/report-generation-training/tests/test_report_writer.py index 1f6cf59..dfcefac 100644 --- a/report-generation-training/tests/test_report_writer.py +++ b/report-generation-training/tests/test_report_writer.py @@ -56,7 +56,7 @@ def test_prefix_runs_once_and_report_adapter_receives_gradient(): torch.tensor([[0.2, -0.4, 1.1]]), torch.tensor([0.5, 0.5, 0.5]), make_report_target( - "x", ["There is no hemorrhage.", "A small infarct is present."] + "x", "There is no hemorrhage.\nA small infarct is present." ), ) output["loss"].backward() @@ -64,4 +64,3 @@ def test_prefix_runs_once_and_report_adapter_receives_gradient(): assert len(calls) == 1 assert llm.lora_report.grad is not None assert torch.isfinite(llm.lora_report.grad) - diff --git a/report-generation-training/tests/test_targets.py b/report-generation-training/tests/test_targets.py index 7ccf880..3bfce1c 100644 --- a/report-generation-training/tests/test_targets.py +++ b/report-generation-training/tests/test_targets.py @@ -2,16 +2,16 @@ def test_complete_report_preserves_statement_order(): - statements = [ - "There is no acute intracranial abnormality.", - "A small chronic infarct is present.", - "Cannot exclude a tiny focus of hemorrhage.", - ] - target = make_report_target("study", statements) - assert target.statements == tuple(statements) - assert target.text == " ".join(statements) + findings = ( + "There is no acute intracranial abnormality.\n" + "A small chronic infarct is present.\n" + "Cannot exclude a tiny focus of hemorrhage." + ) + target = make_report_target("study", findings) + assert target.findings == findings + assert target.text == findings -def test_empty_report_has_explicit_none(): - assert make_report_target("study", []).text == "" - +def test_inline_whitespace_is_normalized_but_lines_are_preserved(): + target = make_report_target("study", " First line. \n\n Second line. ") + assert target.text == "First line.\nSecond line." From 602c3277c1cc8353d9dd0a735166a5f9be4d402d Mon Sep 17 00:00:00 2001 From: Sezgin Er Date: Thu, 30 Jul 2026 15:18:45 +0200 Subject: [PATCH 3/4] Document reproducible training workflow Adds explicit configuration, preflight, cache, smoke, Slurm, and resume commands; forwards smoke limits through launchers; removes the unused num_workers setting. --- report-generation-training/README.md | 167 ++++++++++++++---- report-generation-training/configs/base.yaml | 1 - .../scripts/slurm_train.sh | 4 +- .../scripts/train_cached.sh | 3 +- .../scripts/train_node.sh | 6 + .../scripts/train_online.sh | 3 +- 6 files changed, 150 insertions(+), 34 deletions(-) diff --git a/report-generation-training/README.md b/report-generation-training/README.md index d7c8c7a..ff3a1c9 100644 --- a/report-generation-training/README.md +++ b/report-generation-training/README.md @@ -42,61 +42,168 @@ are soft conditioning context; they are not report targets. ## Required artifacts -Set paths in `configs/base.yaml`: +The files have distinct roles: -1. MR-RATE pretraining checkpoint used by the trained MIL model. -2. The corresponding `mil_head.pt`. -3. A local Gemma 3 model. -4. `all_reports.csv`, plus the MR-RATE JSONL, labels CSV, and splits CSV used - to reconstruct the frozen encoder/MIL cohort. -5. For cached mode, `token_features_{split}.json` plus its ragged memmap files. +| Input | Purpose | +| --- | --- | +| `all_reports.csv` | Writer supervision from `study_uid, findings` | +| `findings_sentences.jsonl` | Reproduce the encoder/MIL cohort only | +| labels CSV | Exact 74-class MIL schema | +| splits CSV | Train/validation/test membership | +| MR-RATE checkpoint | Frozen visual encoder and projection | +| `mil_head.pt` | Frozen Classify-Then-Aggregate MIL head and provenance | +| Gemma snapshot | Frozen language model plus trainable report LoRA | +| exact token cache | Required only for cached mode | -Preflight rejects pooled features, token-capped caches, mismatched -dimensions/classes, missing reports, or mismatched MIL/encoder provenance: +The real encoder checkpoint and `mil_head.pt` are intentionally placeholders +in `configs/base.yaml`; replace them before running. + +## 1. Configure + +```bash +cd /hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training +cp configs/base.yaml configs/my_run.yaml +``` + +Edit at least: + +```yaml +encoder_checkpoint: /absolute/path/to/the/mr-rate-checkpoint.pt +mil_checkpoint: /absolute/path/to/the/corresponding/mil_head.pt +llm_path: /absolute/path/to/gemma-3-12b-it + +data: + data_folder: /absolute/path/to/mri + reports_csv: /absolute/path/to/all_reports.csv + jsonl_file: /absolute/path/to/findings_sentences.jsonl + labels_file: /absolute/path/to/mrrate_labels.csv + splits_csv: /absolute/path/to/splits.csv + cached_tokens_dir: /absolute/path/to/exact_tokens +``` + +Do not pair an arbitrary MIL head and encoder. Preflight checks the recorded +encoder SHA-256, architecture, label order, and cache fingerprint. + +## 2. Run preflight + +Inside the training container: ```bash -python -m mrrate_report_training.preflight --config configs/base.yaml --mode cached +export PROJECT=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training +export PYTHONPATH=/hnvme/workspace/b180dc51-sezgin/extra-pip:$PROJECT/src +cd "$PROJECT" + +python -m mrrate_report_training.preflight \ + --config configs/my_run.yaml \ + --mode cached ``` -## Build exact caches +Use `--mode online` for online training. Preflight rejects pooled features, +token-capped caches, mismatched dimensions/classes, missing findings, and +mismatched MIL/encoder provenance. + +## 3. Prepare an exact cache, if needed ```bash +export MRRATE_REPORT_CONFIG="$PROJECT/configs/my_run.yaml" bash scripts/build_exact_cache.sh train bash scripts/build_exact_cache.sh val ``` -The cache builder always uses full projected token bags with -`max_tokens_per_study=0`. +Run cache generation on a CUDA node. It is an expensive frozen-encoder pass, +not part of every training epoch. The builder always uses +`max_tokens_per_study=0`. If the verified exact cache already exists, skip +this step. + +## 4. Run a small real-data smoke + +Use a separate output directory so the smoke cannot overwrite production +checkpoints: + +```bash +cp configs/my_run.yaml configs/my_smoke.yaml +# Edit output_dir in configs/my_smoke.yaml, for example: +# output_dir: runs/mrrate_single_writer_smoke +``` -## Train +From an allocated GPU node, inside the container: ```bash -bash scripts/train_online.sh configs/base.yaml -bash scripts/train_cached.sh configs/base.yaml +GPUS_PER_NODE=1 bash scripts/train_cached.sh \ + configs/my_smoke.yaml \ + --max-studies 8 \ + --max-updates 2 ``` -For a quick real-data test, add `--max-studies 2 --max-updates 1` to the Python -command in either launcher. Checkpoints contain the query -resampler/connector, report LoRA, optimizer/scheduler state, configuration, -data position, and per-rank RNG state. +For the online path: -For multi-node Slurm training, use `scripts/slurm_train.sh`. It launches one -`torchrun` agent per node, stages Gemma plus the encoder/MIL checkpoints once -per node, and places CUDA, Triton, and TorchInductor caches under node-local -`/tmp`. +```bash +GPUS_PER_NODE=1 bash scripts/train_online.sh \ + configs/my_smoke.yaml \ + --max-studies 2 \ + --max-updates 1 +``` + +Online mode is slower because it loads/resamples MRI volumes and runs the +frozen encoder during every epoch. The online reader supports +`batchXX/.zip`: it extracts only the current study under node-local +`/tmp` and removes it after encoding. + +## 5. Submit Slurm training + +The launcher defaults to two nodes with four GPUs per node. Command-line +`sbatch` options can override the node count. + +```bash +export PROJECT=/hnvme/workspace/b180dc51-sezgin/MR-RATE-report-training +export MODE=cached +export CONFIG="$PROJECT/configs/my_run.yaml" +export LLM_PATH=/absolute/path/to/gemma-3-12b-it +export ENCODER_CHECKPOINT=/absolute/path/to/the/mr-rate-checkpoint.pt +export MIL_CHECKPOINT=/absolute/path/to/the/corresponding/mil_head.pt +export GPUS_PER_NODE=4 +export MAX_STUDIES=0 +export MAX_UPDATES=0 + +sbatch --nodes=2 "$PROJECT/scripts/slurm_train.sh" +``` + +For online training, set `MODE=online`. The node launcher stages Gemma and +both checkpoints once per node, creates node-local CUDA/Triton/Inductor +caches, then launches one `torchrun` rank per GPU. + +The production settings in `configs/my_run.yaml` are: + +```yaml +epochs: 1 +batch_size: 1 +gradient_accumulation: 1 +learning_rate: 0.0001 +``` + +Every real study is seen exactly once. No study is duplicated to fill a +distributed batch. + +## 6. Resume + +```bash +export RESUME=/absolute/path/to/checkpoint-00000500.pt +sbatch --nodes=2 "$PROJECT/scripts/slurm_train.sh" +``` -The online reader supports the released `batchXX/.zip` layout. It -extracts only the current study to node-local `/tmp`, applies upstream MR-RATE -preprocessing, encodes it, and removes the extraction. +Keep `MODE`, `CONFIG`, node count, and artifact variables the same. Checkpoints +contain the resampler/connector, report LoRA, optimizer, scheduler, next data +position, and per-rank RNG state. -## Tests +## 7. Integration tests ```bash python tests/smoke_checks.py +sbatch scripts/synthetic_gpu_e2e.sbatch ``` The tests cover natural findings targets, strict cache validation, online/cached numerical equivalence, frozen MIL behavior, single-prefix construction, optimizer updates, strict weight provenance, and checkpoint -resume. `scripts/synthetic_gpu_e2e.sbatch` runs the integration gate on a -Slurm GPU with a deterministic dummy dataset. +resume. The Slurm test uses a deterministic dummy dataset and does not replace +the required real-data smoke. diff --git a/report-generation-training/configs/base.yaml b/report-generation-training/configs/base.yaml index d66e6d2..e35972a 100644 --- a/report-generation-training/configs/base.yaml +++ b/report-generation-training/configs/base.yaml @@ -48,7 +48,6 @@ training: learning_rate: 0.0001 weight_decay: 0.01 warmup_ratio: 0.03 - num_workers: 2 checkpoint_every: 500 shuffle: true replacement_sampling: false diff --git a/report-generation-training/scripts/slurm_train.sh b/report-generation-training/scripts/slurm_train.sh index d142a27..4a7fb77 100644 --- a/report-generation-training/scripts/slurm_train.sh +++ b/report-generation-training/scripts/slurm_train.sh @@ -26,9 +26,11 @@ MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1)" export MASTER_PORT="${MASTER_PORT:-29541}" export MODE CONFIG LLM_PATH ENCODER_CHECKPOINT MIL_CHECKPOINT export GPUS_PER_NODE="${GPUS_PER_NODE:-4}" +export MAX_STUDIES="${MAX_STUDIES:-0}" +export MAX_UPDATES="${MAX_UPDATES:-0}" +export RESUME="${RESUME:-}" exec srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ --kill-on-bad-exit=1 \ singularity exec --nv -B /hnvme:/hnvme,/tmp:/tmp "$sif" \ bash "$project_dir/scripts/train_node.sh" - diff --git a/report-generation-training/scripts/train_cached.sh b/report-generation-training/scripts/train_cached.sh index aea2a51..8104d69 100755 --- a/report-generation-training/scripts/train_cached.sh +++ b/report-generation-training/scripts/train_cached.sh @@ -4,4 +4,5 @@ project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$project_dir" export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" exec python -m torch.distributed.run --standalone --nproc-per-node="${GPUS_PER_NODE:-1}" \ - -m mrrate_report_training.train --mode cached --config "${1:-configs/base.yaml}" + -m mrrate_report_training.train --mode cached \ + --config "${1:-configs/base.yaml}" "${@:2}" diff --git a/report-generation-training/scripts/train_node.sh b/report-generation-training/scripts/train_node.sh index b021e08..ea01677 100644 --- a/report-generation-training/scripts/train_node.sh +++ b/report-generation-training/scripts/train_node.sh @@ -61,4 +61,10 @@ command=(python -m torch.distributed.run \ if [[ -n "${RESUME:-}" ]]; then command+=(--resume "$RESUME") fi +if [[ "${MAX_STUDIES:-0}" -gt 0 ]]; then + command+=(--max-studies "$MAX_STUDIES") +fi +if [[ "${MAX_UPDATES:-0}" -gt 0 ]]; then + command+=(--max-updates "$MAX_UPDATES") +fi exec "${command[@]}" diff --git a/report-generation-training/scripts/train_online.sh b/report-generation-training/scripts/train_online.sh index fc04779..5d808d4 100755 --- a/report-generation-training/scripts/train_online.sh +++ b/report-generation-training/scripts/train_online.sh @@ -4,4 +4,5 @@ project_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$project_dir" export PYTHONPATH="/hnvme/workspace/b180dc51-sezgin/extra-pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" exec python -m torch.distributed.run --standalone --nproc-per-node="${GPUS_PER_NODE:-1}" \ - -m mrrate_report_training.train --mode online --config "${1:-configs/base.yaml}" + -m mrrate_report_training.train --mode online \ + --config "${1:-configs/base.yaml}" "${@:2}" From 21e86dc49bcd085a7a5e66835cc61d554b6768bb Mon Sep 17 00:00:00 2001 From: Sezgin Er Date: Thu, 30 Jul 2026 15:21:25 +0200 Subject: [PATCH 4/4] Streamline report training README --- report-generation-training/README.md | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/report-generation-training/README.md b/report-generation-training/README.md index ff3a1c9..c44467c 100644 --- a/report-generation-training/README.md +++ b/report-generation-training/README.md @@ -21,25 +21,6 @@ the full bag, a trainable query resampler makes 512 language-prefix tokens, and one Gemma LoRA writer learns the complete findings text. The MIL probabilities are soft conditioning context; they are not report targets. -## Intentional training policy - -- One source-grounded raw findings target per study, preserving line order. -- Raw findings cover all 97,896 split IDs. The encoder/MIL-compatible cohort is - the same 97,887 studies retained by upstream MR-RATE's JSONL loader. -- Natural one-pass coverage: every train study occurs exactly once per epoch. -- No replacement sampler and no pathology oversampling. -- No MIL proposal dropout. -- No localization target, localization token, or localization loss. -- No MR-specific disease loss. -- The 1,536-token target ceiling preserves every normal findings report. It - bounds one corrupted 32k-token outlier (`5NIUCVXWHA`) containing embedded - editing dialogue. -- The encoder and MIL head are frozen. -- Exact cached training rejects `max_tokens_per_study != 0`. -- Startup requires the encoder SHA-256 and configuration recorded by MIL - training to match. Cached MIL additionally requires the training-cache - fingerprint to match the report token cache. - ## Required artifacts The files have distinct roles: