From c11d2ed90ba9cc6bdde6b784a35e0021fa5e0fae Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Sat, 26 Sep 2026 23:52:37 +0200 Subject: [PATCH 1/7] Port the FORA CT-DINO recipe fixes to MR DINOv3 The MR pipeline was adapted from the CT DINOv3-7B continuation recipe that failed in FORA. Bring over every fix the working FORA ViT-L run uses, on by default: - cosine prototypes on both heads (plain std-0.02 prototypes kept Sinkhorn targets flat) - DINOv2 ViT-L recipe: lr 1e-3 @1024 sqrt-scaled, warmup 2.5k, teacher temp warmup 10k, wd 0.04->0.4, EMA 0.992, beta2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr x0.2, prototype layers frozen for 1,250 steps (GroupedAdamW param groups) - ViT-L default (16k/16k prototypes, 2048/256 heads), batch 8/GPU, FP8 and activation checkpointing off; 7B/H+ launchers kept as legacy - global crops >=25% overlap snapped to the patch lattice, locals inside one global - cross-crop twin-token patch objective against last-block position leakage (pairs only same-sequence globals by default, preserving the no cross-sequence patch target rule) - position-binned (2x2x2) Sinkhorn for patch targets - target entropy / top-1 / usage / KL / cross metrics in metrics.jsonl Tests: tests/test_fora_recipe.py (8 new); all 18 pass on CPU. --- dino-pretraining/README.md | 37 ++- dino-pretraining/mr_dino/data.py | 46 ++- dino-pretraining/mr_dino/objective.py | 102 ++++++- dino-pretraining/mr_dino/recipe.py | 263 ++++++++++++++++++ dino-pretraining/mr_dino/train_ddp.py | 58 ++-- dino-pretraining/mr_dino/train_fsdp.py | 82 +++--- dino-pretraining/scripts/train_32n_7b.sbatch | 1 + .../scripts/train_32n_hplus.sbatch | 1 + .../scripts/train_32n_vitl.sbatch | 75 +++++ dino-pretraining/tests/test_fora_recipe.py | 174 ++++++++++++ dino-pretraining/tests/test_mr_dino.py | 9 +- 11 files changed, 786 insertions(+), 62 deletions(-) create mode 100644 dino-pretraining/mr_dino/recipe.py create mode 100644 dino-pretraining/scripts/train_32n_vitl.sbatch create mode 100644 dino-pretraining/tests/test_fora_recipe.py diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index 0cecd9e..9ce5c63 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -53,6 +53,39 @@ masking, distributed KoLeo, optional Gram anchoring, EMA teacher, atomic full state checkpoints, exact sampler/RNG resume, FSDP2, BF16, optional H200 FP8, and compile caches on node-local `/tmp`. +## Recipe fixes ported from FORA CT-DINO (Sep 2026) + +The CT pipeline this module was adapted from started with the same DINOv3-7B continuation +settings and failed; each item below is a diagnosed failure and the fix the working FORA +ViT-L run uses (code: `mr_dino/recipe.py`, `objective.py`, `data.py`; tests: +`tests/test_fora_recipe.py`). All are on by default. + +| Problem seen in FORA | Fix (default) | +|---|---| +| DINOv3 plain-linear prototype layer (std 0.02) never sharpened at lr 5e-5 → flat Sinkhorn targets for 10k steps (iBOT target entropy 8.0/11.5 nats, top-1 0.006) | cosine prototypes on both heads (`--normalize-prototypes`) | +| 1k warmup + 7B settings drove positional drift in the last blocks | DINOv2 ViT-L recipe: lr 1e-3 at batch 1024 (sqrt-scaled), warmup 2,500, teacher-temperature warmup 0.04→0.07 over 10k, weight decay 0.04→0.4, EMA 0.992→1, AdamW β2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr ×0.2, prototype layers frozen for 1,250 steps | +| 7B continuation was worse than a from-scratch ViT-L at the same budget | `--arch large` (ViT-L, 16k/16k prototypes, heads 2048/256), batch 8/GPU, no FP8, no activation checkpointing | +| global crops overlapping ≥75% were near-duplicates | ≥25% overlap per axis (`--global-overlap`), locals inside one of the two globals | +| head-facing block encoded where a token sat in its crop (last-block position leakage 0.20, crop-quadrant prototype fields) | second global snapped to the patch lattice + **cross-crop twin-token objective** (`--cross-view-weight 1.0`): the student's patch prediction in one crop matches the teacher target of the same voxels in the other crop; leakage fell to 0.09 and stayed flat | +| patch targets balanced over all masked tokens at once | Sinkhorn within 2×2×2 spatial bins of the crop (`--position-bins`, CAPI) | +| raw loss hid whether the student learned (CE = H + KL) | `metrics.jsonl` logs `dino/ibot_target_entropy`, `_top1`, `_prototype_usage`, `dino_kl`, `ibot_kl`, `cross`, `cross_kl`, `cross_pairs` | + +Already fixed here before the port: the distributed Sinkhorn uses the actual global token count. + +**MR-specific choice:** twin tokens are paired only when both global crops show the same +sequence (`--cross-view-sequences same`), preserving this module's rule that cross-sequence +images are never patch-level targets. With `--cross-sequence-probability 0.75` that is ~25% of +samples; `any` would also pair co-registered voxels across contrasts (contrast-invariant patch +features, at the risk of suppressing lesions visible in only one sequence). + +Production: `scripts/train_32n_vitl.sbatch`. `train_32n_7b.sbatch` / `train_32n_hplus.sbatch` +keep the old recipe for reference only. + +**Healthy early training (FORA reference, ViT-L, global batch 128–1024):** iBOT target top-1 +rises from ~0.003 to >0.05 within 2.5k steps, patch target entropy falls, prototype usage stays +≳0.4, `cross_kl` ~1.5–2.5 without a rising trend. Gram anchoring (stage `gram`) gave no gain in +FORA and stays optional. + ## Optional input-volume cache From `contrastive-pretraining/`: @@ -113,12 +146,12 @@ cd dino-pretraining DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ SPLITS_CSV=/path/to/splits.csv \ OUTPUT=/path/to/mrdino3d_7b/pretrain \ -sbatch scripts/train_32n_7b.sbatch +sbatch scripts/train_32n_vitl.sbatch # FORA recipe (ViT-L); 7B launcher kept as legacy ``` Important overrides are `STEPS`, `BATCH_SIZE`, `GRAD_ACCUM_STEPS`, `WORKERS`, `LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `WARMUP_STEPS`, `RESUME`, `FP8`, -and `COMPILE`. The default warmup is 1,000 optimizer steps. `RESUME=latest` +and `COMPILE`. The default warmup is 2,500 optimizer steps (FORA recipe). `RESUME=latest` loads the latest checkpoint only when its `COMPLETE` marker exists. For later stages: diff --git a/dino-pretraining/mr_dino/data.py b/dino-pretraining/mr_dino/data.py index 696f5b9..87154a0 100644 --- a/dino-pretraining/mr_dino/data.py +++ b/dino-pretraining/mr_dino/data.py @@ -25,6 +25,8 @@ import torch.nn.functional as F from torch.utils.data import Dataset, Sampler +from .recipe import overlap_pairs, snap_second_start, token_offset + CACHE_MANIFEST_NAME = "_manifest.json" EXPECTED_CACHE_LAYOUT = "per_subject_stack" @@ -376,6 +378,8 @@ def __init__( cross_sequence_probability: float = 0.75, candidate_trials: int = 12, seed: int = 3407, + global_overlap: float = 0.25, + patch_size: tuple[int, int, int] = (2, 16, 16), ) -> None: if (preprocessed_dir is None) == (data_folder is None): raise ValueError("Pass exactly one of preprocessed_dir or data_folder") @@ -383,6 +387,13 @@ def __init__( raise ValueError("Only atlas_space is valid for atlas-registered MR DINO") if not 0 <= cross_sequence_probability <= 1: raise ValueError("cross_sequence_probability must be in [0, 1]") + if not 0 <= global_overlap <= 1: + raise ValueError("global_overlap must be in [0, 1]") + # FORA fix: the two global crops only need >= 25% overlap per axis (0.75 made them near + # duplicates), and the second crop is snapped to the patch lattice so that every token in + # the overlap has an exact twin for the cross-crop patch objective. + self.global_overlap = float(global_overlap) + self.patch_size = tuple(int(v) for v in patch_size) self.crop_spec = crop_spec self.target_shape = tuple(int(x) for x in target_shape) self.target_spacing = tuple(float(x) for x in target_spacing) @@ -500,10 +511,13 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict: first_start = _informative_start( stack[first_sequence], self.crop_spec.global_shape, rng, self.candidate_trials ) - global_starts = _overlapping_global_starts( - tuple(stack.shape[1:]), self.crop_spec.global_shape, rng, first=first_start + container = tuple(stack.shape[1:]) + first, second = _overlapping_global_starts( + container, self.crop_spec.global_shape, rng, first=first_start, + min_overlap=self.global_overlap, ) - lo, hi = _intersection_box(global_starts, self.crop_spec.global_shape) + second = snap_second_start(first, second, container, self.crop_spec.global_shape, self.patch_size) + global_starts = (first, second) teacher_globals, student_globals = [], [] for sequence, start in zip(global_sequences, global_starts): @@ -514,6 +528,11 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict: local_views, local_sequences, local_starts = [], [], [] for _ in range(self.crop_spec.local_crops): sequence = rng.randrange(n_sequences) + # FORA fix: locals lie inside ONE of the globals (the intersection of two crops with + # 25% overlap can be thinner than a local crop). + parent = global_starts[rng.randrange(2)] + lo = parent + hi = tuple(min(p + g, c) for p, g, c in zip(parent, self.crop_spec.global_shape, container)) start = _contained_start(lo, hi, self.crop_spec.local_shape, rng) clean = _crop_at_start(stack[sequence], self.crop_spec.local_shape, start) local_views.append(_student_distortion(clean, rng).unsqueeze(0)) @@ -573,6 +592,7 @@ def collate_dino3d( samples: list[dict], patch_size: tuple[int, int, int] = (2, 16, 16), mask_ratio: tuple[float, float] = (0.1, 0.5), + cross_view_sequences: str = "same", ) -> dict: teacher = torch.stack([sample["teacher_global"] for sample in samples], dim=1) student = torch.stack([sample["student_global"] for sample in samples], dim=1) @@ -590,7 +610,27 @@ def collate_dino3d( masks = torch.stack(masks) indices = masks.flatten().nonzero().flatten() weights = (1.0 / masks.sum(-1).clamp(min=1)).unsqueeze(-1).expand_as(masks)[masks] + # Twin tokens of the two global crops (cross-crop patch objective). By default only when both + # globals show the SAME sequence, preserving this module's contract that cross-sequence images + # are never patch-level targets; "any" also pairs co-registered voxels across contrasts. + if cross_view_sequences not in {"same", "any", "none"}: + raise ValueError("cross_view_sequences must be 'same', 'any' or 'none'") + n = math.prod(grid) + pa, pb, ps = [], [], [] + for index, sample in enumerate(samples): + seq0, seq1 = sample["global_sequences"] + if cross_view_sequences == "none" or (cross_view_sequences == "same" and seq0 != seq1): + continue + g0, g1 = sample["global_starts"] + a, b = overlap_pairs(token_offset(g0, g1, patch_size), grid) + pa.append(a + index * n) + pb.append(b + index * n) + ps.append(torch.full((a.numel(),), index, dtype=torch.long)) + empty = torch.empty(0, dtype=torch.long) return { + "cross_a": torch.cat(pa) if pa else empty, + "cross_b": torch.cat(pb) if pb else empty, + "cross_sample": torch.cat(ps) if ps else empty, "teacher_global": teacher, "student_global": student, "student_local": local, diff --git a/dino-pretraining/mr_dino/objective.py b/dino-pretraining/mr_dino/objective.py index 917cd30..9680e4a 100644 --- a/dino-pretraining/mr_dino/objective.py +++ b/dino-pretraining/mr_dino/objective.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy +import math from dataclasses import dataclass import torch @@ -12,6 +13,7 @@ from dinov3.layers.dino_head import DINOHead from .model import DinoVisionTransformer3D +from .recipe import binned_sinkhorn, entropy, normalize_prototype_layer, padded_head, position_bin_ids, target_stats class SSLNetwork3D(nn.Module): @@ -24,6 +26,7 @@ def __init__( ibot_hidden_dim: int = 4096, dino_bottleneck_dim: int = 384, ibot_bottleneck_dim: int = 384, + normalize_prototypes: bool = True, ) -> None: super().__init__() self.backbone = backbone @@ -41,6 +44,11 @@ def __init__( bottleneck_dim=ibot_bottleneck_dim, nlayers=3, ) + if normalize_prototypes: + # FORA fix: cosine prototypes let the teacher commit from step 0 (plain std-0.02 + # prototypes kept the Sinkhorn targets flat for thousands of steps). + normalize_prototype_layer(self.dino_head) + normalize_prototype_layer(self.ibot_head) def init_weights(self) -> None: self.backbone.init_weights() @@ -184,6 +192,9 @@ def __init__( ibot_loss_chunk_size: int = 1024, with_gram_anchor: bool = False, initialize_weights: bool = True, + normalize_prototypes: bool = True, + position_bins: tuple[int, int, int] | None = (2, 2, 2), + cross_view_weight: float = 1.0, ) -> None: super().__init__() self.student = SSLNetwork3D( @@ -194,7 +205,11 @@ def __init__( ibot_head_hidden_dim or head_hidden_dim, dino_bottleneck_dim, ibot_bottleneck_dim, + normalize_prototypes=normalize_prototypes, ) + self.patch_size = tuple(int(v) for v in backbone.patch_size) + self.position_bins = tuple(int(v) for v in position_bins) if position_bins else None + self.cross_view_weight = float(cross_view_weight) if initialize_weights: self.student.init_weights() self.teacher = copy.deepcopy(self.student).requires_grad_(False) @@ -248,6 +263,49 @@ def reset_gram_anchor_from_teacher(self) -> None: self.gram_anchor.load_state_dict(self.teacher.state_dict()) self.gram_anchor.eval() + def _cross_view(self, batch, t_patch, s_patch, teacher_temperature, grid, n_tokens, batch_size, weighted_mean): + """Cross-crop patch objective (FORA v14): for every token present in both global crops + (crops aligned on the patch lattice; ``cross_a``/``cross_b`` are flat indices of the twins + in view 0 / view 1), the student's patch prediction in one crop must match the teacher's + Sinkhorn target for the same voxels in the other crop, in both directions. Without it the + iBOT term only compares a token with the teacher's view of the SAME crop, and the + head-facing block learns where a token sits in the crop (position leakage).""" + device = s_patch.device + a = batch["cross_a"].to(device) + b = batch["cross_b"].to(device) + batch_size * n_tokens # view 1 follows view 0 + s_flat = s_patch.flatten(0, 1) + t_flat = t_patch.flatten(0, 1) + n_pairs = int(a.numel()) + total_pairs = torch.tensor(n_pairs, device=device) + if dist.is_initialized(): + dist.all_reduce(total_pairs) + zero = s_patch.sum() * 0.0 + if int(total_pairs.item()) == 0: # no sample on any rank has twins (all ranks agree) + return zero, {"cross": zero, "cross_kl": zero, "cross_pairs": zero} + with torch.no_grad(): + t_tokens = torch.cat([t_flat.index_select(0, b), t_flat.index_select(0, a)]) + t_logits, _ = padded_head(self.teacher.ibot_head, t_tokens) + teacher_rows = torch.cat([b, a]) + if self.position_bins is not None: + q = binned_sinkhorn(distributed_sinkhorn, t_logits, teacher_temperature, + position_bin_ids(teacher_rows, n_tokens, grid, self.position_bins), + math.prod(self.position_bins)) + else: + q = distributed_sinkhorn(t_logits, teacher_temperature) + s_tokens = torch.cat([s_flat.index_select(0, a), s_flat.index_select(0, b)]) + s_logits, anchor = padded_head(self.student.ibot_head, s_tokens) + if n_pairs == 0: + return zero + anchor, {"cross": zero, "cross_kl": zero, "cross_pairs": zero} + logp = F.log_softmax(s_logits.float() / self.student_temperature, dim=-1) + ce = -(q * logp).sum(-1) + sample_ids = torch.cat([batch["cross_sample"], batch["cross_sample"]]).to(device) + counts = torch.bincount(sample_ids, minlength=batch_size).clamp_min(1).float() + per_sample = torch.zeros(batch_size, device=device).scatter_add_(0, sample_ids, ce) / counts + per_sample_h = torch.zeros(batch_size, device=device).scatter_add_(0, sample_ids, entropy(q)) / counts + cross = weighted_mean(per_sample) + anchor + return cross, {"cross": cross.detach(), "cross_kl": (cross - weighted_mean(per_sample_h)).detach(), + "cross_pairs": torch.tensor(float(n_pairs), device=device)} + def forward(self, batch: dict, teacher_temperature: float, step: int) -> tuple[Tensor, dict[str, Tensor]]: teacher_images = batch["teacher_global"] # [2,B,1,D,H,W] student_images = batch["student_global"] @@ -284,8 +342,19 @@ def weighted_mean(per_sample: Tensor) -> Tensor: t_patch = tout["x_norm_patchtokens"] t_masked = t_patch.flatten(0, 1).index_select(0, mask_indices) t_patch_logits = self.teacher.ibot_head(t_masked) - t_cls_prob = distributed_sinkhorn(t_cls_logits, teacher_temperature).unflatten(0, (n_global, batch_size)) - t_patch_prob = distributed_sinkhorn(t_patch_logits, teacher_temperature) + t_cls_flat = distributed_sinkhorn(t_cls_logits, teacher_temperature) + t_cls_prob = t_cls_flat.unflatten(0, (n_global, batch_size)) + n_tokens = int(masks.shape[-1]) + grid = tuple(n // p for n, p in zip(teacher_images.shape[-3:], self.patch_size)) + if self.position_bins is not None: + # FORA fix (CAPI): balance patch targets within spatial bins of the crop. + t_patch_prob = binned_sinkhorn( + distributed_sinkhorn, t_patch_logits, teacher_temperature, + position_bin_ids(mask_indices, n_tokens, grid, self.position_bins), + math.prod(self.position_bins), + ) + else: + t_patch_prob = distributed_sinkhorn(t_patch_logits, teacher_temperature) del t_patch_logits, t_masked global_out, local_out = self.student.backbone( @@ -375,14 +444,43 @@ def weighted_mean(per_sample: Tensor) -> Tensor: ).unflatten(0, (n_global, batch_size)).mean(0) gram = weighted_mean(gram_each) + # Teacher-target diagnostics (FORA fix): CE = H(target) + KL, so the raw loss hides + # whether the student learns; log H, KL = CE - H (weighted like the losses), top-1, usage. + with torch.no_grad(): + cls_stats = target_stats(t_cls_flat) + patch_stats = target_stats(t_patch_prob) + h_cls = weighted_mean(entropy(t_cls_prob).mean(0)) + h_rows = entropy(t_patch_prob) + h_mask = torch.zeros(masks.shape[0], device=h_rows.device, dtype=h_rows.dtype) + h_mask.scatter_add_(0, mask_ids, h_rows * mask_weights) + h_patch = weighted_mean(h_mask.unflatten(0, (n_global, batch_size)).mean(0)) + + cross = dino.new_zeros(()) + cross_metrics = {"cross": cross, "cross_kl": cross, "cross_pairs": cross} + if self.cross_view_weight > 0 and "cross_a" in batch: + cross, cross_metrics = self._cross_view( + batch, t_patch, global_out["x_norm_patchtokens"], teacher_temperature, + grid, n_tokens, batch_size, weighted_mean, + ) + total = ( self.weights.dino * dino + self.weights.ibot * ibot + self.weights.koleo * koleo + self.weights.gram * gram + + self.cross_view_weight * cross ) return total, { "loss": total.detach(), + "dino_target_entropy": cls_stats["entropy"], + "ibot_target_entropy": patch_stats["entropy"], + "dino_target_top1": cls_stats["top1"], + "ibot_target_top1": patch_stats["top1"], + "dino_prototype_usage": cls_stats["usage"], + "ibot_prototype_usage": patch_stats["usage"], + "dino_kl": (weighted_mean(dino_global) - h_cls).detach(), + "ibot_kl": (ibot - h_patch).detach(), + **{k: v.detach() for k, v in cross_metrics.items()}, "dino": dino.detach(), "dino_global": weighted_mean(dino_global).detach(), "dino_local": weighted_mean(dino_local).detach(), diff --git a/dino-pretraining/mr_dino/recipe.py b/dino-pretraining/mr_dino/recipe.py new file mode 100644 index 0000000..bfda780 --- /dev/null +++ b/dino-pretraining/mr_dino/recipe.py @@ -0,0 +1,263 @@ +"""Recipe fixes ported from the FORA CT DINOv3 pipeline (v13/v14, Sep 2026). + +The CT pipeline started from the same DINOv3-7B recipe as this module and failed in ways we +diagnosed and fixed one by one; every piece below is what the working CT ViT-L run uses: + +* cosine prototypes on both heads (``CosineLinear``): DINOv3's plain linear prototype layer + (std 0.02) never sharpened at small learning rates, so the Sinkhorn targets stayed flat + (iBOT target entropy 8.0/11.5 nats, top-1 0.006) for thousands of steps; +* position-binned Sinkhorn for the patch head (CAPI, arXiv 2502.08769): the patch targets are + balanced within spatial bins of the crop instead of across all masked tokens at once; +* DINOv2 optimizer groups: layer-wise lr decay, patch-embed lr x0.2, no weight decay on biases, + 1-D parameters and tokens, and the prototype layers frozen for the first steps; +* cross-crop twin tokens: with the two global crops aligned on the patch lattice every token in + their overlap has an exact twin, which the objective uses to stop the head-facing block from + encoding the token's position inside the crop (measured in CT: last-block position leakage + 0.201 -> 0.094, crop-quadrant prototype fields gone). +""" + +from __future__ import annotations + +import json +import math +import os + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import Tensor, nn + +TOKEN_NAMES = ("cls_token", "mask_token", "storage_tokens") + + +# ----------------------------------------------------------------------------- prototypes +class CosineLinear(nn.Linear): + """Bias-free prototype layer computing cosine logits (weights L2-normalized per prototype). + + DINOv3's ``DINOHead`` already L2-normalizes the bottleneck features before ``last_layer``, + so the logits are true cosine similarities. The parameter name stays ``last_layer.weight``, + keeping checkpoints of the plain layer loadable.""" + + def forward(self, x: Tensor) -> Tensor: + return F.linear(x, F.normalize(self.weight, dim=1, eps=1e-12)) + + +def normalize_prototype_layer(head: nn.Module) -> None: + layer = head.last_layer + if not isinstance(layer, nn.Linear) or layer.bias is not None: + raise ValueError("Expected a bias-free linear prototype layer") + if isinstance(layer, CosineLinear): + return + cosine = CosineLinear(layer.in_features, layer.out_features, bias=False, + device=layer.weight.device, dtype=layer.weight.dtype) + with torch.no_grad(): + cosine.weight.copy_(layer.weight) + head.last_layer = cosine + + +# ----------------------------------------------------------------------------- Sinkhorn bins +def position_bin_ids(token_indices: Tensor, n_tokens: int, grid, bins) -> Tensor: + """Spatial bin id in [0, prod(bins)) of flat token indices (modulo the crop's token count).""" + gz, gy, gx = (int(v) for v in grid) + if gz * gy * gx != int(n_tokens): + raise ValueError(f"position grid {grid} does not match {n_tokens} tokens per crop") + bz, by, bx = (int(v) for v in bins) + if not (1 <= bz <= gz and 1 <= by <= gy and 1 <= bx <= gx): + raise ValueError(f"position bins {bins} must fit the token grid {grid}") + pos = token_indices % int(n_tokens) + z = pos // (gy * gx) + y = (pos // gx) % gy + x = pos % gx + return ((z * bz) // gz) * (by * bx) + ((y * by) // gy) * bx + (x * bx) // gx + + +def binned_sinkhorn(base, logits: Tensor, temperature: float, bin_ids: Tensor, n_bins: int, + iterations: int = 3) -> Tensor: + """Run the distributed Sinkhorn ``base`` once per spatial bin; every rank visits every + non-empty bin in the same order (bin counts are all-reduced), so collectives stay matched.""" + if bin_ids.shape[0] != logits.shape[0]: + raise ValueError("bin ids must align with the logit rows") + out = torch.empty(logits.shape, dtype=torch.float32, device=logits.device) + counts = torch.bincount(bin_ids, minlength=n_bins).to(torch.int64).to(logits.device) + if dist.is_initialized(): + dist.all_reduce(counts, op=dist.ReduceOp.SUM) + for b in range(n_bins): + if int(counts[b]) == 0: + continue + rows = (bin_ids == b).nonzero().flatten() + q = base(logits.index_select(0, rows), temperature, iterations) + out.index_copy_(0, rows, q) + return out + + +# ----------------------------------------------------------------------------- optimizer +def _block_index(name: str): + parts = name.split(".") + for i, part in enumerate(parts[:-1]): + if part == "blocks" and parts[i + 1].isdigit(): + return int(parts[i + 1]) + return None + + +def build_param_groups(student: nn.Module, *, layerwise_decay: float, patch_embed_lr_mult: float, + n_blocks: int) -> list[dict]: + """DINOv2 ``get_params_groups_with_decay`` for the 3-D student (backbone + two heads). + + lr_mult: backbone blocks[i] get decay^(n_blocks - i); patch embedding, tokens and RoPE get + decay^(n_blocks + 1) (times ``patch_embed_lr_mult`` for the patch embedding); final norms + get 1; heads get 1. wd_mult: 0 for biases, 1-D parameters and tokens. Parameters that share + (lr_mult, wd_mult) are merged into one group.""" + if not 0.0 < layerwise_decay <= 1.0 or patch_embed_lr_mult <= 0 or n_blocks < 1: + raise ValueError("Require 0 < layerwise_decay <= 1, patch_embed_lr_mult > 0, n_blocks >= 1") + merged: dict[tuple[float, float], dict] = {} + for name, param in student.named_parameters(): + if not param.requires_grad: + continue + clean = name.replace("_fsdp_wrapped_module.", "").replace("_checkpoint_wrapped_module.", "").replace("_orig_mod.", "") + lr_mult, wd_mult = 1.0, 1.0 + if clean.startswith("backbone."): + block = _block_index(clean) + if block is not None: + layer_id = block + 1 + elif "patch_embed" in clean or any(t in clean for t in TOKEN_NAMES) or "rope" in clean: + layer_id = 0 + else: + layer_id = n_blocks + 1 + lr_mult = float(layerwise_decay) ** (n_blocks + 1 - layer_id) + if "patch_embed" in clean: + lr_mult *= float(patch_embed_lr_mult) + if param.ndim <= 1 or clean.endswith(".bias") or any(t in clean for t in TOKEN_NAMES): + wd_mult = 0.0 + key = (round(lr_mult, 12), wd_mult) + group = merged.setdefault(key, dict(params=[], lr_mult=float(lr_mult), wd_mult=float(wd_mult), names=[])) + group["params"].append(param) + group["names"].append(clean) + groups = [merged[k] for k in sorted(merged, key=lambda k: (-k[0], -k[1]))] + if os.environ.get("RANK", "0") == "0": + print(json.dumps(dict(event="param_groups", groups=[ + dict(lr_mult=g["lr_mult"], wd_mult=g["wd_mult"], n=len(g["params"]), example=g["names"][0]) + for g in groups])), flush=True) + return groups + + +class GroupedAdamW(torch.optim.AdamW): + """AdamW applying each group's ``lr_mult``/``wd_mult`` transiently inside ``step`` so the + trainer keeps writing the plain scheduled lr / weight decay into every group.""" + + @torch.no_grad() + def step(self, closure=None): + saved = [(g["lr"], g["weight_decay"]) for g in self.param_groups] + for g in self.param_groups: + g["lr"] = g["lr"] * float(g.get("lr_mult", 1.0)) + g["weight_decay"] = g["weight_decay"] * float(g.get("wd_mult", 1.0)) + try: + return super().step(closure) + finally: + for g, (lr, wd) in zip(self.param_groups, saved): + g["lr"], g["weight_decay"] = lr, wd + + +def make_optimizer(student: nn.Module, *, lr: float, weight_decay: float, betas, layerwise_decay: float, + patch_embed_lr_mult: float, n_blocks: int, fused: bool = True) -> GroupedAdamW: + groups = build_param_groups(student, layerwise_decay=layerwise_decay, + patch_embed_lr_mult=patch_embed_lr_mult, n_blocks=n_blocks) + for g in groups: + g.pop("names") + return GroupedAdamW(groups, lr=lr, betas=tuple(betas), weight_decay=weight_decay, fused=fused) + + +def cancel_last_layer_gradients(student: nn.Module) -> int: + """DINO/DINOv2 ``freeze_last_layer``: drop the prototype-layer gradients of both heads.""" + count = 0 + for name, param in student.named_parameters(): + if "last_layer" in name and param.grad is not None: + param.grad = None + count += 1 + return count + + +# ----------------------------------------------------------------------------- twin tokens +def snap_second_start(origin, second, container_shape, crop_shape, patch): + """Move ``second`` so that (second - origin) is a multiple of the patch on every axis, + staying inside the container.""" + out = [] + for o, s, have, want, p in zip(origin, second, container_shape, crop_shape, patch): + k = int(round((s - o) / p)) + lo = -(o // p) + hi = (max(have - want, 0) - o) // p + if hi < lo: + hi = lo + out.append(o + min(max(k, lo), hi) * p) + return tuple(out) + + +def token_offset(g0, g1, patch): + d = [] + for a, b, p in zip(g0, g1, patch): + if (b - a) % p: + raise ValueError("global crops are not aligned on the patch lattice") + d.append((b - a) // p) + return tuple(d) + + +def overlap_pairs(offset, grid) -> tuple[Tensor, Tensor]: + """Flat token indices (crop A, crop B) of every token covering the same voxels in both crops.""" + ranges = [] + for d, g in zip(offset, grid): + lo, hi = max(0, d), min(g, g + d) + if hi <= lo: + return torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long) + ranges.append(torch.arange(lo, hi)) + z, y, x = torch.meshgrid(*ranges, indexing="ij") + gz, gy, gx = grid + a = (z * gy + y) * gx + x + b = ((z - offset[0]) * gy + (y - offset[1])) * gx + (x - offset[2]) + return a.flatten(), b.flatten() + + +def padded_head(head: nn.Module, tokens: Tensor, chunk: int = 4096) -> tuple[Tensor, Tensor]: + """Run an FSDP-sharded head on a per-rank-variable token count with rank-identical call shapes. + Returns (logits[:local], anchor); anchor keeps a zero graph edge on ranks without tokens.""" + local = int(tokens.shape[0]) + maximum = torch.tensor(local, device=tokens.device) + if dist.is_initialized(): + dist.all_reduce(maximum, op=dist.ReduceOp.MAX) + maximum = int(maximum.item()) + parts, out = [], None + anchor = tokens.sum() * 0.0 + for start in range(0, maximum, chunk): + stop = min(start + chunk, maximum) + size = stop - start + valid = max(0, min(stop, local) - start) + x = tokens[start:start + valid] + if valid < size: + x = F.pad(x, (0, 0, 0, size - valid)) + out = head(x) + if valid == 0: + anchor = anchor + out.sum() * 0.0 + continue + parts.append(out[:valid]) + if parts: + logits = torch.cat(parts) + elif out is not None: + logits = out[:0] + else: + logits = tokens.new_zeros((0, head.last_layer.weight.shape[0])) + return logits, anchor + + +def entropy(q: Tensor) -> Tensor: + return -(q * q.clamp_min(1e-12).log()).sum(-1) + + +def target_stats(q: Tensor) -> dict[str, Tensor]: + """Teacher-target diagnostics: entropy, mean top-1 probability, fraction of prototypes that win.""" + if q.shape[0] == 0: + z = q.new_zeros(()) + return dict(entropy=z, top1=z, usage=z) + return dict(entropy=entropy(q).mean(), top1=q.max(-1).values.mean(), + usage=torch.tensor(q.argmax(-1).unique().numel() / q.shape[1], device=q.device)) + + +def sqrt_scaled_lr(base_lr: float, global_batch: int, reference: int = 1024) -> float: + return base_lr * math.sqrt(global_batch / reference) diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index fa75c74..296c473 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -32,6 +32,7 @@ stage_crop_spec, ) from .model import DinoVisionTransformer3D, vit_7b_3d, vit_hplus_3d, vit_large_3d +from .recipe import cancel_last_layer_gradients, make_optimizer from .objective import DINO3DLearner, LossWeights @@ -48,22 +49,23 @@ def parse_args() -> argparse.Namespace: p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") - p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="hplus") - p.add_argument("--steps", type=int, default=100_000) + p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="large") + p.add_argument("--steps", type=int, default=125_000) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--workers", type=int, default=2) p.add_argument("--seed", type=int, default=3407) - p.add_argument("--prototypes", type=int, default=65_536) - p.add_argument("--head-hidden-dim", type=int, default=4096) - p.add_argument("--lr", type=float, default=2e-4, help="Base LR before sqrt(global_batch/1024) scaling") + p.add_argument("--prototypes", type=int, default=16_384) + p.add_argument("--head-hidden-dim", type=int, default=2048) + p.add_argument("--bottleneck-dim", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3, help="Base LR before sqrt(global_batch/1024) scaling (DINOv2 ViT-L)") p.add_argument("--min-lr", type=float, default=1e-6) - p.add_argument("--warmup-steps", type=int, default=5_000) + p.add_argument("--warmup-steps", type=int, default=2_500) p.add_argument("--weight-decay", type=float, default=0.04) p.add_argument("--weight-decay-end", type=float, default=0.4) - p.add_argument("--teacher-momentum", type=float, default=0.994) + p.add_argument("--teacher-momentum", type=float, default=0.992) p.add_argument("--teacher-temperature", type=float, default=0.07) p.add_argument("--teacher-warmup-temperature", type=float, default=0.04) - p.add_argument("--teacher-warmup-steps", type=int, default=5_000) + p.add_argument("--teacher-warmup-steps", type=int, default=10_000) p.add_argument("--mask-min", type=float, default=0.1) p.add_argument("--mask-max", type=float, default=0.5) p.add_argument("--ibot-weight", type=float, default=1.0) @@ -80,13 +82,23 @@ def parse_args() -> argparse.Namespace: p.add_argument("--posterior-shift-mm", type=float, default=15.0) p.add_argument("--cross-sequence-probability", type=float, default=0.75) p.add_argument("--candidate-trials", type=int, default=12) - p.add_argument("--save-every", type=int, default=250) - p.add_argument("--keep-checkpoints", type=int, default=3) + p.add_argument("--save-every", type=int, default=500) + p.add_argument("--keep-checkpoints", type=int, default=6) p.add_argument("--log-every", type=int, default=10) p.add_argument("--resume", default="latest", help="latest, none, or a checkpoint path") p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--max-studies", type=int, default=None, help="Smoke only: cap studies assigned to each rank") p.add_argument("--deadline-margin-seconds", type=int, default=900) + # FORA recipe fixes (see mr_dino/recipe.py) + p.add_argument("--global-overlap", type=float, default=0.25) + p.add_argument("--normalize-prototypes", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--position-bins", type=int, nargs=3, default=(2, 2, 2)) + p.add_argument("--cross-view-weight", type=float, default=1.0) + p.add_argument("--cross-view-sequences", choices=("same", "any", "none"), default="same") + p.add_argument("--layerwise-decay", type=float, default=0.9) + p.add_argument("--patch-embed-lr-mult", type=float, default=0.2) + p.add_argument("--freeze-last-layer-steps", type=int, default=1250) + p.add_argument("--clip-grad", type=float, default=3.0) return p.parse_args() @@ -531,6 +543,7 @@ def main() -> int: cross_sequence_probability=args.cross_sequence_probability, candidate_trials=args.candidate_trials, seed=args.seed, + global_overlap=args.global_overlap, ) local_instances = torch.tensor(len(dataset), device=device, dtype=torch.long) total_instances = local_instances.clone() @@ -548,6 +561,7 @@ def main() -> int: collate_dino3d, patch_size=(2, 16, 16), mask_ratio=(args.mask_min, args.mask_max), + cross_view_sequences=args.cross_view_sequences if args.cross_view_weight > 0 else "none", ) loader = DataLoader( dataset, @@ -566,6 +580,11 @@ def main() -> int: backbone, prototypes=args.prototypes, head_hidden_dim=args.head_hidden_dim, + dino_bottleneck_dim=args.bottleneck_dim, + ibot_bottleneck_dim=args.bottleneck_dim, + normalize_prototypes=args.normalize_prototypes, + position_bins=None if tuple(args.position_bins) == (1, 1, 1) else tuple(args.position_bins), + cross_view_weight=args.cross_view_weight, loss_weights=LossWeights( dino=1.0, ibot=args.ibot_weight, @@ -582,7 +601,11 @@ def main() -> int: learner.student.backbone.blocks[i] = checkpoint_wrapper(block) learner.to(device) trainable = [p for p in learner.parameters() if p.requires_grad] - optimizer = torch.optim.AdamW(trainable, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay) + optimizer = make_optimizer( + learner.student, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999), + layerwise_decay=args.layerwise_decay, patch_embed_lr_mult=args.patch_embed_lr_mult, + n_blocks=len(learner.student.backbone.blocks), fused=device.type == "cuda", + ) if world > 1: learner = DDP( learner, @@ -657,12 +680,15 @@ def main() -> int: if not finite.item(): raise FloatingPointError(f"Non-finite distributed loss at step {step}: {loss.detach()}") backward_loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(trainable, 3.0) + if step < args.freeze_last_layer_steps: + cancel_last_layer_gradients(unwrap(learner).student) # DINO/DINOv2 freeze_last_layer + grad_norm = torch.nn.utils.clip_grad_norm_(trainable, args.clip_grad) optimizer.step() unwrap(learner).update_teacher(momentum) if (step + 1) % args.log_every == 0: - values = torch.stack([metrics[k].float() for k in ("loss", "dino", "ibot", "koleo", "gram")]) + metric_keys = sorted(metrics) + values = torch.stack([metrics[k].float() for k in metric_keys]) if dist.is_initialized(): dist.all_reduce(values) values /= world @@ -676,11 +702,7 @@ def main() -> int: "sequence_epoch": ( (step + 1) * global_batch / int(total_instances) ), - "loss": float(values[0]), - "dino": float(values[1]), - "ibot": float(values[2]), - "koleo": float(values[3]), - "gram": float(values[4]), + **{k: float(v) for k, v in zip(metric_keys, values)}, "lr": lr, "weight_decay": wd, "teacher_momentum": momentum, diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index f105e42..fce30db 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -33,6 +33,7 @@ ) from .fp8 import enable_fsdp_mixed_precision_fp8 from .objective import DINO3DLearner, LossWeights +from .recipe import cancel_last_layer_gradients, make_optimizer from .train_ddp import ( StopController, balanced_file_assignment, @@ -62,13 +63,14 @@ def parse_args() -> argparse.Namespace: p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") - p.add_argument("--arch", choices=("tiny", "7b"), default="7b") - p.add_argument("--steps", type=int, default=100_000) - p.add_argument("--batch-size", type=int, default=1) + p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="large", + help="FORA lesson: a from-scratch ViT-L beat the 7B continuation; 7b kept for reference") + p.add_argument("--steps", type=int, default=125_000) + p.add_argument("--batch-size", type=int, default=8) p.add_argument( "--grad-accum-steps", type=int, - default=4, + default=1, help="Fixed study microbatches per optimizer update on every rank", ) p.add_argument("--workers", type=int, default=2) @@ -81,18 +83,18 @@ def parse_args() -> argparse.Namespace: p.add_argument("--target-shape", type=int, nargs=3, default=(256, 384, 384)) p.add_argument("--posterior-shift-mm", type=float, default=15.0) p.add_argument("--seed", type=int, default=3407) - p.add_argument("--dino-prototypes", type=int, default=262_144) - p.add_argument("--ibot-prototypes", type=int, default=98_304) - p.add_argument("--dino-head-hidden-dim", type=int, default=8192) - p.add_argument("--ibot-head-hidden-dim", type=int, default=4096) - p.add_argument("--dino-bottleneck-dim", type=int, default=512) - p.add_argument("--ibot-bottleneck-dim", type=int, default=384) - p.add_argument("--lr", type=float, default=5e-5, help="Peak at global batch 1024 before sqrt scaling") + p.add_argument("--dino-prototypes", type=int, default=16_384) + p.add_argument("--ibot-prototypes", type=int, default=16_384) + p.add_argument("--dino-head-hidden-dim", type=int, default=2048) + p.add_argument("--ibot-head-hidden-dim", type=int, default=2048) + p.add_argument("--dino-bottleneck-dim", type=int, default=256) + p.add_argument("--ibot-bottleneck-dim", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3, help="Peak at global batch 1024 before sqrt scaling (DINOv2 ViT-L)") p.add_argument("--min-lr", type=float, default=1e-6) - p.add_argument("--warmup-steps", type=int, default=10_000) + p.add_argument("--warmup-steps", type=int, default=2_500) p.add_argument("--weight-decay", type=float, default=0.04) - p.add_argument("--weight-decay-end", type=float, default=0.04) - p.add_argument("--teacher-momentum", type=float, default=0.994) + p.add_argument("--weight-decay-end", type=float, default=0.4) + p.add_argument("--teacher-momentum", type=float, default=0.992) p.add_argument("--teacher-temperature", type=float, default=0.07) p.add_argument("--teacher-warmup-temperature", type=float, default=0.04) p.add_argument("--teacher-warmup-steps", type=int, default=10_000) @@ -105,16 +107,26 @@ def parse_args() -> argparse.Namespace: p.add_argument("--ibot-loss-chunk-size", type=int, default=1024) p.add_argument("--cross-sequence-probability", type=float, default=0.75) p.add_argument("--candidate-trials", type=int, default=12) - p.add_argument("--clip-grad", type=float, default=30.0) - p.add_argument("--save-every", type=int, default=250) - p.add_argument("--keep-checkpoints", type=int, default=3) + p.add_argument("--clip-grad", type=float, default=3.0) + p.add_argument("--save-every", type=int, default=500) + p.add_argument("--keep-checkpoints", type=int, default=6) p.add_argument("--log-every", type=int, default=10) p.add_argument("--resume", default="latest", help="latest, none, or a DCP checkpoint directory") - p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=False) p.add_argument("--selective-checkpointing", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True) - p.add_argument("--fp8", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--fp8", action=argparse.BooleanOptionalAction, default=False) p.add_argument("--max-studies", type=int, default=None, help="Smoke only: cap studies per rank") + # FORA recipe fixes (see mr_dino/recipe.py) + p.add_argument("--global-overlap", type=float, default=0.25, help="Minimum per-axis overlap of the two global crops") + p.add_argument("--normalize-prototypes", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--position-bins", type=int, nargs=3, default=(2, 2, 2), help="Spatial Sinkhorn bins for patch targets; 1 1 1 disables") + p.add_argument("--cross-view-weight", type=float, default=1.0, help="Cross-crop twin-token patch objective (0 disables)") + p.add_argument("--cross-view-sequences", choices=("same", "any", "none"), default="same") + p.add_argument("--layerwise-decay", type=float, default=0.9) + p.add_argument("--patch-embed-lr-mult", type=float, default=0.2) + p.add_argument("--freeze-last-layer-steps", type=int, default=1250) + p.add_argument("--adamw-beta2", type=float, default=0.999) args = p.parse_args() # A running Slurm allocation can cleanly restart the worker step without # relinquishing its nodes. Keep schedule changes in a small, auditable @@ -245,6 +257,9 @@ def build_learner(args: argparse.Namespace, device: torch.device) -> DINO3DLearn ibot_loss_chunk_size=args.ibot_loss_chunk_size, with_gram_anchor=args.stage in {"gram", "highres"}, initialize_weights=False, + normalize_prototypes=args.normalize_prototypes, + position_bins=None if tuple(args.position_bins) == (1, 1, 1) else tuple(args.position_bins), + cross_view_weight=args.cross_view_weight, ) if args.activation_checkpointing: wrapper = _checkpoint_wrapper(args.selective_checkpointing) @@ -463,6 +478,7 @@ def main() -> int: cross_sequence_probability=args.cross_sequence_probability, candidate_trials=args.candidate_trials, seed=args.seed, + global_overlap=args.global_overlap, ) local_instances = torch.tensor(len(dataset), device=device, dtype=torch.long) total_instances = local_instances.clone() @@ -477,13 +493,14 @@ def main() -> int: ) learner = build_learner(args, device) - trainable = list(learner.student.parameters()) - optimizer = torch.optim.AdamW( - trainable, + optimizer = make_optimizer( + learner.student, lr=args.lr, - betas=(0.9, 0.99), weight_decay=args.weight_decay, - fused=True, + betas=(0.9, args.adamw_beta2), + layerwise_decay=args.layerwise_decay, + patch_embed_lr_mult=args.patch_embed_lr_mult, + n_blocks=len(learner.student.backbone.blocks), ) start = load_checkpoint(resolve_checkpoint(output, args.resume), learner, optimizer, sampler, args, rank) transition_step = int(getattr(args, "warmup_transition_step", 0)) @@ -497,6 +514,7 @@ def main() -> int: collate_dino3d, patch_size=(2, 16, 16), mask_ratio=(args.mask_min, args.mask_max), + cross_view_sequences=args.cross_view_sequences if args.cross_view_weight > 0 else "none", ) loader_generator = torch.Generator().manual_seed(args.seed + rank * 104729) loader = DataLoader( @@ -566,7 +584,7 @@ def main() -> int: group["weight_decay"] = wd optimizer.zero_grad(set_to_none=True) - metric_accumulator = torch.zeros(5, device=device, dtype=torch.float32) + metric_accumulator = None for microstep in range(args.grad_accum_steps): batch = next(iterator) consumed_samples += args.batch_size @@ -596,9 +614,11 @@ def main() -> int: f"{backward_loss.detach()}; metrics={metric_debug}" ) backward_loss.backward() - metric_accumulator += torch.stack( - [metrics[k].float() for k in ("loss", "dino", "ibot", "koleo", "gram")] - ) + metric_keys = sorted(metrics) + values_now = torch.stack([metrics[k].float() for k in metric_keys]) + metric_accumulator = values_now if metric_accumulator is None else metric_accumulator + values_now + if step < args.freeze_last_layer_steps: + cancel_last_layer_gradients(learner.student) # DINO/DINOv2 freeze_last_layer grad_norms = [ torch.nn.utils.clip_grad_norm_(module.parameters(), args.clip_grad) for module in (learner.student.backbone, learner.student.dino_head, learner.student.ibot_head) @@ -621,11 +641,7 @@ def main() -> int: / int(total_instances) ), "grad_accum_steps": args.grad_accum_steps, - "loss": float(values[0]), - "dino": float(values[1]), - "ibot": float(values[2]), - "koleo": float(values[3]), - "gram": float(values[4]), + **{k: float(v) for k, v in zip(metric_keys, values)}, "lr": lr, "weight_decay": wd, "teacher_momentum": momentum, diff --git a/dino-pretraining/scripts/train_32n_7b.sbatch b/dino-pretraining/scripts/train_32n_7b.sbatch index d7dee8b..847406a 100755 --- a/dino-pretraining/scripts/train_32n_7b.sbatch +++ b/dino-pretraining/scripts/train_32n_7b.sbatch @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_32n_vitl.sbatch. #SBATCH --job-name=mr_dino3_7b #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_32n_hplus.sbatch b/dino-pretraining/scripts/train_32n_hplus.sbatch index 2daf2b5..305169d 100755 --- a/dino-pretraining/scripts/train_32n_hplus.sbatch +++ b/dino-pretraining/scripts/train_32n_hplus.sbatch @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_32n_vitl.sbatch. #SBATCH --job-name=mr_dino3_hplus #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch new file mode 100644 index 0000000..25a62f5 --- /dev/null +++ b/dino-pretraining/scripts/train_32n_vitl.sbatch @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_dino3_vitl +#SBATCH --partition=h200 +#SBATCH --nodes=32 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:h200:4 +#SBATCH --cpus-per-task=128 +#SBATCH --exclusive +#SBATCH --exclude=h21-24,h23-24,h24-17,h23-18 +#SBATCH --time=1-00:00:00 +#SBATCH --requeue +#SBATCH --signal=B:USR1@1200 +#SBATCH --output=mrdino_%x_%j.out +#SBATCH --error=mrdino_%x_%j.err +#SBATCH --mail-type=ALL +#SBATCH --mail-user=sezgin.er@dqbm.uzh.ch +set -uo pipefail + +if [[ -z "${CODE:-}" ]]; then + if [[ -d "$SLURM_SUBMIT_DIR/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR" + elif [[ -d "$SLURM_SUBMIT_DIR/dino-pretraining/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR/dino-pretraining" + else + echo "Cannot locate dino-pretraining; set CODE explicitly" >&2 + exit 2 + fi +fi +SCRIPT_DIR="$CODE/scripts" +: "${DATA_FOLDER:?Set DATA_FOLDER to the MR-RATE-atlas NIfTI root}" +STAGE=${STAGE:-pretrain} +OUTPUT=${OUTPUT:-"$PWD/mrdino3d_vitl/$STAGE"} +mkdir -p "$OUTPUT" +rm -f "$OUTPUT/REQUEUE_REQUESTED" + +head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) +MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') +export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29675} +end_ts=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*EndTime=\([0-9T:-]*\).*/\1/p' | head -1) +if [[ -n "$end_ts" ]]; then + export MRDINO_DEADLINE_EPOCH=$(( $(date -d "$end_ts" +%s) - 1200 )) +fi + +case "$STAGE" in + pretrain) STEPS=${STEPS:-125000}; GRAM_WEIGHT=0 ;; + gram) STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0} ;; + highres) STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5} ;; + *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; +esac +[[ "${COMPILE:-1}" == 1 ]] && COMPILE_ARGS=(--compile) || COMPILE_ARGS=(--no-compile) +[[ "${FP8:-1}" == 1 ]] && FP8_ARGS=(--fp8) || FP8_ARGS=(--no-fp8) +SPLIT_ARGS=() +if [[ -n "${SPLITS_CSV:-}" ]]; then + SPLIT_ARGS=(--splits-csv "$SPLITS_CSV") +fi + +echo "MR DINOv3 ViT-L (FORA recipe): stage=$STAGE nodes=$SLURM_NNODES output=$OUTPUT data=$DATA_FOLDER" +srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ + --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" \ + --data-folder "$DATA_FOLDER" \ + --space atlas_space \ + "${SPLIT_ARGS[@]}" \ + --output-dir "$OUTPUT" --stage "$STAGE" --arch large --steps "$STEPS" \ + --batch-size "${BATCH_SIZE:-8}" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ + --workers "${WORKERS:-8}" --local-crops "${LOCAL_CROPS:-8}" \ + --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.75}" \ + --cross-view-sequences "${CROSS_VIEW_SEQUENCES:-same}" \ + --resume "${RESUME:-latest}" --gram-weight "$GRAM_WEIGHT" \ + "${COMPILE_ARGS[@]}" --no-fp8 --no-activation-checkpointing +rc=$? +if [[ -f "$OUTPUT/REQUEUE_REQUESTED" ]]; then + scontrol requeue "$SLURM_JOB_ID" + exit 0 +fi +exit "$rc" diff --git a/dino-pretraining/tests/test_fora_recipe.py b/dino-pretraining/tests/test_fora_recipe.py new file mode 100644 index 0000000..a895b93 --- /dev/null +++ b/dino-pretraining/tests/test_fora_recipe.py @@ -0,0 +1,174 @@ +"""Tests for the recipe fixes ported from the FORA CT DINOv3 pipeline (mr_dino/recipe.py).""" +import math + +import pytest +import torch + +from mr_dino.data import CropSpec, MRAtlasDINO3DDataset, collate_dino3d +from mr_dino.recipe import ( + CosineLinear, + GroupedAdamW, + binned_sinkhorn, + build_param_groups, + cancel_last_layer_gradients, + normalize_prototype_layer, + overlap_pairs, + position_bin_ids, + snap_second_start, + token_offset, +) +from test_mr_dino import make_dummy_cache + +PATCH = (2, 8, 8) +GLOBAL = (8, 32, 32) + + +def big_dataset(cache, cross_sequence_probability=0.0, seed=5): + return MRAtlasDINO3DDataset( + preprocessed_dir=str(cache), + crop_spec=CropSpec(global_shape=GLOBAL, local_shape=(4, 16, 16), local_crops=3), + cross_sequence_probability=cross_sequence_probability, + candidate_trials=3, + seed=seed, + global_overlap=0.25, + patch_size=PATCH, + ) + + +def test_cosine_prototypes_are_bounded_and_keep_parameter_name(): + head = torch.nn.Module() + head.last_layer = torch.nn.Linear(16, 10, bias=False) + torch.nn.init.normal_(head.last_layer.weight, std=5.0) + normalize_prototype_layer(head) + assert isinstance(head.last_layer, CosineLinear) + assert "last_layer.weight" in dict(head.named_parameters()) + x = torch.nn.functional.normalize(torch.randn(7, 16), dim=-1) + logits = head.last_layer(x) + assert logits.abs().max() <= 1.0 + 1e-5 + expected = x @ torch.nn.functional.normalize(head.last_layer.weight, dim=1).t() + torch.testing.assert_close(logits, expected) + + +def test_global_crops_overlap_quarter_and_sit_on_patch_lattice(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + moved = 0 + for i in range(len(dataset)): + sample = dataset[(i, 1)] + g0, g1 = sample["global_starts"] + token_offset(g0, g1, PATCH) # raises unless patch-aligned + for a, b, n in zip(g0, g1, GLOBAL): + assert min(a + n, b + n) - max(a, b) >= math.floor(0.25 * n) - max(PATCH) + moved += g0 != g1 + for start in sample["local_starts"]: + assert any(all(g <= s and s + l <= g + G for s, l, g, G in zip(start, (4, 16, 16), p, GLOBAL)) + for p in (g0, g1)) + assert moved > 0 # crops are no longer near-duplicates + + +def test_twin_tokens_cover_identical_voxels(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + grid = tuple(n // p for n, p in zip(GLOBAL, PATCH)) + checked = 0 + for i in range(len(dataset)): + sample = dataset[(i, 3)] + g0, g1 = sample["global_starts"] + a, b = overlap_pairs(token_offset(g0, g1, PATCH), grid) + crop0, crop1 = sample["teacher_global"][0, 0], sample["teacher_global"][1, 0] + def patch(crop, flat): + z, rem = divmod(int(flat), grid[1] * grid[2]); y, x = divmod(rem, grid[2]) + return crop[z*PATCH[0]:(z+1)*PATCH[0], y*PATCH[1]:(y+1)*PATCH[1], x*PATCH[2]:(x+1)*PATCH[2]] + for ta, tb in list(zip(a, b))[:20]: + torch.testing.assert_close(patch(crop0, ta), patch(crop1, tb)) # same sequence (p=0) + checked += 1 + assert checked > 0 + + +def test_collate_pairs_only_same_sequence_by_default(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64)), cross_sequence_probability=1.0) + samples = [dataset[(i, 0)] for i in range(4)] + assert all(s["global_sequences"][0] != s["global_sequences"][1] for s in samples) + assert collate_dino3d(samples, patch_size=PATCH)["cross_a"].numel() == 0 + assert collate_dino3d(samples, patch_size=PATCH, cross_view_sequences="any")["cross_a"].numel() > 0 + + +def test_binned_sinkhorn_balances_each_bin_separately(): + torch.manual_seed(0) + grid, bins = (4, 4, 4), (2, 2, 2) + idx = torch.arange(64) + ids = position_bin_ids(idx, 64, grid, bins) + assert torch.bincount(ids).tolist() == [8] * 8 + logits = torch.randn(64, 16) * 3 + from mr_dino.objective import distributed_sinkhorn + q = binned_sinkhorn(distributed_sinkhorn, logits, 0.1, ids, 8) + torch.testing.assert_close(q.sum(-1), torch.ones(64)) + for b in range(8): # identical to balancing each bin on its own + rows = (ids == b).nonzero().flatten() + torch.testing.assert_close(q[rows], distributed_sinkhorn(logits[rows], 0.1)) + assert not torch.allclose(q, distributed_sinkhorn(logits, 0.1)) # and differs from one global balance + + +def test_param_groups_layerwise_decay_and_grouped_adamw(): + class Student(torch.nn.Module): + def __init__(self): + super().__init__() + self.backbone = torch.nn.Module() + self.backbone.patch_embed = torch.nn.Linear(4, 4) + self.backbone.blocks = torch.nn.ModuleList([torch.nn.Linear(4, 4) for _ in range(3)]) + self.backbone.norm = torch.nn.LayerNorm(4) + self.backbone.cls_token = torch.nn.Parameter(torch.zeros(1, 1, 4)) + self.dino_head = torch.nn.Module() + self.dino_head.last_layer = torch.nn.Linear(4, 8, bias=False) + student = Student() + groups = build_param_groups(student, layerwise_decay=0.5, patch_embed_lr_mult=0.2, n_blocks=3) + lr_of = {name: g["lr_mult"] for g in groups for name in g["names"]} + wd_of = {name: g["wd_mult"] for g in groups for name in g["names"]} + assert lr_of["backbone.patch_embed.weight"] == pytest.approx(0.5 ** 4 * 0.2) + assert lr_of["backbone.blocks.0.weight"] == pytest.approx(0.5 ** 3) + assert lr_of["backbone.blocks.2.weight"] == pytest.approx(0.5) + assert lr_of["backbone.norm.weight"] == 1.0 and lr_of["dino_head.last_layer.weight"] == 1.0 + assert wd_of["backbone.blocks.0.bias"] == 0.0 and wd_of["backbone.cls_token"] == 0.0 and wd_of["backbone.blocks.0.weight"] == 1.0 + for g in groups: + g.pop("names") + opt = GroupedAdamW(groups, lr=1.0, weight_decay=0.1) + for g in opt.param_groups: + g["lr"] = 1e-2 + student.backbone.patch_embed.weight.grad = torch.ones(4, 4) + before = student.backbone.patch_embed.weight.detach().clone() + opt.step() + assert all(g["lr"] == 1e-2 for g in opt.param_groups) # schedule untouched after step + step = (before - student.backbone.patch_embed.weight.detach()).abs().max() + assert step == pytest.approx(1e-2 * 0.5 ** 4 * 0.2, rel=0.05) # Adam step ~ lr * lr_mult + + +def test_cancel_last_layer_gradients(): + head = torch.nn.Module() + head.last_layer = torch.nn.Linear(3, 3, bias=False) + head.mlp = torch.nn.Linear(3, 3) + (head.last_layer(head.mlp(torch.randn(2, 3)))).sum().backward() + assert cancel_last_layer_gradients(head) == 1 + assert head.last_layer.weight.grad is None and head.mlp.weight.grad is not None + + +def test_cross_view_term_trains_the_backbone(tmp_path): + pytest.importorskip("dinov3") + from mr_dino.model import DinoVisionTransformer3D + from mr_dino.objective import DINO3DLearner, LossWeights + + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + batch = collate_dino3d([dataset[(0, 1)], dataset[(1, 1)]], patch_size=PATCH) + assert batch["cross_a"].numel() > 0 + for key in ("teacher_global", "student_global", "student_local"): + batch[key] = batch[key].float() + batch["loss_weights"] = batch["sample_weights"] + torch.manual_seed(0) + backbone = DinoVisionTransformer3D(volume_size=GLOBAL, patch_size=PATCH, voxel_spacing_mm=(1.0, 0.5, 0.5), + embed_dim=96, depth=2, num_heads=3, ffn_ratio=2, n_storage_tokens=2, drop_path_rate=0) + learner = DINO3DLearner(backbone, prototypes=32, head_hidden_dim=64, dino_bottleneck_dim=32, ibot_bottleneck_dim=32, + loss_weights=LossWeights(dino=0, ibot=0, koleo=0, gram=0), position_bins=(2, 2, 2), cross_view_weight=1.0) + loss, metrics = learner(batch, teacher_temperature=0.07, step=0) + assert metrics["cross_pairs"] > 0 and torch.isfinite(metrics["cross"]) and metrics["cross"] > 0 + for key in ("dino_target_entropy", "ibot_target_entropy", "dino_kl", "ibot_kl", "ibot_prototype_usage"): + assert torch.isfinite(metrics[key]) + loss.backward() # only the cross term is weighted: it must reach the backbone blocks + grads = [p.grad for p in learner.student.backbone.blocks.parameters() if p.grad is not None] + assert grads and any(g.abs().sum() > 0 for g in grads) diff --git a/dino-pretraining/tests/test_mr_dino.py b/dino-pretraining/tests/test_mr_dino.py index 0d91cf8..df50620 100644 --- a/dino-pretraining/tests/test_mr_dino.py +++ b/dino-pretraining/tests/test_mr_dino.py @@ -124,10 +124,11 @@ def test_aligned_cross_sequence_views_and_determinism(tmp_path): assert torch.equal(first["teacher_global"], again["teacher_global"]) assert first["teacher_global"].shape == (2, 1, 8, 32, 32) assert first["student_local"].shape == (2, 1, 4, 16, 16) - lo, hi = _intersection_box(first["global_starts"], (8, 32, 32)) - for start in first["local_starts"]: - for s, n, left, right in zip(start, (4, 16, 16), lo, hi): - assert left <= s and s + n <= right + for start in first["local_starts"]: # each local crop lies inside one of the two globals + assert any( + all(g <= s and s + n <= g + G for s, n, g, G in zip(start, (4, 16, 16), parent, (8, 32, 32))) + for parent in first["global_starts"] + ) def test_collate_masks_match_patch_grid(tmp_path): From 984addb2ed18d1af4562509940842c9d5698c53d Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 03:40:07 +0200 Subject: [PATCH 2/7] Resume checkpoints saved during freeze_last_layer The frozen prototype layers get no gradient during the first --freeze-last-layer-steps steps, so AdamW holds no state for them and checkpoints from that phase lack their optimizer keys. The strict DCP resume then aborted (GPU smoke 941555). Allow exactly those keys to be missing; any other missing model/optimizer key still raises. --- dino-pretraining/mr_dino/recipe.py | 19 +++++ dino-pretraining/mr_dino/train_fsdp.py | 10 ++- .../tests/test_resume_frozen_last_layer.py | 70 +++++++++++++++++++ 3 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 dino-pretraining/tests/test_resume_frozen_last_layer.py diff --git a/dino-pretraining/mr_dino/recipe.py b/dino-pretraining/mr_dino/recipe.py index bfda780..08c65d2 100644 --- a/dino-pretraining/mr_dino/recipe.py +++ b/dino-pretraining/mr_dino/recipe.py @@ -176,6 +176,25 @@ def cancel_last_layer_gradients(student: nn.Module) -> int: return count +def missing_frozen_optimizer_keys(ckpt_dir, model: nn.Module, optimizer: torch.optim.Optimizer) -> list[str]: + """Checkpoint keys the resume would miss that are legitimately absent: optimizer state of the ``last_layer`` + prototype weights, which never receive a gradient during ``freeze_last_layer`` (``cancel_last_layer_gradients`` + sets ``grad=None``, so AdamW never creates their state and checkpoints saved in the first + ``--freeze-last-layer-steps`` steps do not contain it). Any OTHER missing model/optimizer key raises, so a + partial load can never silently skip real weights. Returns the allowed missing keys (empty = strict load).""" + from torch.distributed.checkpoint import FileSystemReader + from torch.distributed.checkpoint._nested_dict import flatten_state_dict + from torch.distributed.checkpoint.state_dict import get_model_state_dict, get_optimizer_state_dict + saved = set(FileSystemReader(str(ckpt_dir)).read_metadata().state_dict_metadata) + want, _ = flatten_state_dict({"model": get_model_state_dict(model), + "optimizer": get_optimizer_state_dict(model, optimizer)}) + missing = [k for k in want if k not in saved] + bad = [k for k in missing if not (k.startswith("optimizer.state.") and ".last_layer." in k)] + if bad: + raise RuntimeError(f"Checkpoint {ckpt_dir} lacks {len(bad)} required keys, e.g. {bad[:5]}") + return missing + + # ----------------------------------------------------------------------------- twin tokens def snap_second_start(origin, second, container_shape, crop_shape, patch): """Move ``second`` so that (second - origin) is a multiple of the patch on every axis, diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index fce30db..7e37dd0 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -33,7 +33,7 @@ ) from .fp8 import enable_fsdp_mixed_precision_fp8 from .objective import DINO3DLearner, LossWeights -from .recipe import cancel_last_layer_gradients, make_optimizer +from .recipe import cancel_last_layer_gradients, make_optimizer, missing_frozen_optimizer_keys from .train_ddp import ( StopController, balanced_file_assignment, @@ -378,11 +378,17 @@ def load_checkpoint( "refusing an inexact sampler resume" ) source_stage = metadata["stage"] + same_stage = source_stage == args.stage + # checkpoints written during freeze_last_layer have no optimizer state for the frozen prototype layers + frozen_missing = missing_frozen_optimizer_keys(path, learner, optimizer) if same_stage else [] + if frozen_missing and rank == 0: + print(f"[resume] {len(frozen_missing)} optimizer-state keys of frozen last_layer params absent " + "(saved during freeze_last_layer); they start fresh", flush=True) loaded_step = int(dcp_load( path, model=learner, optimizer=optimizer, - strict_loading=(source_stage == args.stage), + strict_loading=same_stage and not frozen_missing, )) state = torch.load(path / f"runtime_rank{rank:04d}.pt", map_location="cpu", weights_only=False) sampler.set_offset(int(state["sampler_offset"])) diff --git a/dino-pretraining/tests/test_resume_frozen_last_layer.py b/dino-pretraining/tests/test_resume_frozen_last_layer.py new file mode 100644 index 0000000..c98f472 --- /dev/null +++ b/dino-pretraining/tests/test_resume_frozen_last_layer.py @@ -0,0 +1,70 @@ +"""Resume from a checkpoint saved during freeze_last_layer (GPU smoke 941555 failed on this): +the frozen prototype layer never gets a gradient, so AdamW has no state for it and the checkpoint lacks +`optimizer.state.<...>.last_layer.*`. Only those keys may be missing; anything else must still fail loudly.""" +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from dinov3.checkpointer import load_checkpoint, save_checkpoint +from mr_dino.recipe import cancel_last_layer_gradients, missing_frozen_optimizer_keys + + +@pytest.fixture(autouse=True) +def single_process_group(tmp_path): + """dinov3's DCP checkpointer needs a process group; a 1-rank gloo group is enough on CPU.""" + if not dist.is_initialized(): + dist.init_process_group("gloo", init_method=f"file://{tmp_path}/pg", rank=0, world_size=1) + yield + dist.destroy_process_group() + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.backbone = nn.Linear(4, 4) + self.dino_head = nn.ModuleDict({"mlp": nn.Linear(4, 4), "last_layer": nn.Linear(4, 3, bias=False)}) + + def forward(self, x): + return self.dino_head["last_layer"](self.dino_head["mlp"](self.backbone(x))) + + +def _train_step(model, opt, freeze=True, drop=None): + opt.zero_grad(set_to_none=True) + model(torch.randn(2, 4)).sum().backward() + if freeze: + cancel_last_layer_gradients(model) + if drop is not None: # simulate another parameter without optimizer state + dict(model.named_parameters())[drop].grad = None + opt.step() + + +def test_resume_during_freeze_allows_only_last_layer_state(tmp_path): + torch.manual_seed(0) + model = Tiny() + opt = torch.optim.AdamW(model.parameters(), lr=1e-2) + _train_step(model, opt) + save_checkpoint(tmp_path / "ckpt", iteration=1, model=model, optimizer=opt) + + fresh = Tiny() + fresh_opt = torch.optim.AdamW(fresh.parameters(), lr=1e-2) + missing = missing_frozen_optimizer_keys(tmp_path / "ckpt", fresh, fresh_opt) + assert missing and all(k.startswith("optimizer.state.") and ".last_layer." in k for k in missing) + with pytest.raises(BaseException, match="Missing key in checkpoint state_dict.*last_layer"): # old strict resume = GPU smoke 941555 failure + strict_model = Tiny() + load_checkpoint(tmp_path / "ckpt", model=strict_model, optimizer=torch.optim.AdamW(strict_model.parameters()), strict_loading=True) + it = load_checkpoint(tmp_path / "ckpt", model=fresh, optimizer=fresh_opt, strict_loading=not missing) + assert it == 1 + for (n, a), (_, b) in zip(model.state_dict().items(), fresh.state_dict().items()): + assert torch.equal(a, b), n + + +def test_other_missing_state_still_raises(tmp_path): + torch.manual_seed(0) + model = Tiny() + opt = torch.optim.AdamW(model.parameters(), lr=1e-2) + _train_step(model, opt, drop="dino_head.mlp.weight") + save_checkpoint(tmp_path / "ckpt", iteration=1, model=model, optimizer=opt) + fresh = Tiny() + with pytest.raises(RuntimeError, match="lacks"): + missing_frozen_optimizer_keys(tmp_path / "ckpt", fresh, torch.optim.AdamW(fresh.parameters(), lr=1e-2)) From ff5eefa19677d5ed3bdb7f47a3d5313dac80b993 Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:42:53 +0200 Subject: [PATCH 3/7] Default cross-sequence global pairs to 0.25 A DINO target shared by a T1 and a FLAIR crop rewards dropping contrast-specific findings, and in atlas space the cheapest shared signal is the crop position. Keep a small rate so MIL bags that mix sequences share one embedding space; 0 vs 0.25 to be settled by the validation MIL probe. Also fix the stale 75% global-overlap line in the README. --- dino-pretraining/README.md | 15 +++++++++++---- dino-pretraining/mr_dino/data.py | 2 +- dino-pretraining/mr_dino/train_ddp.py | 2 +- dino-pretraining/mr_dino/train_fsdp.py | 2 +- dino-pretraining/scripts/train_32n_vitl.sbatch | 2 +- 5 files changed, 15 insertions(+), 8 deletions(-) diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index 9ce5c63..4c43f0e 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -34,9 +34,9 @@ keeps only its current study stack in memory, avoiding repeated NIfTI decoding. For each anchor: 1. Global view 1 uses the anchor sequence. -2. Global view 2 uses a different atlas-aligned sequence with probability 0.75 - (or the anchor when no second sequence exists). The two crops overlap by at - least 75% on every axis. +2. Global view 2 uses a different atlas-aligned sequence with probability 0.25 + (`--cross-sequence-probability`; otherwise, or when no second sequence exists, + the anchor). The two crops overlap by at least 25% on every axis. 3. Local crops stay inside the global intersection and may come from any aligned sequence. 4. DINO and KoLeo therefore learn anatomy shared across MR contrasts. @@ -74,10 +74,17 @@ Already fixed here before the port: the distributed Sinkhorn uses the actual glo **MR-specific choice:** twin tokens are paired only when both global crops show the same sequence (`--cross-view-sequences same`), preserving this module's rule that cross-sequence -images are never patch-level targets. With `--cross-sequence-probability 0.75` that is ~25% of +images are never patch-level targets. With `--cross-sequence-probability 0.25` that is ~75% of samples; `any` would also pair co-registered voxels across contrasts (contrast-invariant patch features, at the risk of suppressing lesions visible in only one sequence). +**Cross-sequence global pairs default to 0.25 (was 0.75).** A DINO target shared by a T1 and a +FLAIR crop rewards discarding contrast-specific findings (FLAIR-only gliosis, SWI-only bleeds), +and in atlas space the cheapest shared signal is the crop's atlas position. Many "different" +sequences are the same contrast in another plane, so the true cross-contrast share is lower still. +A small non-zero rate keeps one embedding space across contrasts for MIL bags that mix sequences; +0 vs 0.25 is to be settled with the validation MIL probe. + Production: `scripts/train_32n_vitl.sbatch`. `train_32n_7b.sbatch` / `train_32n_hplus.sbatch` keep the old recipe for reference only. diff --git a/dino-pretraining/mr_dino/data.py b/dino-pretraining/mr_dino/data.py index 87154a0..78fba24 100644 --- a/dino-pretraining/mr_dino/data.py +++ b/dino-pretraining/mr_dino/data.py @@ -375,7 +375,7 @@ def __init__( posterior_shift_mm: float = 15.0, cache_files: list[str] | None = None, raw_samples: list[dict] | None = None, - cross_sequence_probability: float = 0.75, + cross_sequence_probability: float = 0.25, candidate_trials: int = 12, seed: int = 3407, global_overlap: float = 0.25, diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index 296c473..6b3f41d 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -80,7 +80,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--target-spacing", type=float, nargs=3, default=(1.0, 0.5, 0.5)) p.add_argument("--target-shape", type=int, nargs=3, default=(256, 384, 384)) p.add_argument("--posterior-shift-mm", type=float, default=15.0) - p.add_argument("--cross-sequence-probability", type=float, default=0.75) + p.add_argument("--cross-sequence-probability", type=float, default=0.25) p.add_argument("--candidate-trials", type=int, default=12) p.add_argument("--save-every", type=int, default=500) p.add_argument("--keep-checkpoints", type=int, default=6) diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index 7e37dd0..4bf1352 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -105,7 +105,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--gram-weight", type=float, default=1.0) p.add_argument("--gram-max-tokens", type=int, default=1024) p.add_argument("--ibot-loss-chunk-size", type=int, default=1024) - p.add_argument("--cross-sequence-probability", type=float, default=0.75) + p.add_argument("--cross-sequence-probability", type=float, default=0.25) p.add_argument("--candidate-trials", type=int, default=12) p.add_argument("--clip-grad", type=float, default=3.0) p.add_argument("--save-every", type=int, default=500) diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch index 25a62f5..d8c0cf0 100644 --- a/dino-pretraining/scripts/train_32n_vitl.sbatch +++ b/dino-pretraining/scripts/train_32n_vitl.sbatch @@ -63,7 +63,7 @@ srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ --output-dir "$OUTPUT" --stage "$STAGE" --arch large --steps "$STEPS" \ --batch-size "${BATCH_SIZE:-8}" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ --workers "${WORKERS:-8}" --local-crops "${LOCAL_CROPS:-8}" \ - --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.75}" \ + --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.25}" \ --cross-view-sequences "${CROSS_VIEW_SEQUENCES:-same}" \ --resume "${RESUME:-latest}" --gram-weight "$GRAM_WEIGHT" \ "${COMPILE_ARGS[@]}" --no-fp8 --no-activation-checkpointing From 19a72a211bb3e59db5ac7eeabcd1b1d2df777749 Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:03:28 +0200 Subject: [PATCH 4/7] Train MR DINO on co-registered MRI and add the validation MIL probe Coreg (default) keeps every sequence on the study's T1w center grid; atlas resamples into the 1 mm MNI box. The model input is 1.0x0.5x0.5 mm, so from atlas the 0.5 mm in-plane detail is interpolation, and the MNI box crops skull base and neck (-12% head volume on a checked study). Both spaces are rigidly aligned across sequences; native_space is rejected. Checkpoints record their space and resume refuses to cross spaces (older checkpoints are atlas-space). mr_dino.probe scores a checkpoint like FORA CT-DINO: export the EMA teacher, extract frozen last-block tokens of every labelled val study (phase-1-size tiles, 2x2x2 foreground pooling, fp16 shards on node-local /tmp), patient-level 5-fold CV of MR-RATE's ClassifyThenAggregate head on the 14 merged groups, out-of-fold macro/per-class AUROC and AUPRC. Launchers: probe_node.sh (one node or inside a training allocation), probe_checkpoint.sbatch, probe_watch.sh (probes listed steps as checkpoints appear). The GPU smoke now also runs the probe on its own checkpoint. --- dino-pretraining/README.md | 97 +++- dino-pretraining/mr_dino/__init__.py | 6 +- dino-pretraining/mr_dino/data.py | 77 ++- dino-pretraining/mr_dino/probe.py | 513 ++++++++++++++++++ dino-pretraining/mr_dino/train_ddp.py | 77 ++- dino-pretraining/mr_dino/train_fsdp.py | 29 +- .../scripts/probe_checkpoint.sbatch | 15 + dino-pretraining/scripts/probe_node.sh | 56 ++ dino-pretraining/scripts/probe_watch.sh | 36 ++ dino-pretraining/scripts/smoke_dummy.sbatch | 21 +- .../scripts/train_32n_vitl.sbatch | 4 +- ...taset.py => make_dummy_aligned_dataset.py} | 7 +- dino-pretraining/tests/make_dummy_cache.py | 9 +- dino-pretraining/tests/test_fora_recipe.py | 4 +- dino-pretraining/tests/test_mr_dino.py | 58 +- dino-pretraining/tests/test_probe.py | 182 +++++++ 16 files changed, 1062 insertions(+), 129 deletions(-) create mode 100644 dino-pretraining/mr_dino/probe.py create mode 100755 dino-pretraining/scripts/probe_checkpoint.sbatch create mode 100755 dino-pretraining/scripts/probe_node.sh create mode 100755 dino-pretraining/scripts/probe_watch.sh rename dino-pretraining/tests/{make_dummy_atlas_dataset.py => make_dummy_aligned_dataset.py} (82%) create mode 100644 dino-pretraining/tests/test_probe.py diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index 4c43f0e..c077e2c 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -1,22 +1,38 @@ -# Atlas-space 3-D DINOv3 for MR-RATE +# Co-registered 3-D DINOv3 for MR-RATE This module adapts FORA's CT DINOv3 training strategy to MR-RATE's -**atlas-registered multi-sequence MRI**. It is self-supervised: reports and -pathology labels are not used. +**co-registered multi-sequence MRI** (`coreg_space`, default; `atlas_space` remains +selectable). It is self-supervised: reports and pathology labels are not used for +training; they are used only by the validation probe (below) to score checkpoints. ## Data and view contract -Production training reads the existing MR-RATE-atlas NIfTI tree directly: +Production training reads an extracted MR-RATE-coreg NIfTI tree directly: ```text -/batchXX//atlas_img/*.nii.gz +/batchXX//coreg_img/*.nii.gz # --space coreg_space (default) +/batchXX//atlas_img/*.nii.gz # --space atlas_space ``` +**Why coreg, not atlas.** Both spaces use the same rigid (ANTs, Mattes MI) registration, so +both are voxel-aligned across sequences, which the cross-sequence views need. They differ +in the target grid: +- **Coreg** keeps every sequence on the study's T1w center scan's native grid (e.g. + 0.47×0.47×1.2 mm, 190×240×240 mm field of view). +- **Atlas** resamples into the 1 mm MNI box (193×229×193 mm). + +The model input is 1.0×0.5×0.5 mm, so from atlas the 0.5 mm in-plane detail is +interpolation, and the MNI box crops the skull base and neck (−12% head volume on a +checked study), which matters for the spinal labels. `native_space` is rejected +(its sequences are not aligned). + It imports the canonical discovery and per-volume preprocessing functions from `contrastive-pretraining/scripts/data.py`, the loader used by previous MR-RATE MIL training. Therefore atlas selection, canonical RAS orientation, physical resampling, z-score normalization, posterior shift, and crop/padding are exactly -shared rather than reimplemented. `--space atlas_space` maps to `atlas_img`. +shared rather than reimplemented. `--space coreg_space` maps to `coreg_img`, +`--space atlas_space` to `atlas_img`. Checkpoints record their space, and resume refuses +to cross spaces (checkpoints written before the coreg switch are atlas-space). An optional preprocessed volume cache remains available with `--preprocessed-dir`. This is separate from cached MIL: cached MIL stores the @@ -99,9 +115,9 @@ From `contrastive-pretraining/`: ```bash python scripts/preprocess_volumes.py \ - --data_folder /path/to/MR-RATE-atlas/mri \ + --data_folder /path/to/MR-RATE-coreg/mri \ --out_dir /path/to/mrrate_preprocessed \ - --space atlas_space \ + --space coreg_space \ --normalizer zscore \ --num_workers 8 ``` @@ -146,11 +162,12 @@ The job succeeds only after `step_00000002/COMPLETE` exists. ## Production training -The production launcher uses the existing MR-RATE-atlas dataset directly: +The production launcher reads the extracted MR-RATE-coreg tree directly (`SPACE=atlas_space` +switches back to the atlas tree): ```bash cd dino-pretraining -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ +DATA_FOLDER=/path/to/MR-RATE-coreg/mri \ SPLITS_CSV=/path/to/splits.csv \ OUTPUT=/path/to/mrdino3d_7b/pretrain \ sbatch scripts/train_32n_vitl.sbatch # FORA recipe (ViT-L); 7B launcher kept as legacy @@ -161,15 +178,21 @@ Important overrides are `STEPS`, `BATCH_SIZE`, `GRAD_ACCUM_STEPS`, `WORKERS`, and `COMPILE`. The default warmup is 2,500 optimizer steps (FORA recipe). `RESUME=latest` loads the latest checkpoint only when its `COMPLETE` marker exists. -For later stages: +Phases 2 and 3 resume from the previous stage's checkpoint (the schedule restarts; +the Gram anchor is set to the incoming teacher): ```bash -STAGE=gram RESUME=/path/to/pretrain/checkpoints/step_00100000 \ -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ -OUTPUT=/path/to/mrdino3d_7b/gram \ -sbatch scripts/train_32n_7b.sbatch +STAGE=gram RESUME=/path/to/pretrain/checkpoints/step_00020000 \ +DATA_FOLDER=/path/to/MR-RATE-coreg/mri OUTPUT=/path/to/mrdino3d_vitl/gram \ +sbatch scripts/train_32n_vitl.sbatch # 10,000 steps, Gram weight 1.0 +STAGE=highres RESUME=/path/to/gram/checkpoints/step_00010000 \ +DATA_FOLDER=/path/to/MR-RATE-coreg/mri OUTPUT=/path/to/mrdino3d_vitl/highres \ +sbatch scripts/train_32n_vitl.sbatch # 5,000 steps, Gram 1.5, 64x384x384 globals ``` +Probe the phase-2 end against the phase-1 end first. If Gram lost AUROC, start +phase 3 from the phase-1 checkpoint instead (FORA CT did this; there Gram gave no gain). + The launcher catches the Slurm wall-time signal, writes a complete distributed checkpoint with per-rank sampler and RNG state, and requeues the same job. Do not start production before `scripts/smoke_dummy.sbatch` passes in the same @@ -188,13 +211,55 @@ mail notifications, 24-hour requeue policy, deadline checkpointing, node-local compiler caches, InfiniBand/NCCL settings, and optional split filtering. The 7B path is the production target; H+ is a smaller DDP fallback. +## Validation probe (checkpoint scoring) + +This is the FORA CT-DINO protocol: frozen features, MIL on the validation split only, and +patient-level cross-validation. Train and test are never touched. + +1. **Export** (`python -m mr_dino.probe export`): copies the EMA-teacher backbone out of a + DCP checkpoint, so checkpoint rotation cannot delete it. +2. **Extract** (`python -m mr_dino.probe extract`, one worker per GPU): + - Runs every sequence of every labelled `val` study (3,764 studies, 3,413 patients) + through the training preprocessing of the checkpoint's own space/spacing/shape. + - Cuts each volume into phase-1-size tiles (64×192×192 voxels) and keeps the last + block's normalised patch tokens. + - Average-pools the tokens 2×2×2 over head foreground and drops air. + - Writes one fp16 token file plus one index per worker to node-local `/tmp`. +3. **CV** (`python -m mr_dino.probe cv --fold k`): trains MR-RATE's NeuroVFM + `ClassifyThenAggregate` head per patient-level fold (5 folds, 8 epochs, batch 8, + lr 2e-3, pos_weight clipped to [1, 100]). One bag holds all sequences of a study. +4. **Summarize**: out-of-fold macro and per-class AUROC/AUPRC → `results.json`, + `per_class.csv`, `summary.txt`. + +Labels default to the 14 merged neuroradiology groups +(`contrastive-pretraining/scripts/eval_labels/splits_merged_majority`). Set +`LABELS_CSV`/`SPLITS_CSV` to use the 32/37-pathology labels instead. + +```bash +# one checkpoint, one node (extract -> 5 folds -> summary; ~1-2 h on 4 H200): +DATA_FOLDER=/path/to/MR-RATE-coreg/mri \ +sbatch scripts/probe_checkpoint.sbatch /path/to/pretrain/checkpoints/step_00005000 /path/to/pretrain/probe/step_5000 +# follow a running stage: export each listed step as it appears, then submit its probe: +DATA_FOLDER=/path/to/MR-RATE-coreg/mri nohup scripts/probe_watch.sh /path/to/pretrain 2500 5000 10000 15000 20000 & +# inside a running training allocation instead of a new job: +srun --overlap --jobid JOB -N1 -n1 -w NODE scripts/probe_node.sh CKPT OUT_DIR +``` + +Use the probe to: +- decide how long to run phase 1; +- check that phase 2 (Gram) does not lose AUROC before phase 3 starts from it; +- settle open recipe questions, e.g. `--cross-sequence-probability` 0 vs 0.25. + +Compare per-class rows, not just the macro AUROC: FLAIR-, SWI- and DWI-dependent groups +are what cross-sequence choices can hurt. + ## Lightweight/debug run `mr_dino.train_ddp` supports one or more GPUs and a tiny backbone: ```bash torchrun --standalone --nproc-per-node=1 -m mr_dino.train_ddp \ - --data-folder /path/to/MR-RATE-atlas/mri \ + --data-folder /path/to/MR-RATE-coreg/mri \ --output-dir /tmp/mrdino_debug \ --arch tiny --steps 10 --workers 0 --local-crops 2 \ --prototypes 64 --head-hidden-dim 128 \ diff --git a/dino-pretraining/mr_dino/__init__.py b/dino-pretraining/mr_dino/__init__.py index eee5563..378ae02 100644 --- a/dino-pretraining/mr_dino/__init__.py +++ b/dino-pretraining/mr_dino/__init__.py @@ -1,5 +1,5 @@ -"""Volumetric DINOv3 training for atlas-registered MR-RATE studies.""" +"""Volumetric DINOv3 training for aligned (coreg/atlas) MR-RATE studies.""" -from .data import CropSpec, MRAtlasDINO3DDataset, collate_dino3d +from .data import CropSpec, MRAlignedDINO3DDataset, collate_dino3d -__all__ = ["CropSpec", "MRAtlasDINO3DDataset", "collate_dino3d"] +__all__ = ["CropSpec", "MRAlignedDINO3DDataset", "collate_dino3d"] diff --git a/dino-pretraining/mr_dino/data.py b/dino-pretraining/mr_dino/data.py index 78fba24..f018d3c 100644 --- a/dino-pretraining/mr_dino/data.py +++ b/dino-pretraining/mr_dino/data.py @@ -1,6 +1,10 @@ -"""Atlas-registered multi-sequence MR sampling for volumetric DINOv3. +"""Aligned multi-sequence MR sampling for volumetric DINOv3. -The default path discovers and preprocesses atlas NIfTIs with the canonical +Training uses MR-RATE's co-registered space by default (``coreg_space``: every +sequence rigidly registered to the study's T1w center scan and kept on that +scan's native grid, so no MNI field-of-view crop and no 1 mm resampling); +``atlas_space`` remains selectable. Both are voxel-aligned across sequences. +The default path discovers and preprocesses NIfTIs with the canonical loader used by the previous MR-RATE training code. The optional NPZ path reads the volume cache produced by that loader's ``preprocess_volumes.py`` helper. Global and local DINO views share a physical region, but may use different @@ -29,6 +33,8 @@ CACHE_MANIFEST_NAME = "_manifest.json" +ALIGNED_SPACES = ("coreg_space", "atlas_space") +DEFAULT_SPACE = "coreg_space" EXPECTED_CACHE_LAYOUT = "per_subject_stack" _PREVIOUS_DATA_MODULE = None @@ -211,21 +217,30 @@ def _cache_space_dir(preprocessed_dir: str, space: str) -> Path: return nested if nested.is_dir() else root -def validate_atlas_cache( +def check_aligned_space(space: str) -> str: + """Cross-sequence views need voxel alignment: native_space is never accepted.""" + if space not in ALIGNED_SPACES: + raise ValueError( + f"MR DINO requires an aligned space {ALIGNED_SPACES}; got {space!r} " + "(native_space sequences are not voxel-aligned)" + ) + return space + + +def validate_aligned_cache( preprocessed_dir: str, - space: str = "atlas_space", + space: str = DEFAULT_SPACE, expected_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), ) -> dict: - """Fail early unless this is a compatible MR-RATE atlas-space cache.""" - if space != "atlas_space": - raise ValueError("MR DINO requires space='atlas_space'; native/coreg inputs are not accepted") + """Fail early unless this is a compatible MR-RATE aligned-space cache.""" + check_aligned_space(space) space_dir = _cache_space_dir(preprocessed_dir, space) manifest_path = space_dir / CACHE_MANIFEST_NAME if not manifest_path.is_file(): raise FileNotFoundError(f"Missing MR-RATE preprocessing manifest: {manifest_path}") manifest = json.loads(manifest_path.read_text()) - if manifest.get("space") != "atlas_space": - raise ValueError(f"Cache is {manifest.get('space')!r}, expected 'atlas_space'") + if manifest.get("space") != space: + raise ValueError(f"Cache is {manifest.get('space')!r}, expected {space!r}") if manifest.get("layout") != EXPECTED_CACHE_LAYOUT: raise ValueError( f"Unsupported cache layout {manifest.get('layout')!r}; expected {EXPECTED_CACHE_LAYOUT!r}" @@ -236,13 +251,13 @@ def validate_atlas_cache( return manifest -def discover_atlas_cache( +def discover_aligned_cache( preprocessed_dir: str, splits_csv: str | None = None, split: str = "train", - space: str = "atlas_space", + space: str = DEFAULT_SPACE, ) -> list[dict]: - manifest = validate_atlas_cache(preprocessed_dir, space) + manifest = validate_aligned_cache(preprocessed_dir, space) space_dir = _cache_space_dir(preprocessed_dir, space) selected = _load_split_ids(splits_csv, split) samples = [] @@ -258,7 +273,7 @@ def discover_atlas_cache( "volume_shape": shape[1:], }) if not samples: - raise RuntimeError(f"No {split!r} atlas-registered studies found under {space_dir}") + raise RuntimeError(f"No {split!r} {space} studies found under {space_dir}") return samples @@ -308,14 +323,16 @@ def previous_mr_data_module(): return module -def discover_raw_atlas( +def discover_raw_aligned( data_folder: str, selected: set[str] | None = None, + space: str = DEFAULT_SPACE, ) -> list[dict]: - """Discover atlas NIfTIs through the previous training dataloader.""" + """Discover aligned NIfTIs (coreg_img/ or atlas_img/) through the previous training dataloader.""" + check_aligned_space(space) samples = [] previous = previous_mr_data_module() - for subject in previous.discover_subjects(data_folder, "atlas_space"): + for subject in previous.discover_subjects(data_folder, space): study_uid = str(subject["subject_id"]) if selected is not None and study_uid not in selected: continue @@ -326,17 +343,18 @@ def discover_raw_atlas( "n_sequences": len(paths), }) if not samples: - raise RuntimeError(f"No atlas-registered NIfTIs found under {data_folder}") + raise RuntimeError(f"No {space} NIfTIs found under {data_folder}") return samples -def discover_raw_atlas_split( +def discover_raw_aligned_split( data_folder: str, splits_csv: str | None = None, split: str = "train", + space: str = DEFAULT_SPACE, ) -> list[dict]: - """Discover a split using the previous loader's atlas-space convention.""" - return discover_raw_atlas(data_folder, _load_split_ids(splits_csv, split)) + """Discover a split using the previous loader's space convention.""" + return discover_raw_aligned(data_folder, _load_split_ids(splits_csv, split), space) def _raw_volume( @@ -358,8 +376,8 @@ def _raw_volume( return torch.from_numpy(np.ascontiguousarray(array)).to(torch.bfloat16) -class MRAtlasDINO3DDataset(Dataset): - """Study-level SSL samples from atlas-aligned, variable-count MR sequences.""" +class MRAlignedDINO3DDataset(Dataset): + """Study-level SSL samples from voxel-aligned, variable-count MR sequences.""" def __init__( self, @@ -368,7 +386,7 @@ def __init__( data_folder: str | None = None, splits_csv: str | None = None, split: str = "train", - space: str = "atlas_space", + space: str = DEFAULT_SPACE, crop_spec: CropSpec = CropSpec(), target_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), target_shape: tuple[int, int, int] = (256, 384, 384), @@ -383,8 +401,7 @@ def __init__( ) -> None: if (preprocessed_dir is None) == (data_folder is None): raise ValueError("Pass exactly one of preprocessed_dir or data_folder") - if space != "atlas_space": - raise ValueError("Only atlas_space is valid for atlas-registered MR DINO") + self.space = check_aligned_space(space) if not 0 <= cross_sequence_probability <= 1: raise ValueError("cross_sequence_probability must be in [0, 1]") if not 0 <= global_overlap <= 1: @@ -406,10 +423,10 @@ def __init__( self._cached_names: list[str] | None = None self.preprocessed = preprocessed_dir is not None if self.preprocessed: - manifest = validate_atlas_cache(preprocessed_dir, space, target_spacing) + manifest = validate_aligned_cache(preprocessed_dir, space, target_spacing) self.target_shape = tuple(int(x) for x in manifest["target_shape"]) if cache_files is None: - self.samples = discover_atlas_cache(preprocessed_dir, splits_csv, split, space) + self.samples = discover_aligned_cache(preprocessed_dir, splits_csv, split, space) else: selected = _load_split_ids(splits_csv, split) self.samples = [ @@ -423,18 +440,18 @@ def __init__( if selected is None or Path(path).stem in selected ] if not self.samples: - raise RuntimeError("This rank received no atlas-registered cache studies") + raise RuntimeError(f"This rank received no {space} cache studies") else: selected = _load_split_ids(splits_csv, split) if raw_samples is None: - self.samples = discover_raw_atlas(data_folder, selected) + self.samples = discover_raw_aligned(data_folder, selected, space) else: self.samples = [ sample for sample in raw_samples if selected is None or sample["study_uid"] in selected ] if not self.samples: - raise RuntimeError("This rank received no atlas-registered NIfTI studies") + raise RuntimeError(f"This rank received no {space} NIfTI studies") for sample in self.samples: if "volume_shape" in sample and tuple(sample["volume_shape"]) != self.target_shape: raise ValueError( diff --git a/dino-pretraining/mr_dino/probe.py b/dino-pretraining/mr_dino/probe.py new file mode 100644 index 0000000..9108072 --- /dev/null +++ b/dino-pretraining/mr_dino/probe.py @@ -0,0 +1,513 @@ +"""Validation MIL probe for MR DINO checkpoints (the FORA CT-DINO checkpoint protocol). + +A checkpoint is scored without touching train or test data: + +1. ``export``: copy the EMA-teacher backbone out of a DCP training checkpoint into + one small ``.pt`` file, so later steps do not depend on checkpoint rotation. +2. ``extract``: run the frozen backbone over every sequence of every labelled + study of one split (default ``val``). + - Each volume goes through the training preprocessing (same space, spacing and + shape as the checkpoint), cut into training-size tiles (default: the phase-1 + global crop, 64x192x192 voxels). + - Output: the last block's normalised patch tokens. + - Tokens are average-pooled ``--pool`` cells (default 2x2x2) over head-foreground + tokens. Air-only cells are dropped. + - Each worker writes one raw fp16 token file plus one index, so the inode use is + tiny. +3. ``cv``: patient-level K-fold cross-validation of NeuroVFM's ClassifyThenAggregate + head (the MR-RATE MIL head) on the frozen bags. One bag holds all sequences of a + study. Only out-of-fold predictions are kept. +4. ``summarize``: per-class and macro AUROC / AUPRC over all studies. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import importlib.util +import json +import math +import os +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .data import ( + ALIGNED_SPACES, + _load_split_ids, + _raw_volume, + discover_raw_aligned, +) + +TEACHER_PREFIX = "model.teacher.backbone." +EXPORT_FORMAT = "mrdino_teacher_backbone_v1" +INDEX_FORMAT = "mrdino_probe_tokens_v1" +REPO_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_LABELS = REPO_ROOT / "contrastive-pretraining/scripts/eval_labels/splits_merged_majority/mrrate_merged_labels.csv" +DEFAULT_SPLITS = REPO_ROOT / "contrastive-pretraining/scripts/eval_labels/splits_merged_majority/splits.csv" + + +# --------------------------------------------------------------------------- checkpoint + +def read_checkpoint_metadata(checkpoint: Path) -> dict: + """Training args/step of a DCP checkpoint directory or an exported teacher file.""" + checkpoint = Path(checkpoint) + if checkpoint.is_dir(): + meta = json.loads((checkpoint / "metadata.json").read_text()) + return {"step": int(meta["step"]), "stage": meta["stage"], "args": meta["args"], "source": str(checkpoint)} + state = torch.load(checkpoint, map_location="cpu", weights_only=False) + if state.get("format") != EXPORT_FORMAT: + raise ValueError(f"{checkpoint} is neither a DCP checkpoint directory nor a {EXPORT_FORMAT} file") + return {k: state[k] for k in ("step", "stage", "args", "source")} + + +def load_teacher_backbone(checkpoint: Path, device: torch.device | str = "cpu"): + """Rebuild the EMA-teacher backbone from a DCP checkpoint (no process group needed) or an export.""" + from .train_ddp import build_backbone + + checkpoint = Path(checkpoint) + meta = read_checkpoint_metadata(checkpoint) + backbone = build_backbone(meta["args"]["arch"]) + if checkpoint.is_dir(): + if not (checkpoint / "COMPLETE").exists(): + raise RuntimeError(f"Incomplete checkpoint: {checkpoint}") + import torch.distributed.checkpoint as dcp + + wanted = {TEACHER_PREFIX + k: torch.empty_like(v) for k, v in backbone.state_dict().items()} + dcp.load(wanted, storage_reader=dcp.FileSystemReader(str(checkpoint)), no_dist=True) + state = {k[len(TEACHER_PREFIX):]: v for k, v in wanted.items()} + else: + state = torch.load(checkpoint, map_location="cpu", weights_only=False)["state_dict"] + backbone.load_state_dict(state, strict=True) + return backbone.to(device).eval().requires_grad_(False), meta + + +def export_teacher(checkpoint: Path, out: Path) -> Path: + backbone, meta = load_teacher_backbone(checkpoint) + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(out.suffix + ".partial") + torch.save({"format": EXPORT_FORMAT, **meta, "state_dict": backbone.state_dict()}, tmp) + os.replace(tmp, out) + return out + + +# --------------------------------------------------------------------------- tokens + +def _pad_to_multiple(x: torch.Tensor, multiple: tuple[int, int, int], value: float) -> torch.Tensor: + pads = [] + for n, m in reversed(list(zip(x.shape[-3:], multiple))): + pads.extend((0, (-n) % m)) + return F.pad(x, pads, value=value) if any(pads) else x + + +@torch.inference_mode() +def embed_sequence( + backbone, + volume: torch.Tensor, + tile: tuple[int, int, int], + pool: tuple[int, int, int], + patch: tuple[int, int, int] = (2, 16, 16), + fg_margin: float = 0.05, + tiles_per_batch: int = 8, + device: torch.device | str = "cpu", +) -> tuple[torch.Tensor, torch.Tensor]: + """Return pooled foreground tokens ``[M, C]`` (fp16) and their pooled-grid coords ``[M, 3]``. + + Background is the volume minimum (air after MR-RATE's z-score; padding uses it too). + A token counts as foreground when some voxel of its patch is ``fg_margin`` above it. + """ + if any(t % p for t, p in zip(tile, patch)): + raise ValueError(f"Tile {tile} must be a multiple of the patch {patch}") + device = torch.device(device) + background = float(volume.float().min()) + vol = _pad_to_multiple(volume.float(), tile, background) + grid = tuple(n // p for n, p in zip(vol.shape, patch)) + tile_grid = tuple(t // p for t, p in zip(tile, patch)) + starts = [ + (z, y, x) + for z in range(0, vol.shape[0], tile[0]) + for y in range(0, vol.shape[1], tile[1]) + for x in range(0, vol.shape[2], tile[2]) + ] + feats = torch.zeros(*grid, backbone.embed_dim, device=device, dtype=torch.float32) + fg = torch.zeros(grid, device=device, dtype=torch.bool) + autocast = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else torch.autocast("cpu", enabled=False) + for i in range(0, len(starts), tiles_per_batch): + chunk = starts[i : i + tiles_per_batch] + x = torch.stack([vol[z : z + tile[0], y : y + tile[1], xx : xx + tile[2]] for z, y, xx in chunk]) + x = x.to(device)[:, None] + tile_fg = F.max_pool3d(x - background, kernel_size=patch, stride=patch)[:, 0] > fg_margin + keep = tile_fg.flatten(1).any(1) + if not bool(keep.any()): + continue + with autocast: + out = backbone.forward_features(x[keep].to(torch.bfloat16 if device.type == "cuda" else torch.float32)) + tokens = out["x_norm_patchtokens"].float().reshape(-1, *tile_grid, backbone.embed_dim) + for (z, y, xx), t, m in zip([s for s, k in zip(chunk, keep.tolist()) if k], tokens, tile_fg[keep]): + gz, gy, gx = z // patch[0], y // patch[1], xx // patch[2] + feats[gz : gz + tile_grid[0], gy : gy + tile_grid[1], gx : gx + tile_grid[2]] = t + fg[gz : gz + tile_grid[0], gy : gy + tile_grid[1], gx : gx + tile_grid[2]] = m + # foreground-weighted average pooling on the token grid + weight = _pad_to_multiple(fg.float()[None], pool, 0.0)[0] + feats = _pad_to_multiple(feats.permute(3, 0, 1, 2), pool, 0.0) + count = F.avg_pool3d(weight[None, None], pool, pool)[0, 0] + summed = F.avg_pool3d((feats * weight)[None], pool, pool)[0] + cells = count > 0 + pooled = (summed[:, cells] / count[cells]).T + coords = cells.nonzero().to(torch.int16) + return pooled.to(torch.float16).cpu(), coords.cpu() + + +def worker_identity() -> tuple[int, int]: + rank = int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0))) + world = int(os.environ.get("WORLD_SIZE", os.environ.get("SLURM_NTASKS", 1))) + return rank, world + + +def read_labels(labels_csv: Path) -> tuple[list[str], dict[str, np.ndarray]]: + with open(labels_csv, newline="") as handle: + reader = csv.DictReader(handle) + names = [c for c in reader.fieldnames if c not in {"study_uid", "subject_id"}] + key = "study_uid" if "study_uid" in reader.fieldnames else "subject_id" + labels = {row[key].strip(): np.array([float(row[c]) for c in names], dtype=np.float32) for row in reader} + return names, labels + + +def probe_studies(args) -> list[dict]: + """Labelled studies of the split, sorted by uid (identical on every worker).""" + _, labels = read_labels(Path(args.labels_csv)) + selected = _load_split_ids(str(args.splits_csv), args.split) + selected = {uid for uid in selected if uid in labels} + studies = sorted(discover_raw_aligned(args.data_folder, selected, args.space), key=lambda s: s["study_uid"]) + if args.max_studies: + studies = studies[: args.max_studies] + if not studies: + raise RuntimeError(f"No labelled {args.split} studies under {args.data_folder}") + return studies + + +class _StudyVolumes(torch.utils.data.Dataset): + def __init__(self, studies, target_shape, target_spacing, posterior_shift_mm): + self.studies, self.shape, self.spacing, self.shift = studies, target_shape, target_spacing, posterior_shift_mm + + def __len__(self): + return len(self.studies) + + def __getitem__(self, i): + study = self.studies[i] + volumes = [_raw_volume(p, self.shape, self.spacing, self.shift) for p in study["image_paths"]] + return study["study_uid"], [Path(p).name for p in study["image_paths"]], volumes + + +def extract(args) -> Path: + rank, world = worker_identity() + device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) + backbone, meta = load_teacher_backbone(Path(args.checkpoint), device) + train = meta["args"] + args.space = args.space or train.get("space", "atlas_space") # pre-coreg checkpoints were atlas + if args.space != train.get("space", "atlas_space") and not args.allow_space_mismatch: + raise ValueError(f"Checkpoint was trained on {train.get('space')!r}, probe asked for {args.space!r}") + target_shape = tuple(train["target_shape"]) + target_spacing = tuple(train["target_spacing"]) + out = Path(args.features_dir) + out.mkdir(parents=True, exist_ok=True) + index_path = out / f"index_rank{rank:04d}.npz" + if index_path.exists(): + print(f"[probe] rank {rank}: {index_path} exists, skipping", flush=True) + return index_path + studies = probe_studies(args)[rank::world] + loader = torch.utils.data.DataLoader( + _StudyVolumes(studies, target_shape, target_spacing, float(train["posterior_shift_mm"])), + batch_size=None, shuffle=False, num_workers=args.workers, + prefetch_factor=1 if args.workers else None, persistent_workers=False, + ) + token_path = out / f"tokens_rank{rank:04d}.f16" + uids, starts, counts, seq_ids, coords, names = [], [], [], [], [], [] + total, t0 = 0, time.time() + with open(token_path, "wb") as sink: + for i, (uid, seq_names, volumes) in enumerate(loader): + start = total + for s, volume in enumerate(volumes): + tokens, xyz = embed_sequence( + backbone, volume, tuple(args.tile), tuple(args.pool), + fg_margin=args.fg_margin, tiles_per_batch=args.tiles_per_batch, device=device, + ) + sink.write(tokens.numpy().tobytes()) + total += len(tokens) + seq_ids.append(np.full(len(tokens), s, dtype=np.int16)) + coords.append(xyz.numpy()) + if total == start: + raise RuntimeError(f"Study {uid} produced no foreground tokens") + uids.append(uid) + starts.append(start) + counts.append(total - start) + names.append(json.dumps(seq_names)) + if rank == 0 and (i % 20 == 0 or i + 1 == len(studies)): + rate = (i + 1) / max(1e-6, time.time() - t0) + print(f"[probe] rank 0: {i + 1}/{len(studies)} studies, {total:,} tokens, {rate:.2f} studies/s", flush=True) + tmp = out / f"index_rank{rank:04d}.partial.npz" + np.savez( + tmp, + format=INDEX_FORMAT, dim=backbone.embed_dim, token_file=token_path.name, + study_uid=np.array(uids), start=np.array(starts, dtype=np.int64), count=np.array(counts, dtype=np.int64), + seq_id=np.concatenate(seq_ids) if seq_ids else np.zeros(0, np.int16), + coords=np.concatenate(coords) if coords else np.zeros((0, 3), np.int16), + sequence_names=np.array(names), + checkpoint_step=meta["step"], checkpoint_stage=meta["stage"], checkpoint_source=meta["source"], + space=args.space, tile=np.array(args.tile), pool=np.array(args.pool), + ) + os.replace(tmp, index_path) + return index_path + + +# --------------------------------------------------------------------------- MIL head + +def mil_head_class(): + """NeuroVFM ClassifyThenAggregate from MR-RATE's MIL probe (single source of truth).""" + path = Path(os.environ.get("MRRATE_MIL_PROBE_MODULE", REPO_ROOT / "contrastive-pretraining/scripts/mil_probe.py")) + spec = importlib.util.spec_from_file_location("mrrate_mil_probe", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ClassifyThenAggregate + + +class TokenBags: + """All workers' token files as memory maps, one bag per study.""" + + def __init__(self, features_dir: Path): + self.uids, self.bags, self.dim = [], [], None + indexes = sorted(Path(features_dir).glob("index_rank*.npz")) + if not indexes: + raise FileNotFoundError(f"No index_rank*.npz under {features_dir}") + self.meta = {} + for path in indexes: + with np.load(path, allow_pickle=False) as index: + if str(index["format"]) != INDEX_FORMAT: + raise ValueError(f"{path} is not a {INDEX_FORMAT} index") + dim = int(index["dim"]) + if self.dim not in (None, dim): + raise ValueError("Workers disagree on the token width") + self.dim = dim + tokens = np.memmap(path.parent / str(index["token_file"]), dtype=np.float16, mode="r").reshape(-1, dim) + for uid, start, count in zip(index["study_uid"], index["start"], index["count"]): + self.uids.append(str(uid)) + self.bags.append((tokens, int(start), int(count))) + self.meta = {k: index[k].tolist() for k in ("checkpoint_step", "checkpoint_stage", "space", "tile", "pool")} + if len(set(self.uids)) != len(self.uids): + raise ValueError("A study appears in more than one worker shard") + + def __len__(self): + return len(self.uids) + + def tokens(self, i: int) -> np.ndarray: + array, start, count = self.bags[i] + return np.asarray(array[start : start + count]) + + +def patient_folds(uids: list[str], patient_of: dict[str, str], k: int, seed: int) -> np.ndarray: + """Deterministic patient-level fold id per study (all studies of a patient share a fold).""" + patients = sorted({patient_of.get(uid, uid) for uid in uids}) + order = sorted(patients, key=lambda p: hashlib.sha256(f"{seed}:{p}".encode()).hexdigest()) + fold_of = {p: i % k for i, p in enumerate(order)} + return np.array([fold_of[patient_of.get(uid, uid)] for uid in uids]) + + +def read_patients(splits_csv: Path) -> dict[str, str]: + with open(splits_csv, newline="") as handle: + return {row["study_uid"].strip(): row.get("patient_uid", row["study_uid"]).strip() for row in csv.DictReader(handle)} + + +def _batch(bags: TokenBags, items, device): + arrays = [bags.tokens(i) for i in items] + lengths = torch.tensor([0] + [len(a) for a in arrays]) + tokens = torch.from_numpy(np.concatenate(arrays)).to(device, non_blocking=True) + return tokens, torch.cumsum(lengths, 0).to(device) + + +def run_fold(args) -> Path: + device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) + bags = TokenBags(Path(args.features_dir)) + names, labels = read_labels(Path(args.labels_csv)) + y = np.stack([labels[uid] for uid in bags.uids]) + folds = patient_folds(bags.uids, read_patients(Path(args.splits_csv)), args.folds, args.seed) + train = np.flatnonzero(folds != args.fold) + held = np.flatnonzero(folds == args.fold) + torch.manual_seed(args.seed + args.fold) + rng = np.random.default_rng(args.seed + args.fold) + + head = mil_head_class()(bags.dim, len(names), hidden_dim=args.hidden_dim, mlp_hidden_dims=(args.mlp_hidden_dim,)).to(device) + positives = y[train].sum(0) + pos_weight = np.clip((len(train) - positives) / np.clip(positives, 1.0, None), 1.0, 100.0) + loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor(pos_weight, dtype=torch.float32, device=device)) + optimizer = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay) + steps = args.epochs * math.ceil(len(train) / args.batch_size) + warmup = max(1, int(0.1 * steps)) + schedule = torch.optim.lr_scheduler.LambdaLR( + optimizer, lambda s: (s + 1) / warmup if s < warmup else 0.5 * (1 + math.cos(math.pi * (s - warmup) / max(1, steps - warmup))) + ) + autocast = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else torch.autocast("cpu", enabled=False) + head.train() + for epoch in range(args.epochs): + order = rng.permutation(train) + for b in range(0, len(order), args.batch_size): + items = order[b : b + args.batch_size] + tokens, cu = _batch(bags, items, device) + with autocast: + logits = head(tokens.float(), cu) + loss = loss_fn(logits.float(), torch.from_numpy(y[items]).to(device)) + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(head.parameters(), 1.0) + optimizer.step() + schedule.step() + print(f"[probe] fold {args.fold} epoch {epoch + 1}/{args.epochs} loss {float(loss.detach()):.4f}", flush=True) + head.eval() + out_logits = [] + with torch.inference_mode(): + for b in range(0, len(held), args.batch_size): + items = held[b : b + args.batch_size] + tokens, cu = _batch(bags, items, device) + with autocast: + out_logits.append(head(tokens.float(), cu).float().cpu().numpy()) + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + path = out / f"oof_fold{args.fold}.npz" + np.savez(path, study_uid=np.array([bags.uids[i] for i in held]), logits=np.concatenate(out_logits), labels=y[held]) + return path + + +# --------------------------------------------------------------------------- metrics + +def auroc(target: np.ndarray, score: np.ndarray) -> float | None: + """Mann-Whitney AUROC with average ranks for ties.""" + positive = target == 1 + n1, n0 = int(positive.sum()), int((~positive).sum()) + if n1 == 0 or n0 == 0: + return None + _, inverse, counts = np.unique(score, return_inverse=True, return_counts=True) + average_rank = np.cumsum(counts) - (counts - 1) / 2.0 + ranks = average_rank[inverse] + return float((ranks[positive].sum() - n1 * (n1 + 1) / 2.0) / (n1 * n0)) + + +def average_precision(target: np.ndarray, score: np.ndarray) -> float | None: + if not (target == 1).any(): + return None + hits = target[np.argsort(-score, kind="mergesort")] == 1 + precision = np.cumsum(hits) / np.arange(1, len(hits) + 1) + return float(precision[hits].mean()) + + +def summarize(args) -> dict: + out = Path(args.out_dir) + names, _ = read_labels(Path(args.labels_csv)) + parts = [np.load(out / f"oof_fold{f}.npz") for f in range(args.folds)] + uids = np.concatenate([p["study_uid"] for p in parts]) + if len(set(uids.tolist())) != len(uids): + raise ValueError("A study was predicted by more than one fold") + logits = np.concatenate([p["logits"] for p in parts]) + labels = np.concatenate([p["labels"] for p in parts]) + rows = [] + for c, name in enumerate(names): + rows.append({ + "label": name, "positives": int(labels[:, c].sum()), "studies": int(len(labels)), + "auroc": auroc(labels[:, c], logits[:, c]), "auprc": average_precision(labels[:, c], logits[:, c]), + }) + valid = [r for r in rows if r["auroc"] is not None] + meta = TokenBags(Path(args.features_dir)).meta if args.features_dir and Path(args.features_dir).exists() else {} + fold_auroc = [] + for p in parts: + per = [auroc(p["labels"][:, c], p["logits"][:, c]) for c in range(len(names))] + per = [v for v in per if v is not None] + fold_auroc.append(float(np.mean(per)) if per else None) + result = { + "macro_auroc": float(np.mean([r["auroc"] for r in valid])), + "macro_auprc": float(np.mean([r["auprc"] for r in valid])), + "fold_macro_auroc": fold_auroc, + "studies": int(len(labels)), + "folds": args.folds, + "labels_csv": str(args.labels_csv), + "per_class": rows, + **meta, + } + (out / "results.json").write_text(json.dumps(result, indent=2)) + with open(out / "per_class.csv", "w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + print( + f"[probe] step {meta.get('checkpoint_step', '?')}: macro AUROC {result['macro_auroc']:.4f} " + f"AUPRC {result['macro_auprc']:.4f} over {result['studies']} studies " + f"(folds {', '.join(f'{v:.3f}' for v in fold_auroc if v is not None)})", + flush=True, + ) + return result + + +# --------------------------------------------------------------------------- CLI + +def parse_args(argv=None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="command", required=True) + + e = sub.add_parser("export", help="Copy the EMA-teacher backbone out of a DCP checkpoint") + e.add_argument("--checkpoint", required=True) + e.add_argument("--out", required=True) + + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--labels-csv", default=str(DEFAULT_LABELS), help="study_uid + binary label columns (default: 14 merged groups)") + common.add_argument("--splits-csv", default=str(DEFAULT_SPLITS), help="study_uid, patient_uid, split") + common.add_argument("--features-dir", required=True) + common.add_argument("--device", default=None) + + x = sub.add_parser("extract", parents=[common], help="Frozen tokens of every labelled study of a split") + x.add_argument("--checkpoint", required=True, help="DCP checkpoint directory or exported teacher .pt") + x.add_argument("--data-folder", required=True, help="Extracted MR-RATE tree: batchXX//coreg_img (or atlas_img)") + x.add_argument("--space", default=None, choices=ALIGNED_SPACES, help="Default: the checkpoint's training space") + x.add_argument("--allow-space-mismatch", action="store_true") + x.add_argument("--split", default="val") + x.add_argument("--tile", type=int, nargs=3, default=(64, 192, 192), help="Voxels; default = phase-1 global crop") + x.add_argument("--pool", type=int, nargs=3, default=(2, 2, 2), help="Token-grid pooling cell") + x.add_argument("--fg-margin", type=float, default=0.05, help="Foreground: patch max above the volume minimum") + x.add_argument("--tiles-per-batch", type=int, default=8) + x.add_argument("--workers", type=int, default=6) + x.add_argument("--max-studies", type=int, default=None, help="Smoke only") + + c = sub.add_parser("cv", parents=[common], help="Train/predict one patient-level fold") + c.add_argument("--out-dir", required=True) + c.add_argument("--fold", type=int, required=True) + c.add_argument("--folds", type=int, default=5) + c.add_argument("--seed", type=int, default=42) + c.add_argument("--epochs", type=int, default=8) + c.add_argument("--batch-size", type=int, default=8) + c.add_argument("--lr", type=float, default=2e-3) + c.add_argument("--weight-decay", type=float, default=0.05) + c.add_argument("--hidden-dim", type=int, default=512) + c.add_argument("--mlp-hidden-dim", type=int, default=384) + + s = sub.add_parser("summarize", parents=[common], help="Out-of-fold macro/per-class AUROC and AUPRC") + s.add_argument("--out-dir", required=True) + s.add_argument("--folds", type=int, default=5) + return p.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv) + if args.command == "export": + print(export_teacher(Path(args.checkpoint), Path(args.out))) + elif args.command == "extract": + print(extract(args)) + elif args.command == "cv": + print(run_fold(args)) + else: + summarize(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index 6b3f41d..fa86792 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -1,4 +1,4 @@ -"""Distributed, resumable 3-D DINOv3 training for atlas-registered MR-RATE.""" +"""Distributed, resumable 3-D DINOv3 training for aligned (coreg/atlas) MR-RATE.""" from __future__ import annotations @@ -23,11 +23,13 @@ from torch.utils.data import DataLoader from .data import ( - MRAtlasDINO3DDataset, + ALIGNED_SPACES, + DEFAULT_SPACE, + MRAlignedDINO3DDataset, InfiniteStudySampler, collate_dino3d, - discover_atlas_cache, - discover_raw_atlas_split, + discover_aligned_cache, + discover_raw_aligned_split, npz_volume_shape, stage_crop_spec, ) @@ -36,17 +38,37 @@ from .objective import DINO3DLearner, LossWeights +CHECKPOINT_FORMAT = "mrrate_aligned_dinov3d_full_v1" +# Checkpoints written before the coreg switch: same contents, always atlas space. +_LEGACY_ATLAS_FORMATS = { + "mrrate_aligned_dinov3d_full_v1": "mrrate_atlas_dinov3d_full_v1", + "mrrate_aligned_dinov3d_fsdp2_v1": "mrrate_atlas_dinov3d_fsdp2_v1", +} + + +def check_checkpoint_space(fmt: str | None, saved_args: dict, space: str, expected: str) -> None: + """Refuse to resume across MR spaces.""" + if fmt == expected: + saved_space = saved_args.get("space") + elif fmt == _LEGACY_ATLAS_FORMATS.get(expected): + saved_space = "atlas_space" + else: + raise RuntimeError(f"Checkpoint format {fmt!r} is not an MR DINO checkpoint ({expected})") + if saved_space != space: + raise RuntimeError(f"Checkpoint was trained on {saved_space!r}; this run uses {space!r}") + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) source = p.add_mutually_exclusive_group(required=True) - source.add_argument("--preprocessed-dir", help="MR-RATE cache root containing atlas_space/*.npz") + source.add_argument("--preprocessed-dir", help="MR-RATE cache root containing /*.npz") source.add_argument( "--data-folder", - help="MR-RATE-atlas NIfTI tree; uses the same live loader as previous MIL training", + help="MR-RATE-coreg (or -atlas) NIfTI tree; uses the same live loader as previous MIL training", ) p.add_argument("--splits-csv") p.add_argument("--split", default="train") - p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) + p.add_argument("--space", default=DEFAULT_SPACE, choices=ALIGNED_SPACES) p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="large") @@ -119,9 +141,9 @@ def read_train_files( preprocessed_dir: str, splits_csv: str | None, split: str, - space: str = "atlas_space", + space: str = DEFAULT_SPACE, ) -> list[str]: - return [sample["cache_path"] for sample in discover_atlas_cache( + return [sample["cache_path"] for sample in discover_aligned_cache( preprocessed_dir, splits_csv=splits_csv, split=split, space=space )] @@ -148,6 +170,7 @@ def raw_fingerprint( target_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), target_shape: tuple[int, int, int] = (256, 384, 384), posterior_shift_mm: float = 15.0, + space: str = DEFAULT_SPACE, ) -> str: """Fingerprint raw studies without tying exact resume to an absolute mount path.""" records = [] @@ -163,7 +186,7 @@ def raw_fingerprint( "world": world, "max_studies_per_rank": max_studies_per_rank, "preprocessing": { - "space": "atlas_space", + "space": space, "normalizer": "zscore", "target_spacing": list(target_spacing), "target_shape": list(target_shape), @@ -228,16 +251,17 @@ def distributed_raw_assignment( target_spacing: tuple[float, float, float], target_shape: tuple[int, int, int], posterior_shift_mm: float, + space: str = DEFAULT_SPACE, ) -> tuple[list[dict], str]: """Discover once on rank 0 and publish one small assignment file per rank.""" assignment_dir = output / "dataset_assignments" metadata_path = assignment_dir / "metadata.json" if rank == 0: - samples = discover_raw_atlas_split( - data_folder, splits_csv=splits_csv, split=split + samples = discover_raw_aligned_split( + data_folder, splits_csv=splits_csv, split=split, space=space ) if len(samples) < world: - raise RuntimeError(f"{len(samples)} atlas studies cannot supply {world} ranks") + raise RuntimeError(f"{len(samples)} {space} studies cannot supply {world} ranks") fingerprint = raw_fingerprint( samples, world, @@ -245,12 +269,13 @@ def distributed_raw_assignment( target_spacing, target_shape, posterior_shift_mm, + space, ) assignments = balanced_sample_assignment(samples, world) if max_studies_per_rank: assignments = [items[:max_studies_per_rank] for items in assignments] if any(not items for items in assignments): - raise RuntimeError("Raw atlas assignment left at least one distributed rank empty") + raise RuntimeError("Raw assignment left at least one distributed rank empty") assignment_dir.mkdir(parents=True, exist_ok=True) for assigned_rank, items in enumerate(assignments): destination = assignment_dir / f"rank_{assigned_rank:04d}.json" @@ -258,7 +283,8 @@ def distributed_raw_assignment( temporary.write_text(json.dumps(items, separators=(",", ":"))) os.replace(temporary, destination) metadata = { - "format": "mrrate_atlas_raw_assignment_v1", + "format": "mrrate_aligned_raw_assignment_v1", + "space": space, "world_size": world, "dataset_fingerprint": fingerprint, "split": split, @@ -269,13 +295,13 @@ def distributed_raw_assignment( if dist.is_initialized(): dist.barrier() metadata = json.loads(metadata_path.read_text()) - if metadata.get("format") != "mrrate_atlas_raw_assignment_v1": - raise RuntimeError(f"Invalid raw atlas assignment metadata: {metadata_path}") + if metadata.get("format") != "mrrate_aligned_raw_assignment_v1" or metadata.get("space") != space: + raise RuntimeError(f"Invalid or other-space raw assignment metadata: {metadata_path}") if int(metadata.get("world_size", -1)) != world: - raise RuntimeError("Raw atlas assignment world size changed during startup") + raise RuntimeError("Raw assignment world size changed during startup") assigned = json.loads((assignment_dir / f"rank_{rank:04d}.json").read_text()) if not assigned: - raise RuntimeError(f"Rank {rank} received no raw atlas studies") + raise RuntimeError(f"Rank {rank} received no raw {space} studies") return assigned, str(metadata["dataset_fingerprint"]) @@ -370,7 +396,7 @@ def save_checkpoint( rng_by_rank = [local_rng] if rank == 0: state = { - "format": "mrrate_atlas_dinov3d_full_v1", + "format": CHECKPOINT_FORMAT, "step": step, "stage": args.stage, "world_size": dist.get_world_size() if dist.is_initialized() else 1, @@ -407,16 +433,14 @@ def load_checkpoint( current_stage: str, dataset_fingerprint: str | None, world: int, + current_space: str = DEFAULT_SPACE, ) -> int: if not path.exists(): if rank == 0: print(f"[resume] {path} not found; starting from scratch", flush=True) return 0 state = torch.load(path, map_location="cpu", weights_only=False) - if state.get("format") != "mrrate_atlas_dinov3d_full_v1": - raise RuntimeError( - f"Checkpoint format {state.get('format')!r} is not atlas-space MR DINO" - ) + check_checkpoint_space(state.get("format"), state.get("args", {}), current_space, CHECKPOINT_FORMAT) if int(state.get("world_size", world)) != world: raise RuntimeError("Exact DDP resume requires the checkpoint's original world size") saved_fingerprint = state.get("args", {}).get("dataset_fingerprint") @@ -517,6 +541,7 @@ def main() -> int: target_spacing=tuple(args.target_spacing), target_shape=tuple(args.target_shape), posterior_shift_mm=args.posterior_shift_mm, + space=args.space, ) dataset_kwargs = { "data_folder": args.data_folder, @@ -531,7 +556,7 @@ def main() -> int: local_shape=tuple(args.local_shape) if args.local_shape else base_crop_spec.local_shape, local_crops=args.local_crops, ) - dataset = MRAtlasDINO3DDataset( + dataset = MRAlignedDINO3DDataset( **dataset_kwargs, splits_csv=args.splits_csv, split=args.split, @@ -618,7 +643,7 @@ def main() -> int: start = load_checkpoint( checkpoint_path(output, args.resume), learner, optimizer, sampler, device, rank, - args.stage, args.dataset_fingerprint, world, + args.stage, args.dataset_fingerprint, world, args.space, ) if checkpoint_path(output, args.resume) else 0 consumed_samples = sampler.offset # Cross-stage resume fixes the new Gram target to the incoming EMA teacher. diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index 4bf1352..e1aaf90 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -1,4 +1,4 @@ -"""FSDP2 training for the literal DINOv3 ViT-7B on atlas-registered MR-RATE.""" +"""FSDP2 3-D DINOv3 training (ViT-L default) on aligned (coreg/atlas) MR-RATE.""" from __future__ import annotations @@ -26,7 +26,9 @@ from dinov3.layers.fp8_linear import convert_linears_to_fp8 from .data import ( - MRAtlasDINO3DDataset, + ALIGNED_SPACES, + DEFAULT_SPACE, + MRAlignedDINO3DDataset, InfiniteStudySampler, collate_dino3d, stage_crop_spec, @@ -37,6 +39,7 @@ from .train_ddp import ( StopController, balanced_file_assignment, + check_checkpoint_space, build_backbone, cache_fingerprint, cosine, @@ -47,20 +50,24 @@ ) +CHECKPOINT_FORMAT = "mrrate_aligned_dinov3d_fsdp2_v1" + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) source = p.add_mutually_exclusive_group(required=True) source.add_argument( "--data-folder", - help="MR-RATE-atlas NIfTI tree; uses the same live loader as previous MIL training", + help="MR-RATE-coreg (or -atlas) NIfTI tree; uses the same live loader as previous MIL training", ) source.add_argument( "--preprocessed-dir", - help="Optional MR-RATE volume cache root containing atlas_space/*.npz", + help="Optional MR-RATE volume cache root containing /*.npz", ) p.add_argument("--splits-csv") p.add_argument("--split", default="train") - p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) + p.add_argument("--space", default=DEFAULT_SPACE, choices=ALIGNED_SPACES, + help="coreg_space (default): sequences on the center T1's native grid; atlas_space: MNI 1 mm box") p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="large", @@ -325,7 +332,7 @@ def save_checkpoint( os.replace(tmp, path / f"runtime_rank{rank:04d}.pt") if rank == 0: (path / "metadata.json").write_text(json.dumps({ - "format": "mrrate_atlas_dinov3d_fsdp2_v1", + "format": CHECKPOINT_FORMAT, "step": step, "stage": args.stage, "sampler_offset_per_rank": sampler_offset, @@ -361,10 +368,7 @@ def load_checkpoint( if not (path / "COMPLETE").exists(): raise RuntimeError(f"Incomplete checkpoint: {path}") metadata = json.loads((path / "metadata.json").read_text()) - if metadata.get("format") != "mrrate_atlas_dinov3d_fsdp2_v1": - raise RuntimeError( - f"Checkpoint format {metadata.get('format')!r} is not atlas-space MR DINO" - ) + check_checkpoint_space(metadata.get("format"), metadata.get("args", {}), args.space, CHECKPOINT_FORMAT) saved_world = int(metadata.get("world_size", dist.get_world_size())) if saved_world != dist.get_world_size(): raise RuntimeError( @@ -374,7 +378,7 @@ def load_checkpoint( saved_fingerprint = metadata.get("args", {}).get("dataset_fingerprint") if saved_fingerprint and saved_fingerprint != args.dataset_fingerprint: raise RuntimeError( - "Atlas-space training cache/split changed since the checkpoint; " + "Training cache/split/space changed since the checkpoint; " "refusing an inexact sampler resume" ) source_stage = metadata["stage"] @@ -459,6 +463,7 @@ def main() -> int: target_spacing=tuple(args.target_spacing), target_shape=tuple(args.target_shape), posterior_shift_mm=args.posterior_shift_mm, + space=args.space, ) dataset_kwargs = { "data_folder": args.data_folder, @@ -472,7 +477,7 @@ def main() -> int: local_shape=tuple(args.local_shape) if args.local_shape else base_crop_spec.local_shape, local_crops=args.local_crops, ) - dataset = MRAtlasDINO3DDataset( + dataset = MRAlignedDINO3DDataset( **dataset_kwargs, splits_csv=args.splits_csv, split=args.split, diff --git a/dino-pretraining/scripts/probe_checkpoint.sbatch b/dino-pretraining/scripts/probe_checkpoint.sbatch new file mode 100755 index 0000000..7a76b0a --- /dev/null +++ b/dino-pretraining/scripts/probe_checkpoint.sbatch @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_dino_probe +#SBATCH --partition=h200 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:h200:4 +#SBATCH --cpus-per-task=128 +#SBATCH --time=06:00:00 +#SBATCH --output=mrdino_%x_%j.out +#SBATCH --error=mrdino_%x_%j.err +# Usage: DATA_FOLDER= sbatch probe_checkpoint.sbatch CHECKPOINT OUT_DIR [extra extract args] +set -euo pipefail +SCRIPT_DIR=${SCRIPT_DIR:-"$SLURM_SUBMIT_DIR/scripts"} +[[ -x "$SCRIPT_DIR/probe_node.sh" ]] || SCRIPT_DIR="$SLURM_SUBMIT_DIR/dino-pretraining/scripts" +exec "$SCRIPT_DIR/probe_node.sh" "$@" diff --git a/dino-pretraining/scripts/probe_node.sh b/dino-pretraining/scripts/probe_node.sh new file mode 100755 index 0000000..fe10e77 --- /dev/null +++ b/dino-pretraining/scripts/probe_node.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Score one MR DINO checkpoint on THIS node with the validation MIL probe: +# extract (one worker per GPU) -> patient-level K folds in parallel -> out-of-fold summary. +# Usage: DATA_FOLDER= probe_node.sh CHECKPOINT OUT_DIR [extra extract args] +# CHECKPOINT: DCP step directory or the teacher .pt written by `python -m mr_dino.probe export`. +# Tokens go to node-local /tmp (FEATURE_DIR) and are deleted afterwards unless KEEP_FEATURES=1. +# Inside a running training allocation: srun --overlap --jobid JOB -N1 -n1 -w NODE probe_node.sh ... +set -euo pipefail + +CKPT=${1:?checkpoint}; OUT=${2:?output dir}; shift 2 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CODE=${CODE:-"$(dirname "$SCRIPT_DIR")"} +DINOV3_ROOT=${DINOV3_ROOT:-/hnvme/workspace/b180dc51-sezgin/dinov3} +SIF=${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate.sif} +EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} +: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE tree (batchXX//coreg_img)}" +GPUS=${GPUS_PER_NODE:-4} +FOLDS=${FOLDS:-5} +FEAT=${FEATURE_DIR:-/tmp/mrdino_probe_${SLURM_JOB_ID:-$$}_$(basename "$CKPT" .pt)} +LABEL_ARGS=() +[[ -n "${LABELS_CSV:-}" ]] && LABEL_ARGS+=(--labels-csv "$LABELS_CSV") +[[ -n "${SPLITS_CSV:-}" ]] && LABEL_ARGS+=(--splits-csv "$SPLITS_CSV") +export PYTHONPATH="${EXTRA_PIP}:${DINOV3_ROOT}:${CODE}:${PYTHONPATH:-}" +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-8} + +probe() { + singularity exec --nv -B /hnvme/workspace:/hnvme/workspace,/tmp:/tmp "$SIF" \ + python3 -m mr_dino.probe "$@" +} +wait_all() { # fail if any background worker failed + local status=0 + for pid in "$@"; do wait "$pid" || status=1; done + return $status +} + +mkdir -p "$OUT" "$FEAT" +echo "[probe_node] $(date '+%F %T') checkpoint=$CKPT features=$FEAT out=$OUT host=$(hostname)" +pids=() +for g in $(seq 0 $((GPUS - 1))); do + CUDA_VISIBLE_DEVICES=$g RANK=$g WORLD_SIZE=$GPUS probe extract --checkpoint "$CKPT" \ + --data-folder "$DATA_FOLDER" --features-dir "$FEAT" "${LABEL_ARGS[@]}" "$@" \ + > "$OUT/extract_rank$g.log" 2>&1 & + pids+=($!) +done +wait_all "${pids[@]}" || { echo "[probe_node] extraction failed; see $OUT/extract_rank*.log" >&2; exit 1; } +echo "[probe_node] $(date '+%F %T') extraction done ($(du -sh "$FEAT" | cut -f1))" + +pids=() +for f in $(seq 0 $((FOLDS - 1))); do + CUDA_VISIBLE_DEVICES=$((f % GPUS)) probe cv --features-dir "$FEAT" --out-dir "$OUT" \ + --fold "$f" --folds "$FOLDS" "${LABEL_ARGS[@]}" > "$OUT/cv_fold$f.log" 2>&1 & + pids+=($!) +done +wait_all "${pids[@]}" || { echo "[probe_node] CV failed; see $OUT/cv_fold*.log" >&2; exit 1; } +probe summarize --features-dir "$FEAT" --out-dir "$OUT" --folds "$FOLDS" "${LABEL_ARGS[@]}" | tee "$OUT/summary.txt" +[[ "${KEEP_FEATURES:-0}" == 1 ]] || rm -rf "$FEAT" diff --git a/dino-pretraining/scripts/probe_watch.sh b/dino-pretraining/scripts/probe_watch.sh new file mode 100755 index 0000000..b3ea6b2 --- /dev/null +++ b/dino-pretraining/scripts/probe_watch.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Probe chosen steps of a running MR DINO stage as soon as their checkpoints appear. +# Usage (login node, e.g. under nohup): DATA_FOLDER= probe_watch.sh RUN_DIR STEP [STEP ...] +# RUN_DIR is a stage output (…/mrdino3d_vitl/pretrain). For each step the EMA-teacher backbone is +# exported at once to RUN_DIR/probe/step_N/teacher.pt (checkpoint rotation cannot delete it), then +# probe_checkpoint.sbatch is submitted. Results: RUN_DIR/probe/step_N/results.json + summary.txt. +set -euo pipefail +RUN=${1:?run dir}; shift +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CODE=${CODE:-"$(dirname "$SCRIPT_DIR")"} +SIF=${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate.sif} +EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} +DINOV3_ROOT=${DINOV3_ROOT:-/hnvme/workspace/b180dc51-sezgin/dinov3} +: "${DATA_FOLDER:?Set DATA_FOLDER}" +export DATA_FOLDER PYTHONPATH="${EXTRA_PIP}:${DINOV3_ROOT}:${CODE}:${PYTHONPATH:-}" +pending=("$@") +while ((${#pending[@]})); do + left=() + for step in "${pending[@]}"; do + ckpt=$(printf '%s/checkpoints/step_%08d' "$RUN" "$step") + out="$RUN/probe/step_$step" + if [[ -f "$ckpt/COMPLETE" ]]; then + mkdir -p "$out" + singularity exec -B /hnvme/workspace:/hnvme/workspace "$SIF" \ + python3 -m mr_dino.probe export --checkpoint "$ckpt" --out "$out/teacher.pt" + job=$(cd "$CODE" && sbatch --parsable --chdir="$CODE" --output="$out/slurm_%j.out" \ + --error="$out/slurm_%j.err" "$SCRIPT_DIR/probe_checkpoint.sbatch" "$out/teacher.pt" "$out") + echo "[probe_watch] $(date '+%F %T') step $step exported, probe job $job" + else + left+=("$step") + fi + done + pending=("${left[@]+"${left[@]}"}") + ((${#pending[@]})) && sleep 300 +done +echo "[probe_watch] all steps submitted" diff --git a/dino-pretraining/scripts/smoke_dummy.sbatch b/dino-pretraining/scripts/smoke_dummy.sbatch index 59081c5..b0f7761 100755 --- a/dino-pretraining/scripts/smoke_dummy.sbatch +++ b/dino-pretraining/scripts/smoke_dummy.sbatch @@ -28,14 +28,14 @@ EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} export PYTHONPATH="${EXTRA_PIP}:${PYTHONPATH:-}" singularity exec -B /hnvme/workspace:/hnvme/workspace,/tmp:/tmp \ "$SIF" \ - python3 "$CODE/tests/make_dummy_atlas_dataset.py" \ + python3 "$CODE/tests/make_dummy_aligned_dataset.py" \ --out "$DUMMY" --studies 8 --shape 8 32 32 head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29672} COMMON_ARGS=( - --data-folder "$DUMMY" --space atlas_space --output-dir "$OUTPUT" \ + --data-folder "$DUMMY" --space coreg_space --output-dir "$OUTPUT" \ --stage pretrain --arch tiny --batch-size 1 --grad-accum-steps 1 \ --workers 0 --max-studies 2 --local-crops 2 \ --target-shape 8 32 32 --target-spacing 1.0 0.5 0.5 --posterior-shift-mm 0 \ @@ -57,4 +57,19 @@ srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 \ test -f "$OUTPUT/checkpoints/step_00000002/COMPLETE" grep -q '"step": 2' "$OUTPUT/metrics.jsonl" -echo "Synthetic raw-atlas MR DINO FSDP2 integration and resume passed: $OUTPUT" +echo "Synthetic raw-coreg MR DINO FSDP2 integration and resume passed: $OUTPUT" + +# Validation probe on the step-2 checkpoint: all dummy studies labelled "val", 2 folds. +python3 - "$DUMMY" <<'PY' +import csv, pathlib, sys +root = pathlib.Path(sys.argv[1]) +uids = sorted(p.name for p in (root / "batch00").iterdir()) +with open(root / "labels.csv", "w", newline="") as f: + w = csv.writer(f); w.writerow(["study_uid", "a", "b"]); w.writerows([u, i % 2, (i // 2) % 2] for i, u in enumerate(uids)) +with open(root / "splits.csv", "w", newline="") as f: + w = csv.writer(f); w.writerow(["study_uid", "patient_uid", "split"]); w.writerows([u, u, "val"] for u in uids) +PY +DATA_FOLDER="$DUMMY" LABELS_CSV="$DUMMY/labels.csv" SPLITS_CSV="$DUMMY/splits.csv" FOLDS=2 \ + "$SCRIPT_DIR/probe_node.sh" "$OUTPUT/checkpoints/step_00000002" "$OUTPUT/probe" --tile 8 32 32 --workers 0 +test -f "$OUTPUT/probe/results.json" +echo "Validation probe smoke passed: $OUTPUT/probe/summary.txt" diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch index d8c0cf0..a77a0a8 100644 --- a/dino-pretraining/scripts/train_32n_vitl.sbatch +++ b/dino-pretraining/scripts/train_32n_vitl.sbatch @@ -27,7 +27,7 @@ if [[ -z "${CODE:-}" ]]; then fi fi SCRIPT_DIR="$CODE/scripts" -: "${DATA_FOLDER:?Set DATA_FOLDER to the MR-RATE-atlas NIfTI root}" +: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE-coreg NIfTI root (batchXX//coreg_img)}" STAGE=${STAGE:-pretrain} OUTPUT=${OUTPUT:-"$PWD/mrdino3d_vitl/$STAGE"} mkdir -p "$OUTPUT" @@ -58,7 +58,7 @@ echo "MR DINOv3 ViT-L (FORA recipe): stage=$STAGE nodes=$SLURM_NNODES output=$OU srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" \ --data-folder "$DATA_FOLDER" \ - --space atlas_space \ + --space "${SPACE:-coreg_space}" \ "${SPLIT_ARGS[@]}" \ --output-dir "$OUTPUT" --stage "$STAGE" --arch large --steps "$STEPS" \ --batch-size "${BATCH_SIZE:-8}" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ diff --git a/dino-pretraining/tests/make_dummy_atlas_dataset.py b/dino-pretraining/tests/make_dummy_aligned_dataset.py similarity index 82% rename from dino-pretraining/tests/make_dummy_atlas_dataset.py rename to dino-pretraining/tests/make_dummy_aligned_dataset.py index 81ab6d9..a14e069 100644 --- a/dino-pretraining/tests/make_dummy_atlas_dataset.py +++ b/dino-pretraining/tests/make_dummy_aligned_dataset.py @@ -1,4 +1,4 @@ -"""Create a tiny raw MR-RATE-atlas NIfTI tree for integration tests.""" +"""Create a tiny raw MR-RATE coreg/atlas NIfTI tree for integration tests.""" import argparse from pathlib import Path @@ -13,13 +13,14 @@ def main() -> None: parser.add_argument("--studies", type=int, default=8) parser.add_argument("--sequences", type=int, default=3) parser.add_argument("--shape", type=int, nargs=3, default=(8, 32, 32)) + parser.add_argument("--space", choices=("coreg_space", "atlas_space"), default="coreg_space") args = parser.parse_args() d, h, w = args.shape z, y, x = np.indices((d, h, w), dtype=np.float32) affine = np.diag([0.5, 0.5, 1.0, 1.0]).astype(np.float32) for study in range(args.studies): - image_dir = Path(args.out) / "batch00" / f"dummy_{study:03d}" / "atlas_img" + image_dir = Path(args.out) / "batch00" / f"dummy_{study:03d}" / {"coreg_space": "coreg_img", "atlas_space": "atlas_img"}[args.space] image_dir.mkdir(parents=True, exist_ok=True) for sequence in range(args.sequences): center = np.array((d, h, w), dtype=np.float32) / 2 @@ -37,7 +38,7 @@ def main() -> None: nib.Nifti1Image(nifti_array, affine), image_dir / f"sequence_{sequence:02d}.nii.gz", ) - print(f"Wrote {args.studies} raw atlas studies under {args.out}") + print(f"Wrote {args.studies} raw {args.space} studies under {args.out}") if __name__ == "__main__": diff --git a/dino-pretraining/tests/make_dummy_cache.py b/dino-pretraining/tests/make_dummy_cache.py index 37ab876..faf1363 100644 --- a/dino-pretraining/tests/make_dummy_cache.py +++ b/dino-pretraining/tests/make_dummy_cache.py @@ -1,4 +1,4 @@ -"""Create a small MR-RATE-compatible atlas-space cache for GPU smoke tests.""" +"""Create a small MR-RATE-compatible aligned-space (coreg default) cache for GPU smoke tests.""" import argparse import json @@ -13,14 +13,15 @@ def main() -> None: parser.add_argument("--studies", type=int, default=8) parser.add_argument("--sequences", type=int, default=3) parser.add_argument("--shape", type=int, nargs=3, default=(64, 192, 192)) + parser.add_argument("--space", choices=("coreg_space", "atlas_space"), default="coreg_space") args = parser.parse_args() - space = Path(args.out) / "atlas_space" + space = Path(args.out) / args.space space.mkdir(parents=True, exist_ok=True) manifest = { "version": 1, "layout": "per_subject_stack", - "space": "atlas_space", + "space": args.space, "target_spacing": [1.0, 0.5, 0.5], "target_shape": list(args.shape), "posterior_shift_mm": 15.0, @@ -44,7 +45,7 @@ def main() -> None: volume += np.sin((x + sequence * 3) / 9) * 0.04 volumes.append(volume.astype(np.float16)) np.savez(space / f"dummy_{study:03d}.npz", volumes=np.stack(volumes)) - print(f"Wrote {args.studies} dummy atlas-registered studies to {space}") + print(f"Wrote {args.studies} dummy {args.space} studies to {space}") if __name__ == "__main__": diff --git a/dino-pretraining/tests/test_fora_recipe.py b/dino-pretraining/tests/test_fora_recipe.py index a895b93..3a4e4e2 100644 --- a/dino-pretraining/tests/test_fora_recipe.py +++ b/dino-pretraining/tests/test_fora_recipe.py @@ -4,7 +4,7 @@ import pytest import torch -from mr_dino.data import CropSpec, MRAtlasDINO3DDataset, collate_dino3d +from mr_dino.data import CropSpec, MRAlignedDINO3DDataset, collate_dino3d from mr_dino.recipe import ( CosineLinear, GroupedAdamW, @@ -24,7 +24,7 @@ def big_dataset(cache, cross_sequence_probability=0.0, seed=5): - return MRAtlasDINO3DDataset( + return MRAlignedDINO3DDataset( preprocessed_dir=str(cache), crop_spec=CropSpec(global_shape=GLOBAL, local_shape=(4, 16, 16), local_crops=3), cross_sequence_probability=cross_sequence_probability, diff --git a/dino-pretraining/tests/test_mr_dino.py b/dino-pretraining/tests/test_mr_dino.py index df50620..324fc62 100644 --- a/dino-pretraining/tests/test_mr_dino.py +++ b/dino-pretraining/tests/test_mr_dino.py @@ -9,24 +9,24 @@ CACHE_MANIFEST_NAME, CropSpec, InfiniteStudySampler, - MRAtlasDINO3DDataset, + MRAlignedDINO3DDataset, _raw_volume, _contained_start, _intersection_box, collate_dino3d, - discover_raw_atlas, + discover_raw_aligned, previous_mr_data_module, - validate_atlas_cache, + validate_aligned_cache, ) -def make_dummy_cache(root, studies=4, shape=(8, 32, 32)): - space = root / "atlas_space" +def make_dummy_cache(root, studies=4, shape=(8, 32, 32), space_name="coreg_space"): + space = root / space_name space.mkdir(parents=True) manifest = { "version": 1, "layout": "per_subject_stack", - "space": "atlas_space", + "space": space_name, "target_spacing": [1.0, 0.5, 0.5], "target_shape": list(shape), "posterior_shift_mm": 15.0, @@ -49,7 +49,7 @@ def make_dummy_cache(root, studies=4, shape=(8, 32, 32)): def tiny_dataset(cache, seed=17): - return MRAtlasDINO3DDataset( + return MRAlignedDINO3DDataset( preprocessed_dir=str(cache), crop_spec=CropSpec( global_shape=(8, 32, 32), @@ -62,40 +62,42 @@ def tiny_dataset(cache, seed=17): ) -def test_cache_contract_rejects_non_atlas(tmp_path): +def test_cache_contract_rejects_other_space(tmp_path): cache = make_dummy_cache(tmp_path / "cache") - assert validate_atlas_cache(str(cache))["space"] == "atlas_space" - manifest_path = cache / "atlas_space" / CACHE_MANIFEST_NAME + assert validate_aligned_cache(str(cache))["space"] == "coreg_space" + with pytest.raises(FileNotFoundError): + validate_aligned_cache(str(cache), space="atlas_space") # coreg cache, atlas requested + manifest_path = cache / "coreg_space" / CACHE_MANIFEST_NAME manifest = json.loads(manifest_path.read_text()) manifest["space"] = "native_space" manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="expected 'atlas_space'"): - validate_atlas_cache(str(cache)) + with pytest.raises(ValueError, match="expected 'coreg_space'"): + validate_aligned_cache(str(cache)) -def test_cache_contract_rejects_coreg_space_argument(tmp_path): +def test_native_space_is_rejected(tmp_path): cache = make_dummy_cache(tmp_path / "cache") - with pytest.raises(ValueError, match="requires space='atlas_space'"): - validate_atlas_cache(str(cache), space="coreg_space") + with pytest.raises(ValueError, match="requires an aligned space"): + validate_aligned_cache(str(cache), space="native_space") + with pytest.raises(ValueError, match="requires an aligned space"): + discover_raw_aligned(str(tmp_path), space="native_space") -def test_raw_discovery_selects_atlas_img_not_coreg_img(tmp_path): - atlas_dir = tmp_path / "batch00" / "study_atlas" / "atlas_img" - atlas_dir.mkdir(parents=True) - (atlas_dir / "t1.nii.gz").touch() - (atlas_dir / "flair.nii.gz").touch() - coreg_dir = tmp_path / "batch00" / "study_coreg" / "coreg_img" - coreg_dir.mkdir(parents=True) - (coreg_dir / "t1.nii.gz").touch() - - samples = discover_raw_atlas(str(tmp_path), selected=None) - assert [sample["study_uid"] for sample in samples] == ["study_atlas"] - assert samples[0]["n_sequences"] == 2 +def test_raw_discovery_selects_requested_space_folder(tmp_path): + for name, sub, files in (("study_a", "atlas_img", 1), ("study_c", "coreg_img", 2), ("study_n", "img", 3)): + folder = tmp_path / "batch00" / name / sub + folder.mkdir(parents=True) + for i in range(files): + (folder / f"seq{i}.nii.gz").touch() + coreg = discover_raw_aligned(str(tmp_path), selected=None) # default space is coreg + assert [(s["study_uid"], s["n_sequences"]) for s in coreg] == [("study_c", 2)] + atlas = discover_raw_aligned(str(tmp_path), selected=None, space="atlas_space") + assert [(s["study_uid"], s["n_sequences"]) for s in atlas] == [("study_a", 1)] def test_raw_transform_is_exact_previous_mil_transform(tmp_path): nib = pytest.importorskip("nibabel") - image_dir = tmp_path / "batch00" / "study" / "atlas_img" + image_dir = tmp_path / "batch00" / "study" / "coreg_img" image_dir.mkdir(parents=True) x, y, z = np.indices((12, 14, 6), dtype=np.float32) array = np.sin(x / 3) + np.cos(y / 4) + z / 7 diff --git a/dino-pretraining/tests/test_probe.py b/dino-pretraining/tests/test_probe.py new file mode 100644 index 0000000..afe04e8 --- /dev/null +++ b/dino-pretraining/tests/test_probe.py @@ -0,0 +1,182 @@ +"""Validation MIL probe: checkpoint loading, token extraction, patient-level CV, metrics.""" +import csv +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +from mr_dino.probe import ( + TEACHER_PREFIX, + TokenBags, + auroc, + average_precision, + embed_sequence, + export_teacher, + load_teacher_backbone, + main, + patient_folds, +) +from mr_dino.train_ddp import build_backbone, check_checkpoint_space + +TESTS = Path(__file__).resolve().parent + + +def write_dcp_checkpoint(path: Path, backbone, args: dict, step: int = 7) -> Path: + import torch.distributed.checkpoint as dcp + + state = {TEACHER_PREFIX + k: v.clone() for k, v in backbone.state_dict().items()} + state["model.student.backbone.cls_token"] = torch.zeros_like(backbone.cls_token) # ignored by the loader + dcp.save(state, storage_writer=dcp.FileSystemWriter(str(path)), no_dist=True) + (path / "metadata.json").write_text(json.dumps({"step": step, "stage": "pretrain", "args": args})) + (path / "COMPLETE").touch() + return path + + +def tiny_args(**extra) -> dict: + return {"arch": "tiny", "space": "coreg_space", "target_shape": [8, 32, 32], + "target_spacing": [1.0, 0.5, 0.5], "posterior_shift_mm": 0.0, **extra} + + +def test_teacher_backbone_loads_from_dcp_and_export(tmp_path): + torch.manual_seed(0) + reference = build_backbone("tiny").eval() + reference.init_weights() + ckpt = write_dcp_checkpoint(tmp_path / "step_00000007", reference, tiny_args()) + loaded, meta = load_teacher_backbone(ckpt) + assert meta["step"] == 7 and meta["args"]["arch"] == "tiny" + x = torch.randn(1, 1, 8, 32, 32) + expected = reference.forward_features(x)["x_norm_patchtokens"] + torch.testing.assert_close(loaded.forward_features(x)["x_norm_patchtokens"], expected) + exported = export_teacher(ckpt, tmp_path / "teacher.pt") + again, meta2 = load_teacher_backbone(exported) + assert meta2["step"] == 7 + torch.testing.assert_close(again.forward_features(x)["x_norm_patchtokens"], expected) + + +def test_embed_sequence_keeps_pooled_foreground_only(): + torch.manual_seed(0) + backbone = build_backbone("tiny").eval() + backbone.init_weights() + volume = torch.full((16, 64, 64), -0.4) + volume[2:10, 8:40, 8:40] = 0.3 # "head" occupies part of the volume + tokens, coords = embed_sequence(backbone, volume, tile=(8, 32, 32), pool=(2, 2, 2)) + # token grid 8x4x4, pooled 4x2x2; head covers token z 1..4, y/x 0..2 -> pooled z 0..2, y/x 0..1 + assert tokens.dtype == torch.float16 and tokens.shape[1] == backbone.embed_dim + assert {tuple(c) for c in coords.tolist()} == {(z, y, x) for z in range(3) for y in range(2) for x in range(2)} + assert torch.isfinite(tokens.float()).all() + empty, _ = embed_sequence(backbone, torch.full((8, 32, 32), -0.4), tile=(8, 32, 32), pool=(2, 2, 2)) + assert len(empty) == 0 + + +def test_patient_folds_keep_patients_together(): + uids = [f"s{i}" for i in range(40)] + patient_of = {u: f"p{i // 3}" for i, u in enumerate(uids)} + folds = patient_folds(uids, patient_of, 5, seed=1) + for p in set(patient_of.values()): + assert len({f for u, f in zip(uids, folds) if patient_of[u] == p}) == 1 + assert sorted(set(folds.tolist())) == [0, 1, 2, 3, 4] + assert (folds == patient_folds(uids, patient_of, 5, seed=1)).all() + + +def test_auroc_and_ap_match_definitions(): + rng = np.random.default_rng(0) + y = rng.integers(0, 2, 200) + s = np.round(rng.normal(size=200) + y, 1) # ties on purpose + pos, neg = s[y == 1], s[y == 0] + brute = ((pos[:, None] > neg[None]).sum() + 0.5 * (pos[:, None] == neg[None]).sum()) / (len(pos) * len(neg)) + assert auroc(y, s) == pytest.approx(brute) + assert auroc(np.zeros(5), np.arange(5.0)) is None + assert average_precision(np.array([1, 0, 1]), np.array([0.9, 0.8, 0.1])) == pytest.approx((1 + 2 / 3) / 2) + + +def write_token_shards(root: Path, n=60, dim=16, seed=0): + """Bags whose label-0 positives carry a signal token direction; two worker shards.""" + rng = np.random.default_rng(seed) + uids = [f"study{i:03d}" for i in range(n)] + labels = rng.integers(0, 2, (n, 2)) + root.mkdir(parents=True) + for rank in range(2): + mine = list(range(rank, n, 2)) + chunks, starts, counts, total = [], [], [], 0 + for i in mine: + bag = rng.normal(size=(int(rng.integers(5, 12)), dim)).astype(np.float16) + if labels[i, 0]: + bag[0, 0] += 6.0 + chunks.append(bag) + starts.append(total) + counts.append(len(bag)) + total += len(bag) + (root / f"tokens_rank{rank:04d}.f16").write_bytes(np.concatenate(chunks).tobytes()) + np.savez(root / f"index_rank{rank:04d}.npz", format="mrdino_probe_tokens_v1", dim=dim, + token_file=f"tokens_rank{rank:04d}.f16", study_uid=np.array([uids[i] for i in mine]), + start=np.array(starts), count=np.array(counts), checkpoint_step=7, checkpoint_stage="pretrain", + space="coreg_space", tile=np.array([8, 32, 32]), pool=np.array([2, 2, 2])) + with open(root.parent / "labels.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "signal", "noise"]) + w.writerows([u, *labels[i]] for i, u in enumerate(uids)) + with open(root.parent / "splits.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "patient_uid", "split"]) + w.writerows([u, f"p{i // 2}", "val"] for i, u in enumerate(uids)) + return uids + + +def test_cv_finds_signal_and_predicts_every_study_once(tmp_path): + uids = write_token_shards(tmp_path / "feat") + bags = TokenBags(tmp_path / "feat") + assert sorted(bags.uids) == uids + common = ["--features-dir", str(tmp_path / "feat"), "--labels-csv", str(tmp_path / "labels.csv"), + "--splits-csv", str(tmp_path / "splits.csv"), "--device", "cpu"] + for fold in range(3): + main(["cv", *common, "--out-dir", str(tmp_path / "out"), "--fold", str(fold), "--folds", "3", + "--epochs", "15", "--batch-size", "4", "--hidden-dim", "16", "--mlp-hidden-dim", "16"]) + main(["summarize", *common, "--out-dir", str(tmp_path / "out"), "--folds", "3"]) + result = json.loads((tmp_path / "out" / "results.json").read_text()) + assert result["studies"] == 60 and result["checkpoint_step"] == 7 + signal = next(r for r in result["per_class"] if r["label"] == "signal") + assert signal["auroc"] > 0.9 + + +def test_extract_end_to_end_on_raw_coreg_tree(tmp_path): + pytest.importorskip("nibabel") + data = tmp_path / "data" + subprocess.run([sys.executable, str(TESTS / "make_dummy_aligned_dataset.py"), "--out", str(data), + "--studies", "5", "--shape", "8", "32", "32"], check=True) + uids = sorted(p.name for p in (data / "batch00").iterdir()) + with open(tmp_path / "labels.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "a"]) + w.writerows([u, i % 2] for i, u in enumerate(uids)) + with open(tmp_path / "splits.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "patient_uid", "split"]) + w.writerows([u, u, "val" if i < 4 else "train"] for i, u in enumerate(uids)) + torch.manual_seed(0) + backbone = build_backbone("tiny").eval() + backbone.init_weights() + ckpt = write_dcp_checkpoint(tmp_path / "ckpt", backbone, tiny_args()) + args = ["extract", "--checkpoint", str(ckpt), "--data-folder", str(data), "--features-dir", str(tmp_path / "feat"), + "--labels-csv", str(tmp_path / "labels.csv"), "--splits-csv", str(tmp_path / "splits.csv"), + "--tile", "8", "32", "32", "--workers", "0", "--device", "cpu"] + main(args) + bags = TokenBags(tmp_path / "feat") + assert bags.uids == uids[:4] # val only, sorted + assert all(len(bags.tokens(i)) > 0 for i in range(len(bags))) + with np.load(tmp_path / "feat" / "index_rank0000.npz") as index: + assert set(index["seq_id"].tolist()) == {0, 1, 2} # every sequence contributes + with pytest.raises(ValueError, match="trained on 'coreg_space'"): + main([*args[:-2], "--device", "cpu", "--space", "atlas_space", "--features-dir", str(tmp_path / "f2")]) + + +def test_checkpoint_space_guard(): + check_checkpoint_space("mrrate_aligned_dinov3d_fsdp2_v1", {"space": "coreg_space"}, "coreg_space", + "mrrate_aligned_dinov3d_fsdp2_v1") + with pytest.raises(RuntimeError, match="trained on 'atlas_space'"): # pre-coreg checkpoints are atlas + check_checkpoint_space("mrrate_atlas_dinov3d_fsdp2_v1", {}, "coreg_space", "mrrate_aligned_dinov3d_fsdp2_v1") + with pytest.raises(RuntimeError, match="not an MR DINO checkpoint"): + check_checkpoint_space("something_else", {}, "coreg_space", "mrrate_aligned_dinov3d_fsdp2_v1") From c7447ab6d5f282c2bfc8efa14087cdcdac38ae9e Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:08:29 +0200 Subject: [PATCH 5/7] Default phase 1 to 30k steps for the 4-node run At 4 nodes (global batch 128) 30k steps is ~6 passes over the ~635k training sequences; the same ViT-L recipe on 4 nodes peaked at ~23.5k steps in FORA. 125k (~25 epochs at 4 nodes, ~200 at 32) was far past that. The cosine decay ends at STEPS, so the default is the real run length; the probe confirms it. --- dino-pretraining/README.md | 10 ++++++++++ dino-pretraining/mr_dino/train_ddp.py | 2 +- dino-pretraining/mr_dino/train_fsdp.py | 2 +- dino-pretraining/scripts/train_32n_vitl.sbatch | 2 +- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index c077e2c..e8d79f3 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -173,6 +173,16 @@ OUTPUT=/path/to/mrdino3d_7b/pretrain \ sbatch scripts/train_32n_vitl.sbatch # FORA recipe (ViT-L); 7B launcher kept as legacy ``` +**Nodes and length.** The launcher defaults to 32 nodes; run it on 4 with `sbatch --nodes=4 …`. +Nothing else changes: the global batch becomes 4×4×8 = 128 and the peak learning rate is +sqrt-scaled from batch 1024 automatically. +- **Phase 1 = 30,000 steps.** At batch 128 that is ~6 passes over the ~635k training + sequences. +- **FORA reference:** the same ViT-L recipe on 4 nodes (also batch 128) reached its best + validation MIL AUROC at ~23.5k steps. +- **Set `STEPS` before the run starts.** The cosine learning-rate decay ends at `STEPS`, so it + is the real run length; use the probe to confirm the curve has flattened. + Important overrides are `STEPS`, `BATCH_SIZE`, `GRAD_ACCUM_STEPS`, `WORKERS`, `LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `WARMUP_STEPS`, `RESUME`, `FP8`, and `COMPILE`. The default warmup is 2,500 optimizer steps (FORA recipe). `RESUME=latest` diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index fa86792..7c47bc6 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -72,7 +72,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="large") - p.add_argument("--steps", type=int, default=125_000) + p.add_argument("--steps", type=int, default=30_000) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--workers", type=int, default=2) p.add_argument("--seed", type=int, default=3407) diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index e1aaf90..231f73c 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -72,7 +72,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="large", help="FORA lesson: a from-scratch ViT-L beat the 7B continuation; 7b kept for reference") - p.add_argument("--steps", type=int, default=125_000) + p.add_argument("--steps", type=int, default=30_000) p.add_argument("--batch-size", type=int, default=8) p.add_argument( "--grad-accum-steps", diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch index a77a0a8..ed02a83 100644 --- a/dino-pretraining/scripts/train_32n_vitl.sbatch +++ b/dino-pretraining/scripts/train_32n_vitl.sbatch @@ -42,7 +42,7 @@ if [[ -n "$end_ts" ]]; then fi case "$STAGE" in - pretrain) STEPS=${STEPS:-125000}; GRAM_WEIGHT=0 ;; + pretrain) STEPS=${STEPS:-30000}; GRAM_WEIGHT=0 ;; # ~6 sequence-epochs at 4 nodes (batch 128); cosine ends here gram) STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0} ;; highres) STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5} ;; *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; From a935b97a67ac57d7916965458ee54cd018422660 Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:10:52 +0200 Subject: [PATCH 6/7] Default phase 1 to 50k steps ~10 passes over the training sequences at 4 nodes (batch 128). --- dino-pretraining/README.md | 2 +- dino-pretraining/mr_dino/train_ddp.py | 2 +- dino-pretraining/mr_dino/train_fsdp.py | 2 +- dino-pretraining/scripts/train_32n_vitl.sbatch | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index e8d79f3..b4808a3 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -176,7 +176,7 @@ sbatch scripts/train_32n_vitl.sbatch # FORA recipe (ViT-L); 7B launcher kept a **Nodes and length.** The launcher defaults to 32 nodes; run it on 4 with `sbatch --nodes=4 …`. Nothing else changes: the global batch becomes 4×4×8 = 128 and the peak learning rate is sqrt-scaled from batch 1024 automatically. -- **Phase 1 = 30,000 steps.** At batch 128 that is ~6 passes over the ~635k training +- **Phase 1 = 50,000 steps.** At batch 128 that is ~10 passes over the ~635k training sequences. - **FORA reference:** the same ViT-L recipe on 4 nodes (also batch 128) reached its best validation MIL AUROC at ~23.5k steps. diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index 7c47bc6..7f1accb 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -72,7 +72,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="large") - p.add_argument("--steps", type=int, default=30_000) + p.add_argument("--steps", type=int, default=50_000) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--workers", type=int, default=2) p.add_argument("--seed", type=int, default=3407) diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index 231f73c..4d3ee5c 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -72,7 +72,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="large", help="FORA lesson: a from-scratch ViT-L beat the 7B continuation; 7b kept for reference") - p.add_argument("--steps", type=int, default=30_000) + p.add_argument("--steps", type=int, default=50_000) p.add_argument("--batch-size", type=int, default=8) p.add_argument( "--grad-accum-steps", diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch index ed02a83..9026254 100644 --- a/dino-pretraining/scripts/train_32n_vitl.sbatch +++ b/dino-pretraining/scripts/train_32n_vitl.sbatch @@ -42,7 +42,7 @@ if [[ -n "$end_ts" ]]; then fi case "$STAGE" in - pretrain) STEPS=${STEPS:-30000}; GRAM_WEIGHT=0 ;; # ~6 sequence-epochs at 4 nodes (batch 128); cosine ends here + pretrain) STEPS=${STEPS:-50000}; GRAM_WEIGHT=0 ;; # ~10 sequence-epochs at 4 nodes (batch 128); cosine ends here gram) STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0} ;; highres) STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5} ;; *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; From 487344fb4ce7cf8ba79d486f585cfb8633235ece Mon Sep 17 00:00:00 2001 From: Sezgin Er <81181477+sezginerr@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:09:27 +0200 Subject: [PATCH 7/7] Train MR DINO as ViT-H+ on 4 nodes with per-phase settings - train_mrdino.sbatch (replaces train_32n_vitl): ViT-H+ 0.84B default, 4 nodes. Batch/checkpointing from a 1-node H+ benchmark at real crop sizes: phase-1 crops batch 8 OOM without checkpointing, 0.39 steps/s at 32 GiB with it; highres crops batch 4 + checkpointing 0.10 steps/s at 81 GiB. - Phases 2/3 restart the schedule, so they get FORA CT-DINO's tested settings (fixed teacher temperature 0.07, EMA 0.996, short warmup, no prototype freeze; gram lr 5e-4, highres lr 4e-4 batch 4 with 4 locals) instead of re-running the phase-1 warmups. - A requeued phase-2/3 job continued from the previous phase (RESUME stays in its environment); its own complete checkpoint now wins. Phases 2/3 refuse to start without RESUME. - GPU smoke also runs pretrain -> gram -> highres; README describes the full training (model, phase table, memory/speed, time estimates, commands, rules). --- dino-pretraining/README.md | 150 +++++++++++------- dino-pretraining/mr_dino/train_fsdp.py | 20 ++- dino-pretraining/scripts/probe_watch.sh | 2 +- dino-pretraining/scripts/smoke_dummy.sbatch | 13 ++ dino-pretraining/scripts/train_32n_7b.sbatch | 2 +- .../scripts/train_32n_hplus.sbatch | 2 +- .../scripts/train_32n_vitl.sbatch | 75 --------- dino-pretraining/scripts/train_mrdino.sbatch | 102 ++++++++++++ dino-pretraining/tests/test_stage_resume.py | 30 ++++ 9 files changed, 255 insertions(+), 141 deletions(-) delete mode 100644 dino-pretraining/scripts/train_32n_vitl.sbatch create mode 100644 dino-pretraining/scripts/train_mrdino.sbatch create mode 100644 dino-pretraining/tests/test_stage_resume.py diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index b4808a3..60600c4 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -80,7 +80,7 @@ ViT-L run uses (code: `mr_dino/recipe.py`, `objective.py`, `data.py`; tests: |---|---| | DINOv3 plain-linear prototype layer (std 0.02) never sharpened at lr 5e-5 → flat Sinkhorn targets for 10k steps (iBOT target entropy 8.0/11.5 nats, top-1 0.006) | cosine prototypes on both heads (`--normalize-prototypes`) | | 1k warmup + 7B settings drove positional drift in the last blocks | DINOv2 ViT-L recipe: lr 1e-3 at batch 1024 (sqrt-scaled), warmup 2,500, teacher-temperature warmup 0.04→0.07 over 10k, weight decay 0.04→0.4, EMA 0.992→1, AdamW β2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr ×0.2, prototype layers frozen for 1,250 steps | -| 7B continuation was worse than a from-scratch ViT-L at the same budget | `--arch large` (ViT-L, 16k/16k prototypes, heads 2048/256), batch 8/GPU, no FP8, no activation checkpointing | +| 7B continuation was worse than a from-scratch ViT-L at the same budget | train from scratch at a trainable size: FORA used `--arch large` (ViT-L); MR uses `--arch hplus` (ViT-H+, 0.84B). 16k/16k prototypes, heads 2048/256, no FP8 | | global crops overlapping ≥75% were near-duplicates | ≥25% overlap per axis (`--global-overlap`), locals inside one of the two globals | | head-facing block encoded where a token sat in its crop (last-block position leakage 0.20, crop-quadrant prototype fields) | second global snapped to the patch lattice + **cross-crop twin-token objective** (`--cross-view-weight 1.0`): the student's patch prediction in one crop matches the teacher target of the same voxels in the other crop; leakage fell to 0.09 and stayed flat | | patch targets balanced over all masked tokens at once | Sinkhorn within 2×2×2 spatial bins of the crop (`--position-bins`, CAPI) | @@ -96,13 +96,13 @@ features, at the risk of suppressing lesions visible in only one sequence). **Cross-sequence global pairs default to 0.25 (was 0.75).** A DINO target shared by a T1 and a FLAIR crop rewards discarding contrast-specific findings (FLAIR-only gliosis, SWI-only bleeds), -and in atlas space the cheapest shared signal is the crop's atlas position. Many "different" +and in an aligned space the cheapest shared signal is the crop's position. Many "different" sequences are the same contrast in another plane, so the true cross-contrast share is lower still. A small non-zero rate keeps one embedding space across contrasts for MIL bags that mix sequences; 0 vs 0.25 is to be settled with the validation MIL probe. -Production: `scripts/train_32n_vitl.sbatch`. `train_32n_7b.sbatch` / `train_32n_hplus.sbatch` -keep the old recipe for reference only. +Production: `scripts/train_mrdino.sbatch` (ViT-H+, 4 nodes, 3 phases; see below). +`train_32n_7b.sbatch` / `train_32n_hplus.sbatch` keep the old recipe for reference only. **Healthy early training (FORA reference, ViT-L, global batch 128–1024):** iBOT target top-1 rises from ~0.003 to >0.05 within 2.5k steps, patch target entropy falls, prototype usage stays @@ -160,66 +160,106 @@ sbatch scripts/smoke_dummy.sbatch The job succeeds only after `step_00000002/COMPLETE` exists. -## Production training - -The production launcher reads the extracted MR-RATE-coreg tree directly (`SPACE=atlas_space` -switches back to the atlas tree): - -```bash -cd dino-pretraining -DATA_FOLDER=/path/to/MR-RATE-coreg/mri \ -SPLITS_CSV=/path/to/splits.csv \ -OUTPUT=/path/to/mrdino3d_7b/pretrain \ -sbatch scripts/train_32n_vitl.sbatch # FORA recipe (ViT-L); 7B launcher kept as legacy -``` - -**Nodes and length.** The launcher defaults to 32 nodes; run it on 4 with `sbatch --nodes=4 …`. -Nothing else changes: the global batch becomes 4×4×8 = 128 and the peak learning rate is -sqrt-scaled from batch 1024 automatically. -- **Phase 1 = 50,000 steps.** At batch 128 that is ~10 passes over the ~635k training - sequences. -- **FORA reference:** the same ViT-L recipe on 4 nodes (also batch 128) reached its best - validation MIL AUROC at ~23.5k steps. -- **Set `STEPS` before the run starts.** The cosine learning-rate decay ends at `STEPS`, so it - is the real run length; use the probe to confirm the curve has flattened. +## Production training (3 phases, ViT-H+, 4 nodes) + +`scripts/train_mrdino.sbatch` runs every phase. Set `STAGE`; everything else has a phase +default. It reads the extracted MR-RATE-coreg tree directly (`SPACE=atlas_space` switches to +the atlas tree). + +**Model: ViT-H+ (0.84B).** This is the official DINOv3 H+ shape with 3-D patch embedding +and RoPE: 32 blocks, width 1280, 20 heads, SwiGLU ×6, 2×16×16 patch = 2×8×8 mm tokens. The +heads are 2048 → 256 with 16,384 cosine prototypes each. +- **Why not 7B:** the 7B (the previous MR default) failed in FORA with the old recipe, and + 4 nodes cannot train it. +- **Why not ViT-L:** ViT-L (0.30B, `ARCH=large`) is the size proven with this recipe in + FORA CT, but it is smaller than wanted. +- **Learning rate:** it was tuned on ViT-L. Watch the target metrics in the first + 2–3k steps (see "Healthy early training" above). + +| | Phase 1 `pretrain` | Phase 2 `gram` | Phase 3 `highres` | +|---|---|---|---| +| steps | **50,000** (~10 passes over ~635k sequences) | 10,000 | 5,000 | +| starts from | scratch | phase-1 end | phase-2 end (or phase-1 end, see below) | +| global crops | 2 × 64×192×192 vox (64×96×96 mm, 4,608 tokens) | same | 2 × 64×384×384 (whole axial plane, 18,432 tokens) | +| local crops | 8 × 32×96×96 | 8 × 32×96×96 | 4 × 32×192×192 | +| batch / GPU (global on 4 nodes) | 8 (128) | 8 (128) | 4 (64) | +| activation checkpointing | on (selective) | on | on | +| peak lr (at batch 1024, sqrt-scaled) | 1e-3 → 3.5e-4 | 5e-4 → 1.8e-4 | 4e-4 → 1e-4 | +| warmup / end lr | 2,500 / cosine to 1e-6 at step end | 250 / same | 500 / same | +| weight decay | 0.04 → 0.4 | 0.1 → 0.4 | 0.04 → 0.2 | +| teacher temperature | 0.04 → 0.07 over 10k | 0.07 fixed | 0.07 fixed | +| teacher EMA | 0.992 → 1 | 0.996 → 1 | 0.996 → 1 | +| prototype freeze | first 1,250 steps | none | none | +| Gram weight | 0 | 1.0 | 1.5 | + +Phases 2 and 3 use FORA CT-DINO's tested settings (`phase2_gram_recipe.args`, +`phase3_hires_recipe.args`). A phase restarts its schedule at step 0, so reusing phase-1 +values would redo the warmups: teacher temperature back to 0.04, prototypes frozen again. + +**Settings shared by all phases:** +- **Losses:** DINO 1.0 + iBOT 1.0 (block masks 10–50%) + twin-token 1.0 + KoLeo 0.1. + Sinkhorn targets use 2×2×2 spatial bins. +- **Sequences:** cross-sequence global pairs 25%, twin tokens within the same sequence. +- **Optimiser:** AdamW β2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr ×0.2, BF16. +- **Infrastructure:** FSDP2, compiled blocks, checkpoint every 500 steps (last 6 kept). + +**Memory and speed per H200** (1-node benchmark, H+, compiled, real crop sizes): + +| Configuration | Result | +|---|---| +| phase-1 crops, batch 8, no checkpointing | **out of memory** | +| phase-1 crops, batch 8 + checkpointing | 0.39 steps/s, 32 GiB (default) | +| phase-1 crops, batch 4, no checkpointing | 0.95 steps/s, 83 GiB | +| highres crops, batch 2, no checkpointing | 128 GiB (too tight) | +| highres crops, batch 2 + checkpointing | 0.23 steps/s, 43 GiB | +| highres crops, batch 4 + checkpointing | 0.10 steps/s, 81 GiB (default) | -Important overrides are `STEPS`, `BATCH_SIZE`, `GRAD_ACCUM_STEPS`, `WORKERS`, -`LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `WARMUP_STEPS`, `RESUME`, `FP8`, -and `COMPILE`. The default warmup is 2,500 optimizer steps (FORA recipe). `RESUME=latest` -loads the latest checkpoint only when its `COMPLETE` marker exists. +So on 4 nodes, as long as data loading keeps up, expect roughly: +- phase 1: ~36 h; +- phase 2: ~10 h (the Gram anchor adds a forward pass); +- phase 3: ~14 h. -Phases 2 and 3 resume from the previous stage's checkpoint (the schedule restarts; -the Gram anchor is set to the incoming teacher): +Each phase is longer than one 24 h job; the launcher checkpoints at the wall-time signal +and requeues itself. ```bash -STAGE=gram RESUME=/path/to/pretrain/checkpoints/step_00020000 \ -DATA_FOLDER=/path/to/MR-RATE-coreg/mri OUTPUT=/path/to/mrdino3d_vitl/gram \ -sbatch scripts/train_32n_vitl.sbatch # 10,000 steps, Gram weight 1.0 -STAGE=highres RESUME=/path/to/gram/checkpoints/step_00010000 \ -DATA_FOLDER=/path/to/MR-RATE-coreg/mri OUTPUT=/path/to/mrdino3d_vitl/highres \ -sbatch scripts/train_32n_vitl.sbatch # 5,000 steps, Gram 1.5, 64x384x384 globals -``` - -Probe the phase-2 end against the phase-1 end first. If Gram lost AUROC, start -phase 3 from the phase-1 checkpoint instead (FORA CT did this; there Gram gave no gain). +cd dino-pretraining +export DATA_FOLDER=/path/to/MR-RATE-coreg/mri SPLITS_CSV=/path/to/splits.csv +R=/path/to/mrdino3d_hplus -The launcher catches the Slurm wall-time signal, writes a complete distributed -checkpoint with per-rank sampler and RNG state, and requeues the same job. Do -not start production before `scripts/smoke_dummy.sbatch` passes in the same -container and DINOv3 checkout. +# phase 1 (+ validation probe of chosen steps, run alongside on the login node) +STAGE=pretrain OUTPUT=$R/pretrain sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/pretrain 5000 10000 20000 30000 40000 50000 & -The same 32-node H+ fallback used by CT-DINO is also available: +# phase 2 from the phase-1 end +STAGE=gram RESUME=$R/pretrain/checkpoints/step_00050000 OUTPUT=$R/gram sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/gram 5000 10000 & -```bash -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ -OUTPUT=/path/to/mrdino3d_hplus/pretrain \ -sbatch scripts/train_32n_hplus.sbatch +# phase 3 from the phase-2 end, only if the gram probe is >= the phase-1 probe; +# otherwise RESUME=$R/pretrain/checkpoints/step_00050000 (FORA CT: Gram gave no gain) +STAGE=highres RESUME=$R/gram/checkpoints/step_00010000 OUTPUT=$R/highres sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/highres 2500 5000 & ``` -Both production launchers include the established bad-node exclusions, Slurm -mail notifications, 24-hour requeue policy, deadline checkpointing, node-local -compiler caches, InfiniBand/NCCL settings, and optional split filtering. The -7B path is the production target; H+ is a smaller DDP fallback. +**Rules:** +- **`RESUME` for phases 2/3:** it names the previous phase's checkpoint only for the first + start. Once the phase has its own complete checkpoint, a requeued job continues from that + one. A phase-2/3 launch without `RESUME` and without its own checkpoints is refused. +- **Resume does not cross spaces or arches:** checkpoints record their space, and the + probe rebuilds the arch from the checkpoint metadata. +- **`STEPS` is the real run length:** the cosine decay ends there, so set it before + launching, not mid-run. +- **Overrides:** `ARCH`, `STEPS`, `BATCH_SIZE`, `ACT_CKPT`, `GRAD_ACCUM_STEPS`, `WORKERS`, + `LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `CROSS_VIEW_SEQUENCES`, `GRAM_WEIGHT`, + `RESUME`, `SPACE`, `COMPILE`. Container and paths: `SIF`, `DINOV3_ROOT`, `EXTRA_PIP`, + `ZIG_ARCHIVE` (see `scripts/train_node.sh`). +- **Site-specific `#SBATCH` lines:** edit `--partition`, the `--exclude` bad-node list, + and `--mail-user` before running elsewhere. Slurm refuses unknown node names. +- **Run `scripts/smoke_dummy.sbatch` first** (1 node, ~2 min) in the same container and + DINOv3 checkout. It trains, resumes and runs the probe on synthetic coreg data. + +`train_32n_7b.sbatch` and `train_32n_hplus.sbatch` are the old 7B-continuation recipe +(32-node DDP/FSDP) and are kept for reference only. ## Validation probe (checkpoint scoring) diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index 4d3ee5c..aab63f0 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -70,8 +70,8 @@ def parse_args() -> argparse.Namespace: help="coreg_space (default): sequences on the center T1's native grid; atlas_space: MNI 1 mm box") p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") - p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="large", - help="FORA lesson: a from-scratch ViT-L beat the 7B continuation; 7b kept for reference") + p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="hplus", + help="hplus = ViT-H+ 0.84B (MR default); large = ViT-L 0.30B (FORA CT); 7b kept for reference") p.add_argument("--steps", type=int, default=50_000) p.add_argument("--batch-size", type=int, default=8) p.add_argument( @@ -294,15 +294,19 @@ def checkpoint_root(output: Path) -> Path: def resolve_checkpoint(output: Path, resume: str) -> Path | None: + """``latest``, ``none``, or an explicit checkpoint (normally the previous stage's last one). + + Once this stage has written its own complete checkpoint, that one wins over an explicit + path: a requeued gram/highres job keeps RESUME= in its environment and + must continue, not restart from the previous stage. + """ if not resume or resume.lower() == "none": return None - if resume != "latest": - return Path(resume) marker = checkpoint_root(output) / "latest.txt" - if not marker.exists(): - return None - path = checkpoint_root(output) / marker.read_text().strip() - return path if (path / "COMPLETE").exists() else None + own = checkpoint_root(output) / marker.read_text().strip() if marker.exists() else None + if own is not None and (own / "COMPLETE").exists(): + return own + return None if resume == "latest" else Path(resume) def runtime_state(sampler_offset: int) -> dict: diff --git a/dino-pretraining/scripts/probe_watch.sh b/dino-pretraining/scripts/probe_watch.sh index b3ea6b2..ba8dbcf 100755 --- a/dino-pretraining/scripts/probe_watch.sh +++ b/dino-pretraining/scripts/probe_watch.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Probe chosen steps of a running MR DINO stage as soon as their checkpoints appear. # Usage (login node, e.g. under nohup): DATA_FOLDER= probe_watch.sh RUN_DIR STEP [STEP ...] -# RUN_DIR is a stage output (…/mrdino3d_vitl/pretrain). For each step the EMA-teacher backbone is +# RUN_DIR is a stage output (…/mrdino3d_hplus/pretrain). For each step the EMA-teacher backbone is # exported at once to RUN_DIR/probe/step_N/teacher.pt (checkpoint rotation cannot delete it), then # probe_checkpoint.sbatch is submitted. Results: RUN_DIR/probe/step_N/results.json + summary.txt. set -euo pipefail diff --git a/dino-pretraining/scripts/smoke_dummy.sbatch b/dino-pretraining/scripts/smoke_dummy.sbatch index b0f7761..3119ece 100755 --- a/dino-pretraining/scripts/smoke_dummy.sbatch +++ b/dino-pretraining/scripts/smoke_dummy.sbatch @@ -59,6 +59,19 @@ test -f "$OUTPUT/checkpoints/step_00000002/COMPLETE" grep -q '"step": 2' "$OUTPUT/metrics.jsonl" echo "Synthetic raw-coreg MR DINO FSDP2 integration and resume passed: $OUTPUT" +# Stage transitions with the phase-2/3 launcher settings (fixed teacher temperature, no freeze). +STAGE_FIXED=(--teacher-temperature 0.07 --teacher-warmup-temperature 0.07 --teacher-warmup-steps 1 + --teacher-momentum 0.996 --freeze-last-layer-steps 0 --warmup-steps 1) +srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" "${COMMON_ARGS[@]}" \ + --output-dir "$OUTPUT/gram" --stage gram --gram-weight 1.0 --steps 2 \ + --resume "$OUTPUT/checkpoints/step_00000002" "${STAGE_FIXED[@]}" +test -f "$OUTPUT/gram/checkpoints/step_00000002/COMPLETE" +srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" "${COMMON_ARGS[@]}" \ + --output-dir "$OUTPUT/highres" --stage highres --gram-weight 1.5 --steps 2 \ + --resume "$OUTPUT/gram/checkpoints/step_00000002" "${STAGE_FIXED[@]}" +test -f "$OUTPUT/highres/checkpoints/step_00000002/COMPLETE" +echo "Stage transitions pretrain -> gram -> highres passed" + # Validation probe on the step-2 checkpoint: all dummy studies labelled "val", 2 folds. python3 - "$DUMMY" <<'PY' import csv, pathlib, sys diff --git a/dino-pretraining/scripts/train_32n_7b.sbatch b/dino-pretraining/scripts/train_32n_7b.sbatch index 847406a..8a41253 100755 --- a/dino-pretraining/scripts/train_32n_7b.sbatch +++ b/dino-pretraining/scripts/train_32n_7b.sbatch @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_32n_vitl.sbatch. +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_mrdino.sbatch. #SBATCH --job-name=mr_dino3_7b #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_32n_hplus.sbatch b/dino-pretraining/scripts/train_32n_hplus.sbatch index 305169d..3b2c1af 100755 --- a/dino-pretraining/scripts/train_32n_hplus.sbatch +++ b/dino-pretraining/scripts/train_32n_hplus.sbatch @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_32n_vitl.sbatch. +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_mrdino.sbatch. #SBATCH --job-name=mr_dino3_hplus #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_32n_vitl.sbatch b/dino-pretraining/scripts/train_32n_vitl.sbatch deleted file mode 100644 index 9026254..0000000 --- a/dino-pretraining/scripts/train_32n_vitl.sbatch +++ /dev/null @@ -1,75 +0,0 @@ -#!/usr/bin/env bash -#SBATCH --job-name=mr_dino3_vitl -#SBATCH --partition=h200 -#SBATCH --nodes=32 -#SBATCH --ntasks-per-node=1 -#SBATCH --gres=gpu:h200:4 -#SBATCH --cpus-per-task=128 -#SBATCH --exclusive -#SBATCH --exclude=h21-24,h23-24,h24-17,h23-18 -#SBATCH --time=1-00:00:00 -#SBATCH --requeue -#SBATCH --signal=B:USR1@1200 -#SBATCH --output=mrdino_%x_%j.out -#SBATCH --error=mrdino_%x_%j.err -#SBATCH --mail-type=ALL -#SBATCH --mail-user=sezgin.er@dqbm.uzh.ch -set -uo pipefail - -if [[ -z "${CODE:-}" ]]; then - if [[ -d "$SLURM_SUBMIT_DIR/mr_dino" ]]; then - CODE="$SLURM_SUBMIT_DIR" - elif [[ -d "$SLURM_SUBMIT_DIR/dino-pretraining/mr_dino" ]]; then - CODE="$SLURM_SUBMIT_DIR/dino-pretraining" - else - echo "Cannot locate dino-pretraining; set CODE explicitly" >&2 - exit 2 - fi -fi -SCRIPT_DIR="$CODE/scripts" -: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE-coreg NIfTI root (batchXX//coreg_img)}" -STAGE=${STAGE:-pretrain} -OUTPUT=${OUTPUT:-"$PWD/mrdino3d_vitl/$STAGE"} -mkdir -p "$OUTPUT" -rm -f "$OUTPUT/REQUEUE_REQUESTED" - -head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) -MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') -export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29675} -end_ts=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*EndTime=\([0-9T:-]*\).*/\1/p' | head -1) -if [[ -n "$end_ts" ]]; then - export MRDINO_DEADLINE_EPOCH=$(( $(date -d "$end_ts" +%s) - 1200 )) -fi - -case "$STAGE" in - pretrain) STEPS=${STEPS:-50000}; GRAM_WEIGHT=0 ;; # ~10 sequence-epochs at 4 nodes (batch 128); cosine ends here - gram) STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0} ;; - highres) STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5} ;; - *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; -esac -[[ "${COMPILE:-1}" == 1 ]] && COMPILE_ARGS=(--compile) || COMPILE_ARGS=(--no-compile) -[[ "${FP8:-1}" == 1 ]] && FP8_ARGS=(--fp8) || FP8_ARGS=(--no-fp8) -SPLIT_ARGS=() -if [[ -n "${SPLITS_CSV:-}" ]]; then - SPLIT_ARGS=(--splits-csv "$SPLITS_CSV") -fi - -echo "MR DINOv3 ViT-L (FORA recipe): stage=$STAGE nodes=$SLURM_NNODES output=$OUTPUT data=$DATA_FOLDER" -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" \ - --data-folder "$DATA_FOLDER" \ - --space "${SPACE:-coreg_space}" \ - "${SPLIT_ARGS[@]}" \ - --output-dir "$OUTPUT" --stage "$STAGE" --arch large --steps "$STEPS" \ - --batch-size "${BATCH_SIZE:-8}" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ - --workers "${WORKERS:-8}" --local-crops "${LOCAL_CROPS:-8}" \ - --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.25}" \ - --cross-view-sequences "${CROSS_VIEW_SEQUENCES:-same}" \ - --resume "${RESUME:-latest}" --gram-weight "$GRAM_WEIGHT" \ - "${COMPILE_ARGS[@]}" --no-fp8 --no-activation-checkpointing -rc=$? -if [[ -f "$OUTPUT/REQUEUE_REQUESTED" ]]; then - scontrol requeue "$SLURM_JOB_ID" - exit 0 -fi -exit "$rc" diff --git a/dino-pretraining/scripts/train_mrdino.sbatch b/dino-pretraining/scripts/train_mrdino.sbatch new file mode 100644 index 0000000..edd3f2f --- /dev/null +++ b/dino-pretraining/scripts/train_mrdino.sbatch @@ -0,0 +1,102 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_dino3_hplus +#SBATCH --partition=h200 +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:h200:4 +#SBATCH --cpus-per-task=128 +#SBATCH --exclusive +#SBATCH --exclude=h21-24,h23-24,h24-17,h23-18 +#SBATCH --time=1-00:00:00 +#SBATCH --requeue +#SBATCH --signal=B:USR1@1200 +#SBATCH --output=mrdino_%x_%j.out +#SBATCH --error=mrdino_%x_%j.err +#SBATCH --mail-type=ALL +#SBATCH --mail-user=sezgin.er@dqbm.uzh.ch +set -uo pipefail + +if [[ -z "${CODE:-}" ]]; then + if [[ -d "$SLURM_SUBMIT_DIR/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR" + elif [[ -d "$SLURM_SUBMIT_DIR/dino-pretraining/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR/dino-pretraining" + else + echo "Cannot locate dino-pretraining; set CODE explicitly" >&2 + exit 2 + fi +fi +SCRIPT_DIR="$CODE/scripts" +: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE-coreg NIfTI root (batchXX//coreg_img)}" +STAGE=${STAGE:-pretrain} +OUTPUT=${OUTPUT:-"$PWD/mrdino3d_${ARCH:-hplus}/$STAGE"} +mkdir -p "$OUTPUT" +rm -f "$OUTPUT/REQUEUE_REQUESTED" + +head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) +MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') +export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29675} +end_ts=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*EndTime=\([0-9T:-]*\).*/\1/p' | head -1) +if [[ -n "$end_ts" ]]; then + export MRDINO_DEADLINE_EPOCH=$(( $(date -d "$end_ts" +%s) - 1200 )) +fi + +# Phases 2/3 restart the schedule at step 0, so they need their own settings (FORA CT-DINO +# phase2_gram_recipe.args / phase3_hires_recipe.args): short warmup, teacher temperature fixed at +# 0.07, EMA 0.996, no prototype freeze. Phase 3 has 4x the tokens per global crop -> batch 4 + checkpointing. +# --lr is the peak at global batch 1024 and is sqrt-scaled to the real batch. +# H+ memory/speed per H200 (1-node benchmark 943261, real crop sizes, compiled, bf16): +# phase-1 crops: batch 8 no checkpointing OOM; batch 8 + selective checkpointing 0.39 steps/s 32 GiB; +# batch 4 no checkpointing 0.95 steps/s 83 GiB. +# highres crops (Gram on): batch 2 no checkpointing 128 GiB (too tight); batch 2 + checkpointing +# 0.23 steps/s 43 GiB; batch 4 + checkpointing 0.10 steps/s 81 GiB. +# Defaults keep the FORA global batch per phase on 4 nodes: 128 (phases 1-2), 64 (phase 3). +ARCH=${ARCH:-hplus} +FIXED_TEACHER=(--teacher-temperature 0.07 --teacher-warmup-temperature 0.07 --teacher-warmup-steps 1 + --teacher-momentum 0.996 --freeze-last-layer-steps 0) +case "$STAGE" in + pretrain) # ~10 sequence-epochs at 4 nodes (batch 128); the cosine decay ends here + STEPS=${STEPS:-50000}; GRAM_WEIGHT=0; BATCH_SIZE=${BATCH_SIZE:-8}; LOCAL_CROPS=${LOCAL_CROPS:-8} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=() ;; + gram) # Gram anchoring to the incoming teacher + STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0}; BATCH_SIZE=${BATCH_SIZE:-8}; LOCAL_CROPS=${LOCAL_CROPS:-8} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=(--lr 5e-4 --warmup-steps 250 --weight-decay 0.1 --weight-decay-end 0.4 "${FIXED_TEACHER[@]}") ;; + highres) # 64x384x384 globals (whole axial plane, 18,432 tokens), 4 locals 32x192x192 + STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5}; BATCH_SIZE=${BATCH_SIZE:-4}; LOCAL_CROPS=${LOCAL_CROPS:-4} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=(--lr 4e-4 --warmup-steps 500 --weight-decay 0.04 --weight-decay-end 0.2 "${FIXED_TEACHER[@]}") ;; + *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; +esac +if [[ "$STAGE" != pretrain && "${RESUME:-latest}" == latest && ! -f "$OUTPUT/checkpoints/latest.txt" ]]; then + echo "STAGE=$STAGE starts from the previous stage: set RESUME=/checkpoints/step_N" >&2 + exit 2 +fi +[[ "${COMPILE:-1}" == 1 ]] && COMPILE_ARGS=(--compile) || COMPILE_ARGS=(--no-compile) +[[ "$ACT_CKPT" == 1 ]] && CKPT_ARGS=(--activation-checkpointing) || CKPT_ARGS=(--no-activation-checkpointing) +[[ "${FP8:-1}" == 1 ]] && FP8_ARGS=(--fp8) || FP8_ARGS=(--no-fp8) +SPLIT_ARGS=() +if [[ -n "${SPLITS_CSV:-}" ]]; then + SPLIT_ARGS=(--splits-csv "$SPLITS_CSV") +fi + +echo "MR DINOv3 $ARCH (FORA recipe): stage=$STAGE nodes=$SLURM_NNODES output=$OUTPUT data=$DATA_FOLDER" +srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ + --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" \ + --data-folder "$DATA_FOLDER" \ + --space "${SPACE:-coreg_space}" \ + "${SPLIT_ARGS[@]}" \ + --output-dir "$OUTPUT" --stage "$STAGE" --arch "$ARCH" --steps "$STEPS" \ + --batch-size "$BATCH_SIZE" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ + --workers "${WORKERS:-8}" --local-crops "$LOCAL_CROPS" \ + --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.25}" \ + --cross-view-sequences "${CROSS_VIEW_SEQUENCES:-same}" \ + --resume "${RESUME:-latest}" --gram-weight "$GRAM_WEIGHT" \ + "${COMPILE_ARGS[@]}" --no-fp8 "${CKPT_ARGS[@]}" "${STAGE_ARGS[@]}" +rc=$? +if [[ -f "$OUTPUT/REQUEUE_REQUESTED" ]]; then + scontrol requeue "$SLURM_JOB_ID" + exit 0 +fi +exit "$rc" diff --git a/dino-pretraining/tests/test_stage_resume.py b/dino-pretraining/tests/test_stage_resume.py new file mode 100644 index 0000000..9dd7a7b --- /dev/null +++ b/dino-pretraining/tests/test_stage_resume.py @@ -0,0 +1,30 @@ +"""Stage transitions: an explicit RESUME (previous stage) must not override this stage's own progress.""" +from pathlib import Path + +from mr_dino.train_fsdp import resolve_checkpoint + + +def make_ckpt(root: Path, step: int, complete: bool = True) -> Path: + path = root / "checkpoints" / f"step_{step:08d}" + path.mkdir(parents=True) + if complete: + (path / "COMPLETE").touch() + (root / "checkpoints" / "latest.txt").write_text(path.name + "\n") + return path + + +def test_previous_stage_checkpoint_until_own_progress(tmp_path): + previous = make_ckpt(tmp_path / "pretrain", 50_000) + gram = tmp_path / "gram" + assert resolve_checkpoint(gram, str(previous)) == previous # first start of phase 2 + own = make_ckpt(gram, 500) + assert resolve_checkpoint(gram, str(previous)) == own # requeue continues phase 2 + assert resolve_checkpoint(gram, "latest") == own + assert resolve_checkpoint(gram, "none") is None + + +def test_incomplete_own_checkpoint_is_ignored(tmp_path): + previous = make_ckpt(tmp_path / "pretrain", 50_000) + make_ckpt(tmp_path / "gram", 500, complete=False) + assert resolve_checkpoint(tmp_path / "gram", str(previous)) == previous + assert resolve_checkpoint(tmp_path / "fresh", "latest") is None