[DO NOT REVIEW/WIP] Add KumoTabular context-sharded inference

examples/README.md

@@ -7,3 +7,5 @@ This folder contains runnable examples of `structured-data-models`:
 - [**`timeseries/`**](timeseries/): Time series foundation model examples
 
 See [`tabular/quickstart.py`](tabular/quickstart.py) for a minimal runnable `structured-data-models` example with `TabICLv2`.
+
+See [`tabular/context_sharding.py`](tabular/context_sharding.py) for experimental multi-GPU KumoTabular inference with optional FP16 execution. The example includes launch instructions and requires pretrained weights and a finite external job timeout.

examples/tabular/context_sharding.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Reuse a sharded KumoTabular context for held-out synthetic queries.
+
+Run with torchrun --standalone --nproc-per-node=2 --max-restarts=0
+examples/tabular/context_sharding.py --precision float32.
+Use --precision float16 for opt-in FP16 neural execution; model weights and
+numerical inputs remain FP32. Pretrained weights must be available or will be
+downloaded by KumoTabular. This small example demonstrates the API, not a
+speedup or a predictive-quality benchmark. Use a finite external job timeout.
+
+On failure, let torchrun retire the workers. Do not try collective clear/close
+from a one-rank exception handler or reuse a failed executor.
+"""
+
+import argparse
+import os
+from datetime import timedelta
+
+import torch
+import torch.distributed as dist
+
+import sdm
+from sdm.execution import tabular_execution
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--precision", choices=("float32", "float16"), default="float32"
+    )
+    args = parser.parse_args()
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    dist.init_process_group("nccl", timeout=timedelta(minutes=5))
+    control = dist.new_group(backend="gloo", timeout=timedelta(minutes=5))
+    root = dist.get_rank() == 0
+
+    # All ranks load the same stock checkpoint before borrowing the model.
+    model = sdm.models.KumoTabular(
+        task="regression", size="large", device=device
+    ).eval()
+    run = tabular_execution(
+        model,
+        task="regression",
+        group=dist.group.WORLD,
+        control_group=control,
+        precision=args.precision,
+        owner_local_frontend=True,
+        shard_row_workspace=True,
+        fit_head_gather_window=2,
+    )
+
+    x_context = y_context = x_query = batches = None
+    if root:
+        generator = torch.Generator().manual_seed(42)
+        features = torch.randn(576, 8, generator=generator).to(device)
+        # Fit preprocessing only on context; the last 64 rows are held out.
+        x_context, x_query = features[:512], features[512:]
+        y_context = x_context[:, :1] + x_context[:, 1:2].sin()
+        batches = list(x_query.split(16))
+
+    # All ranks call every method. The executor owns the precision scope.
+    # Use no_grad, not inference_mode: fitted caches need version counters.
+    with torch.no_grad():
+        generator = (
+            torch.Generator(device=device).manual_seed(42) if root else None
+        )
+        run.fit(x_context, y_context, generator=generator)
+        predictions = run.predict_batches(batches)
+        run.clear()
+        generator = (
+            torch.Generator(device=device).manual_seed(42) if root else None
+        )
+        run.fit(x_context, y_context, generator=generator)
+        prediction = run.predict(x_query)
+        run.close()
+
+    # Healthy completion only: caller-owned groups outlive the model lease.
+    dist.destroy_process_group(control)
+    dist.destroy_process_group()
+    if root:
+        print("Original batch output shapes:", [p.shape for p in predictions])
+        print("Refitted output shape:", prediction.shape)  # [64, 999]
+
+
+if __name__ == "__main__":
+    main()

sdm/execution/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Opt-in model execution strategies with caller-owned process groups."""
+
+from sdm.execution.tabular import TabularExecution, tabular_execution
+
+__all__ = ["TabularExecution", "tabular_execution"]

sdm/execution/_collectives.py

@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Caller-owned tensor and control collectives."""
+
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+T = TypeVar("T")
+
+
+class Exchange:
+    """Keep allocation/compute failures ahead of the next tensor collective."""
+
+    def __init__(
+        self, group: dist.ProcessGroup, control_group: dist.ProcessGroup
+    ) -> None:
+        self.group, self.control_group = group, control_group
+        self._failure: BaseException | None = None
+        self.rank, self.world = (
+            dist.get_rank(group),
+            dist.get_world_size(group),
+        )
+        self.members = tuple(dist.get_process_group_ranks(group))
+        if self.members != tuple(dist.get_process_group_ranks(control_group)):
+            raise ValueError("Tensor and control group membership must match")
+        if dist.get_backend(control_group) != "gloo":
+            raise ValueError("The control group must use Gloo")
+        # Reused only by the owning thread; nested attempt frames complete
+        # their inner control phase before the outer phase begins.
+        self._status = torch.zeros(1, dtype=torch.int32, device="cpu")
+
+    def records(self, value: T) -> list[T]:
+        return self._records(value)
+
+    def _records(
+        self, value: T, error: BaseException | None = None
+    ) -> list[T]:
+        if self._failure is not None:
+            raise self._failure
+        values: list[Any] = [None] * self.world
+        record = (None if error is None else str(error), value)
+        try:
+            # Every rank takes the same reduction before deciding whether the
+            # full envelope is needed. An inner agree(error) may meet a peer's
+            # records(value), so neither side can skip this phase alone.
+            self._status.fill_(error is not None or value is not None)
+            dist.all_reduce(
+                self._status,
+                op=dist.ReduceOp.MAX,
+                group=self.control_group,
+            )
+            if self._status.item() == 0:
+                return [None] * self.world
+            dist.all_gather_object(values, record, group=self.control_group)
+        except BaseException as collective_error:
+            self._failure = error if error is not None else collective_error
+            if error is None:
+                raise
+            raise error from collective_error
+        failures = [message for message, _ in values if message is not None]
+        if failures:
+            self._failure = (
+                error if error is not None else RuntimeError(failures[0])
+            )
+            raise self._failure
+        return [item for _, item in values]
+
+    def agree(self, error: BaseException | None) -> None:
+        # An inner allocation failure can meet a peer's planning record.
+        # Every control phase consumes the same envelope, then permanently
+        # stops outer attempt frames from issuing additional collectives.
+        self._records(None, error)
+
+    def attempt(self, operation: Callable[[], T]) -> T:
+        error, value = None, None
+        try:
+            value = operation()
+        except BaseException as caught:  # noqa: BLE001
+            error = caught
+        self.agree(error)
+        return cast(T, value)
+
+    def metadata(self, value: T | None) -> T:
+        values = [value]
+        dist.broadcast_object_list(
+            values, src=self.members[0], group=self.control_group
+        )
+        return cast(T, values[0])
+
+    def broadcast(self, value: Tensor, owner: int = 0) -> Tensor:
+        if self.world > 1:
+            dist.broadcast(value, src=self.members[owner], group=self.group)
+        return value
+
+    def gather(
+        self, local: Tensor, lengths: tuple[int, ...], dim: int
+    ) -> Tensor:
+        if self.world == 1:
+            return local.clone(memory_format=torch.contiguous_format)
+        dim %= local.ndim
+        if (
+            dim == 0
+            and lengths[0] > 0
+            and all(size == lengths[0] for size in lengths)
+        ):
+            contiguous, result = self.attempt(
+                lambda: (
+                    local.contiguous(),
+                    local.new_empty((sum(lengths), *local.shape[1:])),
+                )
+            )
+            dist.all_gather_into_tensor(result, contiguous, group=self.group)
+            return result
+        shape = list(local.shape)
+        shape[dim] = max(lengths)
+        padded = self.attempt(lambda: local.new_zeros(shape))
+        padded.narrow(dim, 0, local.size(dim)).copy_(local)
+        received = self.attempt(
+            lambda: [torch.empty_like(padded) for _ in lengths]
+        )
+        dist.all_gather(received, padded, group=self.group)
+        return self.attempt(
+            lambda: torch.cat(
+                [
+                    part.narrow(dim, 0, size)
+                    for part, size in zip(received, lengths)
+                ],
+                dim=dim,
+            )
+        )
+
+    def exchange(
+        self, send: Tensor, sent: list[int], received: list[int]
+    ) -> Tensor:
+        result = self.attempt(lambda: send.new_empty(sum(received)))
+        if self.world == 1:
+            result.copy_(send)
+        else:
+            dist.all_to_all_single(
+                result,
+                send,
+                output_split_sizes=received,
+                input_split_sizes=sent,
+                group=self.group,
+            )
+        return result