From 053b80516451bef7af66ed05979be8dcf8bdc8f5 Mon Sep 17 00:00:00 2001 From: Juan Ovalle Date: Wed, 8 Jul 2026 13:19:58 +0100 Subject: [PATCH] Expose prefill/decode context caching via fit_mode="fit_with_cache" (JAX) Adds an opt-in fit_mode="fit_with_cache" to TabFMClassifier and TabFMRegressor that runs model.prefill once per ensemble-member chunk at fit time and serves predict/predict_proba through model.decode against the cached context, instead of re-encoding [context + queries] on every call. Decode runs eagerly: jit-ing it copies the multi-GB caches into the executable arena and fuses transposes that exhaust memory at real context lengths. Default fit_mode="fit" is unchanged. Caches are dropped on pickle and rebuilt lazily. Raises NotImplementedError under a multi-device data mesh (cached path is single-device for now) and for non-JAX backends. Parity vs the default path: ~1e-7 (fp32 unit tests, incl. the ensemble preset, categorical permutation, row subsampling); bf16 1.6B checkpoint max prob delta ~0.012 on a 10k-row comparison. --- tabfm/src/classifier_and_regressor.py | 496 +++++++++++++++++++-- tabfm/src/classifier_and_regressor_test.py | 376 ++++++++++++++++ 2 files changed, 846 insertions(+), 26 deletions(-) diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index 51dc867..f850d4a 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/tabfm/src/classifier_and_regressor.py @@ -1534,11 +1534,31 @@ def transform_fold( X_test=None, train_fold=train_fold, val_fold=val_fold ) + def context_tensors( + self, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Any]]: + """Build the fitted in-context tensors (all train rows, no query block). + + Produces the same per-member tensors as ``prepare_ensemble_tensors``, but + for the training context alone: no test/query rows are appended, so unlike + ``transform`` no query row (real or dummy) is required. This is the context + that ``fit_mode="fit_with_cache"`` feeds to ``model.prefill``. + + Returns: + The ``prepare_ensemble_tensors`` 5-tuple + ``(Xs, ys, cat_masks, ds, configs_flat)`` with only context rows in + ``Xs`` (so ``Xs.shape[1] == ys.shape[1]``). + """ + check_is_fitted(self, ["ensemble_configs_"]) + data, _ = self._transform_features(X_test=None, context_only=True) + return self.prepare_ensemble_tensors(data) + def _transform_features( self, X_test: Optional[np.ndarray], train_fold: Optional[np.ndarray] = None, val_fold: Optional[np.ndarray] = None, + context_only: bool = False, ) -> Tuple[collections.OrderedDict, List[np.ndarray]]: """Shared helper to construct transformed feature and target dictionaries. @@ -1554,6 +1574,9 @@ def _transform_features( during cross-validation. If None, all active training rows are used. val_fold: Optional array of indices selecting evaluation validation rows during cross-validation. + context_only: If True, build only the training-context tensors (no query + block); ``X_test`` / ``val_fold`` are ignored. Used by + ``context_tensors``. Returns: A tuple containing: @@ -1609,22 +1632,28 @@ def _transform_features( # Note: self.X_ is the fitted training data. X_test is the test data. # We need to construct the full dataset (Train + Test) X_train_to_use = self.X_[train_idx] - X_test_to_use = self.X_[val_idx] if val_idx is not None else X_test - X_full = np.concatenate([X_train_to_use, X_test_to_use], axis=0) + if context_only: + X_full = X_train_to_use + else: + X_test_to_use = self.X_[val_idx] if val_idx is not None else X_test + X_full = np.concatenate([X_train_to_use, X_test_to_use], axis=0) # Apply value permutation _apply_categorical_permutation(X_full, cat_perm) X_variant_instance = preprocessor.transform(X_full) else: X_train_trans = preprocessor.X_transformed_[train_idx] - X_test_trans = ( - preprocessor.X_transformed_[val_idx] - if val_idx is not None - else preprocessor.transform(X_test) - ) - X_variant_instance = np.concatenate( - [X_train_trans, X_test_trans], axis=0 - ) + if context_only: + X_variant_instance = X_train_trans + else: + X_test_trans = ( + preprocessor.X_transformed_[val_idx] + if val_idx is not None + else preprocessor.transform(X_test) + ) + X_variant_instance = np.concatenate( + [X_train_trans, X_test_trans], axis=0 + ) # Apply feature shuffling shuffled_cols = X_variant_instance[:, shuffle_pattern] @@ -1826,12 +1855,55 @@ def _predict_step_pytorch( return out_t.float().cpu().numpy() # upcast: numpy has no bfloat16 -# Compiled predict step functions memoized on the estimators by -# _batch_forward. They close over nnx.jit state and cannot be pickled. +# Compiled predict step functions memoized on the estimators by _batch_forward. +# They close over nnx.jit state and cannot be pickled. (The cached decode path +# runs model.decode eagerly, so it memoizes nothing here.) _COMPILED_PREDICT_CACHE_ATTRS = ( "_predict_step_compiled_with_cat", "_predict_step_compiled_no_cat", ) + +# Accepted values for the ``fit_mode`` constructor argument (TabPFN parity). +_VALID_FIT_MODES = ("fit", "fit_with_cache") + + +def _active_data_shards() -> int: + """Return the size of the active mesh's ``"data"`` axis (1 if none).""" + if not HAS_JAX: + return 1 + mesh = jax.sharding.get_mesh() + if mesh and "data" in mesh.axis_names: + return mesh.axis_sizes[mesh.axis_names.index("data")] + return 1 + + +def _check_fit_with_cache_supported(model: Any) -> None: + """Validate that ``fit_mode="fit_with_cache"`` can run for ``model``. + + Args: + model: The wrapped foundation model. + + Raises: + NotImplementedError: If the model does not expose the JAX prefill/decode + API, or if a data-parallel mesh (a ``"data"`` axis of size > 1) is active. + Unlike the default ``_batch_forward`` path, the cached prefill/decode path + is not sharded and runs on a single device. + """ + if not hasattr(model, "prefill"): + raise NotImplementedError( + 'fit_mode="fit_with_cache" requires a JAX TabFM model exposing' + " prefill/decode; the provided model does not." + ) + num_data_shards = _active_data_shards() + if num_data_shards > 1: + raise NotImplementedError( + 'fit_mode="fit_with_cache" does not support data-parallel sharding' + f' (found a "data" mesh axis of size {num_data_shards}); the cached' + " prefill/decode path runs on a single device. Use the default" + ' fit_mode="fit" under a data-parallel mesh.' + ) + + def _check_classifier_output_dim(output_dim: int, n_classes: int) -> None: """Validates that the model produces logits for all target classes. @@ -1939,6 +2011,7 @@ def __init__( nnls_beta: float = 0.75, calibration_lambda: float = 1e-2, min_rows_for_single_val_split: int = 2000, + fit_mode: str = "fit", ): """Initialises the classifier. @@ -1982,6 +2055,16 @@ def __init__( min_rows_for_single_val_split: Minimum validation rows required to allow learning ensemble/calibration weights on a single train/val split instead of full CV. 0 means always doing full CV. + fit_mode: Either ``"fit"`` (default) or ``"fit_with_cache"``. With + ``"fit_with_cache"`` the training-context KV cache is precomputed at fit + time via ``model.prefill`` so that ``predict`` reuses it through + ``model.decode`` instead of re-encoding the context on every call. This + trades extra accelerator memory (proportional to context length times + ``n_estimators``) for cheaper predictions and requires a JAX model. + Unlike the default path, the cached path is not data-parallel sharded + and runs on a single device: fitting with it raises + ``NotImplementedError`` if a mesh with a ``"data"`` axis of size > 1 is + active. """ self.model = model self.n_estimators = n_estimators @@ -2009,6 +2092,7 @@ def __init__( self.nnls_beta = nnls_beta self.calibration_lambda = calibration_lambda self.min_rows_for_single_val_split = min_rows_for_single_val_split + self.fit_mode = fit_mode if self.average_logits and self.enable_nnls: raise ValueError("average_logits and enable_nnls cannot both be True.") if self.max_num_rows is not None and self.enable_nnls: @@ -2074,8 +2158,16 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": Raises: ValueError - If the number of classes exceeds the model's maximum supported classes. + If the number of classes exceeds the model's maximum supported classes, + or if ``fit_mode`` is not one of ``"fit"`` / ``"fit_with_cache"``. """ + if self.fit_mode not in _VALID_FIT_MODES: + raise ValueError( + f"fit_mode must be one of {_VALID_FIT_MODES}, got {self.fit_mode!r}." + ) + # Fail fast, before any fitting work, if the cached path is unsupported. + if self.fit_mode == "fit_with_cache": + _check_fit_with_cache_supported(self.model) check_classification_targets(y) # Encode class labels @@ -2191,6 +2283,9 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier": ) self._fit_calibration(P, y_fit) + if self.fit_mode == "fit_with_cache": + self._build_context_cache() + return self @jt.typed @@ -2297,11 +2392,10 @@ def _batch_forward( raise ImportError("JAX is required to run a JAX model.") # --- JAX execution path --- mesh = jax.sharding.get_mesh() + num_data_shards = _active_data_shards() if mesh and "data" in mesh.axis_names: - num_data_shards = mesh.axis_sizes[mesh.axis_names.index("data")] data_sharding = NamedSharding(mesh, PartitionSpec("data")) else: - num_data_shards = 1 data_sharding = None num_classes = self.n_classes_ @@ -2465,15 +2559,18 @@ def _predict_step_fn(model, X, y, train_size, d): return np.concatenate(outputs, axis=0) def __getstate__(self): - """Drops memoized compiled predict functions from the pickled state. + """Drops unpicklable prediction caches from the pickled state. The first predict memoizes nnx.jit-compiled step functions on the - estimator (see _batch_forward). Those closures cannot be pickled; they - are pure caches and are rebuilt lazily on the next predict. + estimator (see _batch_forward / _decode_batch_forward), and + fit_mode="fit_with_cache" stores the prefill KV cache in _context_cache_. + Neither can be pickled; both are pure caches, rebuilt lazily on the next + predict (the context cache via _ensure_context_cache). """ state = dict(super().__getstate__()) for attr in _COMPILED_PREDICT_CACHE_ATTRS: state.pop(attr, None) + state.pop("_context_cache_", None) return state @jt.typed @@ -2658,6 +2755,154 @@ def _apply_calibration( return P + def _context_only_tensors( + self, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return the fitted context-only ensemble tensors (no query rows). + + Delegates to ``EnsembleGenerator.context_tensors``, which builds the + per-member in-context tensors from the generator's fitted state without a + query block (so no dummy query row is transformed). + + Returns: + Tuple ``(Xs_ctx, ys_ctx, cat_masks, ds)`` where ``Xs_ctx`` has shape + ``(n_members, context_len, n_features)``, ``ys_ctx`` has shape + ``(n_members, context_len)``, ``cat_masks`` has shape + ``(n_members, n_features)`` and ``ds`` has shape ``(n_members,)``. + """ + Xs, ys, cat_masks, ds, _ = self.ensemble_generator_.context_tensors() + return Xs, ys, cat_masks, ds + + def _build_context_cache(self) -> None: + """Precompute the per-chunk training-context KV caches via ``model.prefill``. + + Runs ``model.prefill`` once per ensemble-member chunk (chunked by + ``self.batch_size``, mirroring ``_batch_forward``) and stores the resulting + caches for reuse by ``_decode_batch_forward``. Members are chunked here so + ``_decode_batch_forward`` can pair each query chunk with the cache built + from the same members' context; column caches are flattened over the + member*feature axis internally, so per-chunk storage avoids slicing them by + member. Called from ``fit`` when ``fit_mode == "fit_with_cache"``. + """ + # Also guards the lazy rebuild path (_ensure_context_cache after unpickle), + # where the active mesh may differ from fit time. + _check_fit_with_cache_supported(self.model) + Xs_ctx, ys_ctx, cat_masks, ds = self._context_only_tensors() + self._cache_ctx_len_ = int(ys_ctx.shape[1]) + + has_cat = cat_masks is not None and hasattr(self.model, "cell_embedder") + batch_size_per_process = self.batch_size or Xs_ctx.shape[0] + n_batches = math.ceil(Xs_ctx.shape[0] / batch_size_per_process) + Xs_split = np.array_split(Xs_ctx, n_batches) + ys_split = np.array_split(ys_ctx, n_batches) + cat_masks_split = ( + np.array_split(cat_masks, n_batches) if has_cat else [None] * n_batches + ) + ds_split = ( + np.array_split(ds, n_batches) if ds is not None else [None] * n_batches + ) + + caches = [] + for X_batch, y_batch, cat_mask_batch, ds_batch_val in zip( + Xs_split, ys_split, cat_masks_split, ds_split + ): + X_batch = jnp.asarray(X_batch, dtype=jnp.float32) + y_batch = jnp.asarray(y_batch, dtype=jnp.float32) + if ds_batch_val is not None: + d_batch = jnp.asarray(ds_batch_val, dtype=jnp.int32) + else: + d_batch = jnp.full((X_batch.shape[0],), X_batch.shape[-1], jnp.int32) + # prefill is run once at fit time, so it is left uncompiled. + if cat_mask_batch is not None: + cat_mask_batch = jnp.asarray(cat_mask_batch, dtype=jnp.bool_) + _, cache = self.model.prefill( + X_batch, y_batch, d=d_batch, cat_mask=cat_mask_batch + ) + else: + _, cache = self.model.prefill(X_batch, y_batch, d=d_batch) + caches.append(cache) + + self._context_cache_ = caches + + def _ensure_context_cache(self) -> None: + """Rebuild the context cache if missing (e.g. after unpickling).""" + if not hasattr(self, "_context_cache_"): + self._build_context_cache() + + @jt.typed + def _decode_batch_forward( + self, + Xs: jt.Float[Array | np.ndarray, "B T_test H"], + cat_masks: Optional[jt.Bool[Array | np.ndarray, "B H"]] = None, + ds: Optional[jt.Int[Array | np.ndarray, "B"]] = None, + ) -> jt.Float[Array | np.ndarray, "B T_test K"]: + """Decode query rows against the precomputed context caches. + + Reuses the per-chunk caches in ``self._context_cache_`` (built by + ``_build_context_cache``), chunking members by ``self.batch_size`` exactly + as ``_batch_forward`` does. ``model.decode`` is called eagerly (no + ``nnx.jit``): passing the multi-gigabyte context caches as jit arguments + makes XLA copy them into the executable's arena (roughly doubling peak + memory) and fuse a large transpose that OOMs or fails autotuning at real + context lengths. Eager decode also matches the supported usage exercised by + the model-level tests (``model_test.py``). ``model.decode`` pads the query + sequence to a multiple of 128 and unpads it internally, so only the query + rows are passed here. + + Args: + Xs: Query features of shape (n_members, n_test, n_features). + cat_masks: Optional categorical mask of shape (n_members, n_features). + ds: Optional active-feature counts of shape (n_members,). + + Returns: + Model outputs of shape (n_members, n_test, n_classes_bins). + """ + if not HAS_JAX: + raise ImportError("JAX is required to run a JAX model.") + + caches = self._context_cache_ + has_cat = cat_masks is not None and hasattr(self.model, "cell_embedder") + batch_size_per_process = self.batch_size or Xs.shape[0] + n_batches = math.ceil(Xs.shape[0] / batch_size_per_process) + if len(caches) != n_batches: + raise RuntimeError( + f"Expected {n_batches} cached context chunks but got {len(caches)};" + " the context cache is stale. Refit the estimator so the cache is" + " rebuilt for the current batch_size (do not mutate batch_size after" + " fit_mode='fit_with_cache')." + ) + Xs_split = np.array_split(Xs, n_batches) + cat_masks_split = ( + np.array_split(cat_masks, n_batches) if has_cat else [None] * n_batches + ) + ds_split = ( + np.array_split(ds, n_batches) if ds is not None else [None] * n_batches + ) + + outputs = [] + for X_batch, cache_batch, cat_mask_batch, ds_batch_val in zip( + Xs_split, caches, cat_masks_split, ds_split + ): + X_batch = jnp.asarray(X_batch, dtype=jnp.float32) + if ds_batch_val is not None: + d_batch = jnp.asarray(ds_batch_val, dtype=jnp.int32) + else: + d_batch = jnp.full((X_batch.shape[0],), X_batch.shape[-1], jnp.int32) + + # Decode eagerly, passing the prefill cache dict straight through; + # model.decode unpads the padded query sequence internally. + if cat_mask_batch is not None: + cat_mask_batch = jnp.asarray(cat_mask_batch, dtype=jnp.bool_) + out = self.model.decode( + X_batch, cache_batch, d=d_batch, cat_mask=cat_mask_batch + ) + else: + out = self.model.decode(X_batch, cache_batch, d=d_batch) + + outputs.append(np.asarray(out)) + + return np.concatenate(outputs, axis=0) + @jt.typed def _predict_proba_internal(self, X: Any) -> jt.Float[Array | np.ndarray, "E T K"]: """Predict class probabilities for test samples.""" @@ -2677,7 +2922,19 @@ def _predict_proba_internal(self, X: Any) -> jt.Float[Array | np.ndarray, "E T K _, ) = self.ensemble_generator_.prepare_ensemble_tensors(data) - outputs = self._batch_forward(Xs_all, ys_all, cat_masks_all, ds=ds_all) + if getattr(self, "fit_mode", "fit") == "fit_with_cache": + self._ensure_context_cache() + # Context rows lead each member view; ys length is the context length. + context_len = self._cache_ctx_len_ + if ys_all.shape[1] != context_len: + raise RuntimeError( + "Context length changed since fit_mode='fit_with_cache' built the" + f" cache ({ys_all.shape[1]} vs {context_len}); refit the estimator." + ) + queries = Xs_all[:, context_len:, :] + outputs = self._decode_batch_forward(queries, cat_masks_all, ds=ds_all) + else: + outputs = self._batch_forward(Xs_all, ys_all, cat_masks_all, ds=ds_all) _check_classifier_output_dim(outputs.shape[-1], self.n_classes_) outputs = outputs[..., :self.n_classes_] @@ -2834,6 +3091,7 @@ def __init__( enable_nnls: bool = False, nnls_beta: float = 0.75, min_rows_for_single_val_split: int = 2000, + fit_mode: str = "fit", ): """Initialises the regressor. @@ -2866,6 +3124,16 @@ def __init__( min_rows_for_single_val_split: Minimum validation rows required to allow learning ensemble weights on a single train/val split instead of full CV. 0 means always doing full CV. + fit_mode: Either ``"fit"`` (default) or ``"fit_with_cache"``. With + ``"fit_with_cache"`` the training-context KV cache is precomputed at fit + time via ``model.prefill`` so that ``predict`` reuses it through + ``model.decode`` instead of re-encoding the context on every call. This + trades extra accelerator memory (proportional to context length times + ``n_estimators``) for cheaper predictions and requires a JAX model. + Unlike the default path, the cached path is not data-parallel sharded + and runs on a single device: fitting with it raises + ``NotImplementedError`` if a mesh with a ``"data"`` axis of size > 1 is + active. """ self.model = model self.n_estimators = n_estimators @@ -2887,6 +3155,7 @@ def __init__( self.enable_nnls = enable_nnls self.nnls_beta = nnls_beta self.min_rows_for_single_val_split = min_rows_for_single_val_split + self.fit_mode = fit_mode if self.max_num_rows is not None and self.enable_nnls: raise ValueError( "max_num_rows and enable_nnls cannot both be set at this time." @@ -2940,7 +3209,17 @@ def fit(self, X: Any, y: Any) -> "TabFMRegressor": Returns: self : TabFMRegressor Fitted regressor instance. + + Raises: + ValueError: If ``fit_mode`` is not one of ``"fit"`` / ``"fit_with_cache"``. """ + if self.fit_mode not in _VALID_FIT_MODES: + raise ValueError( + f"fit_mode must be one of {_VALID_FIT_MODES}, got {self.fit_mode!r}." + ) + # Fail fast, before any fitting work, if the cached path is unsupported. + if self.fit_mode == "fit_with_cache": + _check_fit_with_cache_supported(self.model) y_orig = np.array(y).copy() y = check_array(y, ensure_2d=False, dtype="numeric") self.X_encoder_ = TransformToNumerical( @@ -3003,6 +3282,9 @@ def fit(self, X: Any, y: Any) -> "TabFMRegressor": self.nnls_beta * weights + (1.0 - self.nnls_beta) * avg_weights ) + if self.fit_mode == "fit_with_cache": + self._build_context_cache() + return self @jt.typed @@ -3088,11 +3370,10 @@ def _batch_forward( raise ImportError("JAX is required to run a JAX model.") # --- JAX execution path --- mesh = jax.sharding.get_mesh() + num_data_shards = _active_data_shards() if mesh and "data" in mesh.axis_names: - num_data_shards = mesh.axis_sizes[mesh.axis_names.index("data")] data_sharding = NamedSharding(mesh, PartitionSpec("data")) else: - num_data_shards = 1 data_sharding = None _has_compiled_attr = ( @@ -3250,15 +3531,18 @@ def _inverse_transform_y(self, y_scaled: np.ndarray) -> np.ndarray: return self.y_scaler_.inverse_transform(y_scaled.reshape(-1, 1)).flatten() def __getstate__(self): - """Drops memoized compiled predict functions from the pickled state. + """Drops unpicklable prediction caches from the pickled state. The first predict memoizes nnx.jit-compiled step functions on the - estimator (see _batch_forward). Those closures cannot be pickled; they - are pure caches and are rebuilt lazily on the next predict. + estimator (see _batch_forward / _decode_batch_forward), and + fit_mode="fit_with_cache" stores the prefill KV cache in _context_cache_. + Neither can be pickled; both are pure caches, rebuilt lazily on the next + predict (the context cache via _ensure_context_cache). """ state = dict(super().__getstate__()) for attr in _COMPILED_PREDICT_CACHE_ATTRS: state.pop(attr, None) + state.pop("_context_cache_", None) return state @jt.typed @@ -3339,12 +3623,160 @@ def _compute_oof_preds_scaled( return outputs_oof + def _context_only_tensors( + self, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return the fitted context-only ensemble tensors (no query rows). + + Delegates to ``EnsembleGenerator.context_tensors``, which builds the + per-member in-context tensors from the generator's fitted state without a + query block (so no dummy query row is transformed). + + Returns: + Tuple ``(Xs_ctx, ys_ctx, cat_masks, ds)`` shaped as in the classifier's + counterpart. + """ + Xs, ys, cat_masks, ds, _ = self.ensemble_generator_.context_tensors() + return Xs, ys, cat_masks, ds + + def _build_context_cache(self) -> None: + """Precompute the per-chunk training-context KV caches via ``model.prefill``. + + Runs ``model.prefill`` once per ensemble-member chunk (chunked by + ``self.batch_size``, mirroring ``_batch_forward``) and stores the resulting + caches for reuse by ``_decode_batch_forward``. Members are chunked here so + ``_decode_batch_forward`` can pair each query chunk with the cache built + from the same members' context; column caches are flattened over the + member*feature axis internally, so per-chunk storage avoids slicing them by + member. Called from ``fit`` when ``fit_mode == "fit_with_cache"``. + """ + # Also guards the lazy rebuild path (_ensure_context_cache after unpickle), + # where the active mesh may differ from fit time. + _check_fit_with_cache_supported(self.model) + Xs_ctx, ys_ctx, cat_masks, ds = self._context_only_tensors() + self._cache_ctx_len_ = int(ys_ctx.shape[1]) + + has_cat = cat_masks is not None and hasattr(self.model, "cell_embedder") + batch_size_per_process = self.batch_size or Xs_ctx.shape[0] + n_batches = math.ceil(Xs_ctx.shape[0] / batch_size_per_process) + Xs_split = np.array_split(Xs_ctx, n_batches) + ys_split = np.array_split(ys_ctx, n_batches) + cat_masks_split = ( + np.array_split(cat_masks, n_batches) if has_cat else [None] * n_batches + ) + ds_split = ( + np.array_split(ds, n_batches) if ds is not None else [None] * n_batches + ) + + caches = [] + for X_batch, y_batch, cat_mask_batch, ds_batch_val in zip( + Xs_split, ys_split, cat_masks_split, ds_split + ): + X_batch = jnp.asarray(X_batch, dtype=jnp.float32) + y_batch = jnp.asarray(y_batch, dtype=jnp.float32) + if ds_batch_val is not None: + d_batch = jnp.asarray(ds_batch_val, dtype=jnp.int32) + else: + d_batch = jnp.full((X_batch.shape[0],), X_batch.shape[-1], jnp.int32) + # prefill is run once at fit time, so it is left uncompiled. + if cat_mask_batch is not None: + cat_mask_batch = jnp.asarray(cat_mask_batch, dtype=jnp.bool_) + _, cache = self.model.prefill( + X_batch, y_batch, d=d_batch, cat_mask=cat_mask_batch + ) + else: + _, cache = self.model.prefill(X_batch, y_batch, d=d_batch) + caches.append(cache) + + self._context_cache_ = caches + + def _ensure_context_cache(self) -> None: + """Rebuild the context cache if missing (e.g. after unpickling).""" + if not hasattr(self, "_context_cache_"): + self._build_context_cache() + + @jt.typed + def _decode_batch_forward( + self, + Xs: jt.Float[Array | np.ndarray, "B T_test H"], + cat_masks: Optional[jt.Bool[Array | np.ndarray, "B H"]] = None, + ds: Optional[jt.Int[Array | np.ndarray, "B"]] = None, + ) -> jt.Float[Array | np.ndarray, "B T_test L_out"]: + """Decode query rows against the precomputed context caches. + + Reuses the per-chunk caches in ``self._context_cache_`` (built by + ``_build_context_cache``), chunking members by ``self.batch_size`` exactly + as ``_batch_forward`` does. ``model.decode`` is called eagerly (no + ``nnx.jit``): passing the multi-gigabyte context caches as jit arguments + makes XLA copy them into the executable's arena (roughly doubling peak + memory) and fuse a large transpose that OOMs or fails autotuning at real + context lengths. Eager decode also matches the supported usage exercised by + the model-level tests (``model_test.py``). ``model.decode`` pads the query + sequence to a multiple of 128 and unpads it internally, so only the query + rows are passed here. + + Args: + Xs: Query features of shape (n_members, n_test, n_features). + cat_masks: Optional categorical mask of shape (n_members, n_features). + ds: Optional active-feature counts of shape (n_members,). + + Returns: + Model outputs of shape (n_members, n_test, output_dim). + """ + if not HAS_JAX: + raise ImportError("JAX is required to run a JAX model.") + + caches = self._context_cache_ + has_cat = cat_masks is not None and hasattr(self.model, "cell_embedder") + batch_size_per_process = self.batch_size or Xs.shape[0] + n_batches = math.ceil(Xs.shape[0] / batch_size_per_process) + if len(caches) != n_batches: + raise RuntimeError( + f"Expected {n_batches} cached context chunks but got {len(caches)};" + " the context cache is stale. Refit the estimator so the cache is" + " rebuilt for the current batch_size (do not mutate batch_size after" + " fit_mode='fit_with_cache')." + ) + Xs_split = np.array_split(Xs, n_batches) + cat_masks_split = ( + np.array_split(cat_masks, n_batches) if has_cat else [None] * n_batches + ) + ds_split = ( + np.array_split(ds, n_batches) if ds is not None else [None] * n_batches + ) + + outputs = [] + for X_batch, cache_batch, cat_mask_batch, ds_batch_val in zip( + Xs_split, caches, cat_masks_split, ds_split + ): + X_batch = jnp.asarray(X_batch, dtype=jnp.float32) + if ds_batch_val is not None: + d_batch = jnp.asarray(ds_batch_val, dtype=jnp.int32) + else: + d_batch = jnp.full((X_batch.shape[0],), X_batch.shape[-1], jnp.int32) + + # Decode eagerly, passing the prefill cache dict straight through; + # model.decode unpads the padded query sequence internally. + if cat_mask_batch is not None: + cat_mask_batch = jnp.asarray(cat_mask_batch, dtype=jnp.bool_) + out = self.model.decode( + X_batch, cache_batch, d=d_batch, cat_mask=cat_mask_batch + ) + else: + out = self.model.decode(X_batch, cache_batch, d=d_batch) + + outputs.append(np.asarray(out)) + + return np.concatenate(outputs, axis=0) + @jt.typed def _predict_internal(self, X: Any) -> jt.Float[Array | np.ndarray, "E T"]: """Predict regression target for test samples.""" check_is_fitted(self) if isinstance(X, np.ndarray) and len(X.shape) == 1: - raise ValueError("The provided input X is one-dimensional. Reshape your data.") + raise ValueError( + "The provided input X is one-dimensional. Reshape your data." + ) X_transformed = self.X_encoder_.transform(X) data = self.ensemble_generator_.transform(X_transformed) @@ -3356,7 +3788,19 @@ def _predict_internal(self, X: Any) -> jt.Float[Array | np.ndarray, "E T"]: _, ) = self.ensemble_generator_.prepare_ensemble_tensors(data) - output = self._batch_forward(Xs_all, ys_all, cat_masks_all, ds=ds_all) + if getattr(self, "fit_mode", "fit") == "fit_with_cache": + self._ensure_context_cache() + # Context rows lead each member view; ys length is the context length. + context_len = self._cache_ctx_len_ + if ys_all.shape[1] != context_len: + raise RuntimeError( + "Context length changed since fit_mode='fit_with_cache' built the" + f" cache ({ys_all.shape[1]} vs {context_len}); refit the estimator." + ) + queries = Xs_all[:, context_len:, :] + output = self._decode_batch_forward(queries, cat_masks_all, ds=ds_all) + else: + output = self._batch_forward(Xs_all, ys_all, cat_masks_all, ds=ds_all) _check_regressor_output_dim(output.shape[-1]) predictions = output.squeeze(-1) return predictions diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index 546fbee..8c47630 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/tabfm/src/classifier_and_regressor_test.py @@ -21,6 +21,7 @@ from sklearn.exceptions import NotFittedError try: + import jax.numpy as jnp from flax import nnx from tabfm.src.jax import model as tabfm_model HAS_JAX = True @@ -90,6 +91,45 @@ def test_permute_categorical_structure(self): identity_count, n_estimators, "Categorical permutations should vary." ) + def test_context_tensors_matches_dummy_query_derivation(self): + # context_tensors() must reproduce exactly what the old dummy-query-row + # derivation produced: transform a single dummy query row through the + # public path, then slice it off (ys length is the context length). This + # pins the equivalence the removed dummy-row hack relied on, under both + # categorical permutation and row subsampling. + rng = np.random.RandomState(0) + X = pd.DataFrame({ + "cat": rng.choice(["a", "b", "c"], size=30), + "num1": rng.rand(30), + "num2": rng.rand(30), + }) + y = rng.randint(0, 3, size=30) + X_enc = TransformToNumerical(min_cat_frequency=1).fit_transform(X) + generator = EnsembleGenerator( + n_estimators=6, + norm_methods=["none", "power"], + cat_features=[0], + permute_categorical=True, + max_num_rows=20, + random_state=0, + ) + generator.fit(X_enc, y) + + Xs_ctx, ys_ctx, cat_masks, ds, _ = generator.context_tensors() + + # Old derivation: transform one dummy query row and slice it off. + n_feat = generator.unique_filter_.n_features_in_ + data = generator.transform(np.zeros((1, n_feat))) + Xs_ref, ys_ref, cat_masks_ref, ds_ref, _ = ( + generator.prepare_ensemble_tensors(data) + ) + context_len = ys_ref.shape[1] + + np.testing.assert_array_equal(ys_ctx, ys_ref) + np.testing.assert_array_equal(Xs_ctx, Xs_ref[:, :context_len, :]) + np.testing.assert_array_equal(cat_masks, cat_masks_ref) + np.testing.assert_array_equal(ds, ds_ref) + def test_permute_categorical_application(self): # Test that transform actually changes the data # X has 2 samples, 1 cat feature with values 0 and 1. @@ -1200,5 +1240,341 @@ def test_datetime_column_with_non_string_name(self): self.assertEqual(out.shape, (4, 5)) # unix-ns + 4 derived features +@unittest.skipUnless(HAS_JAX, "JAX is required") +class CachedPredictTest(absltest.TestCase): + """``fit_mode="fit_with_cache"`` must match the default full-forward path. + + With ``fit_mode="fit_with_cache"`` the estimator precomputes the in-context + KV cache at fit time (``model.prefill``) and reuses it at predict time + (``model.decode``) instead of re-encoding the training context on every call. + TabFM conditions queries on the context only (the column embedder's inducing + points are computed from train rows, and ICL attention masks queries to the + train block), so prefill+decode is numerically equivalent to a full forward + pass. These tests pin that equivalence, plus pickle survival and validation. + """ + + def _tiny_model(self, loss): + # float32 (not the bfloat16 default) so prefill+decode agrees with the full + # forward pass to floating-point tolerance rather than bf16 coarseness. + return tabfm_model.TabFM( + loss=loss, + max_classes=3, + embed_dim=8, + col_num_blocks=1, + col_nhead=2, + col_num_inds=8, + row_num_blocks=1, + row_nhead=2, + row_num_cls=1, + icl_num_blocks=1, + icl_nhead=2, + rngs=nnx.Rngs(0), + dtype=jnp.float32, + ) + + def _classification_data(self): + rng = np.random.RandomState(0) + return ( + rng.rand(40, 5).astype(np.float32), + rng.randint(0, 3, size=40), + rng.rand(11, 5).astype(np.float32), + ) + + def _regression_data(self): + rng = np.random.RandomState(1) + return ( + rng.rand(40, 5).astype(np.float32), + rng.rand(40).astype(np.float32), + rng.rand(11, 5).astype(np.float32), + ) + + def _mixed_data(self): + # A categorical (string) column plus numeric columns, to exercise the + # categorical-permutation concat path and non-trivial categorical masks. + rng = np.random.RandomState(4) + X = pd.DataFrame({ + "cat": rng.choice(["a", "b", "c"], size=48), + "num1": rng.rand(48).astype(np.float32), + "num2": rng.rand(48).astype(np.float32), + }) + y = rng.randint(0, 3, size=48) + X_test = pd.DataFrame({ + "cat": rng.choice(["a", "b", "c"], size=10), + "num1": rng.rand(10).astype(np.float32), + "num2": rng.rand(10).astype(np.float32), + }) + return X, y, X_test + + def test_predict_proba_matches_full_forward(self): + X, y, X_test = self._classification_data() + # Share one model instance so both estimators see identical weights; + # batch_size < n_estimators exercises the per-chunk cache path. + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, n_estimators=4, batch_size=2, random_state=0 + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_predict_labels_match_full_forward(self): + X, y, X_test = self._classification_data() + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, n_estimators=4, batch_size=2, random_state=0 + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_array_equal(cached.predict(X_test), base.predict(X_test)) + + def test_cached_classifier_pickle_round_trip(self): + X, y, X_test = self._classification_data() + cached = TabFMClassifier( + model=self._tiny_model("cross_entropy"), + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + cached.fit(X, y) + proba = cached.predict_proba(X_test) + + # The context cache holds jax arrays / nnx pytrees and is dropped on pickle; + # it must rebuild lazily on the next predict and reproduce the outputs. + restored = pickle.loads(pickle.dumps(cached)) + + np.testing.assert_allclose(restored.predict_proba(X_test), proba, atol=1e-6) + + def test_regressor_predict_matches_full_forward(self): + X, y, X_test = self._regression_data() + model = self._tiny_model("rmse") + base = TabFMRegressor( + model=model, n_estimators=4, batch_size=2, random_state=0 + ) + cached = TabFMRegressor( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_allclose( + cached.predict(X_test), base.predict(X_test), atol=1e-5 + ) + + def test_cached_regressor_pickle_round_trip(self): + X, y, X_test = self._regression_data() + cached = TabFMRegressor( + model=self._tiny_model("rmse"), + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + cached.fit(X, y) + preds = cached.predict(X_test) + + restored = pickle.loads(pickle.dumps(cached)) + + np.testing.assert_allclose(restored.predict(X_test), preds, atol=1e-6) + + def test_invalid_fit_mode_raises(self): + X, y, _ = self._classification_data() + clf = TabFMClassifier( + model=self._tiny_model("cross_entropy"), + n_estimators=2, + fit_mode="invalid", + ) + with self.assertRaisesRegex(ValueError, "fit_mode"): + clf.fit(X, y) + + reg = TabFMRegressor( + model=self._tiny_model("rmse"), n_estimators=2, fit_mode="invalid" + ) + with self.assertRaisesRegex(ValueError, "fit_mode"): + reg.fit(X, np.random.RandomState(2).rand(40).astype(np.float32)) + + def test_cached_path_is_used_and_default_is_not(self): + X, y, X_test = self._classification_data() + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, n_estimators=4, batch_size=2, random_state=0 + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + # Only fit_with_cache materializes the context cache. + self.assertFalse(hasattr(base, "_context_cache_")) + self.assertTrue(hasattr(cached, "_context_cache_")) + + # Cached predict must route through decode, not the full forward pass. + with mock.patch.object( + cached, "_decode_batch_forward", wraps=cached._decode_batch_forward + ) as spy_decode: + with mock.patch.object( + cached, "_batch_forward", wraps=cached._batch_forward + ) as spy_full: + cached.predict_proba(X_test) + self.assertTrue(spy_decode.called) + self.assertFalse(spy_full.called) + + def test_permute_categorical_parity(self): + X, y, X_test = self._mixed_data() + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + permute_categorical=True, + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + permute_categorical=True, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_max_num_rows_parity(self): + X, y, X_test = self._classification_data() # 40 rows + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + max_num_rows=25, + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + max_num_rows=25, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + # Per-member row subsampling sets the (constant) context length. + self.assertEqual(cached._cache_ctx_len_, 25) + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_single_chunk_parity(self): + X, y, X_test = self._classification_data() + model = self._tiny_model("cross_entropy") + base = TabFMClassifier( + model=model, n_estimators=4, batch_size=None, random_state=0 + ) + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=None, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_refit_rebuilds_cache(self): + X, y, X_test = self._classification_data() + model = self._tiny_model("cross_entropy") + cached = TabFMClassifier( + model=model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + cached.fit(X, y) + first_cache = cached._context_cache_ + + # Refitting on different data must rebuild the cache and match a fresh fit. + rng = np.random.RandomState(7) + X2 = rng.rand(36, 5).astype(np.float32) + y2 = rng.randint(0, 3, size=36) + cached.fit(X2, y2) + self.assertIsNot(cached._context_cache_, first_cache) + + base = TabFMClassifier( + model=model, n_estimators=4, batch_size=2, random_state=0 + ) + base.fit(X2, y2) + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_ensemble_preset_parity(self): + X, y, X_test = self._classification_data() + model = self._tiny_model("cross_entropy") + base = TabFMClassifier.ensemble( + model, n_estimators=4, batch_size=2, random_state=0 + ) + cached = TabFMClassifier.ensemble( + model, + n_estimators=4, + batch_size=2, + random_state=0, + fit_mode="fit_with_cache", + ) + base.fit(X, y) + cached.fit(X, y) + + np.testing.assert_allclose( + cached.predict_proba(X_test), base.predict_proba(X_test), atol=1e-5 + ) + + def test_fit_with_cache_requires_jax_model(self): + # The JAX-backend capability check runs before any fitting work. + fake_model = mock.Mock(spec=["max_classes"]) # exposes no prefill/decode + clf = TabFMClassifier( + model=fake_model, n_estimators=2, fit_mode="fit_with_cache" + ) + with self.assertRaisesRegex(NotImplementedError, "prefill"): + clf.fit(np.random.RandomState(0).rand(8, 3), np.array([0, 1] * 4)) + + if __name__ == "__main__": absltest.main()