tabicl_baseline: thread offload_mode/batch_size/disk_offload_dir through - #213
Merged
Conversation
Every prior TabICL comparator run in this project scored the basic 17-feature panel, never the strong 609-feature one the GBM baseline uses, because a naive fit at the standard row/estimator settings costs an estimated ~70GB per predict_proba call and reliably OOM-killed the host. TabICLClassifier already exposes offload_mode/batch_size/disk_offload_dir to manage exactly this (the column-wise embedding tensor is the documented memory bottleneck), but none of it was wired through. This threads those three params from fit_tabicl_baselines through to the classifier, and narrows check_inference_cost's RAM-budget gate to skip only when offload_mode="disk" with a real disk_offload_dir set, since that combination changes the actual constraint from resident memory to disk I/O, which the existing GB estimate does not model. docs/tabicl_strong_feature_comparison.md documents the first strong-panel TabICL run this unblocks, including two failed configurations along the way (a disk-offload I/O stall, and a resident-memory OOM from fitting all 12 models up front) before a one-model-at-a-time loop completed cleanly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #213 +/- ##
=======================================
Coverage 89.03% 89.03%
=======================================
Files 57 57
Lines 8491 8493 +2
=======================================
+ Hits 7560 7562 +2
Misses 931 931
🚀 New features to boost your workflow:
|
…igation
Amrit asked whether TabICL could be run at its actual full capability
(8 estimators, 50,000-row context) on the 609-feature panel instead of the
reduced config (1 estimator, 20,000 rows) the existing report used.
Found and worked around a real gap: this project's offload_mode/
disk_offload_dir params only cover TabICL's column-wise embedding stage
(COL_CONFIG), per the library's own docstring. A synthetic-data test with
offload_mode="disk" plus inference_config={"COL_CONFIG": {"disk_dtype":
torch.float16}} confirmed disk offload genuinely works at full scale: a
real ~7.4GB-per-estimator memmap file lands in the configured directory,
one estimator's file is written/used/deleted before the next starts, and
fit completes in 27 seconds without any memory error.
But it is not usable: a predict_proba call for 200 query rows against
that context took 728.9 seconds, and a call for 8,192 rows (this
project's real per-batch size) did not finish inside a 15-minute timeout,
because the entire context is re-read from disk on every call with no
caching. Extrapolated conservatively, a single real (event, horizon) cell
would take multiple hours and the full 12-cell comparison would take at
least a day of sequential compute -- not a memory problem, a throughput
one, and not achievable in a reasonable session on this hardware.
No code changes: the working offload_mode="disk" plumbing already merged
is sufficient to reproduce this finding; the inference_config/disk_dtype
override that made disk offload actually land in the right place did not
translate into a practically usable configuration, so it is documented
rather than wired into the module as unused, unvalidated plumbing.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…PR branch auto-update)
Migrated the GPU host to a2-ultragpu-1g (170GB RAM, A100-80GB, odyssey-cbm-a100-ultra, us-central1-a) to test whether more RAM (not disk) fixes the throughput wall found in the prior disk-offload investigation. It does: RAM-resident full capability (offload_mode="cpu", n_estimators=8, 50,000-row context) fits and scores without crashing or hanging, peak RSS ~98GB, well inside 165GB free RAM. One full cell run end to end (vasopressor_start@8h, fit 109s for all 3 horizons, predict 1854s for this horizon): TabICL(strong, full) scores 0.915 [0.896, 0.936] vs GBM's 0.934 [0.916, 0.950], CIs now overlapping, vs the reduced config's 0.859 [0.832, 0.887] which was a real, separated loss. Full capability closes most of the gap. Extrapolated from the measured per-row rate, the full 12-cell sweep would take on the order of 5 hours of sequential compute. Not run to completion in this session; stopped deliberately rather than run unattended on an hourly-billed A100-80GB host without asking first. VM stopped. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Reproduces the full-capability run (n_estimators=8, TABICL_MAX_ROWS=50,000, this project's real TabICL defaults, not the reduced config the 12-cell report table uses) so the comparison in docs/tabicl_strong_feature_comparison.md can be rerun or extended without recreating the driver from scratch -- the reduced-config comparator script was left disposable, this one isn't, since full capability requires enough RAM (170GB+) that it's worth being able to repeat deliberately rather than by accident. Reuses the existing GBM(strong) scores from alerts_rows_v3.parquet instead of refitting, fits/scores/drops one TabICL model at a time (the discipline that fixed the earlier cross-cell memory accumulation incident), and writes results incrementally so a partial run isn't a lost run. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
All 12 core (event, horizon) cells completed on the migrated a2-ultragpu-1g host (170GB RAM, one A100-80GB): TabICL(strong) at its real full capability (n_estimators=8, 50,000-row context, no reduction), ~6h sequential compute. Headline change from the reduced-config table: TabICL is statistically indistinguishable from the tuned GBM on 9 of 12 cells (all of death, ICU admission, vasopressor start). Only acute kidney injury keeps a real, CI-separated gap across all three horizons. The reduced config (1 estimator, 20,000 rows) was a genuine handicap, not a formality -- it cost real, measurable performance across the board, not just on cells that already looked weak. Also documents a real caveat found after the sweep completed: the AKI label in the reused alerts dump predates the KDIGO staging completeness fix (3d7ecbb, same day, ~9h after the dump was generated), so every AKI number here (both models) is against a known-incomplete label. Flagged, not yet corrected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
predict_probacall and reliably OOM-killed the host (three confirmed kills, 2026-08-23).TabICLClassifieralready exposesoffload_mode/batch_size/disk_offload_dirto manage exactly this memory bottleneck, but none of it was wired through this project'stabicl_baseline.py. This threads those three params throughfit_tabicl_baselines->_fit_one_tabicl-> the classifier, and narrowscheck_inference_cost's RAM-budget gate to skip only foroffload_mode="disk"with a realdisk_offload_dir, since that combination changes the actual constraint from resident memory to disk I/O (which the existing GB estimate doesn't model, and shouldn't be trusted to block).docs/tabicl_strong_feature_comparison.md: the first TabICL run on the strong feature panel this unblocks, with bootstrap CIs against the tuned GBM on all 12 MIMIC alert cells, plus an honest account of two failed configurations before a one-model-at-a-time loop completed cleanly.Test plan
uv run ruff check/uv run ruff format --checkclean on both changed filesuv run mypy odyssey/inference/tabicl_baseline.pycleanuv run pytest tests/odyssey/inference/test_tabicl_baseline.py— 20/20 passed locallyodyssey-cbm-a100): all 12 (event, horizon) cells fit and scored without OOM usingoffload_mode="cpu",n_estimators=1, reduced row cap; seedocs/tabicl_strong_feature_comparison.mdfor full results and caveats🤖 Generated with Claude Code