Batch inference across estimators

benchmark/tabular/model.py

@@ -60,6 +60,7 @@ def _set_default_params(self) -> None:
         self._set_default_param_value("max_context_size", None)
         self._set_default_param_value("max_columns", None)
         self._set_default_param_value("kv_cache", False)
+        self._set_default_param_value("estimator_batch_size", "auto")
 
     def _fit(
         self,
@@ -125,6 +126,7 @@ def _fit(
             y_context = y_context[perm].unflatten(0, shape)
             num_estimators = None
         self._expand_query = num_estimators is None
+        self._context_shape = x_context.shape[-2:]
 
         recipe = self.model.default_recipe()
         if params["max_columns"] is not None:
@@ -144,13 +146,12 @@ def _fit(
                     y=y_context,
                     recipe=recipe,
                     num_estimators=num_estimators,
+                    estimator_batch_size=self._estimator_batch_size(x_context),
                     generator=generator,
                 )
             return
 
         self._recipe_execution = RecipeExecution(recipe)
-        # KumoTabular and TabFM need the column types before preprocessing.
-        self._schema = x_context.schema
         with (
             torch.inference_mode(),
             torch.amp.autocast(self._device.type, enabled=False),
@@ -213,24 +214,32 @@ def _predict_proba(
                 generator = torch.Generator(self._device).set_state(
                     self._rng_state
                 )
-                outputs = []
-                for context, query in zip(self._contexts, queries):
+                outputs: list[sdm.TableTensor] = []
+                size = self._estimator_batch_size(x_query)
+                size = len(self._contexts) if size is None else size
+                for start in range(0, len(self._contexts), size):
+                    batch = [
+                        context._replace(
+                            x=cast(
+                                sdm.TableTensor, context.x.to(self._device)
+                            ),
+                            y=cast(
+                                sdm.TableTensor, context.y.to(self._device)
+                            ),
+                        )
+                        for context in self._contexts[start : start + size]
+                    ]
                     with torch.amp.autocast(
                         self._device.type,
                         self.autocast_dtype,
                         enabled=x_query.is_cuda,
                     ):
-                        out = self.model._forward(
-                            x_context=context.x.to(self._device),
-                            y_context=context.y.to(self._device),
-                            x_query=query.x,
-                            related_context_tables=None,
-                            related_query_tables=None,
-                            cache=None,
+                        outputs += self.model._forward_members(
+                            contexts=batch,
+                            queries=queries[start : start + size],
+                            estimator_batch_size=None,
                             generator=generator,
-                            _schema=self._schema,
                         )
-                    outputs.append(out.to(query.x.dtype))
 
                 if self.problem_type == REGRESSION:
                     outputs = list(
@@ -249,6 +258,28 @@ def _predict_proba(
         probabilities = out.numerical[..., indices].float().cpu().numpy()
         return self._convert_proba_to_unified_form(probabilities)
 
+    def _estimator_batch_size(self, x: torch.Tensor) -> int | None:
+        params = self._get_model_params()
+        estimator_batch_size = params["estimator_batch_size"]
+        if estimator_batch_size != "auto":
+            return estimator_batch_size
+        # Subsampled contexts can produce different cache shapes per estimator.
+        if not x.is_cuda or self._expand_query:
+            return 1
+
+        num_rows, num_cols = self._context_shape
+        if not params["kv_cache"]:
+            # Uncached inference processes context and query rows together.
+            num_rows += x.size(-2)
+            if num_rows > 3_000 or num_rows * num_cols > 50_000:
+                return 1
+            return self._num_estimators
+
+        num_rows = max(num_rows, x.size(-2))
+        if num_rows > 2_000 or num_rows * num_cols >= 50_000:
+            return 1
+        return self._num_estimators
+
     def get_device(self) -> str:
         return str(next(self.model.parameters()).device)
 

sdm/explain/gradient.py

@@ -78,6 +78,7 @@ def on_model_forward_end(
             objective,
             [numerical for _, _, numerical in self._inputs],
             allow_unused=True,
+            retain_graph=True,
         )
         grad_tables: dict[str | None, TableTensor] = {}
         for (table_name, columns, numerical), grad in zip(