tabicl_baseline: thread offload_mode/batch_size/disk_offload_dir through

docs/tabicl_strong_feature_comparison.md

@@ -0,0 +1,117 @@
+# GBM (strong) vs TabICL (strong, full capability) on MIMIC
+
+**Date:** 2026-08-27
+
+## Why this run exists
+
+Every prior TabICL row in this project's registry (`docs/experiments.md`) scored TabICL on the
+**basic** 17-feature panel, while the GBM it was compared against used the **strong** 609-feature
+panel. TabICL never saw the wide panel at all: a naive fit at 50,000 context rows, 609 features,
+and the library's default of 8 estimators costs an estimated ~70 GB per `predict_proba` call, and
+reliably OOM-killed the host (three confirmed kills on 2026-08-23).
+
+This report is the first TabICL run on the strong panel at its actual full capability
+(`n_estimators=8`, 50,000-row context, no reduction). Getting there took three attempts:
+
+1. A reduced config (`n_estimators=1`, 20,000-row context) fit and scored fine on the original
+   82 GB-RAM host, but it isn't TabICL at its documented best, and Amrit asked for the real thing.
+2. Disk offload (`offload_mode="disk"`) was tried next, to fit the full config in less RAM. It
+   works for fitting, but TabICL rereads its entire context from disk on every `predict_proba`
+   call with no caching. Scoring 200 rows took 12 minutes; scoring a normal batch (8,192 rows)
+   did not finish in 15 minutes. Extrapolated, the full comparison would take a day or more. Not
+   a memory problem, a disk I/O throughput problem.
+3. The fix was hardware, not software: the host was migrated to `a2-ultragpu-1g` (170 GB RAM, one
+   A100-80GB, `odyssey-cbm-a100-ultra`, `us-central1-a`). RAM-resident full capability
+   (`offload_mode="cpu"`, no disk) fits comfortably (peak RSS ~102 GB) and scores at normal
+   compute speed, just slower than the reduced config because there is more to compute: roughly
+   1,800 seconds (30 minutes) to score one horizon, versus ~120 seconds for the reduced config.
+   All 12 cells took about 6 hours of sequential GPU compute.
+
+## Setup
+
+- Checkpoint: `subset_run_v8_taskset_v3` (MIMIC, task_set v3, `model_kind=cbm`)
+- GPU host: `odyssey-cbm-a100-ultra` (`a2-ultragpu-1g`, 170 GB RAM, one A100-80GB)
+- GBM (strong) scores: the existing registered alerts run (`alerts_rows_v3.parquet`), fit on 30
+  train shards, scored on 4 held-out shards
+- TabICL (strong, full): `n_estimators=8`, 50,000-row context (this project's real defaults, not
+  reduced), fit on 8 of the 30 train shards, `offload_mode="cpu"`, scored on the same 4 held-out
+  shards. One model fit, scored, and dropped at a time; no more than one model resident at once.
+- Both models are scored on exactly the same rows and labels
+- Script: `scripts/tabicl_strong_compare.py` (committed, reusable; reuses the existing GBM scores
+  instead of refitting)
+
+**Label caveat, found after this sweep completed, not yet corrected:** the acute kidney injury
+label used here, for both the GBM and TabICL scores, comes from `alerts_rows_v3.parquet`,
+generated 2026-08-25 05:07 UTC. The KDIGO AKI staging completeness fix (commit `3d7ecbb`, adding
+the renal-replacement-therapy and urine-output legs that were previously missing, creatinine-only
+staging before that) landed later the same day, at 14:04 UTC. The model checkpoint itself was
+also trained before the fix (03:35 UTC). So every AKI number in this report, GBM and TabICL both,
+uses the old, incomplete label, which under-counts true AKI-3 cases that only qualify through RRT
+or oliguria, not creatinine. This does not explain away the AKI gap below, since the same label
+is used for both models, but it means the true AKI numbers (positive or negative for either model)
+are not yet known. Worth a rerun with a fresh alerts dump once that's a priority.
+
+## The caveat that matters most
+
+Matching the feature set controls how much information each model gets per row. It does not
+control the modeling approach itself. The GBM is gradient-boosted and hyperparameter-tuned per
+task on the full training set (400 rounds, grid search over 4 configs). TabICL is zero-shot
+in-context learning: no gradient descent on this data at all, just one forward pass conditioned
+on a subsampled context. At full capability, that context is closer to the library's own
+validated regime (its authors report strong results up to 50,000 rows), so this asymmetry is
+smaller here than it was in the reduced-config run, but it has not gone away.
+
+## Results
+
+95% subject-clustered bootstrap confidence intervals, 1000 resamples
+(`odyssey.inference.uncertainty.bootstrap_auroc`).
+
+| Event | Horizon | n | GBM (strong) AUROC | TabICL (strong, full) AUROC | Gap | Verdict |
+|---|---|---|---|---|---|---|
+| Acute kidney injury | 8h | 95,471 | 0.894 [0.881, 0.906] | 0.766 [0.746, 0.786] | 0.128 | real gap |
+| Acute kidney injury | 24h | 84,147 | 0.845 [0.827, 0.861] | 0.739 [0.714, 0.761] | 0.106 | real gap |
+| Acute kidney injury | 72h | 57,442 | 0.782 [0.758, 0.805] | 0.685 [0.657, 0.714] | 0.097 | real gap |
+| Death | 8h | 136,850 | 0.953 [0.934, 0.969] | 0.961 [0.941, 0.975] | -0.007 | within noise |
+| Death | 24h | 135,061 | 0.959 [0.947, 0.969] | 0.940 [0.917, 0.958] | 0.020 | within noise |
+| Death | 72h | 130,818 | 0.940 [0.923, 0.954] | 0.925 [0.905, 0.944] | 0.015 | within noise |
+| ICU admission | 8h | 85,838 | 0.968 [0.962, 0.974] | 0.963 [0.956, 0.969] | 0.005 | within noise |
+| ICU admission | 24h | 74,839 | 0.954 [0.945, 0.963] | 0.940 [0.929, 0.951] | 0.014 | within noise |
+| ICU admission | 72h | 48,971 | 0.931 [0.912, 0.946] | 0.914 [0.894, 0.933] | 0.017 | within noise |
+| Vasopressor start | 8h | 111,450 | 0.934 [0.916, 0.950] | 0.916 [0.896, 0.936] | 0.018 | within noise |
+| Vasopressor start | 24h | 98,722 | 0.914 [0.895, 0.933] | 0.893 [0.871, 0.915] | 0.021 | within noise |
+| Vasopressor start | 72h | 67,385 | 0.883 [0.850, 0.913] | 0.868 [0.840, 0.893] | 0.014 | within noise |
+
+"Real gap" means the two confidence intervals do not overlap. "Within noise" means they do.
+
+## What this means
+
+Full capability changes the answer. At the reduced config, TabICL lost on all 12 cells, 7 of them
+a real (CI-separated) loss. At full capability, TabICL is statistically indistinguishable from
+the tuned GBM on 9 of 12 cells: all of death, all of ICU admission, all of vasopressor start.
+Death@8h even slightly favors TabICL, though within noise.
+
+The reduced config was a real handicap, not a formality. Cutting the ensemble from 8 members to 1
+and the context from 50,000 rows to 20,000 cost TabICL real, measurable performance across the
+board, not just on the cells that happened to look weak.
+
+Acute kidney injury is the one place a real gap survives at full capability (0.097 to 0.128,
+separated on all three horizons). This matches an earlier finding in this project
+(`docs/experiments.md`, journal entry 52): the GBM's edge on AKI comes from window aggregates and
+trend statistics it computes explicitly from the raw values, not from having more context to work
+with. TabICL sees the same raw features but has no equivalent way to aggregate them across time
+in a single forward pass, at any context size. But see the label caveat above: this comparison
+used an AKI label known to be incomplete, so the true size of this gap (bigger, smaller, or
+possibly gone) is not yet established.
+
+## Where things live
+
+- Comparison script (committed): `scripts/tabicl_strong_compare.py`
+- Code changes that made this possible: `odyssey/inference/tabicl_baseline.py` (adds
+  `offload_mode`, `batch_size`, and `disk_offload_dir` passthrough to `TabICLClassifier`)
+- Raw per-cell results, including full bootstrap output (mean, std, resample counts): pulled to
+  `/tmp/tabicl_full_sweep.json` this session. Ask if you want it moved somewhere durable.
+- GBM (strong) scores: `~/runs/subset_run_v8_taskset_v3/alerts_rows_v3.parquet`, generated
+  2026-08-25 05:07 UTC (see the label caveat above)
+- GPU host used for this run: `odyssey-cbm-a100-ultra` (`a2-ultragpu-1g`, `us-central1-a`),
+  stopped after this sweep completed. The original host, `odyssey-cbm-a100` (`a2-highgpu-1g`,
+  `us-central1-f`), is stopped and untouched, kept as a fallback.

odyssey/inference/tabicl_baseline.py

@@ -106,9 +106,37 @@ def estimate_peak_gb(n_context_rows: int, n_features: int, n_estimators: int) ->
 
 
 def check_inference_cost(
-    n_context_rows: int, n_features: int, n_estimators: int, *, context: str
+    n_context_rows: int,
+    n_features: int,
+    n_estimators: int,
+    *,
+    context: str,
+    offload_mode: str = "auto",
+    disk_offload_dir: Optional[str] = None,
 ) -> None:
-    """Raise before fitting a TabICL configuration that cannot then be scored."""
+    """Raise before fitting a TabICL configuration that cannot then be scored.
+
+    This GB estimate assumes the column-wise embedding tensor -- the
+    documented ``(n_estimators, n_rows, n_columns, embed_dim)`` memory
+    bottleneck -- stays resident (GPU or CPU RAM). That assumption holds
+    for ``offload_mode`` "gpu"/"cpu"/"auto": "auto"'s own fallback chain
+    ends in "CPU (swap as last resort)" (per tabicl's
+    ``_resolve_offload_mode``), and this host has zero configured swap
+    (measured 2026-08-25, ``free -h``), so a fit that would exceed the
+    budget under those modes is a real, unrecoverable kernel-OOM-kill
+    risk, not a conservative guess -- the gate stays a hard block there.
+
+    ``offload_mode="disk"`` with a working ``disk_offload_dir`` changes
+    the actual constraint from resident memory to disk space and I/O
+    wall-clock, which this GB estimate does not model at all; blocking
+    on it here would be wrong in the other direction (refusing a fit
+    that would actually succeed, just slower). That combination skips
+    this check -- callers are responsible for checking disk headroom
+    themselves (e.g. a ``df`` check before fitting), since there is no
+    equivalent measured constant for disk throughput yet.
+    """
+    if offload_mode == "disk" and disk_offload_dir:
+        return
     peak = estimate_peak_gb(n_context_rows, n_features, n_estimators)
     if peak > _MEMORY_BUDGET_GB:
         raise ValueError(
@@ -120,8 +148,10 @@ def check_inference_cost(
             "OOM-kills, 2026-08-23, the last from a single 2000-row call). Use "
             "the basic feature set (what every completed TabICL run in this "
             "project has used: ~8 GB at the same context size), lower "
-            "TABICL_MAX_ROWS, or raise ODYSSEY_TABICL_MEMORY_BUDGET_GB knowing "
-            "the memory is real."
+            "TABICL_MAX_ROWS, raise ODYSSEY_TABICL_MEMORY_BUDGET_GB knowing "
+            "the memory is real, or pass offload_mode='disk' with a "
+            "disk_offload_dir (this check does not apply to that combination, "
+            "since the constraint becomes disk space/I-O, not RAM)."
         )
 
 
@@ -199,7 +229,7 @@ class TabICLBaselineModel:
 
     feature_set: str = "strong"
     n_features: int = 0
-    params: dict[str, float] = field(default_factory=dict)
+    params: dict[str, Any] = field(default_factory=dict)
     all_nan_cols: np.ndarray | None = None
     """Boolean ``(n_features,)`` mask of columns that were entirely NaN in
     the fit-time context, or ``None`` if none were. tabicl's own
@@ -253,6 +283,9 @@ def _fit_one_tabicl(
     n_estimators: int,
     device: str | None,
     cache: Optional[FitCache] = None,
+    offload_mode: str = "auto",
+    batch_size: Optional[int] = 8,
+    disk_offload_dir: Optional[str] = None,
 ) -> dict[float, TabICLBaselineModel]:
     """Fit one TabICL context per horizon for a single event.
 
@@ -264,6 +297,15 @@ def _fit_one_tabicl(
     stores the (capped, seeded-subsampled) context rather than running
     gradient descent.
 
+    ``offload_mode``/``batch_size``/``disk_offload_dir`` are passed
+    straight through to ``TabICLClassifier`` -- see that class's own
+    docstring for what each controls (in short: where the column-wise
+    embedding tensor, the module's documented memory bottleneck, is
+    materialized, and how many ensemble members are processed per
+    forward pass). Recorded in the fitted model's ``params`` for
+    provenance, the same way ``n_context_rows``/``n_estimators`` already
+    are.
+
     ``cache``, if given, is checked per horizon before fitting and
     written to immediately after -- see :mod:`odyssey.inference.fit_cache`.
     ``_load_tabicl_classifier`` is deferred until the first horizon that
@@ -306,12 +348,17 @@ def _fit_one_tabicl(
             n_estimators=n_estimators,
             device=device,
             random_state=seed,
+            offload_mode=offload_mode,
+            batch_size=batch_size,
+            disk_offload_dir=disk_offload_dir,
         )
         check_inference_cost(
             len(keep),
             int(x_all.shape[1]),
             n_estimators,
             context=f"{event_name}@{h:g}h ({feature_set} features)",
+            offload_mode=offload_mode,
+            disk_offload_dir=disk_offload_dir,
         )
         clf.fit(x_fit, y_fit)
         out[h] = TabICLBaselineModel(
@@ -321,6 +368,9 @@ def _fit_one_tabicl(
             params={
                 "n_context_rows": float(len(keep)),
                 "n_estimators": float(n_estimators),
+                "offload_mode": offload_mode,
+                "batch_size": batch_size,
+                "disk_offload_dir": disk_offload_dir,
             },
             all_nan_cols=all_nan_cols if all_nan_cols.any() else None,
         )
@@ -350,6 +400,9 @@ def fit_tabicl_baselines(
     device: str | None = None,
     cache: Optional[FitCache] = None,
     features: Optional[dict[str, np.ndarray]] = None,
+    offload_mode: str = "auto",
+    batch_size: Optional[int] = 8,
+    disk_offload_dir: Optional[str] = None,
 ) -> dict[tuple[str, float], TabICLBaselineModel]:
     """One TabICLv2 context per (event, horizon), on the same features as the GBM.
 
@@ -372,6 +425,12 @@ def fit_tabicl_baselines(
     ``feature_set`` (see
     :func:`odyssey.inference.baseline_prep.prepare_baseline_data`);
     ``train_events_binned`` is then unused and may be empty.
+
+    ``offload_mode``/``batch_size``/``disk_offload_dir`` pass straight
+    through to every ``TabICLClassifier`` this call fits -- see
+    :func:`_fit_one_tabicl` and ``TabICLClassifier``'s own docstring.
+    Defaults match ``TabICLClassifier``'s own defaults (``"auto"``/``8``/
+    ``None``), so omitting them reproduces prior behavior exactly.
     """
     models: dict[tuple[str, float], TabICLBaselineModel] = {}
     if features is None:
@@ -392,6 +451,9 @@ def fit_tabicl_baselines(
             n_estimators=n_estimators,
             device=device,
             cache=cache,
+            offload_mode=offload_mode,
+            batch_size=batch_size,
+            disk_offload_dir=disk_offload_dir,
         )
         for h, model in per_horizon.items():
             models[(name, h)] = model