From cc56f135586d707c192d9bede65838b8cc058010 Mon Sep 17 00:00:00 2001 From: Weihao Kong Date: Tue, 30 Jun 2026 16:51:19 +0000 Subject: [PATCH 1/2] Run the PyTorch model in bfloat16 to match the JAX compute dtype The JAX release runs in bf16 (load(dtype=jnp.bfloat16)) and casts its input to bf16 at the very start of the model's __call__ (jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)). The native PyTorch estimator path ran everything in float32: load() never cast the float32 checkpoint, and the model never cast its input. This doubled activation memory -- a single forward over a ~33k-row context peaked at 35.6 GB and OOM'd large datasets on a 40 GB GPU -- and diverged from the JAX numerics. Mirror the JAX design: - load(): add a `dtype` arg (default torch.bfloat16) and cast the model after loading the float32 weights. Pass dtype=None to keep float32. - pytorch/model.py forward(): at entry, apply nan_to_num(x, nan=-100.0) and cast x to the compute dtype -- the exact op the JAX model does at the top of __call__ -- so the estimator stays dtype-agnostic. - _predict_step_pytorch: feed float32 input (the model casts it, like JAX) and upcast the output before .numpy() (numpy has no bfloat16). Halves activation memory (35.6 -> 17.8 GB on the kddcup09 forward; large datasets that previously OOM'd now fit, ~24 GB) and brings PyTorch inference in line with the JAX/TPU results. --- tabfm/src/classifier_and_regressor.py | 2 +- tabfm/src/pytorch/model.py | 5 +++++ tabfm/src/pytorch/tabfm_v1_0_0.py | 14 ++++++++++++-- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index a785057..f0d2aa1 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/tabfm/src/classifier_and_regressor.py @@ -1819,7 +1819,7 @@ def _predict_step_pytorch( with torch.no_grad(): out_t = model(X_t, y_t, train_size_t, cat_mask=cat_mask_t, d=d_t) - return out_t.cpu().numpy() + return out_t.float().cpu().numpy() # upcast: numpy has no bfloat16 # --------------------------------------------------------------------------- diff --git a/tabfm/src/pytorch/model.py b/tabfm/src/pytorch/model.py index 0accc02..8c47f7f 100644 --- a/tabfm/src/pytorch/model.py +++ b/tabfm/src/pytorch/model.py @@ -442,6 +442,11 @@ def __init__(self, *, embed_dim=8, max_classes=3, col_num_blocks=2, decoder_hidden or icl_dim * 2, is_classifier) def forward(self, x, y, train_size, cat_mask=None, d=None): + # Mirror the JAX model's entry: replace NaN with the -100 sentinel and cast + # to the compute dtype (JAX: `jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)`). + # NaN is already imputed in the shared preprocessing, so nan_to_num is a + # no-op in the normal flow, but it keeps the model robust + JAX-faithful. + x = torch.nan_to_num(x, nan=-100.0).to(self.cls_tokens.dtype) emb = self.cell_embedder(x, y, train_size, cat_mask, d=d) emb = self.col_embedder(emb, train_size) b, t, _, e = emb.shape diff --git a/tabfm/src/pytorch/tabfm_v1_0_0.py b/tabfm/src/pytorch/tabfm_v1_0_0.py index 18c8258..c3d0993 100644 --- a/tabfm/src/pytorch/tabfm_v1_0_0.py +++ b/tabfm/src/pytorch/tabfm_v1_0_0.py @@ -74,10 +74,17 @@ def load( checkpoint_path: Optional[str] = None, *, device: Optional[str] = None, + dtype: Any = torch.bfloat16, use_cache: bool = True, ) -> TabFM: - """Loads the PyTorch TabFM v1.0.0 model with pre-trained weights.""" - cache_key = (model_type, checkpoint_path, device) + """Loads the PyTorch TabFM v1.0.0 model with pre-trained weights. + + The checkpoint is stored in float32, but the model is designed to run in + bfloat16 (matching the JAX release's ``dtype=jnp.bfloat16`` compute default), + with a few internal fp32 upcasts. ``dtype`` casts the model accordingly; pass + ``None`` to keep the float32 weights. + """ + cache_key = (model_type, checkpoint_path, device, str(dtype)) if use_cache: _LOAD_CACHE_LOCK.acquire() @@ -119,6 +126,9 @@ def load( state_dict = torch.load(checkpoint_file, map_location="cpu") model.load_state_dict(state_dict, strict=True) + if dtype is not None: + model = model.to(dtype) # engage the bf16 compute design (see docstring) + if device is not None: model = model.to(device) From 7bbb9040c6af5a1330e0cda2b41dbbf73e69671c Mon Sep 17 00:00:00 2001 From: Weihao Kong Date: Wed, 1 Jul 2026 23:44:38 +0000 Subject: [PATCH 2/2] Note that the dtype option may be removed in a future release --- tabfm/src/pytorch/tabfm_v1_0_0.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tabfm/src/pytorch/tabfm_v1_0_0.py b/tabfm/src/pytorch/tabfm_v1_0_0.py index c3d0993..88298d1 100644 --- a/tabfm/src/pytorch/tabfm_v1_0_0.py +++ b/tabfm/src/pytorch/tabfm_v1_0_0.py @@ -83,6 +83,9 @@ def load( bfloat16 (matching the JAX release's ``dtype=jnp.bfloat16`` compute default), with a few internal fp32 upcasts. ``dtype`` casts the model accordingly; pass ``None`` to keep the float32 weights. + + ``dtype`` is provided for float32 debugging / quality comparison; the model is + designed for bfloat16 and this option may be removed in a future release. """ cache_key = (model_type, checkpoint_path, device, str(dtype))