Skip to content

Run the PyTorch model in bfloat16 to match the JAX compute dtype #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
weihaokong merged 2 commits into from
Jul 2, 2026
Merged

Run the PyTorch model in bfloat16 to match the JAX compute dtype #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cc56f13
Run the PyTorch model in bfloat16 to match the JAX compute dtype
weihaokong Jun 30, 2026
7bbb904
Note that the dtype option may be removed in a future release
weihaokong Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tabfm/src/classifier_and_regressor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


# ---------------------------------------------------------------------------
Expand Down
5 changes: 5 additions & 0 deletions tabfm/src/pytorch/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
17 changes: 15 additions & 2 deletions tabfm/src/pytorch/tabfm_v1_0_0.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,10 +74,20 @@ def load(
checkpoint_path: Optional[str] = None,
*,
device: Optional[str] = None,
dtype: Any = torch.bfloat16,

@erzel erzel Jun 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the use-case of running the model in float32?

@weihaokong weihaokong Jun 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the model quality is a little bit higher, and I am also using it now for debugging purpose

@erzel erzel Jun 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But during training we compute with bfloat16, right? So there shouldn't be a case for doing inference with float32? I think we should remove it from the JAX code as well.
WDYT?

@erzel erzel Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed offline. Let's keep it for now, but add a comment that "dtype" may not be supported in the future release.

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.

``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))

if use_cache:
_LOAD_CACHE_LOCK.acquire()
Expand Down Expand Up @@ -119,6 +129,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)

Expand Down
Loading