f

configs/model_config.yaml

@@ -11,9 +11,9 @@
 #   - vit      : timm ViT-B/16 ImageNet — generic fallback if above fail.
 image_encoder:
   name: "microsoft/rad-dino"          # informational; backend below drives loading
-  backend: "auto"                      # "auto" | "rad_dino" | "biovilt" | "vit"
+  backend: "rad-dino"                      # "auto" | "rad_dino" | "biovilt" | "vit"
   frozen: true                         # freeze encoder during training
-  img_size: 448                        # input image resolution (RAD-DINO native is 518)
+  img_size: 518                        # input image resolution (RAD-DINO native is 518)
   output_dim: 768                      # patch feature dimension (768 for all backends)
 
 # ── MLP Projection (Alignment Layer) ────────
@@ -30,8 +30,25 @@ llm:
   hidden_size: 4096
   load_in_8bit: false                  # 8-bit quantization (breaks on newer CUDA)
   load_in_4bit: true                   # QLoRA NF4 — recommended for T4 15GB
-  torch_dtype: "float16"
+  torch_dtype: "float16"               # Ampere+ should override to "bfloat16"
   device_map: "auto"
+  # ── Attention backend ──
+  # "auto" → let transformers decide.
+  # "flash_attention_2" → 2–3× speedup, Ampere+ only, needs `pip install flash-attn`.
+  # "sdpa" → Torch 2 native, works on T4 (memory-efficient fallback), good default.
+  # "eager" → legacy slow path.
+  # The Colab notebook auto-picks: FA2 on Ampere+/Ada, SDPA on T4.
+  # A failed flash_attention_2 import auto-falls-back to sdpa in cxr_vlm.py.
+  attn_implementation: "auto"
+  # ── 4-bit (QLoRA) tuning ──
+  # Compute dtype for dequantized matmuls. null → tracks torch_dtype.
+  # On Ampere+/Ada set "bfloat16" for better numerical stability.
+  bnb_4bit_compute_dtype: null
+  bnb_4bit_quant_type:    "nf4"        # nf4 (default) | fp4
+  bnb_4bit_use_double_quant: true      # extra ~0.4 bit/param savings
+  # Gradient checkpointing: saves activation memory at ~25% step-time cost.
+  # On ≥24GB GPUs with 4-bit + LoRA + FA2 you can turn this off.
+  gradient_checkpointing: true
 
 # ── LoRA Config ──────────────────────────────
 lora:

configs/train_config.yaml

@@ -130,6 +130,14 @@ data:
   val_split:   "validate"
   test_split:  "test"
 
+  # Optional: directory containing per-image precomputed patch features
+  # ({feature_cache_dir}/{image_relpath}.pt → (P, 768) tensor). When set
+  # and a cache file exists, dataset.py skips the JPEG/transform and the
+  # frozen image encoder is bypassed (model detects the (P, 768) shape).
+  # Build via:  python -m scripts.precompute_image_features --cache_dir ...
+  # null → disabled (default).
+  feature_cache_dir: null
+
 # ── Task Mix Ratios ───────────────────────────
 # For IU-Xray the `vqa` task is ignored automatically (dataset has no VQA),
 # and remaining task weights are renormalized to sum to 1.
@@ -171,6 +179,14 @@ training:
   bf16: false                          # set true on A100/H100 (better numerical stability, no GradScaler)
   cutoff_len: 512                      # max token length per sample
   dataloader_num_workers: 4            # bump to 8–16 on A100 hosts with more CPU cores
+  dataloader_pin_memory: true          # faster H2D copies; ~free win
+  dataloader_persistent_workers: true  # keep workers alive between epochs
+                                       # (auto-disabled when num_workers=0)
+  # Optimizer choice. "adamw_torch" (fp32 moments) ↔ "paged_adamw_8bit" (bnb
+  # block-wise 8-bit moments, ~4× less VRAM, no measurable quality loss per
+  # Dettmers ICLR'22). 8-bit needs bitsandbytes + Ampere+ for best perf —
+  # the auto-detect cell in scripts/cxrvlm_colab_train.ipynb sets this.
+  optim: "adamw_torch"
   # ── checkpointing ──
   # "steps": save every `save_steps`; best (lowest eval_loss) is kept, older
   # overwritten. `save_total_limit=1` → at most 2 folders live (best + latest).

data/collator.py

@@ -3,6 +3,13 @@ collator.py
 -----------
 Custom DataCollator that handles variable-length sequences and
 stacks images into a batch tensor.
+
+Dynamic padding: input_ids / labels / attention_mask are padded to the
+maximum length WITHIN EACH BATCH, not to `cutoff_len`. Batches drawn
+mostly from short tasks (e.g. VQA) skip the wasted compute on padded
+positions entirely — Llama still runs every matmul on every position,
+so cutting off the padding tail is a direct FLOP saving (~1.5–2× on
+this project's task mix; see commit notes).
 """
 
 from dataclasses import dataclass
@@ -13,23 +20,37 @@ import torch
 @dataclass
 class CXRDataCollator:
     """
-    Collates a list of dataset samples into a batch.
+    Collates a list of dataset samples into a batch with **dynamic padding**.
 
-    Handles:
-      - Padding input_ids, attention_mask, labels to same length
-      - Stacking image tensors
+    Args:
+        pad_token_id: token id used to pad input_ids. labels are padded
+                      with -100 (HF Trainer's ignore index for cross-entropy).
     """
     pad_token_id: int = 0
 
     def __call__(self, samples: List[Dict]) -> Dict[str, torch.Tensor]:
-        images          = torch.stack([s["image"] for s in samples])
-        input_ids       = torch.stack([s["input_ids"] for s in samples])
-        attention_masks = torch.stack([s["attention_mask"] for s in samples])
-        labels          = torch.stack([s["labels"] for s in samples])
+        # Images have fixed shape per dataset config (single-image (C,H,W),
+        # multi-image (N,C,H,W), or cached features (P,D)/(N,P,D)) — torch.stack
+        # works for any of them.
+        images = torch.stack([s["image"] for s in samples])
+
+        # ── Dynamic text padding ─────────────────────────────────────────
+        max_len = max(s["input_ids"].size(0) for s in samples)
+        B       = len(samples)
+
+        input_ids      = torch.full((B, max_len), self.pad_token_id, dtype=torch.long)
+        attention_mask = torch.zeros((B, max_len),                    dtype=torch.long)
+        labels         = torch.full((B, max_len), -100,               dtype=torch.long)
+
+        for i, s in enumerate(samples):
+            L = s["input_ids"].size(0)
+            input_ids[i, :L]      = s["input_ids"]
+            attention_mask[i, :L] = 1
+            labels[i, :L]         = s["labels"]
 
         return {
-            "images":        images,          # (B, C, H, W)
-            "input_ids":     input_ids,       # (B, seq_len)
-            "attention_mask": attention_masks, # (B, seq_len)
-            "labels":        labels,          # (B, seq_len)
+            "images":         images,         # (B, ...) image-shape-dependent
+            "input_ids":      input_ids,      # (B, max_len)
+            "attention_mask": attention_mask, # (B, max_len)
+            "labels":         labels,         # (B, max_len)
         }

data/dataset.py

@@ -112,6 +112,7 @@ class CXRInstructDataset(Dataset):
         cutoff_len:   int = 512,
         task_weights: Optional[Dict[str, float]] = None,
         max_images:   int = 1,           # >1 only useful in multi_image_merged mode
+        feature_cache_dir: Optional[str] = None,
     ):
         self.image_root  = Path(image_root)
         self.tokenizer   = tokenizer
@@ -120,6 +121,12 @@ class CXRInstructDataset(Dataset):
         self.split       = split
         self.cutoff_len  = cutoff_len
         self.max_images  = max(1, int(max_images))
+        # When set, _load_image first checks {feature_cache_dir}/{relpath}.pt
+        # and returns the cached (P, 768) patch-feature tensor instead of the
+        # raw image. The model detects this by tensor last-dim and skips the
+        # frozen encoder entirely. Safe because no random augmentation is
+        # applied (Resize + ToTensor + Normalize are deterministic).
+        self.feature_cache_dir = Path(feature_cache_dir) if feature_cache_dir else None
 
         self.task_weights = task_weights or {
             "findings":   0.4,
@@ -235,10 +242,15 @@ class CXRInstructDataset(Dataset):
             target = training_sample["target"],
         )
 
+        # `input_ids` is un-padded at this point. The collator builds the
+        # final attention_mask (1 for real tokens, 0 for batch-level
+        # padding) after stacking. We pass an all-ones placeholder of the
+        # correct per-sample length so other parts of the pipeline that
+        # peek at the dataset output still see the expected key.
         return {
             "image":          image,
             "input_ids":      input_ids,
-            "attention_mask": (input_ids != self.tokenizer.pad_token_id).long(),
+            "attention_mask": torch.ones_like(input_ids),
             "labels":         labels,
             "task":           sample["task"],   # for per-task logging
         }
@@ -267,12 +279,24 @@ class CXRInstructDataset(Dataset):
 
     def _load_image(self, image_path: str) -> torch.Tensor:
         """
-        Load and transform a chest X-ray image.
+        Load and transform a chest X-ray image, or load pre-computed patch
+        features when a feature cache hits.
 
         Args:
             image_path: relative path from image_root
                         e.g. "files/p10/p10000032/s50414267/02aa804e.jpg"
+
+        Returns either (C, H, W) for raw images or (P, 768) when a cached
+        feature file exists at `{feature_cache_dir}/{image_path}.pt`.
         """
+        # ── Fast path: pre-computed patch features ────────────────────────
+        if self.feature_cache_dir is not None:
+            cache_path = self.feature_cache_dir / (image_path + ".pt")
+            if cache_path.is_file():
+                # Load tensor — should be (P, 768) for a single image.
+                return torch.load(cache_path, map_location="cpu", weights_only=True)
+
+        # ── Slow path: read JPEG + transform ──────────────────────────────
         full_path = self.image_root / image_path
         image = Image.open(full_path).convert("RGB")
 
@@ -300,9 +324,15 @@ class CXRInstructDataset(Dataset):
         Tokenize prompt+target. Labels have -100 for prompt tokens
         (so loss is only computed on target tokens).
 
+        Returns UN-PADDED tensors of variable length ≤ cutoff_len. The
+        CXRDataCollator pads them to the max length within each batch,
+        so short batches (e.g. all-VQA) skip the padded compute entirely.
+        Per-sample padding is no longer applied here; that masking lives
+        in the collator now.
+
         Returns:
-            input_ids: (cutoff_len,)
-            labels:    (cutoff_len,) with -100 for prompt positions
+            input_ids: (L,)    L ≤ cutoff_len
+            labels:    (L,)    with -100 for prompt positions
         """
         full_text      = prompt + " " + target
         prompt_encoded = self.tokenizer.encode(prompt, add_special_tokens=True)
@@ -311,19 +341,18 @@ class CXRInstructDataset(Dataset):
             add_special_tokens  = True,
             max_length          = self.cutoff_len,
             truncation          = True,
-            padding             = "max_length",
+            # No padding — collator pads to max-in-batch (dynamic padding).
         )
 
         input_ids = torch.tensor(full_encoded, dtype=torch.long)
 
-        # Labels: -100 for prompt tokens, actual token ids for target tokens
+        # Labels: -100 for prompt tokens, actual token ids for target tokens.
+        # (Padding masking now happens in the collator since there's no
+        # padding at this stage.)
         labels = input_ids.clone()
         prompt_len = min(len(prompt_encoded), self.cutoff_len)
         labels[:prompt_len] = -100
 
-        # Also mask padding
-        labels[input_ids == self.tokenizer.pad_token_id] = -100
-
         return input_ids, labels
 
 

evaluation/evaluate.py

@@ -247,6 +247,7 @@ def main():
         cutoff_len   = train_cfg.training.cutoff_len,
         task_weights = spec.task_weights,
         max_images   = spec.max_images,
+        feature_cache_dir = getattr(train_cfg.data, "feature_cache_dir", None) or None,
     )
 
     # Build task list, intersected with what's available for this dataset.

model/cxr_vlm.py

@@ -118,19 +118,45 @@ class CXRVisionLanguageModel(nn.Module):
             torch_dtype = llm_dtype,
             device_map  = model_cfg.llm.device_map,
         )
+        # attn_implementation: "flash_attention_2" (Ampere+ with flash-attn pkg),
+        # "sdpa" (Torch 2 native; works on T4), "eager" (legacy), or "auto" → let
+        # transformers pick. Falls back gracefully if FA2 is requested but not
+        # installed.
+        _attn_impl = getattr(model_cfg.llm, "attn_implementation", "auto")
+        if _attn_impl and _attn_impl != "auto":
+            llm_kwargs["attn_implementation"] = _attn_impl
         if getattr(model_cfg.llm, "load_in_8bit", False):
             llm_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
         elif getattr(model_cfg.llm, "load_in_4bit", False):
+            # Allow explicit override of compute dtype; default tracks llm_dtype.
+            _bnb_compute = getattr(model_cfg.llm, "bnb_4bit_compute_dtype", None)
+            if _bnb_compute:
+                bnb_dtype = _DTYPE_MAP.get(_bnb_compute, llm_dtype)
+            else:
+                bnb_dtype = llm_dtype if llm_dtype != torch.float32 else torch.float16
             llm_kwargs["quantization_config"] = BitsAndBytesConfig(
                 load_in_4bit              = True,
-                bnb_4bit_quant_type       = "nf4",
-                bnb_4bit_compute_dtype    = llm_dtype if llm_dtype != torch.float32 else torch.float16,
-                bnb_4bit_use_double_quant = True,
+                bnb_4bit_quant_type       = getattr(model_cfg.llm, "bnb_4bit_quant_type", "nf4"),
+                bnb_4bit_compute_dtype    = bnb_dtype,
+                bnb_4bit_use_double_quant = getattr(model_cfg.llm, "bnb_4bit_use_double_quant", True),
             )
-        self.llm = LlamaForCausalLM.from_pretrained(
-            model_cfg.llm.name,
-            **llm_kwargs,
-        )
+        try:
+            self.llm = LlamaForCausalLM.from_pretrained(
+                model_cfg.llm.name,
+                **llm_kwargs,
+            )
+        except (ImportError, ValueError) as e:
+            # FA2 requested but flash-attn missing / GPU unsupported → fall back to SDPA.
+            if "attn_implementation" in llm_kwargs and llm_kwargs["attn_implementation"] == "flash_attention_2":
+                print(f"[cxr_vlm] flash_attention_2 unavailable ({type(e).__name__}: {e}); "
+                      f"falling back to sdpa")
+                llm_kwargs["attn_implementation"] = "sdpa"
+                self.llm = LlamaForCausalLM.from_pretrained(
+                    model_cfg.llm.name,
+                    **llm_kwargs,
+                )
+            else:
+                raise
         # Resize embeddings to include the newly added <image> special token
         if len(self.tokenizer) > self.llm.get_input_embeddings().weight.size(0):
             self.llm.resize_token_embeddings(len(self.tokenizer))
@@ -151,9 +177,13 @@ class CXRVisionLanguageModel(nn.Module):
         self.llm = get_peft_model(self.llm, lora_cfg)
         self.llm.print_trainable_parameters()
 
-        # Gradient checkpointing to save activations memory
-        self.llm.gradient_checkpointing_enable()
-        self.llm.enable_input_require_grads()
+        # Gradient checkpointing to save activations memory. On big-VRAM GPUs
+        # (≥24GB) with 4-bit + LoRA + FA2 the activations fit without it, so
+        # set `model_cfg.llm.gradient_checkpointing=false` to trade ~25% speed
+        # back for the saved memory.
+        if getattr(model_cfg.llm, "gradient_checkpointing", True):
+            self.llm.gradient_checkpointing_enable()
+            self.llm.enable_input_require_grads()
 
     # ────────────────────────────────────────────────────────────────────────
     # Forward pass
@@ -161,10 +191,20 @@ class CXRVisionLanguageModel(nn.Module):
 
     def _encode_images(self, images: torch.Tensor) -> torch.Tensor:
         """
-        Run the (frozen) image encoder. Supports two input shapes:
+        Run the (frozen) image encoder, or skip it when patch features were
+        precomputed and cached. Supported input shapes:
+
+          Raw images:
+            (B, C, H, W)        → (B, num_patches, 768)         single-image
+            (B, N, C, H, W)     → (B, N*num_patches, 768)       multi-image
 
-            (B, C, H, W)     → (B, num_patches, 768)           single-image mode
-            (B, N, C, H, W)  → (B, N*num_patches, 768)          multi-image mode
+          Cached patch features (encoder is frozen, so this is loss-free):
+            (B, P, 768)         → returned as-is                single-image
+            (B, N, P, 768)      → reshaped to (B, N*P, 768)     multi-image
+
+        Disambiguation rule: a tensor whose LAST dim equals the encoder's
+        feature dim (768) is treated as cached features. Raw images last
+        dim equals the image size (e.g. 448), so the two never collide.
 
         In multi-image mode all N views are encoded independently and their
         patch features are concatenated along the patch axis. MLP Projection's
@@ -172,6 +212,19 @@ class CXRVisionLanguageModel(nn.Module):
         tokens, so the downstream shape (B, 32, 4096) is identical regardless
         of N — no other module needs to know about multi-image.
         """
+        feature_dim = self.image_encoder.output_dim
+        if images.size(-1) == feature_dim:
+            # Pre-encoded path: encoder bypass.
+            if images.dim() == 3:                                       # (B, P, D)
+                return images
+            if images.dim() == 4:                                       # (B, N, P, D)
+                B, N, P, D = images.shape
+                return images.reshape(B, N * P, D)                      # (B, N*P, D)
+            raise ValueError(
+                f"Unexpected cached-feature shape: {tuple(images.shape)} "
+                f"(last dim = {feature_dim}, expected 3D or 4D)"
+            )
+        # Raw-image path: run frozen encoder.
         if images.dim() == 4:
             return self.image_encoder(images)
         if images.dim() == 5:

scripts/_apply_notebook_edits.py

@@ -0,0 +1,350 @@
+"""One-shot helper to surgically edit the Colab training notebook.
+
+Replaces cell-cfg with the GPU auto-profile version and inserts a new
+'pre-compute image features' cell after it. Idempotent — re-running
+replaces the new cell rather than duplicating it.
+
+Run from project root:
+    python scripts/_apply_notebook_edits.py
+"""
+import json
+from pathlib import Path
+
+NB_PATH = Path(__file__).resolve().parent / "cxrvlm_colab_train.ipynb"
+
+
+NEW_CFG_SRC = r'''from omegaconf import OmegaConf
+import torch
+
+train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
+model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')
+
+# ── dataset selector ──
+train_cfg.data.dataset_name = DATASET_NAME
+
+# ── training-scheme switches (thesis ablations) ──
+#   report_mode: 'split'         → 2 tasks (findings + impression separately)
+#                'merged'        → 1 task (full report "Findings: ...\n\nImpression: ...")
+#                'split_cascade' → split, but impression's context = GT findings
+#   image_mode : 'all_views_split' | 'frontal_only_split' | 'multi_image_merged'
+train_cfg.data.report_mode             = 'split'
+train_cfg.data.image_mode              = 'all_views_split'
+train_cfg.data.max_images_per_sample   = 2          # only used in multi_image_merged
+
+# ── dataset-specific paths ──
+if DATASET_NAME == 'MIMIC-CXR':
+    train_cfg.data.mimic_cxr_root = str(CXR_ROOT)
+    # Base path; the resolver suffixes __{report_mode}__{image_mode} and
+    # auto-builds (PNU CheXpert + VQA) via data.mimic_cxr_builder.
+    train_cfg.data.instruct_json  = str(mimic_json_path)
+    train_cfg.data.mimic_auto_build = True
+
+    # RaDialog / U-MultiClass abnormality guidance: locate the CheXpert
+    # label CSV so the builder can bake the PNU structured_findings string.
+    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))
+           or sorted(DATA_SRC.rglob('*chexbert*.csv')))
+    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None
+    print('CheXpert CSV :', train_cfg.data.mimic_chexpert_csv
+          or 'NOT FOUND — PNU abnormality guidance DISABLED!')
+
+    # VQA pairs ({train,valid,test}.json) → abnormality-guided VQA.
+    train_cfg.data.mimic_vqa_root = str(VQA_ROOT) if VQA_ROOT is not None else None
+    print('VQA root     :', train_cfg.data.mimic_vqa_root or '(none — VQA skipped)')
+
+elif DATASET_NAME == 'MIMIC-CXR_resized':
+    # The MIMIC-CXR_resized builder is manifest-driven: it reads
+    # `manifest_{train,val,test}.csv` for split + the 14 chex_* labels
+    # (PNU bucketed directly from the CSV, no separate chexpert.csv needed),
+    # uses `report_relpath` from the manifest to find each .txt, and pulls
+    # VQA from `vqa/{vqa,vqa_val,vqa_test}.json`.
+    train_cfg.data.mimic_cxr_resized.root         = str(MR_ROOT)
+    train_cfg.data.mimic_cxr_resized.manifest_dir = None   # null → defaults to root
+    train_cfg.data.mimic_cxr_resized.vqa_dir      = None   # null → {root}/vqa
+    train_cfg.data.mimic_cxr_resized.reports_root = None   # null → auto-probe {root} then {root}/reports
+    train_cfg.data.mimic_cxr_resized.instruct_json = str(mr_json_path)
+    train_cfg.data.mimic_cxr_resized.auto_build   = True
+
+else:  # IU-Xray
+    train_cfg.data.iu_xray.images_dir    = str(IU_IMAGES_DIR)
+    train_cfg.data.iu_xray.labels_dir    = str(IU_LABELS_DIR)
+    train_cfg.data.iu_xray.instruct_json = str(iu_json_path)
+    train_cfg.data.iu_xray.auto_build    = True
+
+train_cfg.data.train_split = 'train'
+train_cfg.data.val_split   = 'validate'
+train_cfg.data.test_split  = 'test'
+
+# ── checkpoint root (Persistence keeps /content/ckpt/) ──
+CKPT_ROOT = WORK / 'ckpt'
+train_cfg.training.output_root = str(CKPT_ROOT)
+
+
+# ─────────────────────────────────────────────────────────────────────────
+# ── GPU auto-profile ────────────────────────────────────────────────────
+# Pick batch size / precision / attention backend / GC / optimizer based on
+# what the current GPU can actually do. Override anything below this block
+# if you want to force a specific setting.
+#
+# Profile rules (compute capability + total VRAM):
+#   T4 (sm_75, 15GB)             → FP16  + SDPA  + GC ON  + bs=1  accum=16 + fp32 AdamW
+#   3090/L4/A10 (sm_80+, 24GB)   → BF16  + FA2   + GC ON  + bs=8  accum=2  + 8-bit AdamW
+#   A100 40GB (sm_80, 40GB)      → BF16  + FA2   + GC OFF + bs=8  accum=2  + 8-bit AdamW
+#   A100/H100 80GB (sm_80+, 80G) → BF16  + FA2   + GC OFF + bs=8  accum=2  + 8-bit AdamW
+#   unknown                      → conservative T4-style profile
+#
+# Why GC ON for 24GB? Bigger batch amortizes the ~25-30% GC overhead.
+# Math (eff batch = 16):
+#   GC OFF, bs=4, accum=4  →  4 × T               = 4.0T per eff-batch
+#   GC ON,  bs=8, accum=2  →  2 × 1.5T × 1.3      = 3.9T per eff-batch  ✓
+# Sub-linear GPU scaling (time(bs=8) ≈ 1.5 × time(bs=4), not 2×) is what
+# tips the balance. On 40GB+ there's room without GC so we skip it there.
+
+assert torch.cuda.is_available(), 'CUDA not available — refusing to write a CPU profile.'
+_props   = torch.cuda.get_device_properties(0)
+_cap     = (_props.major, _props.minor)
+_vram_gb = _props.total_memory / 1e9
+_bf16_ok = torch.cuda.is_bf16_supported()
+_fa2_ok  = _cap >= (8, 0)        # FA2 needs Ampere+ (sm_80 or newer)
+
+print(f'GPU         : {_props.name}  ({_vram_gb:.1f} GB)')
+print(f'Compute cap : sm_{_cap[0]}{_cap[1]}')
+print(f'BF16 native : {_bf16_ok}')
+print(f'FA2 capable : {_fa2_ok}')
+
+# Try to detect whether flash-attn package is actually importable. If FA2 is
+# requested by the profile but the wheel isn't installed, cxr_vlm.py will
+# auto-fall-back to sdpa, but we surface it here so the user knows.
+_flash_attn_installed = False
+if _fa2_ok:
+    try:
+        import flash_attn  # noqa: F401
+        _flash_attn_installed = True
+    except Exception:
+        _flash_attn_installed = False
+print(f'flash-attn  : {"installed" if _flash_attn_installed else "NOT installed (will fall back to sdpa)"}')
+
+# ── Pick profile ─────────────────────────────────────────────────────────
+if _vram_gb >= 70:                                      # A100/H100 80GB
+    _profile = dict(
+        label='A100/H100 80GB',
+        per_device_train_batch_size=8, per_device_eval_batch_size=8,
+        gradient_accumulation_steps=2, dataloader_num_workers=16,
+        gradient_checkpointing=False,
+    )
+elif _vram_gb >= 35:                                    # A100 40GB
+    _profile = dict(
+        label='A100 40GB',
+        per_device_train_batch_size=8, per_device_eval_batch_size=8,
+        gradient_accumulation_steps=2, dataloader_num_workers=12,
+        gradient_checkpointing=False,
+    )
+elif _vram_gb >= 22:                                    # 3090 / L4 / A10 24GB
+    # GC ON + bigger batch beats GC OFF + smaller batch on throughput here.
+    # Per-eff-batch wall time (eff=16):  4×T (GC OFF, bs=4) vs ~3.9×T (GC ON,
+    # bs=8) — sub-linear scaling means bs=8 step is ~1.5×T, not 2×T, so the
+    # GC overhead (~1.3×) is more than paid back.
+    _profile = dict(
+        label='RTX 3090 / L4 / A10 (24GB)',
+        per_device_train_batch_size=8, per_device_eval_batch_size=8,
+        gradient_accumulation_steps=2, dataloader_num_workers=8,
+        gradient_checkpointing=True,
+    )
+elif _vram_gb >= 14:                                    # T4 / V100 16GB
+    _profile = dict(
+        label='T4 / V100 (15-16GB)',
+        per_device_train_batch_size=1, per_device_eval_batch_size=1,
+        gradient_accumulation_steps=16, dataloader_num_workers=2,
+        gradient_checkpointing=True,
+    )
+else:                                                   # tiny / unknown
+    _profile = dict(
+        label=f'unknown ({_vram_gb:.0f}GB) — conservative',
+        per_device_train_batch_size=1, per_device_eval_batch_size=1,
+        gradient_accumulation_steps=16, dataloader_num_workers=2,
+        gradient_checkpointing=True,
+    )
+
+# Precision: BF16 on Ampere+, FP16 on Turing (T4) and older.
+_profile['bf16'] = bool(_bf16_ok)
+_profile['fp16'] = not _bf16_ok
+
+# Attention backend: FA2 if Ampere+ AND flash-attn wheel present, else SDPA.
+_profile['attn_implementation'] = (
+    'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa'
+)
+
+# 8-bit AdamW: bnb's paged_adamw_8bit cuts optimizer-state VRAM ~4× with no
+# measurable quality loss. Skip on Turing where bnb paged optimizer perf is
+# weaker — keep adamw_torch there.
+_profile['optim'] = 'paged_adamw_8bit' if _cap >= (8, 0) else 'adamw_torch'
+
+# 4-bit compute dtype tracks precision.
+_profile['bnb_4bit_compute_dtype'] = 'bfloat16' if _bf16_ok else 'float16'
+_profile['torch_dtype']            = 'bfloat16' if _bf16_ok else 'float16'
+
+print(f'\n→ Profile   : {_profile["label"]}')
+for k, v in _profile.items():
+    if k == 'label': continue
+    print(f'    {k:<32}= {v}')
+
+# ── Write profile into the configs ───────────────────────────────────────
+train_cfg.training.per_device_train_batch_size = _profile['per_device_train_batch_size']
+train_cfg.training.per_device_eval_batch_size  = _profile['per_device_eval_batch_size']
+train_cfg.training.gradient_accumulation_steps = _profile['gradient_accumulation_steps']
+train_cfg.training.dataloader_num_workers      = _profile['dataloader_num_workers']
+train_cfg.training.fp16                        = _profile['fp16']
+train_cfg.training.bf16                        = _profile['bf16']
+train_cfg.training.dataloader_pin_memory       = True
+train_cfg.training.dataloader_persistent_workers = True
+train_cfg.training.optim                       = _profile['optim']
+# Ensure stage2 still uses the same per-run epoch count we want.
+train_cfg.stage2.num_epochs                    = 5
+
+model_cfg.llm.attn_implementation       = _profile['attn_implementation']
+model_cfg.llm.gradient_checkpointing    = _profile['gradient_checkpointing']
+model_cfg.llm.torch_dtype               = _profile['torch_dtype']
+model_cfg.llm.bnb_4bit_compute_dtype    = _profile['bnb_4bit_compute_dtype']
+model_cfg.llm.bnb_4bit_quant_type       = 'nf4'
+model_cfg.llm.bnb_4bit_use_double_quant = True
+
+# ── task weights (sampling ratio enforced by WeightedRandomSampler) ──
+# Defaults in train_config.yaml: 0.30 / 0.20 / 0.50 (RRG ≈ VQA, impression
+# lower because in split_cascade mode it sees GT findings as input).
+# Resolver auto-renormalizes and drops vqa for IU-Xray. Override here only
+# if you want to experiment per-run, e.g.:
+#   train_cfg.tasks.findings_generation.weight   = 0.30
+#   train_cfg.tasks.impression_generation.weight = 0.20
+#   train_cfg.tasks.vqa.weight                   = 0.50
+
+# ── wandb off ──
+train_cfg.wandb.enabled = False
+
+# ── HuggingFace Hub run tracking ──
+train_cfg.hf_hub.enabled        = True
+train_cfg.hf_hub.repo_id        = 'hieu3636/cxr-vlm-runs'   # <<< EDIT ME
+train_cfg.hf_hub.token_env      = 'HF_TOKEN'
+train_cfg.hf_hub.private        = True
+train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')
+
+# ── 4-bit QLoRA ──
+model_cfg.llm.load_in_8bit = False
+model_cfg.llm.load_in_4bit = True
+# Oracle PNU path does NOT use the CheXpert classifier module (labels come
+# from the GT csv/manifest baked into the prompt). Keep it disabled until
+# you wire the learned classifier for realistic inference.
+model_cfg.chexpert_classifier.enabled = False
+
+OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')
+OmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')
+
+print('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))
+print('--- train_cfg.tasks ---');   print(OmegaConf.to_yaml(train_cfg.tasks))
+print('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))
+print('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))
+print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))
+'''
+
+
+FEATURE_CACHE_SRC = r'''# ─── Optional: pre-compute image patch features (skip frozen encoder forward) ──
+#
+# The image encoder is frozen + the transform is deterministic, so encoding the
+# same image every step is wasted work. Run this ONCE per dataset to cache
+# (P, 768) patch tensors under {WORK}/feature_cache/{DATASET_NAME}/ and the
+# training loop will load them instead of re-encoding.
+#
+# Set CACHE_FEATURES = False to skip (e.g. first time you set up the run, want
+# the smoke test to use the raw path, or you're debugging the encoder).
+#
+# Disk usage:  ~3 MB per image (P=1024 patches × 768 dim × fp16). For ~30k
+# unique images that's ~90 GB — make sure WORK has the room, or set
+# CACHE_FEATURES=False on tight quotas.
+
+CACHE_FEATURES = True
+
+if CACHE_FEATURES:
+    feature_cache_dir = WORK / 'feature_cache' / DATASET_NAME
+    feature_cache_dir.mkdir(parents=True, exist_ok=True)
+    train_cfg.data.feature_cache_dir = str(feature_cache_dir)
+    OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')
+
+    # Re-running this cell is safe: --overwrite is OFF by default so cached
+    # files are skipped. To force a full rebuild, add `--overwrite` below.
+    print(f'feature_cache_dir = {feature_cache_dir}')
+    !python -m scripts.precompute_image_features \
+        --model_config configs/model_config.yaml \
+        --train_config configs/train_config.yaml \
+        --cache_dir "{feature_cache_dir}" \
+        --batch_size 16
+else:
+    train_cfg.data.feature_cache_dir = None
+    OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')
+    print('Feature cache DISABLED. Training will run the image encoder every step.')
+'''
+
+
+def src_to_lines(s: str):
+    """Convert a string into Jupyter's list-of-lines source representation."""
+    lines = s.split("\n")
+    return [ln + "\n" for ln in lines[:-1]] + ([lines[-1]] if lines[-1] else [])
+
+
+def main():
+    with open(NB_PATH, "r", encoding="utf-8") as f:
+        nb = json.load(f)
+
+    # Find cell-cfg index
+    cfg_idx = None
+    for i, c in enumerate(nb["cells"]):
+        if c.get("id") == "cell-cfg":
+            cfg_idx = i
+            break
+    if cfg_idx is None:
+        raise RuntimeError("cell-cfg not found in notebook")
+    print(f"cell-cfg at index {cfg_idx}")
+
+    # Replace cell-cfg
+    nb["cells"][cfg_idx]["source"] = src_to_lines(NEW_CFG_SRC)
+    nb["cells"][cfg_idx]["outputs"] = []
+    nb["cells"][cfg_idx]["execution_count"] = None
+
+    # Remove any pre-existing feature-cache cells (idempotent re-run)
+    nb["cells"] = [
+        c for c in nb["cells"]
+        if c.get("id") not in ("cell-feature-cache", "cell-feature-cache-md")
+    ]
+
+    # Re-find cell-cfg index (may have shifted if we removed earlier ones — but
+    # those would have been after it, so index is stable)
+    for i, c in enumerate(nb["cells"]):
+        if c.get("id") == "cell-cfg":
+            cfg_idx = i
+            break
+
+    # Insert markdown + code cells after cell-cfg
+    md_cell = {
+        "cell_type": "markdown",
+        "id": "cell-feature-cache-md",
+        "metadata": {},
+        "source": ["## 4b. Pre-compute image features (optional speedup)\n"],
+    }
+    code_cell = {
+        "cell_type": "code",
+        "id": "cell-feature-cache",
+        "metadata": {},
+        "execution_count": None,
+        "outputs": [],
+        "source": src_to_lines(FEATURE_CACHE_SRC),
+    }
+    nb["cells"].insert(cfg_idx + 1, md_cell)
+    nb["cells"].insert(cfg_idx + 2, code_cell)
+
+    with open(NB_PATH, "w", encoding="utf-8") as f:
+        json.dump(nb, f, indent=1, ensure_ascii=False)
+        f.write("\n")
+
+    print(f"Wrote {NB_PATH}")
+    print(f"New cell count: {len(nb['cells'])}")
+
+
+if __name__ == "__main__":
+    main()

scripts/cxrvlm_colab_train.ipynb

@@ -90,30 +90,7 @@
     },
     "outputId": "9f59fa97-9392-4d73-da56-b6f791fe1d37"
    },
-   "source": [
-    "!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate\n",
-    "\n",
-    "# Let pip pick latest bnb that matches Colab's CUDA 12.8 + triton 3.x\n",
-    "!pip install -q -U bitsandbytes\n",
-    "\n",
-    "# Transformers ≥4.46 fixed the frozenset bug; don't over-pin\n",
-    "!pip install -q \\\n",
-    "    'transformers>=4.46,<4.50' \\\n",
-    "    'peft>=0.13,<0.15' \\\n",
-    "    'accelerate>=1.0'\n",
-    "\n",
-    "!pip install -q \\\n",
-    "    'huggingface_hub>=0.24,<0.27' \\\n",
-    "    omegaconf sentencepiece 'protobuf>=3.20' \\\n",
-    "    nltk rouge-score bert-score sacrebleu\n",
-    "\n",
-    "import torch, transformers, bitsandbytes, peft, accelerate\n",
-    "print('torch        :', torch.__version__, '| cuda:', torch.cuda.is_available())\n",
-    "print('transformers :', transformers.__version__)\n",
-    "print('bitsandbytes :', bitsandbytes.__version__)\n",
-    "print('peft         :', peft.__version__)\n",
-    "print('accelerate   :', accelerate.__version__)"
-   ],
+   "source": "!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate httpx\n\n# Let pip pick latest bnb that matches Colab's CUDA 12.8 + triton 3.x\n!pip install -q -U bitsandbytes\n\n# httpx ≥0.28 REMOVED the `allow_redirects` kwarg (renamed to `follow_redirects`\n# back in 0.21). huggingface_hub still passes `allow_redirects=...` in some\n# code paths → `Client.head() got an unexpected keyword argument 'allow_redirects'`.\n# Pin httpx<0.28 to keep the legacy kwarg working until hub catches up.\n!pip install -q 'httpx>=0.24,<0.28'\n\n# Transformers ≥4.46 fixed the frozenset bug; don't over-pin\n!pip install -q \\\n    'transformers>=4.46,<4.50' \\\n    'peft>=0.13,<0.15' \\\n    'accelerate>=1.0'\n\n!pip install -q \\\n    'huggingface_hub>=0.27,<2.0' \\\n    omegaconf sentencepiece 'protobuf>=3.20' \\\n    nltk rouge-score bert-score sacrebleu\n\nimport torch, transformers, bitsandbytes, peft, accelerate, huggingface_hub, httpx\nprint('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())\nprint('transformers   :', transformers.__version__)\nprint('bitsandbytes   :', bitsandbytes.__version__)\nprint('peft           :', peft.__version__)\nprint('accelerate     :', accelerate.__version__)\nprint('huggingface_hub:', huggingface_hub.__version__)\nprint('httpx          :', httpx.__version__, '  (must be <0.28)')\nassert tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28), \\\n    'httpx>=0.28 will break huggingface_hub. Re-run this cell or restart runtime.'",
    "execution_count": null,
    "outputs": [
     {
@@ -294,11 +271,292 @@
     },
     "outputId": "80ddabe3-bc8b-4d14-94e2-26ff9e64970c"
    },
-   "source": "from omegaconf import OmegaConf\n\ntrain_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\nmodel_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')\n\n# ── dataset selector ──\ntrain_cfg.data.dataset_name = DATASET_NAME\n\n# ── training-scheme switches (thesis ablations) ──\n#   report_mode: 'split'         → 2 tasks (findings + impression separately)\n#                'merged'        → 1 task (full report \"Findings: ...\\n\\nImpression: ...\")\n#                'split_cascade' → split, but impression's context = GT findings\n#   image_mode : 'all_views_split' | 'frontal_only_split' | 'multi_image_merged'\ntrain_cfg.data.report_mode             = 'split'\ntrain_cfg.data.image_mode              = 'all_views_split'\ntrain_cfg.data.max_images_per_sample   = 2          # only used in multi_image_merged\n\n# ── dataset-specific paths ──\nif DATASET_NAME == 'MIMIC-CXR':\n    train_cfg.data.mimic_cxr_root = str(CXR_ROOT)\n    # Base path; the resolver suffixes __{report_mode}__{image_mode} and\n    # auto-builds (PNU CheXpert + VQA) via data.mimic_cxr_builder.\n    train_cfg.data.instruct_json  = str(mimic_json_path)\n    train_cfg.data.mimic_auto_build = True\n\n    # RaDialog / U-MultiClass abnormality guidance: locate the CheXpert\n    # label CSV so the builder can bake the PNU structured_findings string.\n    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))\n           or sorted(DATA_SRC.rglob('*chexbert*.csv')))\n    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None\n    print('CheXpert CSV :', train_cfg.data.mimic_chexpert_csv\n          or 'NOT FOUND — PNU abnormality guidance DISABLED!')\n\n    # VQA pairs ({train,valid,test}.json) → abnormality-guided VQA.\n    train_cfg.data.mimic_vqa_root = str(VQA_ROOT) if VQA_ROOT is not None else None\n    print('VQA root     :', train_cfg.data.mimic_vqa_root or '(none — VQA skipped)')\n\nelif DATASET_NAME == 'MIMIC-CXR_resized':\n    # The MIMIC-CXR_resized builder is manifest-driven: it reads\n    # `manifest_{train,val,test}.csv` for split + the 14 chex_* labels\n    # (PNU bucketed directly from the CSV, no separate chexpert.csv needed),\n    # uses `report_relpath` from the manifest to find each .txt, and pulls\n    # VQA from `vqa/{vqa,vqa_val,vqa_test}.json`.\n    train_cfg.data.mimic_cxr_resized.root         = str(MR_ROOT)\n    train_cfg.data.mimic_cxr_resized.manifest_dir = None   # null → defaults to root\n    train_cfg.data.mimic_cxr_resized.vqa_dir      = None   # null → {root}/vqa\n    train_cfg.data.mimic_cxr_resized.reports_root = None   # null → auto-probe {root} then {root}/reports\n    train_cfg.data.mimic_cxr_resized.instruct_json = str(mr_json_path)\n    train_cfg.data.mimic_cxr_resized.auto_build   = True\n\nelse:  # IU-Xray\n    train_cfg.data.iu_xray.images_dir    = str(IU_IMAGES_DIR)\n    train_cfg.data.iu_xray.labels_dir    = str(IU_LABELS_DIR)\n    train_cfg.data.iu_xray.instruct_json = str(iu_json_path)\n    train_cfg.data.iu_xray.auto_build    = True\n\ntrain_cfg.data.train_split = 'train'\ntrain_cfg.data.val_split   = 'validate'\ntrain_cfg.data.test_split  = 'test'\n\n# ── checkpoint root (Persistence keeps /content/ckpt/) ──\nCKPT_ROOT = WORK / 'ckpt'\ntrain_cfg.training.output_root = str(CKPT_ROOT)\n\n# ── batching / GPU profile ──────────────────────────────────────────────\n# Default below is tuned for **L4 (24GB)** or **RTX 3090 (24GB)**: BF16 compute,\n# batch 4, 4-bit Vicuna (~3.5GB) + LoRA + projection leaves ~18GB for activations.\n# Effective batch = per_device_train_batch_size * gradient_accumulation_steps.\n#\n# To switch GPU: comment out the active profile and uncomment the target one.\n#\n# ─── L4 / RTX 3090 (24GB)  ← ACTIVE / DEFAULT ────────────────────────────\ntrain_cfg.training.per_device_train_batch_size = 4\ntrain_cfg.training.per_device_eval_batch_size  = 4\ntrain_cfg.training.gradient_accumulation_steps = 4    # effective batch = 16\ntrain_cfg.training.fp16                        = False\ntrain_cfg.training.bf16                        = True  # Ampere+/Ada → native BF16, no GradScaler\ntrain_cfg.training.dataloader_num_workers      = 8\n#\n# ─── T4 (15-16GB, Colab free) ────────────────────────────────────────────\n# Turing architecture: NO BF16 support → must use FP16. Smaller batch.\n# train_cfg.training.per_device_train_batch_size = 1\n# train_cfg.training.per_device_eval_batch_size  = 1\n# train_cfg.training.gradient_accumulation_steps = 16   # effective batch = 16\n# train_cfg.training.fp16                        = True\n# train_cfg.training.bf16                        = False\n# train_cfg.training.dataloader_num_workers      = 2    # Colab free ≈ 2 CPU cores\n#\n# ─── A100 40GB ───────────────────────────────────────────────────────────\n# Plenty of headroom: bigger batch, fewer accumulation steps → faster steps.\n# train_cfg.training.per_device_train_batch_size = 8\n# train_cfg.training.per_device_eval_batch_size  = 8\n# train_cfg.training.gradient_accumulation_steps = 2    # effective batch = 16\n# train_cfg.training.fp16                        = False\n# train_cfg.training.bf16                        = True\n# train_cfg.training.dataloader_num_workers      = 16\n# # Optional: turn off 4-bit quant — Vicuna fp16 (~14GB) fits easily on A100.\n# # Slightly better quality, faster per step. Set\n# # `model_cfg.llm.load_in_4bit = False` in the QLoRA section below.\n\ntrain_cfg.stage2.num_epochs                    = 5\n\n# ── task weights (sampling ratio enforced by WeightedRandomSampler) ──\n# Defaults in train_config.yaml: 0.30 / 0.20 / 0.50 (RRG ≈ VQA, impression\n# lower because in split_cascade mode it sees GT findings as input).\n# Resolver auto-renormalizes and drops vqa for IU-Xray. Override here only\n# if you want to experiment per-run, e.g.:\n#   train_cfg.tasks.findings_generation.weight   = 0.30\n#   train_cfg.tasks.impression_generation.weight = 0.20\n#   train_cfg.tasks.vqa.weight                   = 0.50\n\n# ── wandb off ──\ntrain_cfg.wandb.enabled = False\n\n# ── HuggingFace Hub run tracking ──\ntrain_cfg.hf_hub.enabled        = True\ntrain_cfg.hf_hub.repo_id        = 'hieu3636/cxr-vlm-runs'   # <<< EDIT ME\ntrain_cfg.hf_hub.token_env      = 'HF_TOKEN'\ntrain_cfg.hf_hub.private        = True\ntrain_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')\n\n# ── 4-bit QLoRA ──\nmodel_cfg.llm.load_in_8bit = False\nmodel_cfg.llm.load_in_4bit = True\n# Oracle PNU path does NOT use the CheXpert classifier module (labels come\n# from the GT csv/manifest baked into the prompt). Keep it disabled until\n# you wire the learned classifier for realistic inference.\nmodel_cfg.chexpert_classifier.enabled = False\n\nOmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')\nOmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')\n\nprint('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))\nprint('--- train_cfg.tasks ---');   print(OmegaConf.to_yaml(train_cfg.tasks))\nprint('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))\nprint('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))\nprint('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))",
+   "source": [
+    "from omegaconf import OmegaConf\n",
+    "import torch\n",
+    "\n",
+    "train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')\n",
+    "model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')\n",
+    "\n",
+    "# ── dataset selector ──\n",
+    "train_cfg.data.dataset_name = DATASET_NAME\n",
+    "\n",
+    "# ── training-scheme switches (thesis ablations) ──\n",
+    "#   report_mode: 'split'         → 2 tasks (findings + impression separately)\n",
+    "#                'merged'        → 1 task (full report \"Findings: ...\\n\\nImpression: ...\")\n",
+    "#                'split_cascade' → split, but impression's context = GT findings\n",
+    "#   image_mode : 'all_views_split' | 'frontal_only_split' | 'multi_image_merged'\n",
+    "train_cfg.data.report_mode             = 'split'\n",
+    "train_cfg.data.image_mode              = 'all_views_split'\n",
+    "train_cfg.data.max_images_per_sample   = 2          # only used in multi_image_merged\n",
+    "\n",
+    "# ── dataset-specific paths ──\n",
+    "if DATASET_NAME == 'MIMIC-CXR':\n",
+    "    train_cfg.data.mimic_cxr_root = str(CXR_ROOT)\n",
+    "    # Base path; the resolver suffixes __{report_mode}__{image_mode} and\n",
+    "    # auto-builds (PNU CheXpert + VQA) via data.mimic_cxr_builder.\n",
+    "    train_cfg.data.instruct_json  = str(mimic_json_path)\n",
+    "    train_cfg.data.mimic_auto_build = True\n",
+    "\n",
+    "    # RaDialog / U-MultiClass abnormality guidance: locate the CheXpert\n",
+    "    # label CSV so the builder can bake the PNU structured_findings string.\n",
+    "    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))\n",
+    "           or sorted(DATA_SRC.rglob('*chexbert*.csv')))\n",
+    "    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None\n",
+    "    print('CheXpert CSV :', train_cfg.data.mimic_chexpert_csv\n",
+    "          or 'NOT FOUND — PNU abnormality guidance DISABLED!')\n",
+    "\n",
+    "    # VQA pairs ({train,valid,test}.json) → abnormality-guided VQA.\n",
+    "    train_cfg.data.mimic_vqa_root = str(VQA_ROOT) if VQA_ROOT is not None else None\n",
+    "    print('VQA root     :', train_cfg.data.mimic_vqa_root or '(none — VQA skipped)')\n",
+    "\n",
+    "elif DATASET_NAME == 'MIMIC-CXR_resized':\n",
+    "    # The MIMIC-CXR_resized builder is manifest-driven: it reads\n",
+    "    # `manifest_{train,val,test}.csv` for split + the 14 chex_* labels\n",
+    "    # (PNU bucketed directly from the CSV, no separate chexpert.csv needed),\n",
+    "    # uses `report_relpath` from the manifest to find each .txt, and pulls\n",
+    "    # VQA from `vqa/{vqa,vqa_val,vqa_test}.json`.\n",
+    "    train_cfg.data.mimic_cxr_resized.root         = str(MR_ROOT)\n",
+    "    train_cfg.data.mimic_cxr_resized.manifest_dir = None   # null → defaults to root\n",
+    "    train_cfg.data.mimic_cxr_resized.vqa_dir      = None   # null → {root}/vqa\n",
+    "    train_cfg.data.mimic_cxr_resized.reports_root = None   # null → auto-probe {root} then {root}/reports\n",
+    "    train_cfg.data.mimic_cxr_resized.instruct_json = str(mr_json_path)\n",
+    "    train_cfg.data.mimic_cxr_resized.auto_build   = True\n",
+    "\n",
+    "else:  # IU-Xray\n",
+    "    train_cfg.data.iu_xray.images_dir    = str(IU_IMAGES_DIR)\n",
+    "    train_cfg.data.iu_xray.labels_dir    = str(IU_LABELS_DIR)\n",
+    "    train_cfg.data.iu_xray.instruct_json = str(iu_json_path)\n",
+    "    train_cfg.data.iu_xray.auto_build    = True\n",
+    "\n",
+    "train_cfg.data.train_split = 'train'\n",
+    "train_cfg.data.val_split   = 'validate'\n",
+    "train_cfg.data.test_split  = 'test'\n",
+    "\n",
+    "# ── checkpoint root (Persistence keeps /content/ckpt/) ──\n",
+    "CKPT_ROOT = WORK / 'ckpt'\n",
+    "train_cfg.training.output_root = str(CKPT_ROOT)\n",
+    "\n",
+    "\n",
+    "# ─────────────────────────────────────────────────────────────────────────\n",
+    "# ── GPU auto-profile ────────────────────────────────────────────────────\n",
+    "# Pick batch size / precision / attention backend / GC / optimizer based on\n",
+    "# what the current GPU can actually do. Override anything below this block\n",
+    "# if you want to force a specific setting.\n",
+    "#\n",
+    "# Profile rules (compute capability + total VRAM):\n",
+    "#   T4 (sm_75, 15GB)             → FP16  + SDPA  + GC ON  + bs=1  accum=16 + fp32 AdamW\n",
+    "#   3090/L4/A10 (sm_80+, 24GB)   → BF16  + FA2   + GC ON  + bs=8  accum=2  + 8-bit AdamW\n",
+    "#   A100 40GB (sm_80, 40GB)      → BF16  + FA2   + GC OFF + bs=8  accum=2  + 8-bit AdamW\n",
+    "#   A100/H100 80GB (sm_80+, 80G) → BF16  + FA2   + GC OFF + bs=8  accum=2  + 8-bit AdamW\n",
+    "#   unknown                      → conservative T4-style profile\n",
+    "#\n",
+    "# Why GC ON for 24GB? Bigger batch amortizes the ~25-30% GC overhead.\n",
+    "# Math (eff batch = 16):\n",
+    "#   GC OFF, bs=4, accum=4  →  4 × T               = 4.0T per eff-batch\n",
+    "#   GC ON,  bs=8, accum=2  →  2 × 1.5T × 1.3      = 3.9T per eff-batch  ✓\n",
+    "# Sub-linear GPU scaling (time(bs=8) ≈ 1.5 × time(bs=4), not 2×) is what\n",
+    "# tips the balance. On 40GB+ there's room without GC so we skip it there.\n",
+    "\n",
+    "assert torch.cuda.is_available(), 'CUDA not available — refusing to write a CPU profile.'\n",
+    "_props   = torch.cuda.get_device_properties(0)\n",
+    "_cap     = (_props.major, _props.minor)\n",
+    "_vram_gb = _props.total_memory / 1e9\n",
+    "_bf16_ok = torch.cuda.is_bf16_supported()\n",
+    "_fa2_ok  = _cap >= (8, 0)        # FA2 needs Ampere+ (sm_80 or newer)\n",
+    "\n",
+    "print(f'GPU         : {_props.name}  ({_vram_gb:.1f} GB)')\n",
+    "print(f'Compute cap : sm_{_cap[0]}{_cap[1]}')\n",
+    "print(f'BF16 native : {_bf16_ok}')\n",
+    "print(f'FA2 capable : {_fa2_ok}')\n",
+    "\n",
+    "# Try to detect whether flash-attn package is actually importable. If FA2 is\n",
+    "# requested by the profile but the wheel isn't installed, cxr_vlm.py will\n",
+    "# auto-fall-back to sdpa, but we surface it here so the user knows.\n",
+    "_flash_attn_installed = False\n",
+    "if _fa2_ok:\n",
+    "    try:\n",
+    "        import flash_attn  # noqa: F401\n",
+    "        _flash_attn_installed = True\n",
+    "    except Exception:\n",
+    "        _flash_attn_installed = False\n",
+    "print(f'flash-attn  : {\"installed\" if _flash_attn_installed else \"NOT installed (will fall back to sdpa)\"}')\n",
+    "\n",
+    "# ── Pick profile ─────────────────────────────────────────────────────────\n",
+    "if _vram_gb >= 70:                                      # A100/H100 80GB\n",
+    "    _profile = dict(\n",
+    "        label='A100/H100 80GB',\n",
+    "        per_device_train_batch_size=8, per_device_eval_batch_size=8,\n",
+    "        gradient_accumulation_steps=2, dataloader_num_workers=16,\n",
+    "        gradient_checkpointing=False,\n",
+    "    )\n",
+    "elif _vram_gb >= 35:                                    # A100 40GB\n",
+    "    _profile = dict(\n",
+    "        label='A100 40GB',\n",
+    "        per_device_train_batch_size=8, per_device_eval_batch_size=8,\n",
+    "        gradient_accumulation_steps=2, dataloader_num_workers=12,\n",
+    "        gradient_checkpointing=False,\n",
+    "    )\n",
+    "elif _vram_gb >= 22:                                    # 3090 / L4 / A10 24GB\n",
+    "    # GC ON + bigger batch beats GC OFF + smaller batch on throughput here.\n",
+    "    # Per-eff-batch wall time (eff=16):  4×T (GC OFF, bs=4) vs ~3.9×T (GC ON,\n",
+    "    # bs=8) — sub-linear scaling means bs=8 step is ~1.5×T, not 2×T, so the\n",
+    "    # GC overhead (~1.3×) is more than paid back.\n",
+    "    _profile = dict(\n",
+    "        label='RTX 3090 / L4 / A10 (24GB)',\n",
+    "        per_device_train_batch_size=8, per_device_eval_batch_size=8,\n",
+    "        gradient_accumulation_steps=2, dataloader_num_workers=8,\n",
+    "        gradient_checkpointing=True,\n",
+    "    )\n",
+    "elif _vram_gb >= 14:                                    # T4 / V100 16GB\n",
+    "    _profile = dict(\n",
+    "        label='T4 / V100 (15-16GB)',\n",
+    "        per_device_train_batch_size=1, per_device_eval_batch_size=1,\n",
+    "        gradient_accumulation_steps=16, dataloader_num_workers=2,\n",
+    "        gradient_checkpointing=True,\n",
+    "    )\n",
+    "else:                                                   # tiny / unknown\n",
+    "    _profile = dict(\n",
+    "        label=f'unknown ({_vram_gb:.0f}GB) — conservative',\n",
+    "        per_device_train_batch_size=1, per_device_eval_batch_size=1,\n",
+    "        gradient_accumulation_steps=16, dataloader_num_workers=2,\n",
+    "        gradient_checkpointing=True,\n",
+    "    )\n",
+    "\n",
+    "# Precision: BF16 on Ampere+, FP16 on Turing (T4) and older.\n",
+    "_profile['bf16'] = bool(_bf16_ok)\n",
+    "_profile['fp16'] = not _bf16_ok\n",
+    "\n",
+    "# Attention backend: FA2 if Ampere+ AND flash-attn wheel present, else SDPA.\n",
+    "_profile['attn_implementation'] = (\n",
+    "    'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa'\n",
+    ")\n",
+    "\n",
+    "# 8-bit AdamW: bnb's paged_adamw_8bit cuts optimizer-state VRAM ~4× with no\n",
+    "# measurable quality loss. Skip on Turing where bnb paged optimizer perf is\n",
+    "# weaker — keep adamw_torch there.\n",
+    "_profile['optim'] = 'paged_adamw_8bit' if _cap >= (8, 0) else 'adamw_torch'\n",
+    "\n",
+    "# 4-bit compute dtype tracks precision.\n",
+    "_profile['bnb_4bit_compute_dtype'] = 'bfloat16' if _bf16_ok else 'float16'\n",
+    "_profile['torch_dtype']            = 'bfloat16' if _bf16_ok else 'float16'\n",
+    "\n",
+    "print(f'\\n→ Profile   : {_profile[\"label\"]}')\n",
+    "for k, v in _profile.items():\n",
+    "    if k == 'label': continue\n",
+    "    print(f'    {k:<32}= {v}')\n",
+    "\n",
+    "# ── Write profile into the configs ───────────────────────────────────────\n",
+    "train_cfg.training.per_device_train_batch_size = _profile['per_device_train_batch_size']\n",
+    "train_cfg.training.per_device_eval_batch_size  = _profile['per_device_eval_batch_size']\n",
+    "train_cfg.training.gradient_accumulation_steps = _profile['gradient_accumulation_steps']\n",
+    "train_cfg.training.dataloader_num_workers      = _profile['dataloader_num_workers']\n",
+    "train_cfg.training.fp16                        = _profile['fp16']\n",
+    "train_cfg.training.bf16                        = _profile['bf16']\n",
+    "train_cfg.training.dataloader_pin_memory       = True\n",
+    "train_cfg.training.dataloader_persistent_workers = True\n",
+    "train_cfg.training.optim                       = _profile['optim']\n",
+    "# Ensure stage2 still uses the same per-run epoch count we want.\n",
+    "train_cfg.stage2.num_epochs                    = 5\n",
+    "\n",
+    "model_cfg.llm.attn_implementation       = _profile['attn_implementation']\n",
+    "model_cfg.llm.gradient_checkpointing    = _profile['gradient_checkpointing']\n",
+    "model_cfg.llm.torch_dtype               = _profile['torch_dtype']\n",
+    "model_cfg.llm.bnb_4bit_compute_dtype    = _profile['bnb_4bit_compute_dtype']\n",
+    "model_cfg.llm.bnb_4bit_quant_type       = 'nf4'\n",
+    "model_cfg.llm.bnb_4bit_use_double_quant = True\n",
+    "\n",
+    "# ── task weights (sampling ratio enforced by WeightedRandomSampler) ──\n",
+    "# Defaults in train_config.yaml: 0.30 / 0.20 / 0.50 (RRG ≈ VQA, impression\n",
+    "# lower because in split_cascade mode it sees GT findings as input).\n",
+    "# Resolver auto-renormalizes and drops vqa for IU-Xray. Override here only\n",
+    "# if you want to experiment per-run, e.g.:\n",
+    "#   train_cfg.tasks.findings_generation.weight   = 0.30\n",
+    "#   train_cfg.tasks.impression_generation.weight = 0.20\n",
+    "#   train_cfg.tasks.vqa.weight                   = 0.50\n",
+    "\n",
+    "# ── wandb off ──\n",
+    "train_cfg.wandb.enabled = False\n",
+    "\n",
+    "# ── HuggingFace Hub run tracking ──\n",
+    "train_cfg.hf_hub.enabled        = True\n",
+    "train_cfg.hf_hub.repo_id        = 'hieu3636/cxr-vlm-runs'   # <<< EDIT ME\n",
+    "train_cfg.hf_hub.token_env      = 'HF_TOKEN'\n",
+    "train_cfg.hf_hub.private        = True\n",
+    "train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')\n",
+    "\n",
+    "# ── 4-bit QLoRA ──\n",
+    "model_cfg.llm.load_in_8bit = False\n",
+    "model_cfg.llm.load_in_4bit = True\n",
+    "# Oracle PNU path does NOT use the CheXpert classifier module (labels come\n",
+    "# from the GT csv/manifest baked into the prompt). Keep it disabled until\n",
+    "# you wire the learned classifier for realistic inference.\n",
+    "model_cfg.chexpert_classifier.enabled = False\n",
+    "\n",
+    "OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')\n",
+    "OmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')\n",
+    "\n",
+    "print('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))\n",
+    "print('--- train_cfg.tasks ---');   print(OmegaConf.to_yaml(train_cfg.tasks))\n",
+    "print('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))\n",
+    "print('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))\n",
+    "print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))\n"
+   ],
    "execution_count": null,
    "outputs": [],
    "id": "cell-cfg"
   },
+  {
+   "cell_type": "markdown",
+   "id": "cell-feature-cache-md",
+   "metadata": {},
+   "source": [
+    "## 4b. Pre-compute image features (optional speedup)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "cell-feature-cache",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# ─── Optional: pre-compute image patch features (skip frozen encoder forward) ──\n",
+    "#\n",
+    "# The image encoder is frozen + the transform is deterministic, so encoding the\n",
+    "# same image every step is wasted work. Run this ONCE per dataset to cache\n",
+    "# (P, 768) patch tensors under {WORK}/feature_cache/{DATASET_NAME}/ and the\n",
+    "# training loop will load them instead of re-encoding.\n",
+    "#\n",
+    "# Set CACHE_FEATURES = False to skip (e.g. first time you set up the run, want\n",
+    "# the smoke test to use the raw path, or you're debugging the encoder).\n",
+    "#\n",
+    "# Disk usage:  ~3 MB per image (P=1024 patches × 768 dim × fp16). For ~30k\n",
+    "# unique images that's ~90 GB — make sure WORK has the room, or set\n",
+    "# CACHE_FEATURES=False on tight quotas.\n",
+    "\n",
+    "CACHE_FEATURES = True\n",
+    "\n",
+    "if CACHE_FEATURES:\n",
+    "    feature_cache_dir = WORK / 'feature_cache' / DATASET_NAME\n",
+    "    feature_cache_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    train_cfg.data.feature_cache_dir = str(feature_cache_dir)\n",
+    "    OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')\n",
+    "\n",
+    "    # Re-running this cell is safe: --overwrite is OFF by default so cached\n",
+    "    # files are skipped. To force a full rebuild, add `--overwrite` below.\n",
+    "    print(f'feature_cache_dir = {feature_cache_dir}')\n",
+    "    !python -m scripts.precompute_image_features \\\n",
+    "        --model_config configs/model_config.yaml \\\n",
+    "        --train_config configs/train_config.yaml \\\n",
+    "        --cache_dir \"{feature_cache_dir}\" \\\n",
+    "        --batch_size 16\n",
+    "else:\n",
+    "    train_cfg.data.feature_cache_dir = None\n",
+    "    OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')\n",
+    "    print('Feature cache DISABLED. Training will run the image encoder every step.')\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {

scripts/precompute_image_features.py

@@ -0,0 +1,182 @@
+"""
+precompute_image_features.py
+----------------------------
+One-shot pre-computation of frozen image-encoder patch features.
+
+The image encoder (BioViL-T / RAD-DINO / ViT) is frozen at training time and
+the dataset uses a deterministic transform (Resize + ToTensor + Normalize),
+so the same image always produces the same (P, 768) patch feature tensor.
+Running the encoder every step wastes I/O (re-decoding JPEG) and a small but
+non-trivial slice of GPU compute.
+
+This script walks the unified instruct JSON, encodes each UNIQUE image path
+exactly once, and writes a `.pt` file under `feature_cache_dir` mirroring the
+relative image path. At training time, set `data.feature_cache_dir` in
+train_config.yaml — `data/dataset.py` loads the cached tensor on hit and
+`model/cxr_vlm.py` detects the (P, 768) shape and skips the encoder.
+
+Typical usage:
+
+    python -m scripts.precompute_image_features \
+        --model_config configs/model_config.yaml \
+        --train_config configs/train_config.yaml \
+        --cache_dir cache/image_features \
+        --batch_size 16
+
+After this finishes, edit train_config.yaml:
+    data:
+      feature_cache_dir: "cache/image_features"
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+# project root on path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.rad_dino import BioViLTEncoder
+from utils.dataset_resolver import resolve_dataset_spec
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Pre-compute & cache patch features")
+    p.add_argument("--model_config", default="configs/model_config.yaml")
+    p.add_argument("--train_config", default="configs/train_config.yaml")
+    p.add_argument(
+        "--cache_dir", required=True,
+        help="Output dir. Each image is saved as {cache_dir}/{image_relpath}.pt"
+    )
+    p.add_argument("--batch_size", type=int, default=16)
+    p.add_argument(
+        "--device", default="cuda",
+        help="cuda | cpu (cpu is fine — encoder is small)"
+    )
+    p.add_argument(
+        "--limit", type=int, default=None,
+        help="Cap total images processed (useful for smoke tests)"
+    )
+    p.add_argument(
+        "--overwrite", action="store_true",
+        help="Re-encode and overwrite existing .pt files (default: skip)"
+    )
+    return p.parse_args()
+
+
+def collect_image_paths(instruct_json: str) -> list:
+    """Walk the unified JSON, return sorted list of unique image relpaths
+    (across ALL splits). Multi-image samples contribute each path separately."""
+    with open(instruct_json, "r", encoding="utf-8") as f:
+        samples = json.load(f)
+
+    seen = set()
+    for s in samples:
+        if s.get("image_paths"):
+            for p in s["image_paths"]:
+                seen.add(p)
+        elif s.get("image_path"):
+            seen.add(s["image_path"])
+    return sorted(seen)
+
+
+def main():
+    args = parse_args()
+
+    train_cfg = OmegaConf.load(args.train_config)
+    model_cfg = OmegaConf.load(args.model_config)
+    spec      = resolve_dataset_spec(train_cfg)
+
+    cache_dir = Path(args.cache_dir).resolve()
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    print(f"cache_dir       : {cache_dir}")
+    print(f"image_root      : {spec.image_root}")
+    print(f"instruct_json   : {spec.instruct_json}")
+
+    # ── Build encoder (frozen, in inference dtype) ───────────────────────
+    _DTYPE_MAP = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+    enc_dtype = _DTYPE_MAP.get(model_cfg.llm.torch_dtype, torch.float32)
+    encoder = BioViLTEncoder(
+        frozen   = True,
+        img_size = model_cfg.image_encoder.img_size,
+        backend  = getattr(model_cfg.image_encoder, "backend", "auto"),
+        dtype    = enc_dtype,
+    ).to(args.device).eval()
+
+    transform = BioViLTEncoder.get_transform("val")
+
+    # ── Collect unique image paths ───────────────────────────────────────
+    image_paths = collect_image_paths(spec.instruct_json)
+    print(f"unique images   : {len(image_paths)}")
+    if args.limit:
+        image_paths = image_paths[: args.limit]
+        print(f"(limited to {len(image_paths)})")
+
+    # ── Encode in batches ────────────────────────────────────────────────
+    from PIL import Image
+    import time
+    image_root = Path(spec.image_root)
+    done = 0
+    skipped = 0
+    failed = 0
+    t0 = time.time()
+
+    def _flush_batch(batch_paths, batch_tensors):
+        nonlocal done
+        if not batch_tensors:
+            return
+        x = torch.stack(batch_tensors).to(args.device)
+        with torch.no_grad():
+            with torch.autocast(
+                "cuda" if args.device.startswith("cuda") else "cpu",
+                dtype = enc_dtype if enc_dtype != torch.float32 else torch.float32,
+                enabled = enc_dtype != torch.float32,
+            ):
+                feats = encoder(x)                     # (B, P, 768)
+        feats = feats.to(torch.float16).cpu()           # fp16 .pt → smaller on disk
+        for rel, f in zip(batch_paths, feats):
+            out = cache_dir / (rel + ".pt")
+            out.parent.mkdir(parents=True, exist_ok=True)
+            torch.save(f.contiguous(), out)
+            done += 1
+
+    pending_paths, pending_tensors = [], []
+    for i, rel in enumerate(image_paths):
+        out_path = cache_dir / (rel + ".pt")
+        if out_path.is_file() and not args.overwrite:
+            skipped += 1
+            continue
+        try:
+            img = Image.open(image_root / rel).convert("RGB")
+            t   = transform(img)
+            pending_paths.append(rel)
+            pending_tensors.append(t)
+        except Exception as e:
+            failed += 1
+            print(f"  [skip] {rel}: {type(e).__name__}: {e}")
+            continue
+
+        if len(pending_tensors) >= args.batch_size:
+            _flush_batch(pending_paths, pending_tensors)
+            pending_paths, pending_tensors = [], []
+
+        if (i + 1) % 500 == 0:
+            elapsed = time.time() - t0
+            rate    = (done + skipped) / max(elapsed, 1e-6)
+            print(f"  [{i+1:>6}/{len(image_paths)}]  done={done}  "
+                  f"skipped={skipped}  failed={failed}  ({rate:.1f} img/s)")
+
+    _flush_batch(pending_paths, pending_tensors)
+
+    elapsed = time.time() - t0
+    print(f"\nFinished. encoded={done}  skipped(existing)={skipped}  failed={failed}  "
+          f"elapsed={elapsed/60:.1f} min")
+    print(f"\nNext: set this in configs/train_config.yaml under `data:` ↓")
+    print(f"  feature_cache_dir: \"{cache_dir}\"")
+
+
+if __name__ == "__main__":
+    main()

training/train.py

@@ -335,6 +335,16 @@ def _build_training_args(train_cfg, stage_cfg, out_dir, run_name, *, enable_best
         report_to                   = "wandb" if train_cfg.wandb.enabled else "none",
         run_name                    = run_name,
         dataloader_num_workers      = getattr(tr, "dataloader_num_workers", 4),
+        dataloader_pin_memory       = getattr(tr, "dataloader_pin_memory", True),
+        dataloader_persistent_workers = (
+            getattr(tr, "dataloader_persistent_workers", True)
+            and getattr(tr, "dataloader_num_workers", 4) > 0
+        ),
+        # `paged_adamw_8bit` (bnb) cuts optimizer-state VRAM ~4× with no
+        # measurable quality loss (Dettmers ICLR'22). Default keeps the
+        # legacy fp32 AdamW for backward-compat; the auto-detect cell in
+        # the Colab notebook switches it on for Ampere+ GPUs.
+        optim                       = getattr(tr, "optim", "adamw_torch"),
         remove_unused_columns       = False,
     )
     if save_strategy == "steps":
@@ -434,6 +444,10 @@ class HFBestLastCallback(TrainerCallback):
 
 def _build_datasets(spec: DatasetSpec, train_cfg, model, transform_train, transform_val):
     """Construct train + val CXRInstructDataset instances from a DatasetSpec."""
+    feature_cache_dir = getattr(train_cfg.data, "feature_cache_dir", None) or None
+    if feature_cache_dir:
+        print(f"[_build_datasets] feature_cache_dir = {feature_cache_dir}  "
+              f"(encoder bypass on cache hit)")
     train_ds = CXRInstructDataset(
         data_path    = spec.instruct_json,
         image_root   = spec.image_root,
@@ -444,6 +458,7 @@ def _build_datasets(spec: DatasetSpec, train_cfg, model, transform_train, transf
         cutoff_len   = train_cfg.training.cutoff_len,
         task_weights = spec.task_weights,
         max_images   = spec.max_images,
+        feature_cache_dir = feature_cache_dir,
     )
     val_ds = CXRInstructDataset(
         data_path    = spec.instruct_json,
@@ -455,6 +470,7 @@ def _build_datasets(spec: DatasetSpec, train_cfg, model, transform_train, transf
         cutoff_len   = train_cfg.training.cutoff_len,
         task_weights = spec.task_weights,
         max_images   = spec.max_images,
+        feature_cache_dir = feature_cache_dir,
     )
     return train_ds, val_ds