From 06562b05ed37a9b7f12b129919bc363f464bd9d2 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:35:24 -0700 Subject: [PATCH 01/10] Add opt-in KumoTabular context-sharded execution --- docs/source/index.md | 1 + docs/source/kumotabular-context-sharding.md | 69 ++ sdm/execution/__init__.py | 8 + sdm/execution/_collectives.py | 124 ++++ sdm/execution/_tabular_cache.py | 73 ++ sdm/execution/_tabular_columns.py | 326 ++++++++ sdm/execution/_tabular_icl.py | 149 ++++ sdm/execution/_tabular_rows.py | 334 +++++++++ sdm/execution/tabular.py | 776 ++++++++++++++++++++ sdm/models/base/icl.py | 10 +- sdm/models/base/model.py | 46 +- sdm/models/kumo/tabular/icl.py | 14 +- sdm/models/kumo/tabular/model.py | 14 +- sdm/models/kumo/tabular/row_embedding.py | 35 +- sdm/nn/attention.py | 37 +- test/execution/test_tabular.py | 733 ++++++++++++++++++ test/execution/test_tabular_backpressure.py | 187 +++++ 17 files changed, 2914 insertions(+), 22 deletions(-) create mode 100644 docs/source/kumotabular-context-sharding.md create mode 100644 sdm/execution/__init__.py create mode 100644 sdm/execution/_collectives.py create mode 100644 sdm/execution/_tabular_cache.py create mode 100644 sdm/execution/_tabular_columns.py create mode 100644 sdm/execution/_tabular_icl.py create mode 100644 sdm/execution/_tabular_rows.py create mode 100644 sdm/execution/tabular.py create mode 100644 test/execution/test_tabular.py create mode 100644 test/execution/test_tabular_backpressure.py diff --git a/docs/source/index.md b/docs/source/index.md index feb4079a09..767a5cbd2c 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -30,6 +30,7 @@ overview tensor icl +kumotabular-context-sharding processing nim ``` diff --git a/docs/source/kumotabular-context-sharding.md b/docs/source/kumotabular-context-sharding.md new file mode 100644 index 0000000000..37be4b3e73 --- /dev/null +++ b/docs/source/kumotabular-context-sharding.md @@ -0,0 +1,69 @@ +# KumoTabular context sharding + +`sdm.execution.tabular_execution` is an experimental, opt-in executor for splitting one KumoTabular estimator's context across devices. Ordinary `model.fit` and `model.predict` retain their existing behavior when no executor owns the model. + +## Usage + +Launch one process per CUDA device with `torchrun --standalone --nproc-per-node=4 your_script.py`. Each process loads identical weights and creates the same ordered tensor and control groups: + +```python +import os +from datetime import timedelta + +import torch +import torch.distributed as dist + +from sdm.execution import tabular_execution +from sdm.models import KumoTabular + +local_rank = int(os.environ["LOCAL_RANK"]) +torch.cuda.set_device(local_rank) +dist.init_process_group("nccl", timeout=timedelta(minutes=10)) +control = dist.new_group(backend="gloo", timeout=timedelta(minutes=10)) +root = dist.get_rank() == 0 +model = KumoTabular(task="regression", device=f"cuda:{local_rank}").eval() + +run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=control, + fit_head_gather_window=2, +) +# Only root supplies tensors. All ranks call every method in the same order. +# x_context: [context_rows, features]; y_context: [context_rows, 1]. +run.fit(x_context if root else None, y_context if root else None) +prediction = run.predict(x_query if root else None) +# A list of original batches allows independent batches to run concurrently. +predictions = run.predict_batches(query_batches if root else None) +run.clear() +run.fit(x_context if root else None, y_context if root else None) +run.close() +# The caller owns and destroys groups after healthy completion. +dist.destroy_process_group(control) +dist.destroy_process_group() +``` + +Inputs may be tensors or `TableTensor`s. Root alone supplies the recipe and optional preprocessing generator. Prediction returns the usual public `TableTensor` on root and `None` on peers; `predict_batches` returns a root-only list in the original order. Each original query batch is kept intact, including recipe output processing and regression inverse transforms. Empty batch lists return an empty list. + +## What is distributed + +- Expanded row-embedding cells are stored by contiguous context-row owner, reducing the dominant row-by-feature-by-channel allocation. +- Column attention exchanges row-owned cells into complete native column windows, then returns results to their row owners. It does not replace full-context column attention with independent local fits. +- Independent row-attention windows run on their owner, with one error agreement for the local loop. Windows crossing owner boundaries retain their native shape and arithmetic. +- ICL fit attention heads are divided across owners. Full projections and residuals remain replicated; retained ICL key/value caches are sharded by context row. +- Prediction gathers one cache layer at a time and assigns complete original query batches to owners. This combines context-owned storage with query-batch parallelism, not ensemble parallelism. + +This is not complete tensor parallelism and does not guarantee linear capacity or speed scaling. Raw prepared inputs, some normalization/projection intermediates, full ICL states, one replayed cache layer, and a native column window can still scale with the entire context on a single device. Communication and replicated work can dominate small inputs. `column_staging_bytes` and `attention_workspace_bytes` bound individual planned components, not total allocated, reserved, or physical GPU memory. + +## Supported scope and lifecycle + +The initial scope is eager evaluation, FP32, one unbatched estimator, stock KumoTabular classification or regression, and no callbacks, custom forward hooks, compilation, autocast, or ECOC expansion beyond native class capacity. The number of owners must divide the ICL fit query/key-value head counts. CPU/Gloo is supported for numerical and lifecycle tests; CUDA requires NCCL for tensors plus a separate Gloo control group with identical ordered members. + +The executor borrows its model exclusively. Do not mutate weights, recipes, caches, execution settings, or use the underlying model while the lease is active. `fit` replaces a fitted generation; `clear` releases it while retaining the lease; healthy `close` releases the lease and permits ordinary model reuse. Calls must originate on the constructing thread and cannot overlap. Groups and process lifetime always remain caller-owned. + +Local allocation or compute errors are agreed before the next tensor collective, and the executor retains its first failure. A failed executor cannot be cleared, refitted, or reopened; retire the worker processes, especially after CUDA or collective failures. A process crash or mismatched method order still requires the caller's finite process-group timeout and job supervision. Never attempt a collective `close` in a one-rank-only exception handler. + +## Validation status + +Tests compare complete public classification outputs and all 999 regression quantiles with the same native model on one, two, and four CPU processes. They also exercise original row-window boundaries, refit, clear, close, borrowed process groups, mutation rejection, and collective failure propagation. These CPU checks do not establish CUDA numerical equivalence, trained-model quality, memory scaling, or speedup for this branch. Real-GPU matched native/one/two/four-owner qualification remains required; historical research-prototype performance must not be treated as a measurement of this current-main implementation. diff --git a/sdm/execution/__init__.py b/sdm/execution/__init__.py new file mode 100644 index 0000000000..fb02bdff22 --- /dev/null +++ b/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"] diff --git a/sdm/execution/_collectives.py b/sdm/execution/_collectives.py new file mode 100644 index 0000000000..2a4b5f1067 --- /dev/null +++ b/sdm/execution/_collectives.py @@ -0,0 +1,124 @@ +# 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") + + 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: + 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 + 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 diff --git a/sdm/execution/_tabular_cache.py b/sdm/execution/_tabular_cache.py new file mode 100644 index 0000000000..ec558924e3 --- /dev/null +++ b/sdm/execution/_tabular_cache.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Owning cache slices, with one layer gathered during replay.""" + +from collections.abc import Hashable, Iterator +from typing import Self, cast + +from torch import Tensor + +from sdm.cache import Cache, KVCacheEntry +from sdm.execution._collectives import Exchange + + +class OwnedCache(Cache): + """Store columns for the frontend or context rows for ICL.""" + + def __init__( + self, + exchange: Exchange, + lengths: tuple[int, ...], + dim: int, + prefix: str, + ) -> None: + super().__init__() + self.exchange, self.lengths = exchange, lengths + self.dim, self.prefix = dim, prefix + self.start = sum(lengths[: exchange.rank]) + self.length = lengths[exchange.rank] + + def __setitem__(self, key: Hashable, value: object) -> None: + if isinstance(key, str) and key.startswith(self.prefix): + pair = cast(KVCacheEntry, value) + value = KVCacheEntry( + *( + tensor.narrow(self.dim, self.start, self.length) + .clone() + .contiguous() + for tensor in pair + ) + ) + super().__setitem__(key, value) + + def __getitem__(self, key: Hashable) -> object: + value = super().__getitem__(key) + if ( + self.is_replaying + and isinstance(key, str) + and key.startswith(self.prefix) + ): + pair = cast(KVCacheEntry, value) + return KVCacheEntry( + *( + self.exchange.gather( + tensor, self.lengths, self.dim + ).contiguous() + for tensor in pair + ) + ) + return value + + def _tensors(self) -> Iterator[Tensor]: + # Iteration over Cache.values() would gather replay entries. + for value in self._items.values(): + if isinstance(value, KVCacheEntry): + yield from value + elif isinstance(value, Tensor): + yield value + + def freeze(self) -> Self: + # Entries contain owning tensors only; do not invoke replay getters. + self._mode = Cache.Mode.replay + return self diff --git a/sdm/execution/_tabular_columns.py b/sdm/execution/_tabular_columns.py new file mode 100644 index 0000000000..5c6042779a --- /dev/null +++ b/sdm/execution/_tabular_columns.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Complete native column domains over row-owned cell storage.""" + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor + +from sdm.cache import KVCacheEntry +from sdm.execution._collectives import Exchange +from sdm.nn.attention import _chunk + + +@dataclass(frozen=True) +class Domain: + """Closed inducing/output dependency interval in global column order.""" + + start: int + stop: int + owner: int + + +def domain_rounds( + columns: int, inducing_width: int, output_width: int, world: int +) -> tuple[tuple[Domain, ...], ...]: + """Assign one complete common-boundary domain per rank and round.""" + if min(columns, inducing_width, output_width, world) < 1: + raise ValueError("Column dependency dimensions must be positive") + iw = min(inducing_width, 65535, columns) + ow = min(output_width, 65535, columns) + width = math.lcm(iw, ow) + domains = tuple( + ( + Domain(start, min(start + width, columns), index % world) + for index, start in enumerate(range(0, columns, width)) + ) + ) + return tuple( + domains[start : start + world] + for start in range(0, len(domains), world) + ) + + +def to_column_owners( + local: Tensor, + lengths: tuple[int, ...], + domains: tuple[Domain, ...], + group: Exchange, +) -> dict[Domain, Tensor]: + """Transpose row-owned cells to only assigned global column domains.""" + rank, world = (group.rank, group.world) + channels = local.size(-1) + owned = tuple(domain for domain in domains if domain.owner == rank) + send_splits = [ + sum(d.stop - d.start for d in domains if d.owner == peer) + * lengths[rank] + * channels + for peer in range(world) + ] + receive_splits = [ + sum(d.stop - d.start for d in owned) * length * channels + for length in lengths + ] + + def pack() -> Tensor: + send = local.new_empty(sum(send_splits)) + cursor = 0 + for peer in range(world): + for domain in domains: + if domain.owner == peer: + part = local[domain.start : domain.stop].reshape(-1) + send[cursor : cursor + part.numel()].copy_(part) + cursor += part.numel() + del part + return send + + send = group.attempt(pack) + received = group.exchange(send, send_splits, receive_splits) + del send + + def assemble(received: Tensor = received) -> dict[Domain, Tensor]: + fulls = { + d: local.new_empty(d.stop - d.start, sum(lengths), channels) + for d in owned + } + cursor, offset = (0, 0) + for length in lengths: + for domain in owned: + count = (domain.stop - domain.start) * length * channels + fulls[domain][:, offset : offset + length].copy_( + received[cursor : cursor + count].view( + domain.stop - domain.start, length, channels + ) + ) + cursor += count + offset += length + return fulls + + fulls = group.attempt(assemble) + del received + return fulls + + +def to_row_owners( + outputs: dict[Domain, Tensor], + local_out: Tensor, + lengths: tuple[int, ...], + domains: tuple[Domain, ...], + group: Exchange, +) -> None: + """Return destination rows, releasing global domains before receive.""" + rank, world = (group.rank, group.world) + channels = local_out.size(-1) + owned = tuple(d for d in domains if d.owner == rank) + send_splits = [ + sum(d.stop - d.start for d in owned) * length * channels + for length in lengths + ] + receive_splits = [ + sum(d.stop - d.start for d in domains if d.owner == peer) + * lengths[rank] + * channels + for peer in range(world) + ] + + def pack() -> Tensor: + send = local_out.new_empty(sum(send_splits)) + cursor, offset = (0, 0) + for length in lengths: + for domain in owned: + part = outputs[domain][:, offset : offset + length].reshape(-1) + send[cursor : cursor + part.numel()].copy_(part) + cursor += part.numel() + del part + offset += length + return send + + send = group.attempt(pack) + outputs.clear() + received = group.exchange(send, send_splits, receive_splits) + del send + + def writeback() -> None: + cursor = 0 + for peer in range(world): + for domain in domains: + if domain.owner == peer: + count = ( + (domain.stop - domain.start) * lengths[rank] * channels + ) + local_out[domain.start : domain.stop].copy_( + received[cursor : cursor + count].view( + domain.stop - domain.start, lengths[rank], channels + ) + ) + cursor += count + + group.attempt(writeback) + + +def _native_pitch(full: Tensor, batch: torch.Size, readouts: int) -> Tensor: + columns, rows, channels = full.shape + backing = full.new_empty((*batch[:-1], rows, columns + readouts, channels)) + view = backing[..., readouts:, :].transpose(-2, -3) + view.copy_(full.view(*batch[:-1], columns, rows, channels)) + return view + + +def _column_forward( + column: Any, + inducing: Any, + output: Any, + *, + lengths: tuple[int, ...], + group: Exchange, + readouts: int, +) -> Any: + + def forward( + *, + query: Tensor, + key_value: Tensor, + out: Tensor, + return_key_value: bool = False, + **kwargs: Any, + ) -> Any: + del key_value, kwargs + rank, world = (group.rank, group.world) + batch = query.shape[:-2] + columns, total, channels = ( + math.prod(batch), + sum(lengths), + query.size(-1), + ) + size, points = (query.element_size(), column.inducing_points) + iw = min( + column.inducing_block.auto_batch_size_limit( + query.device, size, points.size(0), total + ), + 65535, + ) + ow = min( + column.output_block.auto_batch_size_limit( + query.device, size, total, points.size(0) + ), + 65535, + ) + rounds = domain_rounds(columns, iw, ow, world) + local = query.reshape(columns, lengths[rank], channels) + local_out = out.reshape(columns, lengths[rank], channels) + kv_shape = ( + columns, + points.size(0), + column.output_block.attn.num_key_value_heads, + column.output_block.attn.head_dim, + ) + cache_pair = group.attempt( + lambda: ( + (query.new_empty(kv_shape), query.new_empty(kv_shape)) + if return_key_value + else None + ) + ) + + def compute(fulls: dict[Domain, Tensor]) -> None: + for domain, full in fulls.items(): + induced = full.new_empty( + domain.stop - domain.start, points.size(0), channels + ) + for start in range(domain.start, domain.stop, iw): + stop = min(start + iw, columns) + part = full[start - domain.start : stop - domain.start] + native = ( + _native_pitch(part, batch, readouts) + if columns <= iw + else part + ) + q = ( + points + if columns <= iw + else _chunk(points, batch, 2, start, stop) + ) + state_out = ( + None + if start == 0 or columns <= iw + else induced[ + start - domain.start : stop - domain.start + ] + ) + state = inducing( + query=q, + key_value=native, + batch_size_limit=iw, + out=state_out, + ) + induced[start - domain.start : stop - domain.start].copy_( + state.reshape(stop - start, points.size(0), channels) + ) + del part, native, q, state_out, state + for start in range(domain.start, domain.stop, ow): + stop = min(start + ow, columns) + part = full[start - domain.start : stop - domain.start] + native = ( + _native_pitch(part, batch, readouts) + if columns <= ow + else part + ) + state = induced[start - domain.start : stop - domain.start] + if columns <= ow: + state = state.view(*batch, points.size(0), channels) + native_has_out = ( + columns <= ow or len(batch) == 1 or total == 1 + ) + result = output( + query=native, + key_value=state, + return_key_value=return_key_value, + batch_size_limit=ow, + out=native if native_has_out else None, + ) + if return_key_value: + assert cache_pair is not None + result, kv = result + cache_pair[0][start:stop].copy_( + kv.key.reshape(stop - start, *kv_shape[1:]) + ) + cache_pair[1][start:stop].copy_( + kv.value.reshape(stop - start, *kv_shape[1:]) + ) + del kv + part.copy_(result.reshape(stop - start, total, channels)) + del part, native, state, result + del induced + + for domains in rounds: + fulls = to_column_owners(local, lengths, domains, group) + group.attempt(lambda fulls=fulls: compute(fulls)) + to_row_owners(fulls, local_out, lengths, domains, group) + if return_key_value: + assert cache_pair is not None + for domain in domains: + for start in range(domain.start, domain.stop, ow): + stop = min(start + ow, columns) + for tensor in cache_pair: + if world > 1: + group.broadcast( + tensor[start:stop], domain.owner + ) + del fulls + if not return_key_value: + return out + assert cache_pair is not None + key, value = cache_pair + return ( + out, + KVCacheEntry( + key.view(*batch, *kv_shape[1:]), + value.view(*batch, *kv_shape[1:]), + ), + ) + + return forward diff --git a/sdm/execution/_tabular_icl.py b/sdm/execution/_tabular_icl.py new file mode 100644 index 0000000000..35fd09cae5 --- /dev/null +++ b/sdm/execution/_tabular_icl.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Head-owned native fit and one-global-layer cached replay.""" + +from collections import deque +from contextlib import nullcontext +from typing import cast + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.nn.attention import SDPBackend, sdpa_kernel + +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_cache import OwnedCache +from sdm.models.kumo.tabular.icl import ICLBlock +from sdm.nn import SDPA, TransformerBlock + + +def fit_icl( + model: ICLBlock, + local: Tensor, + labels: Tensor, + cache: OwnedCache, + exchange: Exchange, + fit_head_gather_window: int | None = None, +) -> Tensor: + """Preserve full projections/residuals; partition only SDPA heads.""" + full = exchange.gather(local, cache.lengths, 0) + pending: deque[torch.cuda.Event] = deque() + + def attention( + sdpa: SDPA, + *, + query: Tensor, + key: Tensor, + value: Tensor, + seqused_key_value: Tensor | None = None, + attn_mask: Tensor | None = None, + is_causal: bool = False, + ) -> Tensor: + if seqused_key_value is not None or attn_mask is not None or is_causal: + raise ValueError("Context-sharded fit requires unmasked attention") + if query.numel() == 0 or exchange.world == 1: + return sdpa(query, key, value) + if ( + fit_head_gather_window is not None + and len(pending) >= fit_head_gather_window + ): + # Prior gathers on this rank are issued before another layer can + # enqueue; peers reach the same calls in model order. + pending.popleft().synchronize() + scaled = exchange.attempt( + lambda: ( + query + if sdpa.query_scaling is None + else sdpa.query_scaling(query, key_len=key.size(-3)) + ) + ) + q_heads, kv_heads = query.size(-2), key.size(-2) + q_width, kv_width = ( + q_heads // exchange.world, + kv_heads // exchange.world, + ) + start = exchange.rank + output = exchange.attempt( + lambda: sdpa( + scaled[..., start * q_width : (start + 1) * q_width, :], + key[..., start * kv_width : (start + 1) * kv_width, :], + value[..., start * kv_width : (start + 1) * kv_width, :], + _already_scaled=True, + ) + ) + source = exchange.attempt(lambda: output.movedim(-2, 0).contiguous()) + received = exchange.attempt( + lambda: source.new_empty( + (exchange.world * source.size(0), *source.shape[1:]) + ) + ) + dist.all_gather_into_tensor(received, source, group=exchange.group) + result = received.movedim(0, -2).contiguous() + if fit_head_gather_window is not None and result.is_cuda: + # The copy of the gathered heads is ordered after NCCL on this + # stream. Bound host enqueue lead without a device-wide fence. + done = torch.cuda.Event() + done.record(torch.cuda.current_stream(result.device)) + pending.append(done) + return result + + # Fit uses the same non-quadratic native backend policy as the accepted + # path. Cached GQA prediction deliberately retains native backend choice. + with ( + sdpa_kernel( + [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] + ) + if full.is_cuda + else nullcontext() + ): + return model(full, labels, cache=cache, _attention=attention) + + +def workspace_bytes(model: ICLBlock, rows: int, cache: OwnedCache) -> int: + """Conservative layer-local logical scratch, not total device memory.""" + context = sum(cache.lengths) + largest = 0 + for module in model.layers: + layer = cast(TransformerBlock, module) + heads = layer.attn.num_query_heads + kv_heads = model.kv_heads or layer.attn.num_key_value_heads + width = layer.attn.head_dim + score = heads * rows * context + q = heads * rows * width + kv = heads * context * width + compact = kv_heads * context * 2 * width + attention = ( + (13 * score + heads * rows + 8 * q + 12 * kv + 4 * compact + 8 * q) + if rows + else 0 + ) + gather = ( + ( + max(cache.lengths) + + len(cache.lengths) * max(cache.lengths) + + 2 * context + ) + * kv_heads + * 2 + * width + * 4 + if len(cache.lengths) > 1 + else compact * 4 + ) + largest = max(largest, attention + gather) + return largest + + +def predict_icl(model: ICLBlock, query: Tensor, cache: OwnedCache) -> Tensor: + """Replay each full original query with one gathered KV layer at a time.""" + state = query.clone() + for index, layer in enumerate(model.layers): + pair = cache[f"icl_block.layer{index}"] + state = cache.exchange.attempt( + lambda layer=layer, pair=pair, state=state: layer( + query=state, key_value=pair, out=state + ) + ) + del pair + return cache.exchange.attempt(lambda: model.head(model.norm(state))) diff --git a/sdm/execution/_tabular_rows.py b/sdm/execution/_tabular_rows.py new file mode 100644 index 0000000000..62b3b82c86 --- /dev/null +++ b/sdm/execution/_tabular_rows.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Original global cell/column/row windows with contiguous row ownership.""" + +from collections.abc import Callable +from functools import partial +from typing import Any, cast + +import torch +from torch import Tensor + +from sdm._memory import chunk_memory_limit +from sdm.cache import Cache +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_columns import _column_forward, domain_rounds +from sdm.models.kumo.tabular.row_embedding import RowEmbedding +from sdm.models.tabfm.cell_embedding import CellEmbedding +from sdm.nn import InducedTransformerBlock, TransformerBlock + + +def bounds(total: int, world: int) -> tuple[int, ...]: + return tuple( + total * (rank + 1) // world - total * rank // world + for rank in range(world) + ) + + +class Rows: + """Bind native AUTO plans once, before allocating expanded cells.""" + + def __init__( + self, + encoder: RowEmbedding, + raw: Tensor, + labels: Tensor, + mask: Tensor, + exchange: Exchange, + staging_bytes: int, + ) -> None: + self.encoder, self.raw, self.labels = encoder, raw, labels + self.exchange, self.mask = exchange, mask + self.total, self.columns = raw.shape + self.lengths = bounds(self.total, exchange.world) + self.start = sum(self.lengths[: exchange.rank]) + self.stop = self.start + self.lengths[exchange.rank] + self.readouts, self.channels = encoder.readout_token.shape + self.mean = raw.nanmean(dim=-2, keepdim=True) + self.mean = self.mean.masked_fill(self.mean.isnan(), 0) + self.cell_rows = self.total + cell = encoder.cell_embedding + if raw.is_cuda: + groups, frequencies = cell.group_size, cell.num_frequencies + fixed = self.columns * ( + groups * frequencies * 4 + + groups * self.channels * 2 * frequencies * 4 + + self.channels * 4 + ) + item = (2 * groups * frequencies + 2 * self.channels) * 4 + items = max((chunk_memory_limit(raw.device) - fixed) // item, 1) + self.cell_rows = max(items // self.columns, 1) + self.handlers: dict[torch.nn.Module, Callable[..., Any]] = { + encoder.cell_embedding: self.cell, + encoder.norm: self.norm, + } + if encoder.y_lin is not None: + self.handlers[encoder.y_lin] = self.target + if encoder.y_emb is not None: + self.handlers[encoder.y_emb] = encoder.y_emb + limits = [] + peak = 0 + for module in encoder.col_blocks: + column = cast(InducedTransformerBlock, module) + points = column.inducing_points.size(0) + iw = min( + column.inducing_block.auto_batch_size_limit( + raw.device, 4, points, self.total + ), + 65_535, + ) + ow = min( + column.output_block.auto_batch_size_limit( + raw.device, 4, self.total, points + ), + 65_535, + ) + limits.append((iw, ow)) + for domains in domain_rounds(self.columns, iw, ow, exchange.world): + width = sum(d.stop - d.start for d in domains) + for peer, length in enumerate(self.lengths): + owned = sum( + d.stop - d.start for d in domains if d.owner == peer + ) + sent = length * width * self.channels * 4 + received = self.total * owned * self.channels * 4 + size = max(2 * (sent + received), 3 * received) + if (self.columns <= iw or self.columns <= ow) and owned: + size += ( + self.total + * (self.columns + self.readouts) + * self.channels + * 4 + ) + peak = max(peak, size) + self.handlers[column] = _column_forward( + column, + column.inducing_block, + column.output_block, + lengths=self.lengths, + group=exchange, + readouts=self.readouts, + ) + row_limits = [] + for index, module in enumerate(encoder.row_blocks): + block = cast(TransformerBlock, module) + last = index == len(encoder.row_blocks) - 1 + tokens = self.readouts if last else self.columns + self.readouts + limit = min( + block.auto_batch_size_limit( + raw.device, 4, tokens, self.columns + self.readouts + ), + 65_535, + ) + row_limits.append(limit) + self.handlers[block] = partial(self.row, block, limit, tokens) + if peak > staging_bytes: + raise ValueError( + "Column staging budget is too small for native windows" + ) + plan = (self.cell_rows, limits, row_limits) + if any(other != plan for other in exchange.records(plan)): + raise ValueError("Native AUTO windows differ between owners") + + def __call__( + self, module: torch.nn.Module, *args: Any, **kwargs: Any + ) -> Any: + return self.handlers[module](*args, **kwargs) + + def cell( + self, + *, + x: Tensor, + categorical_mask: Tensor, + cache: Cache, + out: Tensor, + **kwargs: Any, + ) -> Tensor: + del x, kwargs + cache["cell_impute_mean"] = self.mean + for index, start in enumerate(range(0, self.total, self.cell_rows)): + stop = min(start + self.cell_rows, self.total) + self._cell_window(start, stop, index, categorical_mask, out) + return out + + def _cell_window( + self, start: int, stop: int, index: int, mask: Tensor, out: Tensor + ) -> None: + owner = self.owner(start, stop) + interior = owner is not None + owner = index % self.exchange.world if owner is None else owner + + def allocate() -> Tensor | None: + if interior and self.exchange.rank != owner: + return None + buffer = self.raw.new_empty( + stop - start, self.readouts + self.columns, self.channels + ) + buffer[:, : self.readouts] = self.encoder.readout_token + return buffer + + buffer = self.exchange.attempt(allocate) + + def compute() -> None: + if self.exchange.rank == owner: + assert buffer is not None + raw = self.raw[start:stop] + missing = raw.isnan() + CellEmbedding.forward( + self.encoder.cell_embedding, + torch.where(missing, self.mean, raw), + mask, + batch_size_limit=None, + out=buffer[:, self.readouts :], + missing=missing, + ) + + self.exchange.attempt(compute) + if not interior: + assert buffer is not None + self.exchange.broadcast(buffer, owner) + left, right = max(start, self.start), min(stop, self.stop) + if left < right: + assert buffer is not None + out[left - self.start : right - self.start].copy_( + buffer[left - start : right - start, self.readouts :] + ) + + def owner(self, start: int, stop: int) -> int | None: + offset = 0 + for rank, length in enumerate(self.lengths): + if offset <= start < stop <= offset + length: + return rank + offset += length + return None + + def target(self, labels: Tensor) -> Tensor: + del labels + assert self.encoder.y_lin is not None + return self.encoder.y_lin(self.labels.unsqueeze(-1))[ + self.start : self.stop + ].clone() + + def row( + self, + block: torch.nn.Module, + limit: int, + tokens: int, + *, + key_value: Tensor, + out: Tensor, + **kwargs: Any, + ) -> Tensor: + del kwargs + interior, boundary = [], [] + for index, start in enumerate(range(0, self.total, limit)): + stop = min(start + limit, self.total) + owner = self.owner(start, stop) + if owner is None: + boundary.append((index, start, stop)) + elif owner == self.exchange.rank: + interior.append((start, stop)) + + def local_interiors() -> None: + # Preserve each native window's shape and clone layout, but avoid + # a control collective for every rank-local window. Rows in + # different windows are independent in this attention block. + for start, stop in interior: + full = key_value[ + start - self.start : stop - self.start + ].clone() + query = full[:, :tokens] + block( + query=query, + key_value=full, + out=query, + batch_size_limit=limit, + ) + out[start - self.start : stop - self.start].copy_( + full[:, :tokens] + ) + del query, full + + self.exchange.attempt(local_interiors) + # Boundary collectives start only after all owners have completed + # their local enqueue phase or agreed on its first error. + for index, start, stop in boundary: + self._row_window( + block, limit, tokens, key_value, out, index, start, stop + ) + return out + + def _row_window( + self, + block: torch.nn.Module, + limit: int, + tokens: int, + key_value: Tensor, + out: Tensor, + index: int, + start: int, + stop: int, + ) -> None: + owner = self.owner(start, stop) + interior = owner is not None + owner = index % self.exchange.world if owner is None else owner + left, right = max(start, self.start), min(stop, self.stop) + count = max(right - left, 0) + local = key_value[left - self.start : left - self.start + count] + if interior: + full = self.exchange.attempt( + lambda: local.clone() if self.exchange.rank == owner else None + ) + else: + lengths = tuple( + max( + min(stop, sum(self.lengths[: r + 1])) + - max(start, sum(self.lengths[:r])), + 0, + ) + for r in range(self.exchange.world) + ) + sent = [0] * self.exchange.world + sent[owner] = local.numel() + received = [ + n * (self.columns + self.readouts) * self.channels + if self.exchange.rank == owner + else 0 + for n in lengths + ] + gathered = self.exchange.exchange( + local.contiguous().view(-1), sent, received + ) + full = self.exchange.attempt( + lambda: key_value.new_empty( + stop - start, self.columns + self.readouts, self.channels + ) + ) + if self.exchange.rank == owner: + full.copy_(gathered.view_as(full)) + + def compute() -> None: + if self.exchange.rank == owner: + assert full is not None + query = full[:, :tokens] + block( + query=query, + key_value=full, + out=query, + batch_size_limit=limit, + ) + + self.exchange.attempt(compute) + if not interior: + assert full is not None + self.exchange.broadcast(full, owner) + if count: + assert full is not None + out[left - self.start : left - self.start + count].copy_( + full[left - start : left - start + count, :tokens] + ) + + def norm(self, value: Tensor) -> Tensor: + full = self.exchange.gather(value, self.lengths, 0) + return self.encoder.norm(full)[self.start : self.stop].clone() diff --git a/sdm/execution/tabular.py b/sdm/execution/tabular.py new file mode 100644 index 0000000000..aef78c6c66 --- /dev/null +++ b/sdm/execution/tabular.py @@ -0,0 +1,776 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit owner-fit and original-batch execution for KumoTabular.""" + +import threading +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from types import MethodType +from typing import Any, cast + +import torch +import torch.distributed as dist +from torch import Tensor + +import sdm.processing as sp +from sdm import Recipe, TableTensor, Task, TaskLike +from sdm.cache import Cache +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_cache import OwnedCache +from sdm.execution._tabular_icl import fit_icl, predict_icl, workspace_bytes +from sdm.execution._tabular_rows import Rows, bounds +from sdm.models import KumoTabular +from sdm.models.kumo.tabular.model import _KumoTabular +from sdm.nn import TransformerBlock +from sdm.processing.execution import MemberQuery, RecipeExecution +from sdm.tensor.table import TableSchema + +_RETIRED: list["TabularExecution"] = [] + + +def _tensor_identity(value: Tensor) -> tuple[Any, ...]: + return ( + id(value), + value.data_ptr(), + value._version, + value.shape, + value.stride(), + value.dtype, + value.device, + ) + + +def _model_identity(model: KumoTabular) -> tuple[Any, ...]: + modules = torch.nn.modules.module + if modules._global_forward_hooks or modules._global_forward_pre_hooks: + raise ValueError("Execution does not support global forward hooks") + for module in model.modules(): + if ( + module.training + or module._forward_hooks + or module._forward_pre_hooks + or module._compiled_call_impl is not None + or any( + key in vars(module) + for key in ("forward", "_call_impl", "__call__") + ) + ): + raise ValueError( + "Execution requires unmodified evaluation modules " + "without hooks" + ) + if not cast(MethodType, module.forward).__func__.__module__.startswith( + ("sdm.", "torch.nn.") + ): + raise ValueError( + "Execution does not support custom module dispatch" + ) + return ( + tuple( + (id(module), cast(MethodType, module.forward).__func__) + for module in model.modules() + ), + tuple( + _tensor_identity(value) + for value in (*model.parameters(), *model.buffers()) + ), + ) + + +def _check_output(recipe: Recipe) -> None: + # Waves prepare several inputs before finishing their outputs. Known + # deterministic output steps preserve the original RNG ordering. + allowed = ( + sp.Identity, + sp.Sequential, + sp.EnsembleProcessorAdapter, + sp.AverageEstimators, + sp.Softmax, + sp.SortQuantiles, + sp.TaskDispatch, + torch.nn.ModuleDict, + torch.nn.ModuleList, + ) + if any(type(module) not in allowed for module in recipe.output.modules()): + raise ValueError( + "Execution requires deterministic built-in output processors" + ) + + +class TabularExecution: + """Run an explicitly owned, eager FP32/E1 tabular generation. + + Construct with :func:`tabular_execution`. All ranks call every operation; + only group rank zero supplies inputs and receives public outputs. Groups, + model loading and process lifetime remain the caller's responsibility. + Every rank must load identical weights and use the same operation order. + The model is exclusively borrowed until healthy close. Do not mutate it, + recipes or cached tensors, or launch foreign-stream work during a call. + Operations must run on the constructing thread. Boundary version checks + are not a defense against arbitrary ``.data`` or foreign-stream mutation. + A failed operation permanently retires this executor and retains its + owners; failed CUDA/collectives require process retirement, not recovery. + Original-wave replay requires deterministic built-in output processors. + Fit retains row-owned cells/ICL KV and column-owned frontend KV, but full + ICL state and a native column window can still span the entire context. + """ + + def __init__( + self, + model: KumoTabular, + task: Task, + exchange: Exchange, + column_staging_bytes: int, + attention_workspace_bytes: int, + fit_head_gather_window: int | None = None, + ) -> None: + self.model, self.task, self.exchange = model, task, exchange + self._core = cast(_KumoTabular, model.models[task]) + self.device = next(self._core.parameters()).device + self.column_staging_bytes = column_staging_bytes + self.attention_workspace_bytes = attention_workspace_bytes + self.fit_head_gather_window = fit_head_gather_window + self._lock = threading.Lock() + self._owner_thread = threading.get_ident() + self._thread: int | None = None + self._operation: str | None = None + self._failure: BaseException | None = None + self._closed = False + self._root_cache: Cache | None = None + self._batch_cache: Cache | None = None + self._row_cache: OwnedCache | None = None + self._icl_cache: OwnedCache | None = None + self._result: Tensor | None = None + self._fit_started = False + self._identity = exchange.attempt(self._admit) + signature = tuple( + (value.shape, value.dtype) for value in self._core.parameters() + ) + policies = exchange.records( + ( + task, + column_staging_bytes, + attention_workspace_bytes, + fit_head_gather_window, + signature, + ) + ) + if any(policy != policies[0] for policy in policies): + raise ValueError("Execution policies differ between ranks") + self._cache_identity: tuple[Any, ...] | None = None + self._policy_identity = self._policy() + model._execution = self + + def _policy(self) -> tuple[Any, ...]: + return ( + id(self.exchange.group), + id(self.exchange.control_group), + self.exchange.rank, + self.exchange.world, + self.exchange.members, + self.task, + self.device, + self.column_staging_bytes, + self.attention_workspace_bytes, + self.fit_head_gather_window, + ) + + def _admit(self) -> tuple[Any, ...]: + if self.fit_head_gather_window is not None and ( + type(self.fit_head_gather_window) is not int + or self.fit_head_gather_window < 1 + ): + raise ValueError( + "Fit head gather window must be a positive integer" + ) + if ( + type(self.model) is not KumoTabular + or type(self._core) is not _KumoTabular + ): + raise ValueError( + "Execution requires a stock KumoTabular task core" + ) + if self.model._execution is not None or self.model._cache is not None: + raise ValueError("Execution requires an unfitted, unleased model") + if self.device.type not in ("cpu", "cuda") or dist.get_backend( + self.exchange.group + ) != ("nccl" if self.device.type == "cuda" else "gloo"): + raise ValueError("Execution requires CUDA/NCCL or CPU/Gloo") + if any( + value.dtype != torch.float32 or value.device != self.device + for value in self._core.parameters() + ): + raise ValueError( + "Execution requires one FP32 device per task core" + ) + if torch.is_autocast_enabled(self.device.type): + raise ValueError("Execution does not support autocast") + if any( + cast(TransformerBlock, layer).attn.num_query_heads + % self.exchange.world + or cast(TransformerBlock, layer).attn.num_key_value_heads + % self.exchange.world + for layer in self._core.icl_block.layers + ): + raise ValueError("ICL fit heads must divide evenly across owners") + return _model_identity(self.model) + + def _check_model_call(self, operation: str) -> None: + if self._failure is not None: + raise self._failure + if self._thread != threading.get_ident() or self._operation is None: + raise RuntimeError( + "Use the owning tabular executor while its lease is active" + ) + if operation not in (self._operation, "clear"): + raise RuntimeError( + "The model operation differs from its execution lease" + ) + + def _fence(self) -> None: + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + + def _retire(self, error: BaseException) -> None: + if self._failure is None: + self._failure = error + _RETIRED.append(self) + + def _cache_state(self) -> tuple[Any, ...]: + return ( + id(self.model._cache), + tuple((key, id(value)) for key, value in self.model._cache.items()) + if self.model._cache is not None + else (), + id(self._row_cache), + id(self._icl_cache), + tuple( + _tensor_identity(value) + for cache in (self._row_cache, self._icl_cache) + if cache is not None + for value in cache._tensors() + ), + ) + + @contextmanager + def _call(self, operation: str) -> Iterator[None]: + if self._failure is not None: + raise self._failure + if self._closed: + raise RuntimeError("Tabular execution is closed") + if threading.get_ident() != self._owner_thread: + raise RuntimeError("Tabular execution belongs to another thread") + if not self._lock.acquire(blocking=False): + raise RuntimeError("Tabular execution operations cannot overlap") + self._thread, self._operation = threading.get_ident(), operation + try: + + def validate() -> None: + if ( + self.model._execution is not self + or _model_identity(self.model) != self._identity + ): + raise RuntimeError( + "The borrowed model changed; this execution is stale" + ) + if torch.is_autocast_enabled(self.device.type): + raise ValueError("Execution does not support autocast") + if self._policy() != self._policy_identity: + raise RuntimeError( + "Execution policy changed after admission" + ) + if ( + self._cache_identity is not None + and self._cache_state() != self._cache_identity + ): + raise RuntimeError( + "The fitted cache changed outside execution" + ) + + self.exchange.attempt(validate) + yield + if operation in ("fit", "predict"): + self.exchange.attempt(self._fence) + except BaseException as error: + # Traceback and executor keep submitted inputs, caches and owners + # alive. No cleanup fence can turn uncertain work into reuse. + self._retire(error) + raise + finally: + self._thread, self._operation = None, None + self._lock.release() + + def _prepared_forward( + self, forward: Callable[..., TableTensor], **kwargs: Any + ) -> TableTensor: + cache = kwargs["cache"] + x = ( + kwargs["x_context"] + if kwargs["x_query"] is None + else kwargs["x_query"] + ) + if x.ndim != 2 or cache is None: + raise ValueError( + "Execution supports one unbatched estimator with a cache" + ) + classes = cache["classes"] + actual = Task.regression if classes is None else Task.classification + if actual != self.task: + raise ValueError( + "Prepared task differs from the selected execution task" + ) + if classes is not None and len(classes) > self.model.ecoc.max_classes: + raise ValueError("Execution does not support expanded ECOC tasks") + self._batch_cache = cache + return forward(**kwargs) + + def _run_core( + self, task: Task, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> Tensor: + if self._result is not None: + result, self._result = self._result, None + return result + if task != self.task or model is not self._core: + raise ValueError( + "Prepared core differs from the selected execution core" + ) + x = kwargs.get("x", args[0] if args else None) + y = kwargs.get("y", args[1] if len(args) > 1 else None) + if x.dtype != torch.float32 or x.device != self.device: + raise ValueError("Prepared inputs must match the FP32 core device") + # Finish root-only allocations before announcing the tensor phase. + # Otherwise an allocation error could strand peers in a broadcast. + x, y, mask = ( + x.contiguous(), + y.contiguous(), + kwargs["categorical_mask"].contiguous(), + ) + self._fit_started = True + self.exchange.metadata((None, x.shape, y.dtype)) + self.exchange.agree(None) # Peers allocate the matching destinations. + self._fit_core(x, y, mask) + return x.new_empty((0, model.icl_block.head[-1].out_features)) + + def _fit_core(self, x: Tensor, y: Tensor, mask: Tensor) -> None: + exchange = self.exchange + exchange.broadcast(x) + exchange.broadcast(y) + exchange.broadcast(mask) + rows = exchange.attempt( + lambda: Rows( + self._core.row_embedding, + x, + y, + mask, + exchange, + self.column_staging_bytes, + ) + ) + self._row_cache = OwnedCache( + exchange, bounds(x.size(1), exchange.world), 0, "row_embedding." + ) + embedded = exchange.attempt( + lambda: self._core.row_embedding( + x[rows.start : rows.stop], + y[rows.start : rows.stop], + mask, + cache=self._row_cache, + _runner=rows, + ) + ) + embedded = exchange.attempt(lambda: self._core.row_project(embedded)) + self._icl_cache = OwnedCache(exchange, rows.lengths, 0, "icl_block.") + fit_icl( + self._core.icl_block, + embedded, + y, + self._icl_cache, + exchange, + self.fit_head_gather_window, + ) + self._row_cache.freeze() + self._icl_cache.freeze() + + def fit( + self, + x: Tensor | TableTensor | None, + y: Tensor | TableTensor | None, + *, + recipe: Recipe | None = None, + generator: torch.Generator | None = None, + ) -> None: + """Fit one estimator; only rank zero provides context and recipe.""" + with self._call("fit"), torch.no_grad(): + self.exchange.attempt(self._fence) + self._release() + self._fit_started = False + + def inputs() -> None: + if self.exchange.rank != 0 and ( + x is not None or y is not None + ): + raise ValueError("Only root supplies fit inputs") + if any( + value is not None and value.requires_grad + for value in (x, y) + ): + raise ValueError( + "Execution does not support autograd inputs" + ) + + self.exchange.attempt(inputs) + + def execute() -> None: + if self.exchange.rank == 0: + try: + if ( + x is None + or y is None + or x.ndim != 2 + or y.ndim != 2 + ): + raise ValueError( + "Root must supply two-dimensional fit inputs" + ) + chosen = ( + self.model.default_recipe() + if recipe is None + else recipe + ) + _check_output(chosen) + self.model.fit( + x, + y, + recipe=chosen, + num_estimators=1, + generator=generator, + ) + self._root_cache = self.model._cache + except BaseException as error: + if not self._fit_started: + self.exchange.metadata((str(error), None, None)) + raise + else: + error, shape, dtype = self.exchange.metadata(None) + if error is not None: + raise ValueError(error) + inputs = self.exchange.attempt( + lambda: ( + torch.empty( + shape, dtype=torch.float32, device=self.device + ), + torch.empty( + shape[0], dtype=dtype, device=self.device + ), + torch.empty( + shape[1], dtype=torch.bool, device=self.device + ), + ) + ) + self._fit_core(*inputs) + + self.exchange.attempt(execute) + self._cache_identity = self._cache_state() + + def _prepare( + self, batch: Tensor | TableTensor + ) -> tuple[MemberQuery, Cache]: + root = self._root_cache + assert root is not None + cache = cast(Cache, root[0]) + execution = cast(RecipeExecution, root["recipe_execution"]) + with torch.amp.autocast(batch.device.type, enabled=False): + queries = execution.transform(batch, None) + if len(queries) != 1: + raise ValueError( + "Execution supports exactly one prepared query member" + ) + query = self.model._prepare_query( + queries[0], + cast(tuple[TableSchema, ...], cache["x_schemas"])[0], + None, + (), + ) + if query.x.ndim != 2 or query.x.numerical.size(0) != batch.size(0): + raise ValueError( + "Preparation must preserve each original query batch" + ) + return query, cache + + def _finish( + self, prepared: tuple[MemberQuery, Cache], raw: Tensor + ) -> TableTensor: + query, cache = prepared + self._result = raw + outputs = self.model._forward_batch( + contexts=None, + queries=[query], + cache=cache, + categorical_mask=cast(Tensor, cache["categorical_mask"]), + class_values=cast( + Sequence[tuple[Any, ...] | None], cache["class_values"] + ), + callbacks=(), + requires_grad=False, + generator=None, + ) + if self._result is not None: + raise RuntimeError( + "The public output did not consume its computed batch" + ) + assert self._root_cache is not None + execution = cast(RecipeExecution, self._root_cache["recipe_execution"]) + with torch.amp.autocast(raw.device.type, enabled=False): + if self.task == Task.regression: + outputs = list(execution.inverse_transform_target(outputs)) + return execution.transform_output(outputs) + + def predict(self, x: Tensor | TableTensor | None) -> TableTensor | None: + """Predict one original batch without splitting or rebatching it.""" + values = self.predict_batches( + [cast(Tensor, x)] if self.exchange.rank == 0 else None + ) + return values[0] if values is not None else None + + def predict_batches( + self, + batches: Sequence[Tensor | TableTensor] | None, + ) -> list[TableTensor] | None: + """Predict finite original batches, returning their original order.""" + with self._call("predict"), torch.no_grad(): + + def plan() -> list[list[int]] | None: + if self._icl_cache is None: + raise RuntimeError("Tabular execution is not fitted") + if self.exchange.rank != 0: + if batches is not None: + raise ValueError("Only root supplies query batches") + return None + if not isinstance(batches, Sequence): + raise ValueError( + "Root must provide a finite batch sequence" + ) + waves, previous = [], None + for index, batch in enumerate(batches): + if batch.ndim != 2 or batch.size(0) == 0: + raise ValueError( + "Original query batches must be nonempty " + "and two-dimensional" + ) + geometry = (batch.shape, batch.dtype, batch.device) + if batch.requires_grad: + raise ValueError( + "Execution does not support autograd inputs" + ) + if ( + not waves + or len(waves[-1]) == self.exchange.world + or geometry != previous + ): + waves.append([]) + waves[-1].append(index) + previous = geometry + return waves + + waves = self.exchange.metadata(self.exchange.attempt(plan)) + outputs = [] + for wave in waves: + outputs.extend(self._predict_wave(wave, batches)) + return outputs if self.exchange.rank == 0 else None + + def _predict_wave( + self, wave: list[int], batches: Sequence[Tensor | TableTensor] | None + ) -> list[TableTensor]: + outputs: list[TableTensor] = [] + prepared = self.exchange.attempt( + lambda: ( + [ + self._prepare(cast(Sequence[Tensor], batches)[i]) + for i in wave + ] + if self.exchange.rank == 0 + else None + ) + ) + + def check_geometry() -> None: + if prepared is None: + return + expected = ( + prepared[0][0].x.numerical.shape, + cast(Tensor, prepared[0][1]["categorical_mask"]).shape, + ) + if any( + ( + q.x.numerical.shape, + cast(Tensor, c["categorical_mask"]).shape, + ) + != expected + or q.x.numerical.dtype != torch.float32 + or q.x.device != self.device + for q, c in prepared + ): + raise ValueError( + "Prepared query geometries differ within a wave" + ) + + self.exchange.attempt(check_geometry) + metadata = None + if prepared is not None: + metadata = ( + prepared[0][0].x.numerical.shape, + cast(Tensor, prepared[0][1]["categorical_mask"]).shape, + ) + shape, mask_shape = self.exchange.metadata(metadata) + assert self._icl_cache is not None + if ( + workspace_bytes(self._core.icl_block, shape[0], self._icl_cache) + > self.attention_workspace_bytes + ): + raise ValueError( + "Original query exceeds the attention workspace budget" + ) + local, mask = self.exchange.attempt( + lambda: ( + torch.empty(shape, dtype=torch.float32, device=self.device), + torch.empty(mask_shape, dtype=torch.bool, device=self.device), + ) + ) + sources = ( + prepared + [prepared[-1]] * (self.exchange.world - len(prepared)) + if prepared is not None + else None + ) + packed = self.exchange.attempt( + lambda: ( + ( + [q.x.numerical.contiguous() for q, _ in sources], + [ + cast(Tensor, c["categorical_mask"]).contiguous() + for _, c in sources + ], + ) + if sources is not None + else (None, None) + ) + ) + dist.scatter( + local, + packed[0], + src=self.exchange.members[0], + group=self.exchange.group, + ) + dist.scatter( + mask, + packed[1], + src=self.exchange.members[0], + group=self.exchange.group, + ) + empty = local.new_empty( + (0,), + dtype=torch.int64 + if self.task == Task.classification + else local.dtype, + ) + state = self.exchange.attempt( + lambda: self._core.row_embedding( + local, empty, mask, cache=self._row_cache + ) + ) + state = self.exchange.attempt(lambda: self._core.row_project(state)) + raw = predict_icl(self._core.icl_block, state, self._icl_cache) + raw = self.exchange.attempt(raw.contiguous) + received = self.exchange.attempt( + lambda: ( + [torch.empty_like(raw) for _ in range(self.exchange.world)] + if self.exchange.rank == 0 + else None + ) + ) + dist.gather( + raw, + received, + dst=self.exchange.members[0], + group=self.exchange.group, + ) + + def finish() -> None: + if prepared is not None: + assert received is not None + outputs.extend( + self._finish(item, value) + for item, value in zip(prepared, received) + ) + + self.exchange.attempt(finish) + return outputs + + def _release(self) -> None: + self.model.clear() + self._root_cache = self._batch_cache = None + self._row_cache = self._icl_cache = None + self._cache_identity = None + + def clear(self) -> None: + """Fence and release the fitted generation, keeping the model lease.""" + with self._call("clear"): + self.exchange.attempt(self._fence) + self._release() + + def close(self) -> None: + """Fence and release the lease; never destroy caller-owned groups.""" + if self._closed: + return + with self._call("close"): + self.exchange.attempt(self._fence) + self._release() + self.model._execution = None + self._closed = True + + +def tabular_execution( + model: KumoTabular, + *, + task: TaskLike, + group: dist.ProcessGroup, + control_group: dist.ProcessGroup, + column_staging_bytes: int = 8 << 30, + attention_workspace_bytes: int = 12 << 30, + fit_head_gather_window: int | None = None, +) -> TabularExecution: + """Borrow a loaded model for explicit original-wave execution. + + Args: + model: Caller-owned, unfitted FP32 evaluation KumoTabular model. + task: Selected task; classification must fit native class capacity. + group: Caller-owned tensor process group, one rank per device. + control_group: Caller-owned Gloo group with identical ordered members. + column_staging_bytes: Component bound for native column transport. + attention_workspace_bytes: Logical cached-attention scratch admission. + Neither component budget bounds total allocator/device memory. + fit_head_gather_window: Optional limit on fit head gathers enqueued + ahead of completed gathered outputs. ``None`` preserves the + original asynchronous enqueue behavior. + + All ranks construct the executor and call its methods in the same order:: + + run = tabular_execution(model, task="regression", group=group, + control_group=control_group) + run.fit(x if root else None, y if root else None) + outputs = run.predict_batches(batches if root else None) + run.close() + + Each batch keeps its original shape and public recipe postprocessing. + This first scope is eager FP32 with one unbatched estimator, no callbacks, + and no ECOC expansion beyond the native class capacity. Groups and model + loading are caller-owned; a healthy close permits ordinary model reuse. + """ + return TabularExecution( + model, + Task(task), + Exchange(group, control_group), + column_staging_bytes, + attention_workspace_bytes, + fit_head_gather_window, + ) diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 78734c8892..aef216db67 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -135,6 +135,7 @@ def forward( The processed prediction after applying ``recipe.output`` to the stacked estimator outputs with shape ``[E, ..., R_query, *]``. """ + self._check_execution("forward") callbacks = () if callbacks is None else callbacks requires_grad = self.training requires_grad |= any(callback.requires_grad for callback in callbacks) @@ -257,6 +258,7 @@ def fit( pre-processing and model execution. kwargs: Additional keyword arguments passed to the model. """ + self._check_execution("fit") if context_only_columns and not self.supports_context_only_columns: raise ValueError( f"{self.__class__.__name__!r} does not support context-only " @@ -321,7 +323,7 @@ def fit( class_values=class_values[batch], categorical_mask=categorical_mask, ) - self._forward( + self._call_prepared( x_context=context.x, y_context=context.y, x_query=None, @@ -374,6 +376,7 @@ def predict( The processed prediction after applying ``recipe.output`` to the stacked estimator outputs with shape ``[E, ..., R, *]``. """ + self._check_execution("predict") if self.training: raise RuntimeError( f"{self.__class__.__name__!r}.predict() does not support " @@ -526,6 +529,11 @@ def _forward( # Helpers ################################################################# + def _call_prepared(self, **kwargs: Any) -> TableTensor: + if self._execution is not None: + return self._execution._prepared_forward(self._forward, **kwargs) + return self._forward(**kwargs) + def _forward_members( self, contexts: Sequence[MemberContext], diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py index 1b48d1ecfd..d6c7fc73e2 100644 --- a/sdm/models/base/model.py +++ b/sdm/models/base/model.py @@ -2,8 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import abc -from collections.abc import Iterator, Sequence -from typing import ClassVar, Generic, Self, TypeVar +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import Any, ClassVar, Generic, Protocol, Self, TypeVar import torch from torch import Tensor @@ -13,6 +13,17 @@ from sdm.cache import Cache from sdm.tensor.mixin import DeviceMixin + +class _Execution(Protocol): + def _check_model_call(self, operation: str) -> None: ... + def _prepared_forward( + self, forward: Callable[..., Any], **kwargs: Any + ) -> Any: ... + def _run_core( + self, task: Any, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> Tensor: ... + + T = TypeVar("T", bound=Tensor | DeviceMixin) @@ -30,6 +41,7 @@ def __init__(self) -> None: super().__init__() self._cache: Cache | None = None + self._execution: _Execution | None = None self._transfer_streams: dict[torch.device, torch.cuda.Stream] = {} @classmethod @@ -39,8 +51,36 @@ def default_recipe(cls) -> Recipe: def clear(self) -> None: r"""Clear any cached state.""" + self._check_execution("clear") self._cache = None + def _check_execution(self, operation: str) -> None: + if self._execution is not None: + self._execution._check_model_call(operation) + + def train(self, mode: bool = True) -> Self: + """Set training mode when no executor owns this model.""" + self._check_execution("train") + return super().train(mode) + + def _apply( + self, fn: Callable[[Tensor], Tensor], recurse: bool = True + ) -> Self: + self._check_execution("_apply") + return super()._apply(fn, recurse=recurse) + + def load_state_dict( + self, + state_dict: Mapping[str, Any], + strict: bool = True, + assign: bool = False, + ) -> Any: + """Load weights when no executor owns this model.""" + self._check_execution("load_state_dict") + return super().load_state_dict( + state_dict, strict=strict, assign=assign + ) + # Helpers ################################################################# def _prefetch( @@ -60,6 +100,7 @@ def _prefetch( # Python Builtins ######################################################### def __getstate__(self) -> dict[str, object]: + self._check_execution("serialize") for stream in self._transfer_streams.values(): stream.synchronize() state = super().__getstate__() @@ -69,6 +110,7 @@ def __getstate__(self) -> dict[str, object]: def __setstate__(self, state: dict[str, object]) -> None: super().__setstate__(state) self._transfer_streams = {} + self._execution = None def __repr__(self) -> str: device = next(self.parameters()).device diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py index 067d5f6191..832c6b226f 100644 --- a/sdm/models/kumo/tabular/icl.py +++ b/sdm/models/kumo/tabular/icl.py @@ -3,6 +3,8 @@ # ruff: noqa: D101, D102 +from collections.abc import Callable +from functools import partial from typing import Any, cast import torch @@ -63,6 +65,7 @@ def forward( y: Tensor, # [..., R_train] *, cache: Cache | None = None, + _attention: Callable[..., Tensor] | None = None, ) -> Tensor: # [..., R_test, out_channels] R_train = y.size(-1) @@ -77,11 +80,16 @@ def forward( del y_emb for i, layer in enumerate(self.layers): + call = ( + layer + if _attention is None + else partial(layer, _attention=_attention) + ) cache_key = f"icl_block.layer{i}" last_layer = i == len(self.layers) - 1 if self.kv_heads is None or cache is not None: - result = layer( + result = call( query=x[..., R_train:, :] if last_layer else x, key_value=( cast(KVCacheEntry, cache[cache_key]) @@ -108,7 +116,7 @@ def forward( del result continue - x_context, (key, value) = layer( + x_context, (key, value) = call( query=x[..., :0, :] if last_layer else x[..., :R_train, :], key_value=x[..., :R_train, :], return_key_value=True, @@ -123,7 +131,7 @@ def forward( value=value[..., : self.kv_heads, :].contiguous(), ) del key, value - x_query = layer( + x_query = call( query=x[..., R_train:, :], key_value=key_value, out=None if torch.is_grad_enabled() else x[..., R_train:, :], diff --git a/sdm/models/kumo/tabular/model.py b/sdm/models/kumo/tabular/model.py index b946f963e3..d8a968ff8b 100644 --- a/sdm/models/kumo/tabular/model.py +++ b/sdm/models/kumo/tabular/model.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Iterable +from functools import partial from typing import Any, ClassVar, Literal, cast import torch @@ -209,7 +210,8 @@ def _forward( categorical_mask = categorical_mask.expand(*x.size()[:-2], -1) if classes is None: - out = self.models[Task.regression]( + out = self._call_core( + Task.regression, x=x, y=y, categorical_mask=categorical_mask, @@ -223,7 +225,9 @@ def _forward( ) out = self.ecoc( - model=self.models[Task.classification], + model=self.models[Task.classification] + if self._execution is None + else partial(self._call_core, Task.classification), x=x, y=y, num_classes=len(classes), @@ -245,6 +249,12 @@ def _forward( numerical=out, ) + def _call_core(self, task: Task, *args: Any, **kwargs: Any) -> Tensor: + model = self.models[task] + if self._execution is not None: + return self._execution._run_core(task, model, *args, **kwargs) + return model(*args, **kwargs) + class _KumoTabular(torch.nn.Module): def __init__( diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index 6714c5ce5e..94df1387d2 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -3,6 +3,8 @@ # ruff: noqa: D101, D102 +from collections.abc import Callable +from functools import partial from typing import Any, cast import torch @@ -114,6 +116,7 @@ def forward( *, cache: Cache | None = None, num_columns: Tensor | None = None, + _runner: Callable[..., Any] | None = None, ) -> Tensor: # [..., R, K * D] *B, R, C = x.size() @@ -125,10 +128,13 @@ def forward( if num_columns is None else (K + num_columns).unsqueeze(-1).to(torch.int32) ) + cell = self.cell_embedding + if _runner is not None: + cell = partial(_runner, cell) buffer: Tensor | None = None if torch.is_grad_enabled(): - x = self.cell_embedding( + x = cell( x=x, categorical_mask=categorical_mask, train_size=R_train, @@ -144,7 +150,7 @@ def forward( else x.dtype, ) buffer[..., :K, :] = self.readout_token.to(buffer.dtype) - x = self.cell_embedding( + x = cell( x=x, categorical_mask=categorical_mask, train_size=R_train, @@ -156,10 +162,16 @@ def forward( if y.numel() > 0: if self.y_emb is not None: - y_emb = self.y_emb(y).unsqueeze(-2) # [..., R_train, 1, D] + target = self.y_emb else: assert self.y_lin is not None - y_emb = self.y_lin(y.unsqueeze(-1)).unsqueeze(-2) + target = self.y_lin + target_call = ( + target if _runner is None else partial(_runner, target) + ) + y_emb = target_call( + y if self.y_emb is not None else y.unsqueeze(-1) + ).unsqueeze(-2) # [..., R_train, 1, D] x[..., :R_train, :, :] += y_emb.to(x.dtype) if buffer is not None: @@ -179,7 +191,13 @@ def forward( x = x.transpose(-2, -3) # [..., C, R, D] key = f"row_embedding.col_block{i}" - result = col_block( + col_call = ( + col_block if _runner is None else partial(_runner, col_block) + ) + row_call = ( + row_block if _runner is None else partial(_runner, row_block) + ) + result = col_call( query=x, key_value=cast(KVCacheEntry, cache[key]) if cache is not None and cache.is_replaying @@ -200,7 +218,7 @@ def forward( [readout_token.to(x.dtype), x.transpose(-2, -3)], dim=-2, ) - x = row_block( + x = row_call( query=x[..., :K, :] if i == len(self.row_blocks) - 1 else x, @@ -209,7 +227,7 @@ def forward( batch_size_limit="auto", ) else: - buffer = row_block( + buffer = row_call( query=buffer[..., :K, :] if i == len(self.row_blocks) - 1 else buffer, @@ -221,4 +239,5 @@ def forward( else buffer, ) - return self.norm(buffer if buffer is not None else x).flatten(-2) + norm = self.norm if _runner is None else partial(_runner, self.norm) + return norm(buffer if buffer is not None else x).flatten(-2) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index 3d02751198..9d72277bf7 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -4,6 +4,8 @@ """Attention modules for structured tensor models.""" import math +from collections.abc import Callable +from functools import partial from typing import Any, Literal, overload import torch @@ -70,6 +72,7 @@ def forward( attn_mask: Tensor | None = None, # [..., Hq, Q, KV] *, is_causal: bool = False, + _already_scaled: bool = False, ) -> Tensor: # [..., Q, Hq, C] r"""The forward pass. @@ -113,7 +116,7 @@ def forward( batch_shapes.append(attn_mask.size()[:-3]) batch_shape = torch.broadcast_shapes(*batch_shapes) - if self.query_scaling is not None: + if self.query_scaling is not None and not _already_scaled: if attn_mask is not None and attn_mask.dtype != torch.bool: raise ValueError( "Query scaling requires boolean attention mask" @@ -288,6 +291,7 @@ def forward( *, is_causal: bool = False, return_key_value: Literal[False] = False, + _attention: Callable[..., Tensor] | None = None, ) -> Tensor: ... @overload @@ -302,6 +306,7 @@ def forward( *, is_causal: bool = False, return_key_value: Literal[True], + _attention: Callable[..., Tensor] | None = None, ) -> tuple[Tensor, KVCacheEntry]: ... @overload @@ -316,6 +321,7 @@ def forward( *, is_causal: bool = False, return_key_value: bool, + _attention: Callable[..., Tensor] | None = None, ) -> Tensor | tuple[Tensor, KVCacheEntry]: ... def forward( @@ -329,6 +335,7 @@ def forward( *, is_causal: bool = False, return_key_value: bool = False, + _attention: Callable[..., Tensor] | None = None, ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] r"""The forward pass. @@ -421,7 +428,10 @@ def forward( else: key = self.key_transform(key, pos=key_pos) - out = self.sdpa( + attention = ( + self.sdpa if _attention is None else partial(_attention, self.sdpa) + ) + out = attention( query=query, # [..., Q, Hq, C // Hq] key=key, # [..., KV, Hkv, C // Hq] value=value, # [..., KV, Hkv, C // Hq] @@ -527,6 +537,7 @@ def forward( *, is_causal: bool = False, return_key_value: Literal[False] = False, + _attention: Callable[..., Tensor] | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, ) -> Tensor: ... @@ -543,6 +554,7 @@ def forward( *, is_causal: bool = False, return_key_value: Literal[True], + _attention: Callable[..., Tensor] | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, ) -> tuple[Tensor, KVCacheEntry]: ... @@ -559,6 +571,7 @@ def forward( *, is_causal: bool = False, return_key_value: bool, + _attention: Callable[..., Tensor] | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, ) -> Tensor | tuple[Tensor, KVCacheEntry]: ... @@ -574,6 +587,7 @@ def forward( *, is_causal: bool = False, return_key_value: bool = False, + _attention: Callable[..., Tensor] | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] @@ -624,8 +638,13 @@ def forward( "already cached" ) + forward = ( + self._forward + if _attention is None + else partial(self._forward, _attention=_attention) + ) if torch.is_grad_enabled(): - return self._forward( + return forward( query=query, key_value=key_value, seqused_key_value=seqused_key_value, @@ -680,7 +699,7 @@ def forward( disable_chunking = True if disable_chunking: - return self._forward( + return forward( query=query, key_value=key_value, seqused_key_value=seqused_key_value, @@ -702,7 +721,7 @@ def forward( for start in range(0, batch_size, batch_size_limit): end = min(start + batch_size_limit, batch_size) - result = self._forward( + result = forward( query=_chunk(query, batch_shape, 2, start, end), key_value=_chunk(key_value, batch_shape, 2, start, end), seqused_key_value=_chunk( @@ -779,6 +798,7 @@ def _forward( *, is_causal: bool = False, return_key_value: bool = False, + _attention: Callable[..., Tensor] | None = None, out: Tensor | None = None, ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] @@ -792,7 +812,12 @@ def _forward( elif isinstance(key_value, Tensor): key_value = self.key_value_norm(key_value) - result = self.attn( + attention = ( + self.attn + if _attention is None + else partial(self.attn, _attention=_attention) + ) + result = attention( query=query if self.query_norm is None else self.query_norm(query), key_value=key_value, seqused_key_value=seqused_key_value, diff --git a/test/execution/test_tabular.py b/test/execution/test_tabular.py new file mode 100644 index 0000000000..11de09f690 --- /dev/null +++ b/test/execution/test_tabular.py @@ -0,0 +1,733 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import subprocess +import sys +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +import sdm.processing as sp +from sdm import CategoricalTensor, Recipe, Stype, TableTensor, Task +from sdm.execution import tabular_execution +from sdm.models import KumoTabular +from sdm.models.kumo.tabular.model import _KumoTabular +from sdm.nn import SDPA, TransformerBlock + + +def fixture(task: str): + model = KumoTabular(task=task, pretrained=False, device="meta").eval() + core = _KumoTabular( + num_classes=10 if task == "classification" else 0, + num_quantiles=999, + cell_channels=8, + num_embedding_layers=2, + num_embedding_heads=2, + num_inducing_points=3, + num_readout_tokens=2, + group_size=3, + num_frequencies=3, + icl_channels=32, + num_icl_layers=3, + num_icl_heads=8, + num_icl_key_value_heads_for_query=2, + ).eval() + with torch.no_grad(): + for index, parameter in enumerate(core.parameters()): + values = torch.arange(parameter.numel()).reshape(parameter.shape) + parameter.copy_(0.09 * (values * 0.31 + index * 0.17).sin()) + model.models[Task(task)] = core + x = torch.arange(11 * 3, dtype=torch.float32).reshape(11, 3).sin() + query = torch.arange(7 * 3, dtype=torch.float32).reshape(7, 3).cos() + if task == "classification": + y = TableTensor( + columns={Stype.categorical: ["label"]}, + categorical=CategoricalTensor( + code=(torch.arange(11) % 3).view(-1, 1), + categories=(torch.tensor([8, 3, 9]),), + ), + ) + else: + y = torch.arange(11, dtype=torch.float32).sin().view(-1, 1) + return model, x, y, query + + +def journey( + rank: int, + world: int, + directory: str, + task: str, + chunked: bool = False, + fit_head_gather_window: int | None = None, +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=35), + ) + try: + model, x, y, query = fixture(task) + native = copy.deepcopy(model) + recipe = Recipe( + target=sp.StypeDispatch(numerical=sp.Standardize()), + output=[sp.AverageEstimators()], + ) + + def fit(target, *args, **kwargs): + # Exercise original fit windows; query retains its native AUTO. + with ( + patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=2 + ) + if chunked + else nullcontext() + ): + target.fit(*args, **kwargs) + + fit(native, x, y, recipe=recipe, num_estimators=1) + batches = [query[:2], query[2:4], query[4:6], query[6:]] + expected = [native.predict(batch) for batch in batches] + run = tabular_execution( + model, + task=task, + group=dist.group.WORLD, + control_group=dist.group.WORLD, + fit_head_gather_window=fit_head_gather_window, + ) + for _ in range(2): + fit( + run, + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe, + ) + actual = run.predict_batches(batches if rank == 0 else None) + if rank == 0: + for left, right in zip(actual, expected, strict=True): + assert left.columns == right.columns + torch.testing.assert_close( + left.numerical, right.numerical, atol=2e-6, rtol=1e-5 + ) + run.clear() + fit( + run, + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe, + ) + last = run.predict(query if rank == 0 else None) + run.close() + assert dist.is_initialized() + fit(model, x, y, recipe=recipe, num_estimators=1) + if rank == 0: + torch.testing.assert_close( + last.numerical, + model.predict(query).numerical, + atol=2e-6, + rtol=1e-5, + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2, 4]) +def test_public_journey(tmp_path: Path, task: str, world: int) -> None: + if world == 1: + journey(0, world, str(tmp_path), task) + else: + mp.spawn( + journey, args=(world, str(tmp_path), task), nprocs=world, join=True + ) + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +def test_chunked_native_windows(tmp_path: Path, task: str) -> None: + mp.spawn(journey, args=(2, str(tmp_path), task, True), nprocs=2, join=True) + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2, 4]) +@pytest.mark.parametrize("window", [1, 2, 4]) +def test_fit_head_gather_window_journey( + tmp_path: Path, task: str, world: int, window: int +) -> None: + if world == 1: + journey(0, world, str(tmp_path), task, False, window) + else: + mp.spawn( + journey, + args=(world, str(tmp_path), task, False, window), + nprocs=world, + join=True, + ) + + +@pytest.fixture +def group(tmp_path: Path) -> Iterator[dist.ProcessGroup]: + previous = torch.get_num_threads() + torch.set_num_threads(1) + dist.init_process_group( + "gloo", init_method=f"file://{tmp_path}/single", rank=0, world_size=1 + ) + try: + yield dist.group.WORLD + finally: + dist.destroy_process_group() + torch.set_num_threads(previous) + + +def execution(model: KumoTabular, group: dist.ProcessGroup): + return tabular_execution( + model, task="regression", group=group, control_group=group + ) + + +def test_fit_head_gather_window_defaults_to_none(group) -> None: + model, _, _, _ = fixture("regression") + run = execution(model, group) + assert run.fit_head_gather_window is None + run.close() + + +@pytest.mark.parametrize( + "window", + [True, False, 0, -1, 2.0, 2.5, float("nan"), float("inf"), "2"], +) +def test_invalid_fit_head_gather_window_before_lease(group, window) -> None: + model, x, y, query = fixture("regression") + with pytest.raises(ValueError, match="positive integer"): + tabular_execution( + model, + task="regression", + group=group, + control_group=group, + fit_head_gather_window=window, + ) + assert model._execution is None + model.fit(x, y, recipe=Recipe(), num_estimators=1) + assert model.predict(query).shape[-1] == 999 + + +def test_fit_head_gather_window_mutation_retires_first_error(group) -> None: + model, x, y, query = fixture("regression") + run = tabular_execution( + model, + task="regression", + group=group, + control_group=group, + fit_head_gather_window=2, + ) + run.fit(x, y, recipe=Recipe()) + cache = model._cache + run.fit_head_gather_window = 4 + with pytest.raises(RuntimeError, match="policy changed") as caught: + run.predict(query) + # Restoring the policy cannot release uncertain work or renew the lease. + run.fit_head_gather_window = 2 + for operation in (run.clear, run.close, lambda: run.fit(x, y)): + with pytest.raises(RuntimeError) as repeated: + operation() + assert repeated.value is caught.value + assert model._cache is cache + assert model._execution is run + + +def mismatched_fit_head_gather_window(rank: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/mismatch", + rank=rank, + world_size=2, + timeout=timedelta(seconds=15), + ) + try: + model, _, _, _ = fixture("regression") + with pytest.raises(ValueError, match="policies differ"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + fit_head_gather_window=rank + 1, + ) + assert model._execution is None + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_rank_fit_head_gather_window_mismatch(tmp_path: Path) -> None: + mp.spawn( + mismatched_fit_head_gather_window, + args=(str(tmp_path),), + nprocs=2, + join=True, + ) + + +@pytest.mark.parametrize("kind", ["pre", "post", "local", "compiled"]) +def test_unsupported_hooks_before_lease(group, kind: str) -> None: + model, x, y, query = fixture("regression") + handle = None + if kind == "pre": + handle = torch.nn.modules.module.register_module_forward_pre_hook( + lambda module, args: None + ) + elif kind == "post": + handle = torch.nn.modules.module.register_module_forward_hook( + lambda module, args, result: None + ) + elif kind == "local": + handle = model.register_forward_hook(lambda module, args, result: None) + else: + model._compiled_call_impl = lambda *args, **kwargs: None + try: + with pytest.raises(ValueError, match=r"hooks|unmodified"): + execution(model, group) + assert model._execution is None + finally: + if handle is not None: + handle.remove() + model._compiled_call_impl = None + model.fit(x, y, recipe=Recipe()) + assert model.predict(query).shape[-1] == 999 + + +def test_borrowed_model_mutations_rejected_before_change(group) -> None: + model, x, y, query = fixture("regression") + run = execution(model, group) + for operation in ( + model.clear, + model.train, + lambda: model.to(torch.float64), + lambda: model.fit(x, y), + lambda: model.predict(query), + lambda: model.load_state_dict(model.state_dict()), + lambda: copy.deepcopy(model), + ): + with pytest.raises(RuntimeError, match="owning tabular executor"): + operation() + assert not model.training + assert next(model.parameters()).dtype == torch.float32 + run.close() + model.train() + assert model.training + + +@pytest.mark.parametrize("mutation", ["parameter", "cache", "kv", "hook"]) +def test_generation_rejects_stale_state(group, mutation: str) -> None: + model, x, y, query = fixture("regression") + run = execution(model, group) + run.fit(x, y, recipe=Recipe()) + handle = None + if mutation == "parameter": + with torch.no_grad(): + next(model.parameters()).add_(1) + elif mutation == "cache": + model._cache = copy.copy(model._cache) + elif mutation == "kv": + with torch.no_grad(): + next(run._icl_cache._tensors()).add_(1) + else: + handle = model.register_forward_hook(lambda module, args, output: None) + try: + with pytest.raises( + (RuntimeError, ValueError), match=r"changed|hooks" + ) as caught: + run.predict(query) + for operation in (run.clear, run.close, lambda: run.fit(x, y)): + with pytest.raises(type(caught.value)) as again: + operation() + assert again.value is caught.value + finally: + if handle is not None: + handle.remove() + + +def test_failed_fence_retains_generation_and_first_error( + group, monkeypatch +) -> None: + model, x, y, query = fixture("regression") + run = execution(model, group) + run.fit(x, y, recipe=Recipe()) + cache, rows, icl = model._cache, run._row_cache, run._icl_cache + primary = RuntimeError("first failed completion fence") + + def fail() -> None: + raise primary + + monkeypatch.setattr(run, "_fence", fail) + with pytest.raises(RuntimeError, match="first failed") as caught: + run.clear() + assert caught.value is primary + assert model._cache is cache + assert run._row_cache is rows + assert run._icl_cache is icl + monkeypatch.setattr(run, "_fence", lambda: None) + for operation in (run.clear, run.close, lambda: run.predict(query)): + with pytest.raises(RuntimeError) as again: + operation() + assert again.value is primary + + +def test_default_recipe_rng_and_native_output(group) -> None: + model, x, y, query = fixture("regression") + native = copy.deepcopy(model) + initial = torch.Generator().manual_seed(42).get_state() + native_generator = torch.Generator().set_state(initial) + owner_generator = torch.Generator().set_state(initial) + native.fit(x, y, generator=native_generator, num_estimators=1) + expected = native.predict(query) + run = execution(model, group) + run.fit(x, y, generator=owner_generator) + actual = run.predict(query) + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + assert torch.equal( + native_generator.get_state(), owner_generator.get_state() + ) + assert actual.columns == expected.columns + run.close() + + +def test_two_models_have_independent_instance_dispatch(group) -> None: + first, x, y, query = fixture("regression") + second = copy.deepcopy(first) + original = type(first)._forward + a, b = execution(first, group), execution(second, group) + a.fit(x, y, recipe=Recipe()) + b.fit(x, y + 2, recipe=Recipe()) + expected = a.predict(query).clone() + b.close() + torch.testing.assert_close(a.predict(query).numerical, expected.numerical) + assert type(first)._forward is original + assert "forward" not in vars(first.models[Task.regression]) + a.close() + + +def test_import_has_no_research_or_cuda() -> None: + code = """ +import sys +import torch +def reject(*args, **kwargs): + raise AssertionError('CUDA initialized during execution import') +torch.cuda._lazy_init = reject +import sdm.execution +assert not any(name.startswith('research') for name in sys.modules) +assert not torch.cuda.is_initialized() +""" + result = subprocess.run( + [sys.executable, "-B", "-c", code], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_owner_thread_reentry_and_duplicate_lease(group) -> None: + model, x, y, query = fixture("regression") + run = execution(model, group) + with pytest.raises(ValueError, match="unleased"): + execution(model, group) + with ThreadPoolExecutor(max_workers=1) as thread: + future = thread.submit(run.clear) + with pytest.raises(RuntimeError, match="another thread"): + future.result() + with run._call("fit"), pytest.raises(RuntimeError, match="overlap"): + run.clear() + run.fit(x, y, recipe=Recipe()) + assert run.predict(query).shape[-1] == 999 + run.close() + + +@pytest.mark.parametrize("kind", ["missing", "grad", "output", "budget"]) +def test_unsupported_fit_retires_without_silent_fallback(group, kind) -> None: + model, x, y, query = fixture("regression") + run = tabular_execution( + model, + task="regression", + group=group, + control_group=group, + column_staging_bytes=1 if kind == "budget" else 8 << 30, + ) + if kind == "missing": + y = None + if kind == "grad": + x.requires_grad_() + recipe = ( + Recipe(output=sp.Callable(lambda value: value)) + if kind == "output" + else Recipe() + ) + with pytest.raises( + ValueError, match=r"inputs|processors|staging" + ) as caught: + run.fit(x, y, recipe=recipe) + with pytest.raises( + ValueError, match=r"inputs|processors|staging" + ) as again: + run.predict(query) + assert again.value is caught.value + + +def test_expanded_ecoc_is_rejected(group) -> None: + model, x, _, _ = fixture("classification") + y = TableTensor( + columns={Stype.categorical: ["label"]}, + categorical=CategoricalTensor( + code=torch.arange(11).view(-1, 1), + categories=(torch.arange(11),), + ), + ) + run = tabular_execution( + model, task="classification", group=group, control_group=group + ) + with pytest.raises(ValueError, match="ECOC"): + run.fit(x, y, recipe=Recipe()) + + +def test_autocast_and_policy_changes_rejected(group) -> None: + model, _, _, _ = fixture("regression") + with torch.autocast("cpu"), pytest.raises(ValueError, match="autocast"): + execution(model, group) + run = execution(model, group) + run.attention_workspace_bytes += 1 + with pytest.raises(RuntimeError, match="policy changed"): + run.clear() + + +@pytest.mark.parametrize("grad", [False, True]) +@pytest.mark.parametrize("limit", [None, 2]) +def test_attention_dispatch_preserves_native_forward_and_cache(grad, limit): + model, _, _, _ = fixture("regression") + block = model.models[Task.regression].icl_block.layers[0] + query = torch.arange(3 * 4 * 32).view(3, 4, 32).float().sin() + calls = [] + + def attention(sdpa: SDPA, **kwargs): + calls.append(sdpa) + return sdpa(**kwargs) + + with torch.set_grad_enabled(grad): + expected, expected_kv = block( + query=query.clone(), return_key_value=True, batch_size_limit=limit + ) + actual, actual_kv = block( + query=query.clone(), + return_key_value=True, + batch_size_limit=limit, + _attention=attention, + ) + assert calls + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + for left, right in zip(actual_kv, expected_kv, strict=True): + torch.testing.assert_close(left, right, atol=0, rtol=0) + + +def test_explicit_scaling_does_not_mutate_native_sdpa() -> None: + model, _, _, _ = fixture("regression") + sdpa = model.models[Task.regression].icl_block.layers[0].attn.sdpa + query = torch.arange(5 * 8 * 4).view(5, 8, 4).float().sin() + key, value = query.cos(), query.sin() + scaling = sdpa.query_scaling + with torch.no_grad(): + expected = sdpa(query, key, value) + scaled = scaling(query, key_len=key.size(-3)) + actual = sdpa(scaled, key, value, _already_scaled=True) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert sdpa.query_scaling is scaling + + +def test_native_serialization_without_execution_field() -> None: + model, x, y, query = fixture("regression") + state = model.__getstate__() + state.pop("_execution") + restored = object.__new__(KumoTabular) + restored.__setstate__(state) + restored.fit(x, y, recipe=Recipe()) + assert restored.predict(query).shape[-1] == 999 + + +def rejected_fit(rank: int, directory: str, failure: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/failed-group", + rank=rank, + world_size=2, + timeout=timedelta(seconds=15), + ) + try: + model, x, y, _ = fixture("regression") + run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + column_staging_bytes=1 if failure == "budget" else 8 << 30, + ) + original = ( + torch.Tensor.nanmean + if failure == "mean" + else torch.Tensor.clone + if failure == "row" + else torch.Tensor.contiguous + if failure == "root" + else torch.empty + ) + first = RuntimeError("rank-local allocation failed") + + def allocate(*args, **kwargs): + shape = ( + args[0].shape + if failure in ("mean", "row", "root") + else (args[0],) + if len(args) == 1 and isinstance(args[0], int) + else args[0] + if len(args) == 1 + else args + ) + target = ( + (11, 3) + if failure == "mean" + else (2, 5, 8) + if failure == "row" + else (3,) + if failure == "root" + else (6, 5, 8) + ) + failing_rank = 0 if failure == "root" else 1 + if rank == failing_rank and tuple(shape) == target: + raise first + return original(*args, **kwargs) + + injection = ( + patch.object(torch.Tensor, "nanmean", allocate) + if failure == "mean" + else patch.object(torch.Tensor, "clone", allocate) + if failure == "row" + else patch.object(torch.Tensor, "contiguous", allocate) + if failure == "root" + else patch.object(torch, "empty", allocate) + if failure == "buffer" + else nullcontext() + ) + with ( + injection, + patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=2 + ) + if failure == "row" + else nullcontext(), + pytest.raises( + (ValueError, RuntimeError), match=r"staging budget|rank-local" + ) as primary, + ): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=Recipe(), + ) + if rank == (0 if failure == "root" else 1) and failure != "budget": + assert primary.value is first + with pytest.raises( + (ValueError, RuntimeError), match=r"staging budget|rank-local" + ) as repeated: + run.close() + assert repeated.value is primary.value + # The executor must not destroy caller-owned groups even on failure. + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "failure", ["budget", "mean", "buffer", "row", "root"] +) +def test_collective_preallocation_failure( + tmp_path: Path, failure: str +) -> None: + mp.spawn(rejected_fit, args=(str(tmp_path), failure), nprocs=2, join=True) + + +def test_local_primary_survives_control_collective_failure(group, monkeypatch): + model, _, _, _ = fixture("regression") + run = execution(model, group) + first = RuntimeError("original allocation failure") + + def fail_operation(): + raise first + + def fail_collective(*args, **kwargs): + raise RuntimeError("secondary transport failure") + + monkeypatch.setattr(dist, "all_gather_object", fail_collective) + with pytest.raises(RuntimeError, match="original allocation") as caught: + run.exchange.attempt(fail_operation) + assert caught.value is first + with pytest.raises(RuntimeError, match="original allocation") as repeated: + run.exchange.attempt(lambda: None) + assert repeated.value is first + + +def rejected_query_packing(rank: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/query-packing", + rank=rank, + world_size=2, + timeout=timedelta(seconds=15), + ) + try: + model, x, y, query = fixture("regression") + run = execution(model, dist.group.WORLD) + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=Recipe(), + ) + mask = model._cache[0]["categorical_mask"] if rank == 0 else None + original = torch.Tensor.contiguous + failure = RuntimeError("root query packing allocation failed") + + def contiguous(value, *args, **kwargs): + if value is mask: + raise failure + return original(value, *args, **kwargs) + + with ( + patch.object(torch.Tensor, "contiguous", contiguous), + pytest.raises(RuntimeError, match="query packing") as caught, + ): + run.predict(query if rank == 0 else None) + if rank == 0: + assert caught.value is failure + with pytest.raises(RuntimeError, match="query packing") as repeated: + run.close() + assert caught.value is repeated.value + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_root_query_packing_failure_reaches_peers(tmp_path: Path) -> None: + mp.spawn( + rejected_query_packing, args=(str(tmp_path),), nprocs=2, join=True + ) diff --git a/test/execution/test_tabular_backpressure.py b/test/execution/test_tabular_backpressure.py new file mode 100644 index 0000000000..c03ef27e37 --- /dev/null +++ b/test/execution/test_tabular_backpressure.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fit enqueue ordering with fake CUDA events, not CUDA performance.""" + +from collections.abc import Callable +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest + +import sdm.execution._tabular_icl as icl + + +class Value: + is_cuda = True + device = "cuda:3" + shape = (8, 4, 4) + + def __init__(self, calls: list[str], empty: bool = False) -> None: + self.calls = calls + self.empty = empty + + def numel(self) -> int: + return 0 if self.empty else 32 + + def size(self, axis: int) -> int: + return 8 if axis == -2 else 4 + + def __getitem__(self, key: object) -> "Value": + return self + + def movedim(self, source: int, destination: int) -> "Value": + self.calls.append("movedim") + return self + + def contiguous(self) -> "Value": + self.calls.append("copy") + return self + + def new_empty(self, shape: tuple[int, ...]) -> "Value": + return Value(self.calls) + + +def queue( + monkeypatch: pytest.MonkeyPatch, + window: int | None, + world: int, + count: int, + fault: str | None = None, +) -> SimpleNamespace: + calls: list[str] = [] + events: list[object] = [] + outstanding: list[int] = [] + snapshots: list[tuple[str, int, tuple[int, ...]]] = [] + primary = RuntimeError("first fit enqueue failure") + + class Event: + def __init__(self) -> None: + self.index = len(events) + events.append(self) + + def record(self, stream: str) -> None: + assert calls[-1] == "copy" + assert stream == "stream:cuda:3" + if fault == "record": + raise primary + outstanding.append(self.index) + snapshots.append(("record", self.index, tuple(outstanding))) + + def synchronize(self) -> None: + if fault == "wait": + raise primary + assert outstanding.pop(0) == self.index + snapshots.append(("wait", self.index, tuple(outstanding))) + + def gather(received: Value, source: Value, *, group: object) -> None: + if window is not None: + assert len(outstanding) < window + calls.append("gather") + if fault == "gather": + raise primary + + class Attention: + query_scaling = None + + def __call__(self, *args: Value, **kwargs: object) -> Value: + return Value(calls) + + class Model: + def __call__( + self, + full: Value, + labels: Value, + *, + cache: object, + _attention: Callable[..., Value], + ) -> Value: + for _ in range(count): + _attention( + Attention(), + query=Value(calls), + key=Value(calls), + value=Value(calls), + ) + # Empty queries cannot issue a collective or queue an event. + _attention( + Attention(), + query=Value(calls, empty=True), + key=Value(calls), + value=Value(calls), + ) + return Value(calls) + + monkeypatch.setattr( + icl, + "torch", + SimpleNamespace( + cuda=SimpleNamespace( + Event=Event, + current_stream=lambda device: "stream:" + device, + ) + ), + ) + monkeypatch.setattr( + icl, "dist", SimpleNamespace(all_gather_into_tensor=gather) + ) + monkeypatch.setattr(icl, "sdpa_kernel", lambda backends: nullcontext()) + exchange = SimpleNamespace( + world=world, + rank=0, + group=object(), + gather=lambda *args: Value(calls), + attempt=lambda operation: operation(), + ) + + def run() -> Value: + return icl.fit_icl( + Model(), + Value(calls), + Value(calls), + SimpleNamespace(lengths=(4,) * world), + exchange, + window, + ) + + return SimpleNamespace( + run=run, + calls=calls, + events=events, + outstanding=outstanding, + snapshots=snapshots, + primary=primary, + ) + + +@pytest.mark.parametrize("window", [None, 1, 2, 4]) +@pytest.mark.parametrize("world", [1, 2]) +@pytest.mark.parametrize("count", [0, 1, 2, 3, 23]) +def test_fit_head_gather_fifo( + monkeypatch: pytest.MonkeyPatch, + window: int | None, + world: int, + count: int, +) -> None: + state = queue(monkeypatch, window, world, count) + state.run() + expected = count if world > 1 and window is not None else 0 + assert len(state.events) == expected + assert state.calls.count("gather") == (count if world > 1 else 0) + waits = [index for kind, index, _ in state.snapshots if kind == "wait"] + assert waits == list(range(max(0, expected - (window or expected)))) + assert len(state.outstanding) == min(expected, window or expected) + # The fit's outer completion fence owns the remaining tail; no extra + # per-layer wait or device-wide synchronization is required here. + + +@pytest.mark.parametrize("fault", ["gather", "record", "wait"]) +def test_fit_head_gather_stops_at_first_fault( + monkeypatch: pytest.MonkeyPatch, fault: str +) -> None: + state = queue(monkeypatch, window=2, world=2, count=23, fault=fault) + with pytest.raises(RuntimeError, match="first fit enqueue") as caught: + state.run() + assert caught.value is state.primary + assert state.calls.count("gather") == (2 if fault == "wait" else 1) + assert len(state.events) == {"gather": 0, "record": 1, "wait": 2}[fault] From 0483d818179985860400614b50f47d16feefce06 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:41:41 -0700 Subject: [PATCH 02/10] Propagate rank-local allocation failures before collectives --- sdm/execution/_tabular_icl.py | 4 +- sdm/execution/tabular.py | 12 +++-- test/execution/test_tabular.py | 97 ++++++++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 7 deletions(-) diff --git a/sdm/execution/_tabular_icl.py b/sdm/execution/_tabular_icl.py index 35fd09cae5..a602cfc740 100644 --- a/sdm/execution/_tabular_icl.py +++ b/sdm/execution/_tabular_icl.py @@ -79,7 +79,7 @@ def attention( ) ) dist.all_gather_into_tensor(received, source, group=exchange.group) - result = received.movedim(0, -2).contiguous() + result = exchange.attempt(lambda: received.movedim(0, -2).contiguous()) if fit_head_gather_window is not None and result.is_cuda: # The copy of the gathered heads is ordered after NCCL on this # stream. Bound host enqueue lead without a device-wide fence. @@ -137,7 +137,7 @@ def workspace_bytes(model: ICLBlock, rows: int, cache: OwnedCache) -> int: def predict_icl(model: ICLBlock, query: Tensor, cache: OwnedCache) -> Tensor: """Replay each full original query with one gathered KV layer at a time.""" - state = query.clone() + state = cache.exchange.attempt(query.clone) for index, layer in enumerate(model.layers): pair = cache[f"icl_block.layer{index}"] state = cache.exchange.attempt( diff --git a/sdm/execution/tabular.py b/sdm/execution/tabular.py index aef78c6c66..d760ac6830 100644 --- a/sdm/execution/tabular.py +++ b/sdm/execution/tabular.py @@ -667,11 +667,13 @@ def check_geometry() -> None: src=self.exchange.members[0], group=self.exchange.group, ) - empty = local.new_empty( - (0,), - dtype=torch.int64 - if self.task == Task.classification - else local.dtype, + empty = self.exchange.attempt( + lambda: local.new_empty( + (0,), + dtype=torch.int64 + if self.task == Task.classification + else local.dtype, + ) ) state = self.exchange.attempt( lambda: self._core.row_embedding( diff --git a/test/execution/test_tabular.py b/test/execution/test_tabular.py index 11de09f690..67ea26e878 100644 --- a/test/execution/test_tabular.py +++ b/test/execution/test_tabular.py @@ -16,6 +16,8 @@ import torch.distributed as dist import torch.multiprocessing as mp +import sdm.execution._tabular_icl as icl +import sdm.execution.tabular as tabular import sdm.processing as sp from sdm import CategoricalTensor, Recipe, Stype, TableTensor, Task from sdm.execution import tabular_execution @@ -731,3 +733,98 @@ def test_root_query_packing_failure_reaches_peers(tmp_path: Path) -> None: mp.spawn( rejected_query_packing, args=(str(tmp_path),), nprocs=2, join=True ) + + +def rejected_rank_allocation(rank: int, directory: str, stage: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/rank-allocation", + rank=rank, + world_size=2, + timeout=timedelta(seconds=15), + ) + try: + model, x, y, query = fixture("regression") + run = execution(model, dist.group.WORLD) + failure = RuntimeError(f"rank-local {stage} allocation failed") + original_clone = torch.Tensor.clone + original_empty = torch.Tensor.new_empty + original_contiguous = torch.Tensor.contiguous + original_predict = icl.predict_icl + + def inject_clone(block, state, cache): + def clone(value, *args, **kwargs): + if rank == 0 and value is state: + raise failure + return original_clone(value, *args, **kwargs) + + with patch.object(torch.Tensor, "clone", clone): + return original_predict(block, state, cache) + + def new_empty(value, *args, **kwargs): + caller = sys._getframe(1).f_code.co_filename + if ( + rank == 0 + and caller.endswith("/tabular.py") + and args == ((0,),) + ): + raise failure + return original_empty(value, *args, **kwargs) + + def contiguous(value, *args, **kwargs): + caller = sys._getframe(1).f_code.co_filename + if ( + rank == 0 + and caller.endswith("/_tabular_icl.py") + and tuple(value.shape) == (11, 8, 4) + ): + raise failure + return original_contiguous(value, *args, **kwargs) + + def fit(): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=Recipe(), + ) + + if stage != "fit-gather": + fit() + scope = ( + patch.object(tabular, "predict_icl", inject_clone) + if stage == "query-clone" + else patch.object(torch.Tensor, "new_empty", new_empty) + if stage == "empty-labels" + else patch.object(torch.Tensor, "contiguous", contiguous) + ) + operation = ( + fit + if stage == "fit-gather" + else lambda: run.predict(query if rank == 0 else None) + ) + with scope, pytest.raises(RuntimeError, match=stage) as caught: + operation() + if rank == 0: + assert caught.value is failure + with pytest.raises(RuntimeError, match=stage) as repeated: + run.close() + assert caught.value is repeated.value + # A peer timeout is not acceptable error propagation. + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "stage", ["query-clone", "empty-labels", "fit-gather"] +) +def test_rank_local_allocation_reaches_peers( + tmp_path: Path, stage: str +) -> None: + mp.spawn( + rejected_rank_allocation, + args=(str(tmp_path), stage), + nprocs=2, + join=True, + ) From 30efb7bd3b067cc81c1fc4f758d68aa07ecd4bd3 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:58:14 -0700 Subject: [PATCH 03/10] Extend context sharding with local row workspaces and FP16 Keep default FP32 execution while adding explicit owner-local frontend and precision controls. Bound row workspaces per owner, preserve native FP16 inducing state, and harden fitted-state lifecycle checks. --- sdm/execution/_tabular_columns.py | 4 +- sdm/execution/_tabular_icl.py | 25 +- sdm/execution/_tabular_rows.py | 94 ++++- sdm/execution/tabular.py | 89 +++- test/execution/test_lifecycle_hardening.py | 335 +++++++++++++++ test/execution/test_owner_local_frontend.py | 383 ++++++++++++++++++ test/execution/test_sharded_row_workspace.py | 283 +++++++++++++ .../test_sharded_row_workspace_p8.py | 101 +++++ test/execution/test_tabular.py | 72 +++- test/execution/test_tabular_backend_policy.py | 103 +++++ test/execution/test_tabular_precision.py | 345 ++++++++++++++++ .../test_tabular_precision_owners.py | 273 +++++++++++++ .../test_tabular_precision_protocol.py | 127 ++++++ 13 files changed, 2201 insertions(+), 33 deletions(-) create mode 100644 test/execution/test_lifecycle_hardening.py create mode 100644 test/execution/test_owner_local_frontend.py create mode 100644 test/execution/test_sharded_row_workspace.py create mode 100644 test/execution/test_sharded_row_workspace_p8.py create mode 100644 test/execution/test_tabular_backend_policy.py create mode 100644 test/execution/test_tabular_precision.py create mode 100644 test/execution/test_tabular_precision_owners.py create mode 100644 test/execution/test_tabular_precision_protocol.py diff --git a/sdm/execution/_tabular_columns.py b/sdm/execution/_tabular_columns.py index 5c6042779a..9add05f90d 100644 --- a/sdm/execution/_tabular_columns.py +++ b/sdm/execution/_tabular_columns.py @@ -228,7 +228,9 @@ def forward( def compute(fulls: dict[Domain, Tensor]) -> None: for domain, full in fulls.items(): - induced = full.new_empty( + # Native inducing queries and their residual remain FP32 + # under autocast, unlike the expanded cell buffers. + induced = points.new_empty( domain.stop - domain.start, points.size(0), channels ) for start in range(domain.start, domain.stop, iw): diff --git a/sdm/execution/_tabular_icl.py b/sdm/execution/_tabular_icl.py index a602cfc740..86aaabb573 100644 --- a/sdm/execution/_tabular_icl.py +++ b/sdm/execution/_tabular_icl.py @@ -88,21 +88,26 @@ def attention( pending.append(done) return result - # Fit uses the same non-quadratic native backend policy as the accepted - # path. Cached GQA prediction deliberately retains native backend choice. - with ( - sdpa_kernel( - [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] - ) - if full.is_cuda - else nullcontext() + # Preserve the original FP32 policy; FP16 also permits native cuDNN. + # Never enable quadratic math fit attention. Cached GQA prediction + # deliberately retains native backend choice. + backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] + if ( + full.is_cuda + and torch.is_autocast_enabled(full.device.type) + and torch.get_autocast_dtype(full.device.type) == torch.float16 ): + backends.append(SDPBackend.CUDNN_ATTENTION) + with sdpa_kernel(backends) if full.is_cuda else nullcontext(): return model(full, labels, cache=cache, _attention=attention) def workspace_bytes(model: ICLBlock, rows: int, cache: OwnedCache) -> int: """Conservative layer-local logical scratch, not total device memory.""" context = sum(cache.lengths) + # Actual cache dtype controls transport; preserve conservative FP32 + # attention scratch because native math attention can promote inputs. + cache_element_size = next(cache._tensors()).element_size() largest = 0 for module in model.layers: layer = cast(TransformerBlock, module) @@ -127,9 +132,9 @@ def workspace_bytes(model: ICLBlock, rows: int, cache: OwnedCache) -> int: * kv_heads * 2 * width - * 4 + * cache_element_size if len(cache.lengths) > 1 - else compact * 4 + else compact * cache_element_size ) largest = max(largest, attention + gather) return largest diff --git a/sdm/execution/_tabular_rows.py b/sdm/execution/_tabular_rows.py index 62b3b82c86..ef5df1bf6e 100644 --- a/sdm/execution/_tabular_rows.py +++ b/sdm/execution/_tabular_rows.py @@ -37,14 +37,25 @@ def __init__( mask: Tensor, exchange: Exchange, staging_bytes: int, + *, + owner_local_frontend: bool = False, + shard_row_workspace: bool = False, ) -> None: self.encoder, self.raw, self.labels = encoder, raw, labels self.exchange, self.mask = exchange, mask + self.owner_local_frontend = owner_local_frontend + self.shard_row_workspace = shard_row_workspace self.total, self.columns = raw.shape self.lengths = bounds(self.total, exchange.world) self.start = sum(self.lengths[: exchange.rank]) self.stop = self.start + self.lengths[exchange.rank] self.readouts, self.channels = encoder.readout_token.shape + self.dtype = ( + torch.get_autocast_dtype(raw.device.type) + if torch.is_autocast_enabled(raw.device.type) + else raw.dtype + ) + element_size = self.dtype.itemsize self.mean = raw.nanmean(dim=-2, keepdim=True) self.mean = self.mean.masked_fill(self.mean.isnan(), 0) self.cell_rows = self.total @@ -53,10 +64,12 @@ def __init__( groups, frequencies = cell.group_size, cell.num_frequencies fixed = self.columns * ( groups * frequencies * 4 - + groups * self.channels * 2 * frequencies * 4 - + self.channels * 4 + + groups * self.channels * 2 * frequencies * element_size + + self.channels * element_size ) - item = (2 * groups * frequencies + 2 * self.channels) * 4 + item = ( + 2 * groups * frequencies + 2 * self.channels + ) * element_size items = max((chunk_memory_limit(raw.device) - fixed) // item, 1) self.cell_rows = max(items // self.columns, 1) self.handlers: dict[torch.nn.Module, Callable[..., Any]] = { @@ -74,13 +87,13 @@ def __init__( points = column.inducing_points.size(0) iw = min( column.inducing_block.auto_batch_size_limit( - raw.device, 4, points, self.total + raw.device, element_size, points, self.total ), 65_535, ) ow = min( column.output_block.auto_batch_size_limit( - raw.device, 4, self.total, points + raw.device, element_size, self.total, points ), 65_535, ) @@ -91,15 +104,17 @@ def __init__( owned = sum( d.stop - d.start for d in domains if d.owner == peer ) - sent = length * width * self.channels * 4 - received = self.total * owned * self.channels * 4 + sent = length * width * self.channels * element_size + received = ( + self.total * owned * self.channels * element_size + ) size = max(2 * (sent + received), 3 * received) if (self.columns <= iw or self.columns <= ow) and owned: size += ( self.total * (self.columns + self.readouts) * self.channels - * 4 + * element_size ) peak = max(peak, size) self.handlers[column] = _column_forward( @@ -117,7 +132,10 @@ def __init__( tokens = self.readouts if last else self.columns + self.readouts limit = min( block.auto_batch_size_limit( - raw.device, 4, tokens, self.columns + self.readouts + raw.device, + element_size, + tokens, + self.columns + self.readouts, ), 65_535, ) @@ -127,7 +145,13 @@ def __init__( raise ValueError( "Column staging budget is too small for native windows" ) - plan = (self.cell_rows, limits, row_limits) + plan = ( + self.dtype, + self.cell_rows, + limits, + row_limits, + shard_row_workspace, + ) if any(other != plan for other in exchange.records(plan)): raise ValueError("Native AUTO windows differ between owners") @@ -147,11 +171,33 @@ def cell( ) -> Tensor: del x, kwargs cache["cell_impute_mean"] = self.mean + if self.owner_local_frontend: + self.exchange.attempt( + lambda: self._local_cells(categorical_mask, out) + ) + return out for index, start in enumerate(range(0, self.total, self.cell_rows)): stop = min(start + self.cell_rows, self.total) self._cell_window(start, stop, index, categorical_mask, out) return out + def _local_cells(self, mask: Tensor, out: Tensor) -> None: + for start in range(0, self.total, self.cell_rows): + left = max(start, self.start) + right = min(start + self.cell_rows, self.total, self.stop) + if left >= right: + continue + raw = self.raw[left:right] + missing = raw.isnan() + CellEmbedding.forward( + self.encoder.cell_embedding, + torch.where(missing, self.mean, raw), + mask, + batch_size_limit=None, + out=out[left - self.start : right - self.start], + missing=missing, + ) + def _cell_window( self, start: int, stop: int, index: int, mask: Tensor, out: Tensor ) -> None: @@ -163,7 +209,10 @@ def allocate() -> Tensor | None: if interior and self.exchange.rank != owner: return None buffer = self.raw.new_empty( - stop - start, self.readouts + self.columns, self.channels + stop - start, + self.readouts + self.columns, + self.channels, + dtype=self.dtype, ) buffer[:, : self.readouts] = self.encoder.readout_token return buffer @@ -221,6 +270,29 @@ def row( **kwargs: Any, ) -> Tensor: del kwargs + if self.owner_local_frontend: + local_limit = ( + max(1, limit // self.exchange.world) + if self.shard_row_workspace + else limit + ) + + def local_windows() -> None: + for start in range(0, self.total, limit): + left = max(start, self.start) + right = min(start + limit, self.total, self.stop) + if left >= right: + continue + local = key_value[left - self.start : right - self.start] + block( + query=local[:, :tokens], + key_value=local, + out=out[left - self.start : right - self.start], + batch_size_limit=local_limit, + ) + + self.exchange.attempt(local_windows) + return out interior, boundary = [], [] for index, start in enumerate(range(0, self.total, limit)): stop = min(start + limit, self.total) diff --git a/sdm/execution/tabular.py b/sdm/execution/tabular.py index d760ac6830..dbfce356a8 100644 --- a/sdm/execution/tabular.py +++ b/sdm/execution/tabular.py @@ -5,9 +5,9 @@ import threading from collections.abc import Callable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager, nullcontext from types import MethodType -from typing import Any, cast +from typing import Any, Literal, cast import torch import torch.distributed as dist @@ -99,7 +99,7 @@ def _check_output(recipe: Recipe) -> None: class TabularExecution: - """Run an explicitly owned, eager FP32/E1 tabular generation. + """Run an explicitly owned, eager E1 tabular generation. Construct with :func:`tabular_execution`. All ranks call every operation; only group rank zero supplies inputs and receives public outputs. Groups, @@ -124,6 +124,9 @@ def __init__( column_staging_bytes: int, attention_workspace_bytes: int, fit_head_gather_window: int | None = None, + owner_local_frontend: bool = False, + shard_row_workspace: bool = False, + precision: Literal["float32", "float16"] = "float32", ) -> None: self.model, self.task, self.exchange = model, task, exchange self._core = cast(_KumoTabular, model.models[task]) @@ -131,6 +134,9 @@ def __init__( self.column_staging_bytes = column_staging_bytes self.attention_workspace_bytes = attention_workspace_bytes self.fit_head_gather_window = fit_head_gather_window + self.owner_local_frontend = owner_local_frontend + self.shard_row_workspace = shard_row_workspace + self.precision = precision self._lock = threading.Lock() self._owner_thread = threading.get_ident() self._thread: int | None = None @@ -153,6 +159,9 @@ def __init__( column_staging_bytes, attention_workspace_bytes, fit_head_gather_window, + owner_local_frontend, + shard_row_workspace, + precision, signature, ) ) @@ -174,9 +183,44 @@ def _policy(self) -> tuple[Any, ...]: self.column_staging_bytes, self.attention_workspace_bytes, self.fit_head_gather_window, + self.owner_local_frontend, + self.shard_row_workspace, + self.precision, + ) + + def _check_precision(self) -> None: + if self.precision not in ("float32", "float16"): + raise ValueError("Execution precision must be float32 or float16") + if torch.is_autocast_enabled(self.device.type): + if self.precision == "float32": + raise ValueError("Execution does not support autocast") + if torch.get_autocast_dtype(self.device.type) != torch.float16: + raise ValueError( + "Ambient autocast differs from execution precision" + ) + + def _precision_scope(self) -> AbstractContextManager: + # Every worker owns the same scope, even without caller autocast. + # Original recipe and public-processing islands disable it. + return ( + nullcontext() + if self.precision == "float32" + else torch.autocast( + self.device.type, dtype=torch.float16, cache_enabled=True + ) ) def _admit(self) -> tuple[Any, ...]: + if torch.is_inference_mode_enabled(): + raise ValueError("Execution does not support inference_mode") + if self.shard_row_workspace and ( + not self.owner_local_frontend + or self.exchange.world not in (1, 2, 4, 8) + ): + raise ValueError( + "Sharded row workspace requires owner-local frontend " + "and P1/2/4/8" + ) if self.fit_head_gather_window is not None and ( type(self.fit_head_gather_window) is not int or self.fit_head_gather_window < 1 @@ -204,8 +248,7 @@ def _admit(self) -> tuple[Any, ...]: raise ValueError( "Execution requires one FP32 device per task core" ) - if torch.is_autocast_enabled(self.device.type): - raise ValueError("Execution does not support autocast") + self._check_precision() if any( cast(TransformerBlock, layer).attn.num_query_heads % self.exchange.world @@ -267,6 +310,12 @@ def _call(self, operation: str) -> Iterator[None]: try: def validate() -> None: + if operation in ("fit", "predict") and ( + torch.is_inference_mode_enabled() + ): + raise ValueError( + "Execution does not support inference_mode" + ) if ( self.model._execution is not self or _model_identity(self.model) != self._identity @@ -274,8 +323,7 @@ def validate() -> None: raise RuntimeError( "The borrowed model changed; this execution is stale" ) - if torch.is_autocast_enabled(self.device.type): - raise ValueError("Execution does not support autocast") + self._check_precision() if self._policy() != self._policy_identity: raise RuntimeError( "Execution policy changed after admission" @@ -289,7 +337,8 @@ def validate() -> None: ) self.exchange.attempt(validate) - yield + with self._precision_scope(): + yield if operation in ("fit", "predict"): self.exchange.attempt(self._fence) except BaseException as error: @@ -365,6 +414,8 @@ def _fit_core(self, x: Tensor, y: Tensor, mask: Tensor) -> None: mask, exchange, self.column_staging_bytes, + owner_local_frontend=self.owner_local_frontend, + shard_row_workspace=self.shard_row_workspace, ) ) self._row_cache = OwnedCache( @@ -391,6 +442,9 @@ def _fit_core(self, x: Tensor, y: Tensor, mask: Tensor) -> None: ) self._row_cache.freeze() self._icl_cache.freeze() + # Successful work no longer needs callbacks that keep Rows in a cycle. + # Failures retain the original traceback/owners instead of cleaning up. + rows.handlers.clear() def fit( self, @@ -740,6 +794,9 @@ def tabular_execution( column_staging_bytes: int = 8 << 30, attention_workspace_bytes: int = 12 << 30, fit_head_gather_window: int | None = None, + owner_local_frontend: bool = False, + shard_row_workspace: bool = False, + precision: Literal["float32", "float16"] = "float32", ) -> TabularExecution: """Borrow a loaded model for explicit original-wave execution. @@ -754,6 +811,16 @@ def tabular_execution( fit_head_gather_window: Optional limit on fit head gathers enqueued ahead of completed gathered outputs. ``None`` preserves the original asynchronous enqueue behavior. + owner_local_frontend: Experimental local intersections of native cell + and row windows. Changes projection geometry; trained CUDA + numerical equivalence and capacity require separate evidence. + shard_row_workspace: Further bound owner-local row-block calls to + max(1, native row limit // owners). Requires owner-local frontend + and P1/2/4/8. Default False retains the native row-window budget. + precision: Neural execution precision. Default float32 rejects + ambient autocast. Explicit float16 owns matching autocast scopes + on every rank; weights, prepared inputs and original processing + islands remain FP32. This is not half-stored model support. All ranks construct the executor and call its methods in the same order:: @@ -764,7 +831,8 @@ def tabular_execution( run.close() Each batch keeps its original shape and public recipe postprocessing. - This first scope is eager FP32 with one unbatched estimator, no callbacks, + This scope is eager with one unbatched estimator, no callbacks, + no outer ``torch.inference_mode()`` (tensor version fences are required), and no ECOC expansion beyond the native class capacity. Groups and model loading are caller-owned; a healthy close permits ordinary model reuse. """ @@ -775,4 +843,7 @@ def tabular_execution( column_staging_bytes, attention_workspace_bytes, fit_head_gather_window, + owner_local_frontend, + shard_row_workspace, + precision, ) diff --git a/test/execution/test_lifecycle_hardening.py b/test/execution/test_lifecycle_hardening.py new file mode 100644 index 0000000000..e28bbcb989 --- /dev/null +++ b/test/execution/test_lifecycle_hardening.py @@ -0,0 +1,335 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Successful cache release and early rejection of versionless inference.""" + +import copy +import gc +import weakref +from contextlib import nullcontext +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture + +import sdm.execution.tabular as tabular +import sdm.processing as sp +from sdm import Recipe +from sdm.execution import tabular_execution + + +def start(rank: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=2, + timeout=timedelta(seconds=20), + ) + + +def execution(model, task: str, flag: bool): + return tabular_execution( + model, + task=task, + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=flag, + shard_row_workspace=flag, + ) + + +def recipe() -> Recipe: + return Recipe( + target=sp.StypeDispatch(numerical=sp.Standardize()), + output=[sp.AverageEstimators()], + ) + + +def cache_refs(run): + # Owned entries only: replay getters would allocate gathered temporaries. + return [ + weakref.ref(tensor) + for cache in (run._row_cache, run._icl_cache) + for tensor in cache._tensors() + ] + + +def release_journey(rank: int, directory: str, task: str, flag: bool) -> None: + start(rank, directory) + enabled = gc.isenabled() + gc.collect() + gc.disable() + try: + model, x, y, query = fixture(task) + native = copy.deepcopy(model) + native.fit(x, y, recipe=recipe(), num_estimators=1) + expected = native.predict(query) + run = execution(model, task, flag) + + def fit() -> None: + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe() if rank == 0 else None, + ) + + fit() + refs = cache_refs(run) + actual = run.predict(query if rank == 0 else None) + if rank == 0: + assert actual.columns == expected.columns + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + if task == "regression": + assert actual.numerical.shape == (7, 999) + else: + assert torch.equal( + actual.numerical.argmax(-1), + expected.numerical.argmax(-1), + ) + del actual + fit() + assert all(ref() is None for ref in refs) + refs = cache_refs(run) + run.clear() + assert all(ref() is None for ref in refs) + assert run._row_cache is run._icl_cache is None + assert model._execution is run + fit() + refs = cache_refs(run) + run.close() + assert all(ref() is None for ref in refs) + assert model._execution is None + run.close() + # A healthy close leaves native model reuse and borrowed groups intact. + model.fit(x, y, recipe=recipe(), num_estimators=1) + reused = model.predict(query) + torch.testing.assert_close( + reused.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + dist.barrier() + finally: + if enabled: + gc.enable() + dist.destroy_process_group() + + +def inference_journey( + rank: int, directory: str, operation: str, flag: bool +) -> None: + start(rank, directory) + try: + model, x, y, query = fixture("regression") + if operation == "constructor": + # Only one rank enables inference mode; both must reject before + # acquiring a model lease or entering fit/core/cache allocation. + with ( + torch.inference_mode() if rank == 1 else nullcontext(), + pytest.raises( + (ValueError, RuntimeError), match="inference_mode" + ), + ): + execution(model, "regression", flag) + assert model._execution is None + assert model._cache is None + run = execution(model, "regression", flag) + run.close() + dist.barrier() + return + + run = execution(model, "regression", flag) + if operation != "fit": + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe() if rank == 0 else None, + ) + old_identity = run._cache_identity + old_cache = model._cache, run._row_cache, run._icl_cache + + def invoke() -> None: + if operation in ("fit", "refit"): + run.fit(x if rank == 0 else None, y if rank == 0 else None) + elif operation == "predict": + run.predict(query if rank == 0 else None) + else: + getattr(run, operation)() + + with ( + patch.object( + tabular, + "Rows", + side_effect=AssertionError("Unexpected frontend allocation"), + ) as rows, + torch.inference_mode() if rank == 1 else nullcontext(), + pytest.raises( + (ValueError, RuntimeError), match="inference_mode" + ) as caught, + ): + invoke() + assert rows.call_count == 0 + assert run._cache_identity is old_identity + assert all( + actual is original + for actual, original in zip( + (model._cache, run._row_cache, run._icl_cache), old_cache + ) + ) + if operation == "fit": + assert not run._fit_started + assert old_cache == (None, None, None) + assert model._execution is run + for repeat in (run.clear, run.close, invoke): + with pytest.raises(type(caught.value)) as again: + repeat() + assert again.value is caught.value + dist.barrier() + finally: + dist.destroy_process_group() + + +def failed_journey(rank: int, directory: str, point: str, flag: bool) -> None: + start(rank, directory) + try: + model, x, y, _ = fixture("regression") + run = execution(model, "regression", flag) + original_call = tabular.Rows.__call__ + original_icl = tabular.fit_icl + rows_refs = [] + first = RuntimeError("rank-local submitted fit failure") + + def row_call(rows, module, *args, **kwargs): + rows_refs.append(weakref.ref(rows)) + if point == "frontend" and rank == 1: + raise first + return original_call(rows, module, *args, **kwargs) + + def icl(*args, **kwargs): + if point == "icl": + + def fail() -> None: + if rank == 1: + raise first + + run.exchange.attempt(fail) + return original_icl(*args, **kwargs) + + with ( + patch.object(tabular.Rows, "__call__", row_call), + patch.object(tabular, "fit_icl", icl), + pytest.raises( + RuntimeError, match="rank-local submitted fit failure" + ) as caught, + ): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe() if rank == 0 else None, + ) + if rank == 1: + assert caught.value is first + assert rows_refs + gc.collect() + # Do not break cycles or discard owners when fit did not complete. + assert all(ref() is not None and ref().handlers for ref in rows_refs) + assert model._execution is run + for repeat in (run.clear, run.close): + with pytest.raises(RuntimeError) as again: + repeat() + assert again.value is caught.value + dist.barrier() + finally: + dist.destroy_process_group() + + +def inference_cleanup_journey(rank: int, directory: str, flag: bool) -> None: + start(rank, directory) + try: + model, x, y, query = fixture("regression") + native = copy.deepcopy(model) + native.fit(x, y, recipe=recipe(), num_estimators=1) + expected = native.predict(query) + run = execution(model, "regression", flag) + + def fit() -> None: + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe() if rank == 0 else None, + ) + + fit() + with torch.inference_mode() if rank == 1 else nullcontext(): + run.clear() + assert run._failure is None + assert run._cache_identity is None + assert run._row_cache is run._icl_cache is None + assert model._execution is run + fit() + actual = run.predict(query if rank == 0 else None) + if rank == 0: + assert actual.numerical.shape == (7, 999) + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + with torch.inference_mode() if rank == 1 else nullcontext(): + run.close() + run.close() + assert run._failure is None + assert model._execution is None + assert run._row_cache is run._icl_cache is None + model.fit(x, y, recipe=recipe(), num_estimators=1) + reused = model.predict(query) + torch.testing.assert_close( + reused.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("flag", [False, True]) +@pytest.mark.parametrize("task", ["classification", "regression"]) +def test_immediate_cache_release(tmp_path: Path, task: str, flag: bool): + mp.spawn( + release_journey, args=(str(tmp_path), task, flag), nprocs=2, join=True + ) + + +@pytest.mark.parametrize("flag", [False, True]) +@pytest.mark.parametrize( + "operation", ["constructor", "fit", "refit", "predict"] +) +def test_early_inference_rejection(tmp_path: Path, operation: str, flag: bool): + mp.spawn( + inference_journey, + args=(str(tmp_path), operation, flag), + nprocs=2, + join=True, + ) + + +@pytest.mark.parametrize("flag", [False, True]) +@pytest.mark.parametrize("point", ["frontend", "icl"]) +def test_failed_rows_remain_owned(tmp_path: Path, point: str, flag: bool): + mp.spawn( + failed_journey, args=(str(tmp_path), point, flag), nprocs=2, join=True + ) + + +@pytest.mark.parametrize("flag", [False, True]) +def test_healthy_cleanup_under_inference_mode(tmp_path: Path, flag: bool): + mp.spawn( + inference_cleanup_journey, + args=(str(tmp_path), flag), + nprocs=2, + join=True, + ) diff --git a/test/execution/test_owner_local_frontend.py b/test/execution/test_owner_local_frontend.py new file mode 100644 index 0000000000..a07c05ef54 --- /dev/null +++ b/test/execution/test_owner_local_frontend.py @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Owner-local frontend output, window-boundary and failure controls.""" + +from contextlib import ExitStack +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture, journey + +from sdm import Task +from sdm.cache import Cache +from sdm.execution import tabular_execution +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_rows import Rows +from sdm.models.tabfm.cell_embedding import CellEmbedding + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2, 4]) +@pytest.mark.parametrize("count", [1, 11]) +def test_owner_local_public_lifecycle( + tmp_path: Path, task: str, world: int, count: int +) -> None: + args = (world, str(tmp_path), task, count > 1, None, True, count) + if world == 1: + journey(0, *args) + else: + mp.spawn(journey, args=args, nprocs=world, join=True) + + +def component_worker(rank: int, world: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=25), + ) + try: + with torch.no_grad(): + for count in (1, 11): + model, _, _, _ = fixture("regression") + encoder = model.models[Task.regression].row_embedding + raw = ( + torch.arange(count * 3, dtype=torch.float32) + .reshape(count, 3) + .sin() + ) + raw[::2, 0] = torch.nan + raw[:, 2] = torch.nan + mask = torch.tensor([False, True, False]) + labels = torch.arange(count, dtype=torch.float32) + baseline = Rows( + encoder, + raw, + labels, + mask, + Exchange(dist.group.WORLD, dist.group.WORLD), + 8 << 30, + ) + candidate = Rows( + encoder, + raw, + labels, + mask, + Exchange(dist.group.WORLD, dist.group.WORLD), + 8 << 30, + owner_local_frontend=True, + ) + baseline.cell_rows = candidate.cell_rows = 3 + n, features, readouts, channels = ( + candidate.stop - candidate.start, + candidate.columns, + candidate.readouts, + candidate.channels, + ) + before = torch.full((n, features + readouts, channels), -123.0) + expected_back, actual_back = before.clone(), before.clone() + expected, actual = ( + expected_back[:, readouts:], + actual_back[:, readouts:], + ) + expected_cache, actual_cache = Cache(), Cache() + assert ( + baseline.cell( + x=raw[baseline.start : baseline.stop], + categorical_mask=mask, + cache=expected_cache, + out=expected, + ) + is expected + ) + calls = [] + forward = CellEmbedding.forward + + def observed( + module, + x, + categorical_mask, + *, + _calls=calls, + _forward=forward, + **kwargs, + ): + _calls.append((x.size(0), kwargs["out"].stride(0))) + return _forward(module, x, categorical_mask, **kwargs) + + with ExitStack() as stack: + stack.enter_context( + patch.object( + candidate.exchange, + "broadcast", + side_effect=AssertionError( + "NO_BOUNDARY_BROADCAST" + ), + ) + ) + stack.enter_context( + patch.object( + candidate.exchange, + "exchange", + side_effect=AssertionError("NO_BOUNDARY_EXCHANGE"), + ) + ) + stack.enter_context( + patch.object(CellEmbedding, "forward", observed) + ) + assert ( + candidate.cell( + x=raw[candidate.start : candidate.stop], + categorical_mask=mask, + cache=actual_cache, + out=actual, + ) + is actual + ) + expected_rows = [ + min(start + 3, count, candidate.stop) + - max(start, candidate.start) + for start in range(0, count, 3) + if max(start, candidate.start) + < min(start + 3, count, candidate.stop) + ] + assert calls == [ + (rows, (features + readouts) * channels) + for rows in expected_rows + ] + torch.testing.assert_close( + actual_back, expected_back, atol=2e-6, rtol=1e-5 + ) + torch.testing.assert_close( + actual_cache["cell_impute_mean"], + expected_cache["cell_impute_mean"], + atol=0, + rtol=0, + ) + assert actual_cache["cell_impute_mean"][0, 2] == 0 + for tokens in (readouts, features + readouts): + for separate in (False, True): + full = ( + torch.arange( + count * (features + readouts) * channels, + dtype=torch.float32, + ) + .reshape(count, features + readouts, channels) + .cos() + ) + original = full[candidate.start : candidate.stop] + expected_kv, actual_kv = ( + original.clone(), + original.clone(), + ) + expected_out = ( + torch.empty(n, tokens, channels) + if separate + else expected_kv[:, :tokens] + ) + actual_out = ( + torch.empty(n, tokens, channels) + if separate + else actual_kv[:, :tokens] + ) + block = encoder.row_blocks[-1] + assert ( + baseline.row( + block, + 4, + tokens, + key_value=expected_kv, + out=expected_out, + ) + is expected_out + ) + row_calls = [] + block_forward = block.forward + + def observed_row( + *, + _calls=row_calls, + _forward=block_forward, + **kwargs, + ): + _calls.append(kwargs["key_value"].size(0)) + return _forward(**kwargs) + + with ExitStack() as stack: + stack.enter_context( + patch.object( + candidate.exchange, + "broadcast", + side_effect=AssertionError( + "NO_BOUNDARY_BROADCAST" + ), + ) + ) + stack.enter_context( + patch.object( + candidate.exchange, + "exchange", + side_effect=AssertionError( + "NO_BOUNDARY_EXCHANGE" + ), + ) + ) + stack.enter_context( + patch.object(block, "forward", observed_row) + ) + assert ( + candidate.row( + block, + 4, + tokens, + key_value=actual_kv, + out=actual_out, + ) + is actual_out + ) + expected_windows = [ + min(start + 4, count, candidate.stop) + - max(start, candidate.start) + for start in range(0, count, 4) + if max(start, candidate.start) + < min(start + 4, count, candidate.stop) + ] + assert row_calls == expected_windows + torch.testing.assert_close( + actual_out, expected_out, atol=2e-6, rtol=1e-5 + ) + torch.testing.assert_close( + actual_kv, expected_kv, atol=2e-6, rtol=1e-5 + ) + assert actual_out.stride() == expected_out.stride() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [1, 2, 4]) +def test_local_intersections_pitch_cache_and_aliases( + tmp_path: Path, world: int +) -> None: + mp.spawn( + component_worker, args=(world, str(tmp_path)), nprocs=world, join=True + ) + + +def fault_worker(rank: int, directory: str, stage: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=2, + timeout=timedelta(seconds=20), + ) + try: + model, x, _, _ = fixture("regression") + if stage == "policy": + with pytest.raises(ValueError, match="policies differ"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=bool(rank), + ) + assert model._execution is None + elif stage == "mutation": + run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=True, + ) + run.owner_local_frontend = False + with pytest.raises(RuntimeError, match="policy changed") as first: + run.clear() + with pytest.raises( + RuntimeError, match="policy changed" + ) as repeated: + run.clear() + assert repeated.value is first.value + assert model._execution is run + else: + encoder = model.models[Task.regression].row_embedding + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + rows = Rows( + encoder, + x, + torch.zeros(11), + torch.zeros(3, dtype=torch.bool), + exchange, + 8 << 30, + owner_local_frontend=True, + ) + n = rows.stop - rows.start + state = x.new_zeros(n, rows.columns + rows.readouts, rows.channels) + injected = ValueError( + f"CONTROLLED_OWNER_LOCAL_{stage.upper()}_RANK1" + ) + original = ( + CellEmbedding.forward + if stage == "cell" + else encoder.row_blocks[0].forward + ) + + def fail(*args, **kwargs): + if rank == 1: + raise injected + return original(*args, **kwargs) + + target = ( + CellEmbedding if stage == "cell" else encoder.row_blocks[0] + ) + + def execute_local(): + if stage == "cell": + rows.cell( + x=x[rows.start : rows.stop], + categorical_mask=rows.mask, + cache=Cache(), + out=state[:, rows.readouts :], + ) + else: + rows.row( + encoder.row_blocks[0], + 4, + rows.readouts, + key_value=state, + out=state[:, : rows.readouts], + ) + + with patch.object(target, "forward", fail), torch.no_grad(): + with pytest.raises( + (ValueError, RuntimeError), match=str(injected) + ) as first: + execute_local() + with pytest.raises( + (ValueError, RuntimeError), match=str(injected) + ) as repeated: + exchange.attempt(lambda: None) + assert repeated.value is first.value + if rank == 1: + assert first.value is injected + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("stage", ["cell", "row", "policy", "mutation"]) +def test_local_failures_and_rank_policy_agreement( + tmp_path: Path, stage: str +) -> None: + mp.spawn(fault_worker, args=(str(tmp_path), stage), nprocs=2, join=True) diff --git a/test/execution/test_sharded_row_workspace.py b/test/execution/test_sharded_row_workspace.py new file mode 100644 index 0000000000..2040179c59 --- /dev/null +++ b/test/execution/test_sharded_row_workspace.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Actual row-call bounds and unchanged public CPU/Gloo semantics.""" + +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture, journey + +from sdm.execution import tabular_execution +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_rows import Rows + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2, 4]) +@pytest.mark.parametrize("count", [1, 11]) +def test_public_sharded_workspace( + tmp_path: Path, task: str, world: int, count: int +) -> None: + args = (world, str(tmp_path), task, count > 1, None, True, count, True) + mp.spawn(journey, args=args, nprocs=world, join=True) + + +def component_worker( + rank: int, world: int, directory: str, failure: bool = False +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=25), + ) + try: + model, _, _, _ = fixture("regression") + encoder = next(iter(model.models.values())).row_embedding + with torch.no_grad(): + for total in (1, 11, 29): + raw = torch.arange(total * 3).float().reshape(total, 3).sin() + mask = torch.zeros(3, dtype=torch.bool) + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + baseline = Rows( + encoder, + raw, + torch.zeros(total), + mask, + exchange, + 1 << 30, + owner_local_frontend=True, + ) + candidate = Rows( + encoder, + raw, + torch.zeros(total), + mask, + exchange, + 1 << 30, + owner_local_frontend=True, + shard_row_workspace=True, + ) + assert baseline.shard_row_workspace is False + features, readouts, channels = ( + candidate.columns, + candidate.readouts, + candidate.channels, + ) + block = encoder.row_blocks[-1] + for limit in (1, 3, 8, 31): + local_limit = max(1, limit // world) + segments = [] + for start in range(0, total, limit): + left = max(start, candidate.start) + right = min(start + limit, total, candidate.stop) + segments.extend( + (begin, min(begin + local_limit, right)) + for begin in range(left, right, local_limit) + ) + for tokens in (readouts, features + readouts): + for separate in (False, True): + full = ( + torch.arange( + total * (features + readouts) * channels, + dtype=torch.float32, + ) + .reshape(total, features + readouts, channels) + .cos() + ) + original = full[candidate.start : candidate.stop] + expected_kv, actual_kv = ( + original.clone(), + original.clone(), + ) + expected = ( + torch.empty(original.size(0), tokens, channels) + if separate + else expected_kv[:, :tokens] + ) + actual = ( + torch.empty(original.size(0), tokens, channels) + if separate + else actual_kv[:, :tokens] + ) + baseline.row( + block, + limit, + tokens, + key_value=expected_kv, + out=expected, + ) + calls = [] + original_forward = block._forward + error = RuntimeError( + "rank-local subdivided row failure" + ) + + def observed( + *, + query, + key_value, + out, + _calls=calls, + _error=error, + _forward=original_forward, + **kwargs, + ): + _calls.append( + ( + key_value.clone(), + out.stride(), + ) + ) + if failure and rank == world - 1: + raise _error + return _forward( + query=query, + key_value=key_value, + out=out, + **kwargs, + ) + + with patch.object(block, "_forward", observed): + if failure: + with pytest.raises( + RuntimeError, match=str(error) + ) as caught: + candidate.row( + block, + limit, + tokens, + key_value=actual_kv, + out=actual, + ) + if rank == world - 1: + assert caught.value is error + with pytest.raises( + RuntimeError + ) as repeated: + exchange.attempt(lambda: None) + assert repeated.value is caught.value + dist.barrier() + return + assert ( + candidate.row( + block, + limit, + tokens, + key_value=actual_kv, + out=actual, + ) + is actual + ) + assert len(calls) == len(segments) + for (value, strides), ( + left, + right, + ) in zip(calls, segments, strict=True): + assert 0 < value.size(0) <= local_limit + assert strides == actual.stride() + torch.testing.assert_close( + value, full[left:right], atol=0, rtol=0 + ) + torch.testing.assert_close( + actual, expected, atol=2e-6, rtol=1e-5 + ) + torch.testing.assert_close( + actual_kv, expected_kv, atol=2e-6, rtol=1e-5 + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [1, 2, 4]) +def test_actual_max_row_bound_and_outputs(tmp_path: Path, world: int) -> None: + mp.spawn( + component_worker, args=(world, str(tmp_path)), nprocs=world, join=True + ) + + +@pytest.mark.parametrize("world", [2, 4]) +def test_rank_local_row_failure(tmp_path: Path, world: int) -> None: + mp.spawn( + component_worker, + args=(world, str(tmp_path), True), + nprocs=world, + join=True, + ) + + +def policy_worker(rank: int, world: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=25), + ) + try: + model, x, y, query = fixture("regression") + for frontend in [False, True] if world == 3 else [False]: + with pytest.raises( + ValueError, match="requires owner-local frontend and P1/2/4" + ): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=frontend, + shard_row_workspace=True, + ) + assert model._execution is None + if world == 3: + dist.barrier() + return + with pytest.raises(ValueError, match="policies differ"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=True, + shard_row_workspace=bool(rank), + ) + assert model._execution is None + run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=True, + shard_row_workspace=True, + ) + run.fit(x if rank == 0 else None, y if rank == 0 else None) + cache = run._row_cache + run.shard_row_workspace = False + with pytest.raises(RuntimeError, match="policy changed") as caught: + run.predict(query if rank == 0 else None) + run.shard_row_workspace = True + for operation in (run.clear, run.close): + with pytest.raises(RuntimeError) as repeated: + operation() + assert repeated.value is caught.value + assert run._row_cache is cache + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 3]) +def test_workspace_admission_and_policy(tmp_path: Path, world: int) -> None: + mp.spawn( + policy_worker, args=(world, str(tmp_path)), nprocs=world, join=True + ) diff --git a/test/execution/test_sharded_row_workspace_p8.py b/test/execution/test_sharded_row_workspace_p8.py new file mode 100644 index 0000000000..f023fade17 --- /dev/null +++ b/test/execution/test_sharded_row_workspace_p8.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Eight actual Gloo owners, including empty row and column owners.""" + +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_sharded_row_workspace import ( + component_worker, + policy_worker, +) +from test.execution.test_tabular import fixture, journey + +from sdm.execution import tabular_execution +from sdm.nn import Attention + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("count", [1, 11]) +@pytest.mark.parametrize( + "policy", [(False, False), (True, False), (True, True)] +) +def test_p8_public( + tmp_path: Path, task: str, count: int, policy: tuple[bool, bool] +) -> None: + frontend, workspace = policy + mp.spawn( + journey, + args=( + 8, + str(tmp_path), + task, + count > 1, + None, + frontend, + count, + workspace, + ), + nprocs=8, + join=True, + ) + + +@pytest.mark.parametrize("failure", [False, True]) +def test_p8_actual_row_bound_and_failure( + tmp_path: Path, failure: bool +) -> None: + mp.spawn( + component_worker, + args=(8, str(tmp_path), failure), + nprocs=8, + join=True, + ) + + +def test_p8_policy_and_retirement(tmp_path: Path) -> None: + mp.spawn(policy_worker, args=(8, str(tmp_path)), nprocs=8, join=True) + + +def head_worker(rank: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=8, + timeout=timedelta(seconds=35), + ) + try: + # Real, internally consistent attention modules: separately reject + # four query heads and four KV heads before executor ownership. + for query_heads, kv_heads in ((4, 4), (8, 4)): + model, _, _, _ = fixture("regression") + core = next(iter(model.models.values())) + core.icl_block.layers[0].attn = Attention( + channels=32, + num_query_heads=query_heads, + num_key_value_heads=kv_heads, + ).eval() + with pytest.raises(ValueError, match="heads must divide"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + owner_local_frontend=True, + shard_row_workspace=True, + ) + assert model._execution is None + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_p8_fit_head_divisibility(tmp_path: Path) -> None: + mp.spawn(head_worker, args=(str(tmp_path),), nprocs=8, join=True) diff --git a/test/execution/test_tabular.py b/test/execution/test_tabular.py index 67ea26e878..7754e8151a 100644 --- a/test/execution/test_tabular.py +++ b/test/execution/test_tabular.py @@ -20,6 +20,7 @@ import sdm.execution.tabular as tabular import sdm.processing as sp from sdm import CategoricalTensor, Recipe, Stype, TableTensor, Task +from sdm.cache import Cache, KVCacheEntry from sdm.execution import tabular_execution from sdm.models import KumoTabular from sdm.models.kumo.tabular.model import _KumoTabular @@ -70,6 +71,10 @@ def journey( task: str, chunked: bool = False, fit_head_gather_window: int | None = None, + owner_local_frontend: bool = False, + context_rows: int = 11, + shard_row_workspace: bool = False, + default_recipe: bool = False, ) -> None: torch.set_num_threads(1) dist.init_process_group( @@ -81,13 +86,22 @@ def journey( ) try: model, x, y, query = fixture(task) + x, y = x[:context_rows], y[:context_rows] native = copy.deepcopy(model) recipe = Recipe( target=sp.StypeDispatch(numerical=sp.Standardize()), output=[sp.AverageEstimators()], ) + if default_recipe: + recipe = None def fit(target, *args, **kwargs): + generator = None + if default_recipe: + generator = torch.Generator().manual_seed(42) + kwargs["generator"] = ( + generator if target is native or rank == 0 else None + ) # Exercise original fit windows; query retains its native AUTO. with ( patch.object( @@ -97,8 +111,9 @@ def fit(target, *args, **kwargs): else nullcontext() ): target.fit(*args, **kwargs) + return generator - fit(native, x, y, recipe=recipe, num_estimators=1) + native_generator = fit(native, x, y, recipe=recipe, num_estimators=1) batches = [query[:2], query[2:4], query[4:6], query[6:]] expected = [native.predict(batch) for batch in batches] run = tabular_execution( @@ -107,21 +122,62 @@ def fit(target, *args, **kwargs): group=dist.group.WORLD, control_group=dist.group.WORLD, fit_head_gather_window=fit_head_gather_window, + owner_local_frontend=owner_local_frontend, + shard_row_workspace=shard_row_workspace, ) for _ in range(2): - fit( + owner_generator = fit( run, x if rank == 0 else None, y if rank == 0 else None, recipe=recipe, ) + if owner_local_frontend: + for cache in (run._row_cache, run._icl_cache): + assert cache.is_replaying + for key in cache: + actual_cache = Cache.__getitem__(cache, key) + native_cache = native._cache[0][key] + if isinstance(actual_cache, KVCacheEntry): + for actual_tensor, native_tensor in zip( + actual_cache, native_cache, strict=True + ): + torch.testing.assert_close( + actual_tensor, + native_tensor.narrow( + 0, cache.start, cache.length + ), + atol=2e-6, + rtol=1e-5, + ) + elif isinstance(actual_cache, torch.Tensor): + torch.testing.assert_close( + actual_cache, + native_cache, + atol=2e-6, + rtol=1e-5, + ) actual = run.predict_batches(batches if rank == 0 else None) if rank == 0: + if default_recipe: + assert torch.equal( + native_generator.get_state(), + owner_generator.get_state(), + ) for left, right in zip(actual, expected, strict=True): assert left.columns == right.columns + if task == "regression": + assert left.shape[-1] == 999 torch.testing.assert_close( left.numerical, right.numerical, atol=2e-6, rtol=1e-5 ) + if owner_local_frontend and task == "classification": + assert torch.equal( + left.numerical.argmax(-1), + right.numerical.argmax(-1), + ) + else: + assert actual is None run.clear() fit( run, @@ -130,6 +186,8 @@ def fit(target, *args, **kwargs): recipe=recipe, ) last = run.predict(query if rank == 0 else None) + if rank != 0: + assert last is None run.close() assert dist.is_initialized() fit(model, x, y, recipe=recipe, num_estimators=1) @@ -160,6 +218,16 @@ def test_chunked_native_windows(tmp_path: Path, task: str) -> None: mp.spawn(journey, args=(2, str(tmp_path), task, True), nprocs=2, join=True) +@pytest.mark.parametrize("task", ["classification", "regression"]) +def test_two_owner_default_recipe_journey(tmp_path: Path, task: str) -> None: + mp.spawn( + journey, + args=(2, str(tmp_path), task, False, None, False, 11, False, True), + nprocs=2, + join=True, + ) + + @pytest.mark.parametrize("task", ["classification", "regression"]) @pytest.mark.parametrize("world", [1, 2, 4]) @pytest.mark.parametrize("window", [1, 2, 4]) diff --git a/test/execution/test_tabular_backend_policy.py b/test/execution/test_tabular_backend_policy.py new file mode 100644 index 0000000000..8fca73c1f3 --- /dev/null +++ b/test/execution/test_tabular_backend_policy.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fit attention backend policy and ambient restoration without CUDA work.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from sdm.execution._tabular_icl import fit_icl + + +def enabled() -> tuple[bool, bool, bool, bool]: + backend = torch.backends.cuda + return ( + backend.flash_sdp_enabled(), + backend.mem_efficient_sdp_enabled(), + backend.cudnn_sdp_enabled(), + backend.math_sdp_enabled(), + ) + + +@pytest.mark.parametrize("raises", [False, True]) +@pytest.mark.parametrize( + ("autocast", "dtype", "cudnn"), + [ + (False, torch.float16, False), + (True, torch.float16, True), + (True, torch.bfloat16, False), + ], +) +def test_cuda_fit_precision_policy_and_ambient_restoration( + raises: bool, + autocast: bool, + dtype: torch.dtype, + cudnn: bool, +) -> None: + primary = RuntimeError("injected fit failure") + full = SimpleNamespace(is_cuda=True, device=torch.device("cuda")) + exchange = SimpleNamespace(gather=lambda *args: full) + + def model(*args: object, **kwargs: object) -> object: + assert enabled() == (True, True, cudnn, False) + if raises: + raise primary + return full + + before = enabled() + was_enabled = torch.is_autocast_enabled("cuda") + previous_dtype = torch.get_autocast_dtype("cuda") + try: + # Thread-local policy only: no CUDA allocation or initialization. + # torch.autocast("cuda") would disable itself in CPU-only CI. + torch.set_autocast_enabled("cuda", autocast) + torch.set_autocast_dtype("cuda", dtype) + with sdpa_kernel(SDPBackend.MATH): + assert enabled() == (False, False, False, True) + if raises: + with pytest.raises(RuntimeError) as caught: + fit_icl( + model, + full, + full, + SimpleNamespace(lengths=(1,)), + exchange, + ) + assert caught.value is primary + else: + assert ( + fit_icl( + model, + full, + full, + SimpleNamespace(lengths=(1,)), + exchange, + ) + is full + ) + assert enabled() == (False, False, False, True) + finally: + torch.set_autocast_enabled("cuda", was_enabled) + torch.set_autocast_dtype("cuda", previous_dtype) + assert enabled() == before + + +def test_cpu_fit_leaves_ambient_attention_policy_unchanged() -> None: + full = SimpleNamespace(is_cuda=False) + exchange = SimpleNamespace(gather=lambda *args: full) + + def model(*args: object, **kwargs: object) -> object: + assert enabled() == (False, False, False, True) + return full + + before = enabled() + with sdpa_kernel(SDPBackend.MATH): + assert ( + fit_icl(model, full, full, SimpleNamespace(lengths=(1,)), exchange) + is full + ) + assert enabled() == (False, False, False, True) + assert enabled() == before diff --git a/test/execution/test_tabular_precision.py b/test/execution/test_tabular_precision.py new file mode 100644 index 0000000000..9ac26f0739 --- /dev/null +++ b/test/execution/test_tabular_precision.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU dtype/protocol checks, not CUDA FP16 numerical qualification.""" + +import copy +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture +from test.execution.test_tabular import group as group + +import sdm +import sdm.processing as sp +from sdm import Recipe, Task +from sdm.cache import Cache, KVCacheEntry +from sdm.execution import tabular_execution +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_icl import workspace_bytes +from sdm.execution._tabular_rows import Rows +from sdm.nn import TransformerBlock + + +@pytest.fixture(autouse=True) +def no_cuda(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("CPU source test attempted CUDA") + + monkeypatch.setattr(torch.cuda, "_lazy_init", forbidden) + monkeypatch.setattr(torch.cuda, "synchronize", forbidden) + yield + assert not torch.cuda.is_initialized() + + +def make_run(model, group, *, task="regression", precision="float16"): + return tabular_execution( + model, task=task, group=group, control_group=group, precision=precision + ) + + +def recipe(): + return Recipe( + target=sp.StypeDispatch(numerical=sp.Standardize()), + output=[sp.AverageEstimators()], + ) + + +def journey(rank, world, directory, task, mismatch=False): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=15), + ) + try: + model, x, y, query = fixture(task) + if mismatch: + with pytest.raises(ValueError, match="policies differ"): + make_run( + model, + dist.group.WORLD, + task=task, + precision="float32" if rank else "float16", + ) + assert model._execution is None + return + native = copy.deepcopy(model) + weights = {k: v.clone() for k, v in model.state_dict().items()} + batches = [query[:2], query[2:4], query[4:6], query[6:]] + with torch.autocast("cpu", dtype=torch.float16): + native.fit(x, y, recipe=recipe(), num_estimators=1) + expected = [native.predict(q) for q in batches] + run = make_run(model, dist.group.WORLD, task=task) + # Neither owner nor workers require an ambient autocast context. + for _ in range(2): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + assert not torch.is_autocast_enabled("cpu") + for cache in (run._row_cache, run._icl_cache): + for value in cache._items.values(): + if isinstance(value, KVCacheEntry): + assert ( + value.key.dtype + == value.value.dtype + == torch.float16 + ) + elif isinstance(value, torch.Tensor): + assert value.dtype == torch.float32 # Native mean. + actual = run.predict_batches(batches if rank == 0 else None) + if rank == 0: + for left, right in zip(actual, expected, strict=True): + assert left.columns == right.columns + assert ( + left.numerical.dtype + == right.numerical.dtype + == torch.float32 + ) + assert torch.isfinite(left.numerical).all() + torch.testing.assert_close( + left.numerical, right.numerical, atol=2e-6, rtol=1e-5 + ) + if task == "regression": + assert left.numerical.shape[-1] == 999 + run.clear() + assert run._icl_cache is run._row_cache is None + run.close() + assert dist.is_initialized() + assert model._execution is None + for key, value in model.state_dict().items(): + assert value.dtype == torch.float32 + assert torch.equal(value, weights[key]) + # Healthy release permits original native caller use, including AMP. + with torch.autocast("cpu", dtype=torch.float16): + model.fit(x, y, recipe=recipe(), num_estimators=1) + torch.testing.assert_close( + model.predict(batches[0]).numerical, + expected[0].numerical, + atol=2e-6, + rtol=1e-5, + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2]) +def test_actual_cpu_public_amp(tmp_path, task, world): + if world == 1: + journey(0, world, str(tmp_path), task) + else: + mp.spawn(journey, args=(world, str(tmp_path), task), nprocs=world) + + +def test_precision_must_agree_across_ranks(tmp_path): + mp.spawn(journey, args=(2, str(tmp_path), "regression", True), nprocs=2) + + +@pytest.mark.parametrize("limit", [2, 65_535]) +def test_native_amp_inducing_residual_is_fp32(group, limit): + model, x, y, _ = fixture("regression") + y = y.squeeze(-1) # Original task core's scalar target seam. + encoder = model.models[Task.regression].row_embedding + mask = torch.zeros(x.size(1), dtype=torch.bool) + seen = [] + handles = [ + block.inducing_block.register_forward_hook( + lambda module, args, output: seen.append(output.dtype) + ) + for block in encoder.col_blocks + ] + try: + # Mocked CPU AUTO hardware leaf only; genuine modules and raw data. + with ( + torch.no_grad(), + torch.autocast("cpu", dtype=torch.float16), + patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=limit + ), + ): + native_cache, owned_cache = Cache(), Cache() + expected = encoder(x, y, mask, cache=native_cache) + rows = Rows(encoder, x, y, mask, Exchange(group, group), 8 << 30) + assert rows.dtype == torch.float16 + assert rows.mean.dtype == torch.float32 + actual = encoder(x, y, mask, cache=owned_cache, _runner=rows) + assert seen + assert set(seen) == {torch.float32} + assert actual.dtype == expected.dtype == torch.float16 + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=1e-5) + for key, value in native_cache.items(): + other = owned_cache[key] + if isinstance(value, KVCacheEntry): + for left, right in zip(value, other, strict=True): + assert left.dtype == right.dtype == torch.float16 + torch.testing.assert_close( + left, right, atol=2e-6, rtol=1e-5 + ) + finally: + for handle in handles: + handle.remove() + + +@pytest.mark.parametrize("precision", ["bfloat16", "half", None]) +def test_unsupported_precision_rejected_before_lease(group, precision): + model, _, _, _ = fixture("regression") + with pytest.raises(ValueError, match="precision must be"): + make_run(model, group, precision=precision) + assert model._execution is None + + +def test_conflicting_ambient_rejected(group): + model, _, _, _ = fixture("regression") + with ( + torch.autocast("cpu", dtype=torch.bfloat16), + pytest.raises(ValueError, match="Ambient autocast differs"), + ): + make_run(model, group) + assert model._execution is None + + +def test_conflicting_ambient_after_fit_retires(group): + model, x, y, query = fixture("regression") + run = make_run(model, group) + run.fit(x, y, recipe=recipe()) + cache = run._icl_cache + with ( + torch.autocast("cpu", dtype=torch.bfloat16), + pytest.raises(ValueError, match="Ambient autocast differs") as caught, + ): + run.predict(query) + with pytest.raises(ValueError, match="Ambient autocast differs") as again: + run.close() + assert caught.value is again.value + assert run._icl_cache is cache + + +@pytest.mark.parametrize("kind", ["weight", "cache"]) +def test_amp_generation_version_guard_remains(group, kind): + model, x, y, query = fixture("regression") + run = make_run(model, group) + run.fit(x, y, recipe=recipe()) + value = ( + next(model.parameters()) + if kind == "weight" + else next(run._icl_cache._tensors()) + ) + with torch.no_grad(): + value.add_(1) + with pytest.raises(RuntimeError, match="changed") as caught: + run.predict(query) + with pytest.raises(RuntimeError, match="changed") as again: + run.close() + assert caught.value is again.value + assert model._execution is run + + +def test_fp16_does_not_admit_half_stored_model(group): + model, _, _, _ = fixture("regression") + model.half() + with pytest.raises(ValueError, match="FP32 device"): + make_run(model, group) + assert model._execution is None + + +def test_source_origins_are_this_external_copy(): + source = Path(__file__).resolve().parents[2] + assert Path(sdm.__file__).resolve().is_relative_to(source) + assert Path(fixture.__code__.co_filename).resolve().is_relative_to(source) + + +def test_matching_ambient_restored(group): + model, x, y, _ = fixture("regression") + with torch.autocast("cpu", dtype=torch.float16, cache_enabled=False): + run = make_run(model, group) + run.fit(x, y, recipe=recipe()) + assert torch.is_autocast_enabled("cpu") + assert not torch.is_autocast_cache_enabled() + run.close() + assert not torch.is_autocast_enabled("cpu") + + +def test_precision_mutation_retires_fitted_generation(group): + model, x, y, query = fixture("regression") + run = make_run(model, group) + run.fit(x, y, recipe=recipe()) + cache = run._icl_cache + run.precision = "float32" + with pytest.raises(RuntimeError, match="policy changed") as caught: + run.predict(query) + with pytest.raises(RuntimeError, match="policy changed") as repeated: + run.close() + assert caught.value is repeated.value + assert run._icl_cache is cache + assert model._execution is run + assert not torch.is_autocast_enabled("cpu") + + +def test_failure_restores_amp_and_preserves_primary(group): + model, x, y, _ = fixture("regression") + run = make_run(model, group) + primary = RuntimeError("injected neural failure") + original = run._fit_core + + def fail(*args, **kwargs): + assert torch.is_autocast_enabled("cpu") + assert torch.get_autocast_dtype("cpu") == torch.float16 + raise primary + + run._fit_core = fail + with pytest.raises(RuntimeError, match="injected neural") as caught: + run.fit(x, y, recipe=recipe()) + run._fit_core = original + assert caught.value is primary + assert not torch.is_autocast_enabled("cpu") + with pytest.raises(RuntimeError, match="injected neural") as repeated: + run.close() + assert repeated.value is primary + assert model._execution is run + + +def test_transport_budget_uses_cached_dtype_not_halved_attention(group): + model, x, y, _ = fixture("regression") + run = make_run(model, group) + run.fit(x, y, recipe=recipe()) + cache, block = run._icl_cache, run._core.icl_block + half = workspace_bytes(block, 3, cache) + original = dict(cache._items) + cache._items = { + key: KVCacheEntry(*(t.float() for t in pair)) + for key, pair in original.items() + } + full = workspace_bytes(block, 3, cache) + cache._items = original + compact = 11 * block.kv_heads * 2 * block.layers[0].attn.head_dim + assert full - half == compact * 2 + assert half > full // 2 + run.close() + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_fp32_default_matches_explicit_and_native(group, explicit): + model, x, y, query = fixture("regression") + native = copy.deepcopy(model) + native.fit(x, y, recipe=recipe(), num_estimators=1) + kwargs = {"precision": "float32"} if explicit else {} + run = tabular_execution( + model, task="regression", group=group, control_group=group, **kwargs + ) + run.fit(x, y, recipe=recipe()) + actual, expected = run.predict(query), native.predict(query) + assert actual.numerical.dtype == torch.float32 + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=2e-6, rtol=1e-5 + ) + run.close() diff --git a/test/execution/test_tabular_precision_owners.py b/test/execution/test_tabular_precision_owners.py new file mode 100644 index 0000000000..8495d167f8 --- /dev/null +++ b/test/execution/test_tabular_precision_owners.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FP16 native controls for owner-local/subchunk compatibility, CPU only.""" + +import copy +from datetime import timedelta +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture +from test.execution.test_tabular_precision import no_cuda as no_cuda +from test.execution.test_tabular_precision import recipe + +from sdm.cache import Cache, KVCacheEntry +from sdm.execution import tabular_execution +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_rows import Rows +from sdm.nn import TransformerBlock + + +def owner_journey(rank, world, directory, task, policy, count): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=40), + ) + try: + model, x, y, query = fixture(task) + x, y = x[:count], y[:count] + native = copy.deepcopy(model) + before = { + key: value.clone() for key, value in model.state_dict().items() + } + batches = [query[:2], query[2:4], query[4:6], query[6:]] + # Original modules, only CPU AUTO hardware leaf controlled. Native + # and candidate use the same base limit; candidate may subdivide rows. + with torch.autocast("cpu", dtype=torch.float16): + with patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=8 + ): + native.fit(x, y, recipe=recipe(), num_estimators=1) + expected = [native.predict(batch) for batch in batches] + run = tabular_execution( + model, + task=task, + group=dist.group.WORLD, + control_group=dist.group.WORLD, + precision="float16", + owner_local_frontend=policy[0], + shard_row_workspace=policy[1], + ) + for repeat in range(3): + if repeat == 2: + run.clear() + assert run._row_cache is run._icl_cache is None + with patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=8 + ): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + assert not torch.is_autocast_enabled("cpu") + for cache in (run._row_cache, run._icl_cache): + for key in cache: + actual = Cache.__getitem__(cache, key) + full = native._cache[0][key] + if isinstance(actual, KVCacheEntry): + for left, right in zip(actual, full, strict=True): + assert left.dtype == right.dtype == torch.float16 + torch.testing.assert_close( + left, + right.narrow(0, cache.start, cache.length), + atol=2e-6, + rtol=1e-5, + ) + elif isinstance(actual, torch.Tensor): + assert actual.dtype == full.dtype == torch.float32 + torch.testing.assert_close( + actual, full, atol=2e-6, rtol=1e-5 + ) + outputs = run.predict_batches(batches if rank == 0 else None) + if rank == 0: + for actual, full in zip(outputs, expected, strict=True): + assert actual.columns == full.columns + assert ( + actual.numerical.dtype + == full.numerical.dtype + == torch.float32 + ) + torch.testing.assert_close( + actual.numerical, full.numerical, atol=2e-6, rtol=1e-5 + ) + if task == "regression": + assert actual.size(-1) == 999 + else: + assert torch.equal( + actual.numerical.argmax(-1), + full.numerical.argmax(-1), + ) + else: + assert outputs is None + run.close() + assert model._execution is None + assert dist.is_initialized() + for key, value in model.state_dict().items(): + assert value.dtype == torch.float32 + assert torch.equal(value, before[key]) + with torch.autocast("cpu", dtype=torch.float16): + with patch.object( + TransformerBlock, "auto_batch_size_limit", return_value=8 + ): + model.fit(x, y, recipe=recipe(), num_estimators=1) + torch.testing.assert_close( + model.predict(batches[0]).numerical, + expected[0].numerical, + atol=2e-6, + rtol=1e-5, + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2, 4, 8]) +@pytest.mark.parametrize( + "policy", [(False, False), (True, False), (True, True)] +) +def test_fp16_owner_policies(tmp_path, task, world, policy): + mp.spawn( + owner_journey, + args=(world, str(tmp_path), task, policy, 11), + nprocs=world, + ) + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +def test_fp16_empty_owners(tmp_path, task): + mp.spawn( + owner_journey, args=(8, str(tmp_path), task, (True, True), 1), nprocs=8 + ) + + +def row_worker(rank, world, directory, fail): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=30), + ) + try: + model, _, _, _ = fixture("regression") + encoder = next(iter(model.models.values())).row_embedding + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + with torch.no_grad(), torch.autocast("cpu", dtype=torch.float16): + for count in (1, 29): + raw = torch.arange(count * 3).float().reshape(count, 3).sin() + rows = Rows( + encoder, + raw, + torch.zeros(count), + torch.zeros(3, dtype=torch.bool), + exchange, + 1 << 30, + owner_local_frontend=True, + shard_row_workspace=True, + ) + assert rows.dtype == torch.float16 + assert rows.mean.dtype == torch.float32 + block = encoder.row_blocks[-1] + for tokens in (2, 5): + for separate in (False, True): + full = ( + torch.arange(count * 5 * 8) + .float() + .reshape(count, 5, 8) + .cos() + .half() + ) + expected = full.clone() + block( + query=expected[:, :tokens], + key_value=expected, + out=expected[:, :tokens], + batch_size_limit=8, + ) + local = full[rows.start : rows.stop].clone() + out = ( + local.new_empty(local.size(0), tokens, 8) + if separate + else local[:, :tokens] + ) + observed = [] + original = block._forward + error = RuntimeError("FP16 owner-local row failure") + + def checked( + *, + query, + key_value, + out, + _observed=observed, + _error=error, + _original=original, + **kwargs, + ): + _observed.append(key_value.size(0)) + assert ( + key_value.dtype == out.dtype == torch.float16 + ) + assert key_value.size(0) <= max(1, 8 // world) + if fail and rank == world - 1: + raise _error + return _original( + query=query, + key_value=key_value, + out=out, + **kwargs, + ) + + with patch.object(block, "_forward", checked): + if fail: + with pytest.raises( + RuntimeError, + match="FP16 owner-local row failure", + ) as caught: + rows.row( + block, + 8, + tokens, + key_value=local, + out=out, + ) + if rank == world - 1: + assert caught.value is error + dist.barrier() + return + result = rows.row( + block, 8, tokens, key_value=local, out=out + ) + assert result is out + assert bool(observed) == bool(local.size(0)) + torch.testing.assert_close( + out, + expected[rows.start : rows.stop, :tokens], + atol=2e-6, + rtol=1e-5, + ) + if separate: + assert torch.equal( + local, full[rows.start : rows.stop] + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [1, 2, 4, 8]) +@pytest.mark.parametrize("failure", [False, True]) +def test_fp16_actual_inner_row_bound_alias_and_failure( + tmp_path, world, failure +): + mp.spawn(row_worker, args=(world, str(tmp_path), failure), nprocs=world) diff --git a/test/execution/test_tabular_precision_protocol.py b/test/execution/test_tabular_precision_protocol.py new file mode 100644 index 0000000000..dae3eecea5 --- /dev/null +++ b/test/execution/test_tabular_precision_protocol.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rank-asymmetric precision/mutation/failure behavior with borrowed groups.""" + +from datetime import timedelta +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture +from test.execution.test_tabular_precision import no_cuda as no_cuda +from test.execution.test_tabular_precision import recipe + +from sdm.execution import tabular_execution +from sdm.nn import TransformerBlock + + +def worker(rank, world, directory): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=30), + ) + try: + model, _, _, _ = fixture("regression") + with pytest.raises(ValueError, match="policies differ"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + precision="float32" if rank else "float16", + owner_local_frontend=True, + shard_row_workspace=True, + ) + assert model._execution is None + dist.barrier() + + for problem in ("precision", "cache", "row_failure"): + model, x, y, query = fixture("regression") + run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + precision="float16", + owner_local_frontend=True, + shard_row_workspace=True, + ) + if problem != "row_failure": + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + cache = run._icl_cache + if rank == 0: + if problem == "precision": + run.precision = "float32" + else: + with torch.no_grad(): + next(cache._tensors()).add_(1) + with pytest.raises(RuntimeError, match="changed") as caught: + run.predict(query if rank == 0 else None) + assert run._icl_cache is cache + else: + original = torch.add + error = RuntimeError("rank-local FP16 row residual failure") + seen = [] + + def fail( + left, + *args, + _original=original, + _error=error, + _seen=seen, + **kwargs, + ): + if ( + rank == world - 1 + and isinstance(left, torch.Tensor) + and left.dtype == torch.float16 + and left.shape[-2:] == (5, 8) + ): + _seen.append(tuple(left.shape)) + raise _error + return _original(left, *args, **kwargs) + + with ( + patch.object(torch, "add", fail), + patch.object( + TransformerBlock, + "auto_batch_size_limit", + return_value=8, + ), + pytest.raises( + RuntimeError, + match="rank-local FP16 row residual failure", + ) as caught, + ): + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + if rank == world - 1: + assert seen + assert caught.value is error + assert not torch.is_autocast_enabled("cpu") + with pytest.raises(RuntimeError) as repeated: + run.close() + assert repeated.value is caught.value + assert model._execution is run + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 8]) +def test_rank_precision_cache_and_row_failure(tmp_path, world): + mp.spawn(worker, args=(world, str(tmp_path)), nprocs=world) From 835bad9a0e3140c8082e2dd2912dff235a8271db Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:58:20 -0700 Subject: [PATCH 04/10] Document context-sharded inference configuration and scaling Add a runnable precision-selectable example and separate usage, lifecycle limits, and measured predecessor results. Record cold-start and aggregate-memory caveats without implying final-source GPU qualification. --- docs/source/kumotabular-context-sharding.md | 132 +++++++++++++++++--- examples/README.md | 2 + examples/tabular/context_sharding.py | 93 ++++++++++++++ 3 files changed, 213 insertions(+), 14 deletions(-) create mode 100644 examples/tabular/context_sharding.py diff --git a/docs/source/kumotabular-context-sharding.md b/docs/source/kumotabular-context-sharding.md index 37be4b3e73..ec9e6779d5 100644 --- a/docs/source/kumotabular-context-sharding.md +++ b/docs/source/kumotabular-context-sharding.md @@ -2,9 +2,19 @@ `sdm.execution.tabular_execution` is an experimental, opt-in executor for splitting one KumoTabular estimator's context across devices. Ordinary `model.fit` and `model.predict` retain their existing behavior when no executor owns the model. +For large contexts, the measured configuration combines **context-owned row embeddings**, **owner-local frontend computation**, and **sharded row workspaces**. FP16 neural execution is a separate, explicit option. This scales one fitted model's context; it does not train independent models on subsets or ensemble their predictions. + ## Usage -Launch one process per CUDA device with `torchrun --standalone --nproc-per-node=4 your_script.py`. Each process loads identical weights and creates the same ordered tensor and control groups: +Run the self-contained [synthetic regression example](../../examples/tabular/context_sharding.py) from the repository root: + +```sh +torchrun --standalone --nproc-per-node=2 --max-restarts=0 examples/tabular/context_sharding.py --precision float32 +``` + +Use `--precision float16` to exercise the FP16 path. The example loads the stock large regression checkpoint on each device (downloading weights if absent), fits preprocessing only on the context rows, predicts original held-out query batches, then clears, refits and closes. Its small synthetic input is not a performance or quality benchmark. Run under a finite external job timeout, with one process per CUDA device; the measured frontend configuration supports one, two, four or eight processes. + +For integration into another script, each rank loads identical weights and creates the same ordered tensor and control groups. The following fragment assumes caller-provided context and query tensors on rank zero's CUDA device; numerical inputs and regression targets are FP32, while classification labels preserve their integer dtype: ```python import os @@ -17,11 +27,12 @@ from sdm.execution import tabular_execution from sdm.models import KumoTabular local_rank = int(os.environ["LOCAL_RANK"]) -torch.cuda.set_device(local_rank) +device = torch.device("cuda", local_rank) +torch.cuda.set_device(device) dist.init_process_group("nccl", timeout=timedelta(minutes=10)) control = dist.new_group(backend="gloo", timeout=timedelta(minutes=10)) root = dist.get_rank() == 0 -model = KumoTabular(task="regression", device=f"cuda:{local_rank}").eval() +model = KumoTabular(task="regression", size="large", device=device).eval() run = tabular_execution( model, @@ -29,36 +40,61 @@ run = tabular_execution( group=dist.group.WORLD, control_group=control, fit_head_gather_window=2, + owner_local_frontend=True, + shard_row_workspace=True, + precision="float32", # Explicitly choose "float16" for neural autocast. ) # Only root supplies tensors. All ranks call every method in the same order. # x_context: [context_rows, features]; y_context: [context_rows, 1]. -run.fit(x_context if root else None, y_context if root else None) -prediction = run.predict(x_query if root else None) -# A list of original batches allows independent batches to run concurrently. -predictions = run.predict_batches(query_batches if root else None) -run.clear() -run.fit(x_context if root else None, y_context if root else None) -run.close() +with torch.no_grad(): + run.fit(x_context if root else None, y_context if root else None) + prediction = run.predict(x_query if root else None) + # Original query batches can run concurrently without splitting a batch. + predictions = run.predict_batches(query_batches if root else None) + run.clear() + run.fit(x_context if root else None, y_context if root else None) + run.close() # The caller owns and destroys groups after healthy completion. dist.destroy_process_group(control) dist.destroy_process_group() ``` -Inputs may be tensors or `TableTensor`s. Root alone supplies the recipe and optional preprocessing generator. Prediction returns the usual public `TableTensor` on root and `None` on peers; `predict_batches` returns a root-only list in the original order. Each original query batch is kept intact, including recipe output processing and regression inverse transforms. Empty batch lists return an empty list. +Inputs may be tensors or `TableTensor`s, with shapes `[context_rows, features]`, `[context_rows, 1]` for targets, and `[query_rows, features]`. Root alone supplies inputs, the recipe and an optional preprocessing generator; peers supply `None`. Prediction returns the usual public `TableTensor` on root and `None` on peers; `predict_batches` returns a root-only list in the original order. Each original query batch is kept intact, including recipe output processing and regression inverse transforms. Empty batch lists are supported; individual empty batches are not. `predict` does not automatically partition one query batch. + +## Configuration and compatibility + +The original FP32 behavior remains the default. Enable the memory-saving frontend explicitly: + +| Option | Default | Purpose | +| --- | --- | --- | +| `owner_local_frontend` | `False` | Compute each owner's part of cell and row windows directly, avoiding boundary staging. | +| `shard_row_workspace` | `False` | Divide the native row-block chunk limit by the owner count; requires owner-local frontend and 1, 2, 4 or 8 owners. | +| `precision` | `"float32"` | `"float16"` enables per-worker neural autocast; weights, prepared numerical inputs and public processing remain FP32. | +| `fit_head_gather_window` | `None` | Limit outstanding fit head gathers; the measured configuration uses `2`. | +| `column_staging_bytes` | `8 << 30` | Admission budget for column exchange staging. | +| `attention_workspace_bytes` | `12 << 30` | Admission budget for layer-local prediction attention and cache transport. The large-context measurements below use `48 << 30`. | + +These component budgets are not total device-memory limits. The example retains the public 12 GiB attention budget; it does not reproduce the large-context benchmark configuration automatically. Choose budgets for the workload and device. An oversized original query batch is rejected rather than silently split. + +The FP16 option owns matching autocast scopes on every worker. Do not call `model.half()` or cast preprocessing inputs to FP16. Do not wrap the default FP32 executor in autocast as a substitute for selecting `precision="float16"`. BF16 and half-stored model weights are not supported. FP16 can differ from FP32; same-precision distributed/native agreement is a separate question from precision-related quality changes. ## What is distributed - Expanded row-embedding cells are stored by contiguous context-row owner, reducing the dominant row-by-feature-by-channel allocation. - Column attention exchanges row-owned cells into complete native column windows, then returns results to their row owners. It does not replace full-context column attention with independent local fits. -- Independent row-attention windows run on their owner, with one error agreement for the local loop. Windows crossing owner boundaries retain their native shape and arithmetic. +- With the default frontend, independent row-attention windows run on their owner and boundary windows retain their native shape. With owner-local frontend enabled, each rank computes its intersections of the native windows. Sharded row workspace additionally bounds the original whole-row-block chunker to `max(1, native_row_limit // owners)`. These options change projection batch geometry and require numerical validation. - ICL fit attention heads are divided across owners. Full projections and residuals remain replicated; retained ICL key/value caches are sharded by context row. - Prediction gathers one cache layer at a time and assigns complete original query batches to owners. This combines context-owned storage with query-batch parallelism, not ensemble parallelism. This is not complete tensor parallelism and does not guarantee linear capacity or speed scaling. Raw prepared inputs, some normalization/projection intermediates, full ICL states, one replayed cache layer, and a native column window can still scale with the entire context on a single device. Communication and replicated work can dominate small inputs. `column_staging_bytes` and `attention_workspace_bytes` bound individual planned components, not total allocated, reserved, or physical GPU memory. +For the tested large model, the expanded row tensor has shape `[context_rows, prepared_features + 4, 256]`. At 2,097,152 rows and 500 prepared features in FP16, this tensor alone is 504 GiB without sharding, 126 GiB per rank with four balanced owners, and 63 GiB with eight. Those are component sizes, not predictions of whole-workflow peak memory. Chunk limits bound temporary work, not the total number of context rows. + ## Supported scope and lifecycle -The initial scope is eager evaluation, FP32, one unbatched estimator, stock KumoTabular classification or regression, and no callbacks, custom forward hooks, compilation, autocast, or ECOC expansion beyond native class capacity. The number of owners must divide the ICL fit query/key-value head counts. CPU/Gloo is supported for numerical and lifecycle tests; CUDA requires NCCL for tensors plus a separate Gloo control group with identical ordered members. +The supported scope is eager evaluation, FP32 or explicit FP16 neural execution, one unbatched estimator, stock KumoTabular classification or regression, and no autograd inputs, callbacks, custom dispatch/forward hooks, compilation, or ECOC expansion beyond native class capacity. The number of owners must divide both ICL fit query- and key/value-head counts. CPU/Gloo supports numerical and lifecycle tests; CUDA requires NCCL for tensors plus a separate Gloo control group with identical ordered members. Multi-node and heterogeneous-device execution have not been qualified. + +Use normal calls or `torch.no_grad()`, **not an outer `torch.inference_mode()`**: fitted-cache identity checks require tensor version counters. Constructor, fit and prediction reject inference mode. Healthy `clear` and `close` are allowed inside inference mode. The executor borrows its model exclusively. Do not mutate weights, recipes, caches, execution settings, or use the underlying model while the lease is active. `fit` replaces a fitted generation; `clear` releases it while retaining the lease; healthy `close` releases the lease and permits ordinary model reuse. Calls must originate on the constructing thread and cannot overlap. Groups and process lifetime always remain caller-owned. @@ -66,4 +102,72 @@ Local allocation or compute errors are agreed before the next tensor collective, ## Validation status -Tests compare complete public classification outputs and all 999 regression quantiles with the same native model on one, two, and four CPU processes. They also exercise original row-window boundaries, refit, clear, close, borrowed process groups, mutation rejection, and collective failure propagation. These CPU checks do not establish CUDA numerical equivalence, trained-model quality, memory scaling, or speedup for this branch. Real-GPU matched native/one/two/four-owner qualification remains required; historical research-prototype performance must not be treated as a measurement of this current-main implementation. +The CPU execution tests cover complete public classification outputs and all 999 regression quantiles, one/two/four/eight-owner paths, uneven and empty owners, default-recipe RNG preservation, precision-policy agreement, row workspace bounds, refit, clear, close, borrowed groups, mutation rejection and collective failure propagation. Numerical comparisons use `atol=2e-6, rtol=1e-5`. CPU checks alone do not establish trained CUDA equivalence or performance. + +During consolidation, all 252 execution tests passed on the imported measured runtime. After the FP32 backend compatibility adjustment described below, 83 focused backend-policy, precision and lifecycle tests passed on the final runtime. These sets overlap; they are not 335 distinct tests or a complete final-source suite rerun. Both runs had ten CPU Half-input/FP32-weight RMSNorm fallback warnings and no failures or skips. Ruff, formatting, diff checks and the example's CPU-only `--help` also passed. Run `python -m pytest test/execution` from the repository root for the complete execution suite. + +The measured runtime identified below passed trained FP16 native-versus-eight-owner classification (EPSILON, 4,096 context rows) and all-999-quantile regression (YearPrediction, 16,384 rows) checks, plus the larger EPSILON comparison below. These are selected checkpoints and inputs, not general precision guarantees. Earlier FP16 source attempts failed strict native parity; their results are not silently included in the passing cohort. + +Dedicated CUDA lifecycle diagnostics for this runtime remain pending: earlier attempts failed test collection, and the packaging/device fixture corrections were checked locally only. Older-source CUDA lifecycle passes do not qualify the successor. The trained benchmark runs do exercise repeated fit/predict/clear/close, but are not a substitute for those dedicated failure and ownership tests. + +## Measured FP16 scaling + +These measurements use a single eight-B200 host, PyTorch 2.9.1+cu130, cuDNN 9.2, one large classification estimator, TF32 disabled, both frontend options enabled, `fit_head_gather_window=2`, 8 GiB column staging and 48 GiB attention workspace. Each case has 512 fixed held-out validation queries in original batches of eight. Preprocessing and feature selection use training/context data only. The measured predecessor's 176-file SDM Python source aggregate is `2f5314f0efd69cad02135c9a66a502de5eadd12f2a692a388aefd7e9dab67e43`. + +This branch incorporates that runtime with one compatibility adjustment: cuDNN fit attention is eligible only under CUDA FP16 autocast, preserving the original flash/efficient-only policy for default FP32 calls. The enabled backend set in the measured FP16 path is unchanged. The consolidated source has CPU tests but has not been re-run on GPUs; the following measurements belong to the identified predecessor, not a fresh qualification of this branch revision. + +Each GPU count has one clean multiprocess run, with one worker per GPU, three complete synchronized fits and twelve complete predictions. Fit medians include all three fits; prediction medians use nine calls, excluding the first call after each fit. Memory includes setup, warmups, refits, clear and close on every rank. Allocated memory is the busiest rank's allocator high-water; physical memory is sampled device usage. These are not independent-run confidence intervals or full-process cold-start measurements. + +### EPSILON: 262,144 context rows + +The input has 2,000 raw features and 500 prepared features. The one-GPU denominator is the same executor with the same options, not a different native implementation. + +| GPUs | Fit median (s) | Predict 512 rows (s) | Peak allocated / GPU (GiB) | Busiest physical (GiB) | Aligned all-device physical (GiB) | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 20.731 | 11.610 | 77.16 | 79.89 | 79.92 | +| 2 | 11.368 | 6.185 | 41.32 | 49.24 | 95.48 | +| 4 | 6.056 | 4.134 | 25.36 | 36.30 | 130.41 | +| 8 | 3.285 | 3.647 | 17.38 | 29.93 | 214.80 | + +Eight GPUs yield **6.31× fit and 3.18× prediction speedup** over one. First-fit times are 26.442 / 19.890 / 17.765 / 22.531 seconds respectively: the median speedup is not a first-fit speedup. Prediction scaling flattens at this workload, and aggregate memory increases despite lower per-device peaks. + +Complete outputs are bit-identical across all four GPU counts and the native FP16 reference. Selected-512 validation accuracy is 0.765625 and log loss is 0.48790048. A historical native FP32 reference from a different host has accuracy 0.767578125 and log loss 0.48788189, with one changed label relative to FP16. This contrast is descriptive, not a matched current-host precision experiment, a quality improvement or a relaxed distributed parity gate. + +### URL: 2,097,152 context rows + +The input has 500 prepared features: more than one billion context cells. Native one-GPU execution runs out of memory requesting the 504 GiB expanded row buffer; the two-owner executor fails requesting 252 GiB per rank. These failed requests are not observed allocated peaks. + +| GPUs | Outcome | Fit median (s) | Predict 512 rows (s) | Peak allocated / GPU (GiB) | Busiest physical (GiB) | Aligned all-device physical (GiB) | +| ---: | --- | ---: | ---: | ---: | ---: | ---: | +| 1 | Native OOM | — | — | — | — | — | +| 2 | Executor OOM | — | — | — | — | — | +| 4 | Completes | 166.145 | 23.663 | 144.38 | 159.37 | 599.84 | +| 8 | Completes | 85.096 | 11.903 | 80.53 | 97.87 | 695.11 | + +Four to eight GPUs gives **1.95× fit and 1.99× prediction speedup**, with 44% lower busiest-rank allocated memory. First fits are 176.208 and 103.237 seconds. All complete four- and eight-owner predictions are bit-identical; selected-512 validation accuracy is 0.974609375 and log loss is 0.07221999. There is no successful native oracle or selected FP32 quality reference at this size. + +This is near-linear **fixed-context speed scaling from four to eight GPUs**, not a universal linear maximum-context law. It demonstrates a workload that fits with multiple GPUs but not one; it does not locate the largest possible context. The smaller EPSILON and larger URL cases must not be combined into a same-dataset capacity curve. + +### Measurement identities + +The source hash is `sha256(json.dumps(mapping, sort_keys=True).encode())`, where `mapping` contains every recursively found Python file's path relative to `sdm/` (for example, `execution/tabular.py`) and its SHA-256 hex digest. Timing and numerical evidence is tied to that source and the individual workloads, not to every future revision of this branch. + +| Record | SHA-256 | +| --- | --- | +| Consolidated runtime, CPU-validated compatibility successor | `8c00b18e7996a0812c48481bdfe6ac3db4825bf38709c53828f020f098e3f320` | +| EPSILON timing runner | `d8f1b2b3e6f49a5f6cf57639c2d080b881e454fc882578e5e35cf5783aa976f8` | +| EPSILON accepted four-point comparison | `f0d5f989d73b2bddc23daed4e4f542f7fc37a6f4ca83496d85f75e3ab6581a08` | +| URL timing runner | `23a8fa5cd6e293253b0ae219fd045ee93f926d046fdad56b0d4dfc451361a3f9` | +| URL accepted four/eight-GPU comparison | `91adbbae737128724bbe023e7a056066ee1d9727216883bacf99e6eb4ca66880` | + +These identify the retained research records; raw experiment bundles and dataset caches are not shipped in this library branch. The synthetic example is an API example, not a reproduction of these timings. Local-final-normalization, alternative collectives and replicated prediction-KV prototypes are not part of this implementation or its claims. + +### Earlier branch measurements + +The following historical results apply to the original FP32 implementation, not a matched before/after measurement of the current FP16 path. They are retained separately to avoid mixing source versions, precision or timing denominators. + +Separately, published revision `131f794254dccf78e3ea396cc3523d86dfeaacc2` has trained LARGE FP32/E1 CUDA qualification on B200 GPUs. Matched EPSILON C131072, 2000 raw → 500 prepared features, Q512 in original B8 batches completes three fits and twelve predictions on native and one/two/four owners; full saved output, prepared-value, RNG, model-state and axis checks pass with zero output differences or class flips. Separate clean four-owner fit/predict medians are 22.444074/3.771376 seconds versus native 77.864580/9.322108 seconds (3.469×/2.472×). California regression C4097/F8/Q64 passes full all-999-quantile checks across native/one/two/four owners; its single query batch does not establish query-parallel throughput scaling. + +The same production source also has matched EPSILON C262144 native/four/eight-owner qualification under runner `7a40a89b26ad05857f59e9f347fdd2c7436ac0e2c06ce74a4210af48cf891f81`. This workflow uses `column_staging_bytes=8 << 30`, `attention_workspace_bytes=48 << 30`, `fit_head_gather_window=2`, TF32 disabled and the stock default recipe. Clean fit/predict medians are 226.735692/18.492159 seconds native, 61.366311/5.048313 with four owners and 34.625737/3.527150 with eight. The 48-GiB workspace setting is an explicit benchmark component admission override; the example and usage snippet above use the public default of 12 GiB. These results do not establish the same capacity for a default-budget launch. + +At C350000, two/four owners complete where native's first fit OOMs, and their full outputs agree. Clean two→four fit/predict medians are 187.399507→96.843257 and 13.059308→6.636829 seconds. There is no native prediction oracle or native speed denominator at that size. C131072 busiest allocated memory falls from 76.412 GiB native to 30.382 GiB on four owners, while sampled aggregate physical rises from 82.174 to 142.615 GiB. C350000 aggregate physical similarly rises from 211.832 to 252.496 GiB on two→four owners. These sequential, workload-specific measurements are not confidence intervals, universal linear scaling, maximum-context bounds or mixed-precision qualification. Fixed validation subsets are not TEST scores. Later experimental source variants and local documentation/example/test edits do not inherit qualification beyond these recorded source and workflow identities. diff --git a/examples/README.md b/examples/README.md index 467e72d460..5e147e4798 100644 --- a/examples/README.md +++ b/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`. + +For opt-in KumoTabular context sharding across CUDA devices, run `torchrun --standalone --nproc-per-node=2 --max-restarts=0 examples/tabular/context_sharding.py --precision float32`. Select `--precision float16` for experimental FP16 neural execution; weights and numerical inputs remain FP32. The [context-sharding example](tabular/context_sharding.py) demonstrates held-out synthetic regression queries, original query batches and clear/refit/close; see the [scope, memory options and measured scaling](../docs/source/kumotabular-context-sharding.md). Pretrained weights must be available or will be downloaded when the example runs. Use a finite external job timeout. diff --git a/examples/tabular/context_sharding.py b/examples/tabular/context_sharding.py new file mode 100644 index 0000000000..e281efc1fb --- /dev/null +++ b/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() From f34b0e207f8ffa6bda7ced500cefd70387dbf0dd Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 09:04:14 -0700 Subject: [PATCH 05/10] Keep context-sharding research notes outside the library --- docs/source/index.md | 1 - docs/source/kumotabular-context-sharding.md | 173 -------------------- examples/README.md | 2 +- 3 files changed, 1 insertion(+), 175 deletions(-) delete mode 100644 docs/source/kumotabular-context-sharding.md diff --git a/docs/source/index.md b/docs/source/index.md index 767a5cbd2c..feb4079a09 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -30,7 +30,6 @@ overview tensor icl -kumotabular-context-sharding processing nim ``` diff --git a/docs/source/kumotabular-context-sharding.md b/docs/source/kumotabular-context-sharding.md deleted file mode 100644 index ec9e6779d5..0000000000 --- a/docs/source/kumotabular-context-sharding.md +++ /dev/null @@ -1,173 +0,0 @@ -# KumoTabular context sharding - -`sdm.execution.tabular_execution` is an experimental, opt-in executor for splitting one KumoTabular estimator's context across devices. Ordinary `model.fit` and `model.predict` retain their existing behavior when no executor owns the model. - -For large contexts, the measured configuration combines **context-owned row embeddings**, **owner-local frontend computation**, and **sharded row workspaces**. FP16 neural execution is a separate, explicit option. This scales one fitted model's context; it does not train independent models on subsets or ensemble their predictions. - -## Usage - -Run the self-contained [synthetic regression example](../../examples/tabular/context_sharding.py) from the repository root: - -```sh -torchrun --standalone --nproc-per-node=2 --max-restarts=0 examples/tabular/context_sharding.py --precision float32 -``` - -Use `--precision float16` to exercise the FP16 path. The example loads the stock large regression checkpoint on each device (downloading weights if absent), fits preprocessing only on the context rows, predicts original held-out query batches, then clears, refits and closes. Its small synthetic input is not a performance or quality benchmark. Run under a finite external job timeout, with one process per CUDA device; the measured frontend configuration supports one, two, four or eight processes. - -For integration into another script, each rank loads identical weights and creates the same ordered tensor and control groups. The following fragment assumes caller-provided context and query tensors on rank zero's CUDA device; numerical inputs and regression targets are FP32, while classification labels preserve their integer dtype: - -```python -import os -from datetime import timedelta - -import torch -import torch.distributed as dist - -from sdm.execution import tabular_execution -from sdm.models import KumoTabular - -local_rank = int(os.environ["LOCAL_RANK"]) -device = torch.device("cuda", local_rank) -torch.cuda.set_device(device) -dist.init_process_group("nccl", timeout=timedelta(minutes=10)) -control = dist.new_group(backend="gloo", timeout=timedelta(minutes=10)) -root = dist.get_rank() == 0 -model = KumoTabular(task="regression", size="large", device=device).eval() - -run = tabular_execution( - model, - task="regression", - group=dist.group.WORLD, - control_group=control, - fit_head_gather_window=2, - owner_local_frontend=True, - shard_row_workspace=True, - precision="float32", # Explicitly choose "float16" for neural autocast. -) -# Only root supplies tensors. All ranks call every method in the same order. -# x_context: [context_rows, features]; y_context: [context_rows, 1]. -with torch.no_grad(): - run.fit(x_context if root else None, y_context if root else None) - prediction = run.predict(x_query if root else None) - # Original query batches can run concurrently without splitting a batch. - predictions = run.predict_batches(query_batches if root else None) - run.clear() - run.fit(x_context if root else None, y_context if root else None) - run.close() -# The caller owns and destroys groups after healthy completion. -dist.destroy_process_group(control) -dist.destroy_process_group() -``` - -Inputs may be tensors or `TableTensor`s, with shapes `[context_rows, features]`, `[context_rows, 1]` for targets, and `[query_rows, features]`. Root alone supplies inputs, the recipe and an optional preprocessing generator; peers supply `None`. Prediction returns the usual public `TableTensor` on root and `None` on peers; `predict_batches` returns a root-only list in the original order. Each original query batch is kept intact, including recipe output processing and regression inverse transforms. Empty batch lists are supported; individual empty batches are not. `predict` does not automatically partition one query batch. - -## Configuration and compatibility - -The original FP32 behavior remains the default. Enable the memory-saving frontend explicitly: - -| Option | Default | Purpose | -| --- | --- | --- | -| `owner_local_frontend` | `False` | Compute each owner's part of cell and row windows directly, avoiding boundary staging. | -| `shard_row_workspace` | `False` | Divide the native row-block chunk limit by the owner count; requires owner-local frontend and 1, 2, 4 or 8 owners. | -| `precision` | `"float32"` | `"float16"` enables per-worker neural autocast; weights, prepared numerical inputs and public processing remain FP32. | -| `fit_head_gather_window` | `None` | Limit outstanding fit head gathers; the measured configuration uses `2`. | -| `column_staging_bytes` | `8 << 30` | Admission budget for column exchange staging. | -| `attention_workspace_bytes` | `12 << 30` | Admission budget for layer-local prediction attention and cache transport. The large-context measurements below use `48 << 30`. | - -These component budgets are not total device-memory limits. The example retains the public 12 GiB attention budget; it does not reproduce the large-context benchmark configuration automatically. Choose budgets for the workload and device. An oversized original query batch is rejected rather than silently split. - -The FP16 option owns matching autocast scopes on every worker. Do not call `model.half()` or cast preprocessing inputs to FP16. Do not wrap the default FP32 executor in autocast as a substitute for selecting `precision="float16"`. BF16 and half-stored model weights are not supported. FP16 can differ from FP32; same-precision distributed/native agreement is a separate question from precision-related quality changes. - -## What is distributed - -- Expanded row-embedding cells are stored by contiguous context-row owner, reducing the dominant row-by-feature-by-channel allocation. -- Column attention exchanges row-owned cells into complete native column windows, then returns results to their row owners. It does not replace full-context column attention with independent local fits. -- With the default frontend, independent row-attention windows run on their owner and boundary windows retain their native shape. With owner-local frontend enabled, each rank computes its intersections of the native windows. Sharded row workspace additionally bounds the original whole-row-block chunker to `max(1, native_row_limit // owners)`. These options change projection batch geometry and require numerical validation. -- ICL fit attention heads are divided across owners. Full projections and residuals remain replicated; retained ICL key/value caches are sharded by context row. -- Prediction gathers one cache layer at a time and assigns complete original query batches to owners. This combines context-owned storage with query-batch parallelism, not ensemble parallelism. - -This is not complete tensor parallelism and does not guarantee linear capacity or speed scaling. Raw prepared inputs, some normalization/projection intermediates, full ICL states, one replayed cache layer, and a native column window can still scale with the entire context on a single device. Communication and replicated work can dominate small inputs. `column_staging_bytes` and `attention_workspace_bytes` bound individual planned components, not total allocated, reserved, or physical GPU memory. - -For the tested large model, the expanded row tensor has shape `[context_rows, prepared_features + 4, 256]`. At 2,097,152 rows and 500 prepared features in FP16, this tensor alone is 504 GiB without sharding, 126 GiB per rank with four balanced owners, and 63 GiB with eight. Those are component sizes, not predictions of whole-workflow peak memory. Chunk limits bound temporary work, not the total number of context rows. - -## Supported scope and lifecycle - -The supported scope is eager evaluation, FP32 or explicit FP16 neural execution, one unbatched estimator, stock KumoTabular classification or regression, and no autograd inputs, callbacks, custom dispatch/forward hooks, compilation, or ECOC expansion beyond native class capacity. The number of owners must divide both ICL fit query- and key/value-head counts. CPU/Gloo supports numerical and lifecycle tests; CUDA requires NCCL for tensors plus a separate Gloo control group with identical ordered members. Multi-node and heterogeneous-device execution have not been qualified. - -Use normal calls or `torch.no_grad()`, **not an outer `torch.inference_mode()`**: fitted-cache identity checks require tensor version counters. Constructor, fit and prediction reject inference mode. Healthy `clear` and `close` are allowed inside inference mode. - -The executor borrows its model exclusively. Do not mutate weights, recipes, caches, execution settings, or use the underlying model while the lease is active. `fit` replaces a fitted generation; `clear` releases it while retaining the lease; healthy `close` releases the lease and permits ordinary model reuse. Calls must originate on the constructing thread and cannot overlap. Groups and process lifetime always remain caller-owned. - -Local allocation or compute errors are agreed before the next tensor collective, and the executor retains its first failure. A failed executor cannot be cleared, refitted, or reopened; retire the worker processes, especially after CUDA or collective failures. A process crash or mismatched method order still requires the caller's finite process-group timeout and job supervision. Never attempt a collective `close` in a one-rank-only exception handler. - -## Validation status - -The CPU execution tests cover complete public classification outputs and all 999 regression quantiles, one/two/four/eight-owner paths, uneven and empty owners, default-recipe RNG preservation, precision-policy agreement, row workspace bounds, refit, clear, close, borrowed groups, mutation rejection and collective failure propagation. Numerical comparisons use `atol=2e-6, rtol=1e-5`. CPU checks alone do not establish trained CUDA equivalence or performance. - -During consolidation, all 252 execution tests passed on the imported measured runtime. After the FP32 backend compatibility adjustment described below, 83 focused backend-policy, precision and lifecycle tests passed on the final runtime. These sets overlap; they are not 335 distinct tests or a complete final-source suite rerun. Both runs had ten CPU Half-input/FP32-weight RMSNorm fallback warnings and no failures or skips. Ruff, formatting, diff checks and the example's CPU-only `--help` also passed. Run `python -m pytest test/execution` from the repository root for the complete execution suite. - -The measured runtime identified below passed trained FP16 native-versus-eight-owner classification (EPSILON, 4,096 context rows) and all-999-quantile regression (YearPrediction, 16,384 rows) checks, plus the larger EPSILON comparison below. These are selected checkpoints and inputs, not general precision guarantees. Earlier FP16 source attempts failed strict native parity; their results are not silently included in the passing cohort. - -Dedicated CUDA lifecycle diagnostics for this runtime remain pending: earlier attempts failed test collection, and the packaging/device fixture corrections were checked locally only. Older-source CUDA lifecycle passes do not qualify the successor. The trained benchmark runs do exercise repeated fit/predict/clear/close, but are not a substitute for those dedicated failure and ownership tests. - -## Measured FP16 scaling - -These measurements use a single eight-B200 host, PyTorch 2.9.1+cu130, cuDNN 9.2, one large classification estimator, TF32 disabled, both frontend options enabled, `fit_head_gather_window=2`, 8 GiB column staging and 48 GiB attention workspace. Each case has 512 fixed held-out validation queries in original batches of eight. Preprocessing and feature selection use training/context data only. The measured predecessor's 176-file SDM Python source aggregate is `2f5314f0efd69cad02135c9a66a502de5eadd12f2a692a388aefd7e9dab67e43`. - -This branch incorporates that runtime with one compatibility adjustment: cuDNN fit attention is eligible only under CUDA FP16 autocast, preserving the original flash/efficient-only policy for default FP32 calls. The enabled backend set in the measured FP16 path is unchanged. The consolidated source has CPU tests but has not been re-run on GPUs; the following measurements belong to the identified predecessor, not a fresh qualification of this branch revision. - -Each GPU count has one clean multiprocess run, with one worker per GPU, three complete synchronized fits and twelve complete predictions. Fit medians include all three fits; prediction medians use nine calls, excluding the first call after each fit. Memory includes setup, warmups, refits, clear and close on every rank. Allocated memory is the busiest rank's allocator high-water; physical memory is sampled device usage. These are not independent-run confidence intervals or full-process cold-start measurements. - -### EPSILON: 262,144 context rows - -The input has 2,000 raw features and 500 prepared features. The one-GPU denominator is the same executor with the same options, not a different native implementation. - -| GPUs | Fit median (s) | Predict 512 rows (s) | Peak allocated / GPU (GiB) | Busiest physical (GiB) | Aligned all-device physical (GiB) | -| ---: | ---: | ---: | ---: | ---: | ---: | -| 1 | 20.731 | 11.610 | 77.16 | 79.89 | 79.92 | -| 2 | 11.368 | 6.185 | 41.32 | 49.24 | 95.48 | -| 4 | 6.056 | 4.134 | 25.36 | 36.30 | 130.41 | -| 8 | 3.285 | 3.647 | 17.38 | 29.93 | 214.80 | - -Eight GPUs yield **6.31× fit and 3.18× prediction speedup** over one. First-fit times are 26.442 / 19.890 / 17.765 / 22.531 seconds respectively: the median speedup is not a first-fit speedup. Prediction scaling flattens at this workload, and aggregate memory increases despite lower per-device peaks. - -Complete outputs are bit-identical across all four GPU counts and the native FP16 reference. Selected-512 validation accuracy is 0.765625 and log loss is 0.48790048. A historical native FP32 reference from a different host has accuracy 0.767578125 and log loss 0.48788189, with one changed label relative to FP16. This contrast is descriptive, not a matched current-host precision experiment, a quality improvement or a relaxed distributed parity gate. - -### URL: 2,097,152 context rows - -The input has 500 prepared features: more than one billion context cells. Native one-GPU execution runs out of memory requesting the 504 GiB expanded row buffer; the two-owner executor fails requesting 252 GiB per rank. These failed requests are not observed allocated peaks. - -| GPUs | Outcome | Fit median (s) | Predict 512 rows (s) | Peak allocated / GPU (GiB) | Busiest physical (GiB) | Aligned all-device physical (GiB) | -| ---: | --- | ---: | ---: | ---: | ---: | ---: | -| 1 | Native OOM | — | — | — | — | — | -| 2 | Executor OOM | — | — | — | — | — | -| 4 | Completes | 166.145 | 23.663 | 144.38 | 159.37 | 599.84 | -| 8 | Completes | 85.096 | 11.903 | 80.53 | 97.87 | 695.11 | - -Four to eight GPUs gives **1.95× fit and 1.99× prediction speedup**, with 44% lower busiest-rank allocated memory. First fits are 176.208 and 103.237 seconds. All complete four- and eight-owner predictions are bit-identical; selected-512 validation accuracy is 0.974609375 and log loss is 0.07221999. There is no successful native oracle or selected FP32 quality reference at this size. - -This is near-linear **fixed-context speed scaling from four to eight GPUs**, not a universal linear maximum-context law. It demonstrates a workload that fits with multiple GPUs but not one; it does not locate the largest possible context. The smaller EPSILON and larger URL cases must not be combined into a same-dataset capacity curve. - -### Measurement identities - -The source hash is `sha256(json.dumps(mapping, sort_keys=True).encode())`, where `mapping` contains every recursively found Python file's path relative to `sdm/` (for example, `execution/tabular.py`) and its SHA-256 hex digest. Timing and numerical evidence is tied to that source and the individual workloads, not to every future revision of this branch. - -| Record | SHA-256 | -| --- | --- | -| Consolidated runtime, CPU-validated compatibility successor | `8c00b18e7996a0812c48481bdfe6ac3db4825bf38709c53828f020f098e3f320` | -| EPSILON timing runner | `d8f1b2b3e6f49a5f6cf57639c2d080b881e454fc882578e5e35cf5783aa976f8` | -| EPSILON accepted four-point comparison | `f0d5f989d73b2bddc23daed4e4f542f7fc37a6f4ca83496d85f75e3ab6581a08` | -| URL timing runner | `23a8fa5cd6e293253b0ae219fd045ee93f926d046fdad56b0d4dfc451361a3f9` | -| URL accepted four/eight-GPU comparison | `91adbbae737128724bbe023e7a056066ee1d9727216883bacf99e6eb4ca66880` | - -These identify the retained research records; raw experiment bundles and dataset caches are not shipped in this library branch. The synthetic example is an API example, not a reproduction of these timings. Local-final-normalization, alternative collectives and replicated prediction-KV prototypes are not part of this implementation or its claims. - -### Earlier branch measurements - -The following historical results apply to the original FP32 implementation, not a matched before/after measurement of the current FP16 path. They are retained separately to avoid mixing source versions, precision or timing denominators. - -Separately, published revision `131f794254dccf78e3ea396cc3523d86dfeaacc2` has trained LARGE FP32/E1 CUDA qualification on B200 GPUs. Matched EPSILON C131072, 2000 raw → 500 prepared features, Q512 in original B8 batches completes three fits and twelve predictions on native and one/two/four owners; full saved output, prepared-value, RNG, model-state and axis checks pass with zero output differences or class flips. Separate clean four-owner fit/predict medians are 22.444074/3.771376 seconds versus native 77.864580/9.322108 seconds (3.469×/2.472×). California regression C4097/F8/Q64 passes full all-999-quantile checks across native/one/two/four owners; its single query batch does not establish query-parallel throughput scaling. - -The same production source also has matched EPSILON C262144 native/four/eight-owner qualification under runner `7a40a89b26ad05857f59e9f347fdd2c7436ac0e2c06ce74a4210af48cf891f81`. This workflow uses `column_staging_bytes=8 << 30`, `attention_workspace_bytes=48 << 30`, `fit_head_gather_window=2`, TF32 disabled and the stock default recipe. Clean fit/predict medians are 226.735692/18.492159 seconds native, 61.366311/5.048313 with four owners and 34.625737/3.527150 with eight. The 48-GiB workspace setting is an explicit benchmark component admission override; the example and usage snippet above use the public default of 12 GiB. These results do not establish the same capacity for a default-budget launch. - -At C350000, two/four owners complete where native's first fit OOMs, and their full outputs agree. Clean two→four fit/predict medians are 187.399507→96.843257 and 13.059308→6.636829 seconds. There is no native prediction oracle or native speed denominator at that size. C131072 busiest allocated memory falls from 76.412 GiB native to 30.382 GiB on four owners, while sampled aggregate physical rises from 82.174 to 142.615 GiB. C350000 aggregate physical similarly rises from 211.832 to 252.496 GiB on two→four owners. These sequential, workload-specific measurements are not confidence intervals, universal linear scaling, maximum-context bounds or mixed-precision qualification. Fixed validation subsets are not TEST scores. Later experimental source variants and local documentation/example/test edits do not inherit qualification beyond these recorded source and workflow identities. diff --git a/examples/README.md b/examples/README.md index 5e147e4798..c806a52de1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,4 +8,4 @@ This folder contains runnable examples of `structured-data-models`: See [`tabular/quickstart.py`](tabular/quickstart.py) for a minimal runnable `structured-data-models` example with `TabICLv2`. -For opt-in KumoTabular context sharding across CUDA devices, run `torchrun --standalone --nproc-per-node=2 --max-restarts=0 examples/tabular/context_sharding.py --precision float32`. Select `--precision float16` for experimental FP16 neural execution; weights and numerical inputs remain FP32. The [context-sharding example](tabular/context_sharding.py) demonstrates held-out synthetic regression queries, original query batches and clear/refit/close; see the [scope, memory options and measured scaling](../docs/source/kumotabular-context-sharding.md). Pretrained weights must be available or will be downloaded when the example runs. Use a finite external job timeout. +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. From d7777c6f95136fbbc6fdb29cd868f4dce31018cf Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 10:26:56 -0700 Subject: [PATCH 06/10] Preserve padded-column semantics after context-sharding rebase --- sdm/models/kumo/tabular/row_embedding.py | 4 + test/execution/test_tabular_backpressure.py | 8 +- test/execution/test_tabular_padded_columns.py | 77 +++++++++++++++++++ 3 files changed, 86 insertions(+), 3 deletions(-) create mode 100644 test/execution/test_tabular_padded_columns.py diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index 94df1387d2..61fc735b12 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -118,6 +118,10 @@ def forward( num_columns: Tensor | None = None, _runner: Callable[..., Any] | None = None, ) -> Tensor: # [..., R, K * D] + if _runner is not None and num_columns is not None: + raise ValueError( + "Context-sharded row embedding does not support padded columns" + ) *B, R, C = x.size() R_train = y.size(-1) diff --git a/test/execution/test_tabular_backpressure.py b/test/execution/test_tabular_backpressure.py index c03ef27e37..6abf42906a 100644 --- a/test/execution/test_tabular_backpressure.py +++ b/test/execution/test_tabular_backpressure.py @@ -8,13 +8,14 @@ from types import SimpleNamespace import pytest +import torch import sdm.execution._tabular_icl as icl class Value: is_cuda = True - device = "cuda:3" + device = torch.device("cuda:3") shape = (8, 4, 4) def __init__(self, calls: list[str], empty: bool = False) -> None: @@ -118,8 +119,9 @@ def __call__( SimpleNamespace( cuda=SimpleNamespace( Event=Event, - current_stream=lambda device: "stream:" + device, - ) + current_stream=lambda device: f"stream:{device}", + ), + is_autocast_enabled=lambda device_type: False, ), ) monkeypatch.setattr( diff --git a/test/execution/test_tabular_padded_columns.py b/test/execution/test_tabular_padded_columns.py new file mode 100644 index 0000000000..25137c5817 --- /dev/null +++ b/test/execution/test_tabular_padded_columns.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The row executor rejects padding without changing native padded tables.""" + +from unittest.mock import patch + +import pytest +import torch +from test.execution.test_tabular import fixture +from test.execution.test_tabular import group as group + +from sdm import TableTensor, Task +from sdm.cache import Cache +from sdm.execution._collectives import Exchange +from sdm.execution._tabular_rows import Rows + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("owner_local_frontend", [False, True]) +def test_padded_columns_are_native_only(group, task, owner_local_frontend): + model, x, target, _ = fixture(task) + labels = ( + target.categorical.code.squeeze(-1) + if isinstance(target, TableTensor) + else target.squeeze(-1) + ) + encoder = model.models[Task(task)].row_embedding + mask = torch.zeros(x.size(-1), dtype=torch.bool) + valid = x.size(-1) - 1 + + with torch.no_grad(): + native = encoder(x, labels, mask, cache=Cache()) + rows = Rows( + encoder, + x, + labels, + mask, + Exchange(group, group), + 8 << 30, + owner_local_frontend=owner_local_frontend, + ) + sharded = encoder(x, labels, mask, cache=Cache(), _runner=rows) + torch.testing.assert_close(sharded, native, atol=2e-6, rtol=1e-5) + + with ( + patch.object( + torch, + "empty", + side_effect=AssertionError("allocated before padding guard"), + ), + pytest.raises( + ValueError, + match="does not support padded columns", + ), + ): + encoder( + x, + labels, + mask, + cache=Cache(), + num_columns=torch.tensor(valid), + _runner=rows, + ) + + padded = encoder( + x.unsqueeze(0), + labels.unsqueeze(0), + mask.unsqueeze(0), + num_columns=torch.tensor([valid]), + ) + clipped = encoder( + x[None, :, :valid], + labels.unsqueeze(0), + mask[None, :valid], + ) + torch.testing.assert_close(padded, clipped, atol=1e-4, rtol=1e-4) From 2ec572a4cf3500600038ff5679d797c4b5b9070b Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:58:07 -0700 Subject: [PATCH 07/10] Reduce serialization in successful execution control rounds --- sdm/execution/_collectives.py | 14 +++ test/execution/test_control_status.py | 151 ++++++++++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 test/execution/test_control_status.py diff --git a/sdm/execution/_collectives.py b/sdm/execution/_collectives.py index 2a4b5f1067..658b37e1f0 100644 --- a/sdm/execution/_collectives.py +++ b/sdm/execution/_collectives.py @@ -30,6 +30,9 @@ def __init__( 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) @@ -42,6 +45,17 @@ def _records( 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 diff --git a/test/execution/test_control_status.py b/test/execution/test_control_status.py new file mode 100644 index 0000000000..3d0fe0bae9 --- /dev/null +++ b/test/execution/test_control_status.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Control status fast path must preserve mixed nested failure envelopes.""" + +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import group as group + +from sdm.execution._collectives import Exchange + + +def journey(rank: int, world: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=25), + ) + try: + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + original = dist.all_gather_object + with patch.object(dist, "all_gather_object", wraps=original) as gather: + assert exchange.records(None) == [None] * world + exchange.agree(None) + assert exchange.attempt(lambda: None) is None + assert gather.call_count == 0 + + for count, payload in enumerate((0, False, [], ""), start=1): + assert exchange.records(payload if rank == 0 else None) == [ + payload + ] + [None] * (world - 1) + assert gather.call_count == count + assert exchange.metadata({"owner": 0} if rank == 0 else None) == { + "owner": 0 + } + assert gather.call_count == 4 + assert exchange.records(rank) == list(range(world)) + assert gather.call_count == 5 + + blank = Exchange(dist.group.WORLD, dist.group.WORLD) + if rank == 0: + empty_error = ValueError("") + + def fail_empty() -> None: + raise empty_error + + with pytest.raises(ValueError, match=r"^$") as caught_empty: + blank.attempt(fail_empty) + assert caught_empty.value is empty_error + else: + with pytest.raises(RuntimeError) as caught_empty: + blank.records("peer record") + assert str(caught_empty.value) == "" + + failure = ValueError("original inner failure") + + def fail() -> None: + raise failure + + with patch.object(dist, "all_gather_object", wraps=original) as gather: + if rank == 0: + with pytest.raises( + ValueError, match="original inner" + ) as caught: + exchange.attempt(lambda: exchange.attempt(fail)) + assert caught.value is failure + else: + with pytest.raises( + RuntimeError, match="original inner" + ) as caught: + exchange.records("peer record") + assert gather.call_count == 1 + + # An inner failure cannot cause the outer frame or a later call to + # issue another control collective on any rank. + with ( + patch.object( + dist, + "all_reduce", + side_effect=AssertionError("late reduction"), + ), + patch.object( + dist, + "all_gather_object", + side_effect=AssertionError("late object gather"), + ), + ): + with pytest.raises( + (ValueError, RuntimeError), match="original inner" + ) as again: + exchange.records(None) + assert again.value is caught.value + with pytest.raises( + (ValueError, RuntimeError), match="original inner" + ) as once_more: + exchange.agree(None) + assert once_more.value is caught.value + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [1, 2, 4]) +def test_status_and_mixed_failure(tmp_path: Path, world: int) -> None: + if world == 1: + journey(0, world, str(tmp_path)) + else: + mp.spawn(journey, args=(world, str(tmp_path)), nprocs=world) + + +@pytest.mark.parametrize("local_error", [False, True]) +@pytest.mark.parametrize("phase", ["status", "object"]) +def test_transport_failure_preserves_local_error( + group, monkeypatch, local_error, phase +): + exchange = Exchange(group, group) + primary = ValueError("original local failure") + secondary = RuntimeError("status transport failure") + + def fail_transport(*args, **kwargs): + raise secondary + + monkeypatch.setattr( + dist, + "all_reduce" if phase == "status" else "all_gather_object", + fail_transport, + ) + + def operation() -> None: + if local_error: + raise primary + + def invoke() -> object: + if local_error: + return exchange.attempt(operation) + return exchange.records("payload") + + with pytest.raises((ValueError, RuntimeError)) as caught: + invoke() + assert caught.value is (primary if local_error else secondary) + with pytest.raises((ValueError, RuntimeError)) as repeated: + exchange.records(None) + assert repeated.value is caught.value From 06859f5d588035d0f2e2122d0426f205c05b25a8 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:01:01 -0700 Subject: [PATCH 08/10] Gather equal context shards directly into the output tensor --- sdm/execution/_collectives.py | 13 ++ test/execution/test_equal_context_gather.py | 147 ++++++++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 test/execution/test_equal_context_gather.py diff --git a/sdm/execution/_collectives.py b/sdm/execution/_collectives.py index 658b37e1f0..6ab48bf3a7 100644 --- a/sdm/execution/_collectives.py +++ b/sdm/execution/_collectives.py @@ -103,6 +103,19 @@ def gather( 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)) diff --git a/test/execution/test_equal_context_gather.py b/test/execution/test_equal_context_gather.py new file mode 100644 index 0000000000..58c5219f84 --- /dev/null +++ b/test/execution/test_equal_context_gather.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Equal dim-zero context gathering preserves ordered bytes and failures.""" + +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sdm.execution._collectives import Exchange + + +def journey( + rank: int, world: int, directory: str, failure: str | None +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=25), + ) + try: + if failure is not None: + failed_allocation(rank, world, failure) + return + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + gather_into = dist.all_gather_into_tensor + gather_list = dist.all_gather + for dtype in ( + torch.float16, + torch.float32, + torch.float64, + torch.int64, + ): + for lengths, dim in ( + ((3,) * world, 0), + ((3,) * world, -3), + (tuple(range(1, world + 1)), 0), + (tuple(range(world)), 0), + ((0,) * world, 0), + ((3,) * world, 1), + ): + axis = dim % 3 + fast = ( + world > 1 + and axis == 0 + and lengths[0] > 0 + and len(set(lengths)) == 1 + ) + for layout in ("contiguous", "strided", "mixed"): + shape = [2, 2, 4] + shape[axis] = sum(lengths) + expected = torch.arange( + torch.tensor(shape).prod().item(), dtype=dtype + ).reshape(shape) + local = expected.narrow( + axis, sum(lengths[:rank]), lengths[rank] + ).clone() + if layout == "strided" or (layout == "mixed" and rank % 2): + # Preserve values/shape while forcing strided input. + local = ( + local.transpose(0, 2).contiguous().transpose(0, 2) + ) + before = local.clone() + with ( + patch.object( + dist, "all_gather_into_tensor", wraps=gather_into + ) as direct, + patch.object( + dist, "all_gather", wraps=gather_list + ) as general, + ): + actual = exchange.gather(local, lengths, dim) + torch.testing.assert_close( + actual, expected, rtol=0, atol=0 + ) + assert torch.equal( + actual.view(torch.uint8), expected.view(torch.uint8) + ) + torch.testing.assert_close(local, before, rtol=0, atol=0) + assert actual.dtype == dtype + assert actual.is_contiguous() + assert direct.call_count == int(fast) + assert general.call_count == int(world > 1 and not fast) + if world == 1 and local.numel(): + assert actual.data_ptr() != local.data_ptr() + dist.barrier() + finally: + dist.destroy_process_group() + + +def failed_allocation(rank: int, world: int, failure: str) -> None: + exchange = Exchange(dist.group.WORLD, dist.group.WORLD) + local = torch.arange(24).reshape(3, 2, 4).transpose(1, 2) + primary = MemoryError(f"rank-local {failure} failure") + method = "contiguous" if failure == "input" else "new_empty" + original = getattr(torch.Tensor, method) + + def allocate(tensor, *args, **kwargs): + if rank == world - 1 and tensor is local: + raise primary + return original(tensor, *args, **kwargs) + + with ( + patch.object(torch.Tensor, method, allocate), + patch.object( + dist, + "all_gather_into_tensor", + side_effect=AssertionError("unsafe tensor collective"), + ), + pytest.raises((MemoryError, RuntimeError), match=failure) as caught, + ): + exchange.attempt(lambda: exchange.gather(local, (3,) * world, 0)) + if rank == world - 1: + assert caught.value is primary + with ( + patch.object( + dist, "all_reduce", side_effect=AssertionError("late agreement") + ), + pytest.raises((MemoryError, RuntimeError), match=failure) as repeated, + ): + exchange.attempt(lambda: exchange.gather(local, (3,) * world, 0)) + assert repeated.value is caught.value + dist.barrier() + + +@pytest.mark.parametrize("world", [1, 2, 4]) +def test_equal_and_general_gather(tmp_path: Path, world: int) -> None: + if world == 1: + journey(0, world, str(tmp_path), None) + else: + mp.spawn(journey, args=(world, str(tmp_path), None), nprocs=world) + + +@pytest.mark.parametrize("world", [2, 4]) +@pytest.mark.parametrize("failure", ["input", "output"]) +def test_nested_allocation_failure( + tmp_path: Path, world: int, failure: str +) -> None: + mp.spawn(journey, args=(world, str(tmp_path), failure), nprocs=world) From 711a7e2e9139f3ed58ac6fbd92c6e86312158fd5 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:21:00 -0700 Subject: [PATCH 09/10] Add bounded gathered KV reuse for tabular prediction --- sdm/execution/_tabular_cache.py | 71 +++++++- sdm/execution/tabular.py | 47 +++++ test/execution/test_bounded_kv_reuse.py | 233 ++++++++++++++++++++++++ 3 files changed, 349 insertions(+), 2 deletions(-) create mode 100644 test/execution/test_bounded_kv_reuse.py diff --git a/sdm/execution/_tabular_cache.py b/sdm/execution/_tabular_cache.py index ec558924e3..d907f2318e 100644 --- a/sdm/execution/_tabular_cache.py +++ b/sdm/execution/_tabular_cache.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Owning cache slices, with one layer gathered during replay.""" +"""Owning cache slices, optionally retaining bounded ICL replay replicas.""" from collections.abc import Hashable, Iterator -from typing import Self, cast +from math import prod +from typing import Any, Self, cast from torch import Tensor @@ -27,6 +28,8 @@ def __init__( self.dim, self.prefix = dim, prefix self.start = sum(lengths[: exchange.rank]) self.length = lengths[exchange.rank] + self._replicas: dict[str, KVCacheEntry] | None = None + self._building_tensors: list[Tensor] | None = None def __setitem__(self, key: Hashable, value: object) -> None: if isinstance(key, str) and key.startswith(self.prefix): @@ -48,6 +51,8 @@ def __getitem__(self, key: Hashable) -> object: and isinstance(key, str) and key.startswith(self.prefix) ): + if self._replicas is not None: + return self._replicas[key] pair = cast(KVCacheEntry, value) return KVCacheEntry( *( @@ -59,6 +64,63 @@ def __getitem__(self, key: Hashable) -> object: ) return value + def replica_plan(self) -> tuple[tuple[Any, ...], int]: + """Return full-KV geometry and retained bytes on each rank.""" + layout: list[Any] = [] + total = 0 + for key, value in self._items.items(): + if not isinstance(key, str) or not key.startswith(self.prefix): + continue + pair = cast(KVCacheEntry, value) + shapes = [] + for tensor in pair: + shape = list(tensor.shape) + shape[self.dim % tensor.ndim] = sum(self.lengths) + shapes.append((tuple(shape), tensor.dtype)) + total += prod(shape) * tensor.element_size() + layout.append((key, tuple(shapes))) + return tuple(layout), total + + def materialize_replicas(self, layout: tuple[Any, ...]) -> None: + """Gather whole ICL KV once, retaining partial payload on failure.""" + building: dict[str, KVCacheEntry] = {} + building_tensors: list[Tensor] = [] + self._building_tensors = building_tensors + for key, _ in layout: + pair = self.exchange.attempt( + lambda key=key: cast( + KVCacheEntry, Cache.__getitem__(self, key) + ) + ) + full = [] + for tensor in pair: + gathered = self.exchange.gather(tensor, self.lengths, self.dim) + self.exchange.attempt( + lambda gathered=gathered: building_tensors.append(gathered) + ) + full.append(gathered) + self.exchange.attempt( + lambda key=key, full=full: building.__setitem__( + key, KVCacheEntry(*full) + ) + ) + + def install() -> None: + self._replicas = building + self._building_tensors = None + + self.exchange.attempt(install) + + def replica_state(self) -> tuple[Any, ...]: + """Capture mapping identity beside tensor version identities.""" + replicas = self._replicas + return ( + id(replicas), + tuple((key, id(pair)) for key, pair in replicas.items()) + if replicas is not None + else (), + ) + def _tensors(self) -> Iterator[Tensor]: # Iteration over Cache.values() would gather replay entries. for value in self._items.values(): @@ -66,6 +128,11 @@ def _tensors(self) -> Iterator[Tensor]: yield from value elif isinstance(value, Tensor): yield value + if self._replicas is not None: + for pair in self._replicas.values(): + yield from pair + if self._building_tensors is not None: + yield from self._building_tensors def freeze(self) -> Self: # Entries contain owning tensors only; do not invoke replay getters. diff --git a/sdm/execution/tabular.py b/sdm/execution/tabular.py index dbfce356a8..a1046cbb72 100644 --- a/sdm/execution/tabular.py +++ b/sdm/execution/tabular.py @@ -127,6 +127,7 @@ def __init__( owner_local_frontend: bool = False, shard_row_workspace: bool = False, precision: Literal["float32", "float16"] = "float32", + gathered_kv_cache_bytes: int = 0, ) -> None: self.model, self.task, self.exchange = model, task, exchange self._core = cast(_KumoTabular, model.models[task]) @@ -137,6 +138,7 @@ def __init__( self.owner_local_frontend = owner_local_frontend self.shard_row_workspace = shard_row_workspace self.precision = precision + self.gathered_kv_cache_bytes = gathered_kv_cache_bytes self._lock = threading.Lock() self._owner_thread = threading.get_ident() self._thread: int | None = None @@ -149,6 +151,7 @@ def __init__( self._icl_cache: OwnedCache | None = None self._result: Tensor | None = None self._fit_started = False + self._kv_decision_made = False self._identity = exchange.attempt(self._admit) signature = tuple( (value.shape, value.dtype) for value in self._core.parameters() @@ -162,6 +165,7 @@ def __init__( owner_local_frontend, shard_row_workspace, precision, + gathered_kv_cache_bytes, signature, ) ) @@ -186,6 +190,7 @@ def _policy(self) -> tuple[Any, ...]: self.owner_local_frontend, self.shard_row_workspace, self.precision, + self.gathered_kv_cache_bytes, ) def _check_precision(self) -> None: @@ -228,6 +233,13 @@ def _admit(self) -> tuple[Any, ...]: raise ValueError( "Fit head gather window must be a positive integer" ) + if ( + type(self.gathered_kv_cache_bytes) is not int + or self.gathered_kv_cache_bytes < 0 + ): + raise ValueError( + "Gathered KV cache budget must be a nonnegative integer" + ) if ( type(self.model) is not KumoTabular or type(self._core) is not _KumoTabular @@ -294,6 +306,9 @@ def _cache_state(self) -> tuple[Any, ...]: if cache is not None for value in cache._tensors() ), + self._icl_cache.replica_state() + if self._icl_cache is not None + else None, ) @contextmanager @@ -628,11 +643,34 @@ def plan() -> list[list[int]] | None: return waves waves = self.exchange.metadata(self.exchange.attempt(plan)) + if waves: + self._materialize_icl() outputs = [] for wave in waves: outputs.extend(self._predict_wave(wave, batches)) return outputs if self.exchange.rank == 0 else None + def _materialize_icl(self) -> None: + cache = self._icl_cache + assert cache is not None + if ( + self._kv_decision_made + or self.gathered_kv_cache_bytes == 0 + or self.exchange.world == 1 + ): + return + layout, needed = self.exchange.attempt(cache.replica_plan) + decision = (layout, needed <= self.gathered_kv_cache_bytes) + decisions = self.exchange.records(decision) + if any(other != decision for other in decisions): + raise RuntimeError("ICL KV replica plans differ between ranks") + self._kv_decision_made = True + if not decision[1]: + return + cache.materialize_replicas(layout) + # The owned slices and complete replicas become one boundary identity. + self._cache_identity = self._cache_state() + def _predict_wave( self, wave: list[int], batches: Sequence[Tensor | TableTensor] | None ) -> list[TableTensor]: @@ -767,6 +805,7 @@ def _release(self) -> None: self._root_cache = self._batch_cache = None self._row_cache = self._icl_cache = None self._cache_identity = None + self._kv_decision_made = False def clear(self) -> None: """Fence and release the fitted generation, keeping the model lease.""" @@ -797,6 +836,7 @@ def tabular_execution( owner_local_frontend: bool = False, shard_row_workspace: bool = False, precision: Literal["float32", "float16"] = "float32", + gathered_kv_cache_bytes: int = 0, ) -> TabularExecution: """Borrow a loaded model for explicit original-wave execution. @@ -821,6 +861,12 @@ def tabular_execution( ambient autocast. Explicit float16 owns matching autocast scopes on every rank; weights, prepared inputs and original processing islands remain FP32. This is not half-stored model support. + gathered_kv_cache_bytes: Optional per-rank retained-payload budget + for a whole generation's gathered ICL K/V. Default zero keeps + per-layer gathers. A positive budget is used only when all full + layers fit; owned slices remain resident alongside replicas. + This does not bound transient collective workspace or total + allocator/device memory. All ranks construct the executor and call its methods in the same order:: @@ -846,4 +892,5 @@ def tabular_execution( owner_local_frontend, shard_row_workspace, precision, + gathered_kv_cache_bytes, ) diff --git a/test/execution/test_bounded_kv_reuse.py b/test/execution/test_bounded_kv_reuse.py new file mode 100644 index 0000000000..11b7c48911 --- /dev/null +++ b/test/execution/test_bounded_kv_reuse.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Budgeted ICL KV reuse with real CPU owners and public outputs.""" + +import copy +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_tabular import fixture + +import sdm.processing as sp +from sdm import Recipe +from sdm.cache import KVCacheEntry +from sdm.execution import tabular_execution + + +def recipe() -> Recipe: + return Recipe( + target=sp.StypeDispatch(numerical=sp.Standardize()), + output=[sp.AverageEstimators()], + ) + + +def worker( + rank: int, + world: int, + directory: str, + task: str, + precision: str, + budget: int, + mutation: str | None = None, +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=35), + ) + try: + model, x, y, query = fixture(task) + native = copy.deepcopy(model) + with torch.autocast( + "cpu", dtype=torch.float16, enabled=precision == "float16" + ): + native.fit(x, y, recipe=recipe(), num_estimators=1) + expected = native.predict(query[:2]) + run = tabular_execution( + model, + task=task, + group=dist.group.WORLD, + control_group=dist.group.WORLD, + precision=precision, + gathered_kv_cache_bytes=budget, + ) + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + cache = run._icl_cache + assert cache is not None + layout, needed = cache.replica_plan() + assert layout + assert needed > 1 + owned_bytes = cache.size() + assert cache._replicas is None + assert run.predict_batches([] if rank == 0 else None) == ( + [] if rank == 0 else None + ) + assert cache._replicas is None + + with ( + patch.object( + run.exchange, "records", wraps=run.exchange.records + ) as records, + patch.object( + run.exchange, "gather", wraps=run.exchange.gather + ) as gather, + ): + first = run.predict_batches([query[:2]] if rank == 0 else None) + first_gathers = gather.call_count + second = run.predict_batches([query[:2]] if rank == 0 else None) + second_gathers = gather.call_count - first_gathers + assert records.call_count == int(world > 1 and budget > 0) + if world > 1 and budget >= needed: + assert first_gathers == second_gathers + 2 * len(layout) + else: + assert first_gathers == second_gathers + + if rank == 0: + assert first is not None + assert second is not None + for actual in (first[0], second[0]): + assert actual.columns == expected.columns + torch.testing.assert_close( + actual.numerical, + expected.numerical, + atol=2e-6, + rtol=1e-5, + ) + if world > 1 and budget >= needed: + assert cache._replicas is not None + assert cache.size() == owned_bytes + needed + # A hit has no per-layer Gloo handshake or tensor gather. + with ( + patch.object( + run.exchange, + "records", + side_effect=AssertionError("per-hit Gloo record"), + ), + patch.object( + run.exchange, + "gather", + side_effect=AssertionError("per-hit tensor gather"), + ), + ): + for key, _ in layout: + assert isinstance(cache[key], KVCacheEntry) + else: + assert cache._replicas is None + assert cache.size() == owned_bytes + + if mutation is not None: + if rank == 0: + with torch.no_grad(): + if mutation == "owner": + next(cache._tensors()).add_(1) + else: + assert cache._replicas is not None + next(iter(cache._replicas.values())).key.add_(1) + with pytest.raises(RuntimeError, match="fitted cache changed"): + run.predict_batches([query[:2]] if rank == 0 else None) + assert run._icl_cache is cache + return + + if world > 1 and budget >= needed: + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + assert run._icl_cache is not cache + assert run._icl_cache is not None + assert run._icl_cache._replicas is None + run.predict_batches([query[:2]] if rank == 0 else None) + assert run._icl_cache._replicas is not None + run.clear() + assert run._icl_cache is None + if world > 1 and budget >= needed: + run.fit( + x if rank == 0 else None, + y if rank == 0 else None, + recipe=recipe(), + ) + run.predict_batches([query[:2]] if rank == 0 else None) + assert run._icl_cache is not None + assert run._icl_cache._replicas is not None + run.close() + assert run._icl_cache is None + assert model._execution is None + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("task", ["classification", "regression"]) +@pytest.mark.parametrize("world", [1, 2]) +def test_float32_reuse(tmp_path: Path, task: str, world: int) -> None: + args = (world, str(tmp_path), task, "float32", 1 << 20) + if world == 1: + worker(0, *args) + else: + mp.spawn(worker, args=args, nprocs=world) + + +@pytest.mark.parametrize("budget", [0, 1, 1 << 20]) +def test_float16_and_budget(tmp_path: Path, budget: int) -> None: + mp.spawn( + worker, + args=(2, str(tmp_path), "regression", "float16", budget), + nprocs=2, + ) + + +@pytest.mark.parametrize("mutation", ["owner", "replica"]) +def test_mutation_retires(tmp_path: Path, mutation: str) -> None: + mp.spawn( + worker, + args=( + 2, + str(tmp_path), + "classification", + "float32", + 1 << 20, + mutation, + ), + nprocs=2, + ) + + +def mismatch(rank: int, directory: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=2, + timeout=timedelta(seconds=20), + ) + try: + model, _, _, _ = fixture("regression") + with pytest.raises(ValueError, match="policies differ"): + tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + gathered_kv_cache_bytes=1 << 20 if rank == 0 else 0, + ) + assert model._execution is None + finally: + dist.destroy_process_group() + + +def test_rank_budget_mismatch(tmp_path: Path) -> None: + mp.spawn(mismatch, args=(str(tmp_path),), nprocs=2) From 7b37ac52d25c19840eb786d5188734d38ddf2720 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:29:30 -0700 Subject: [PATCH 10/10] Agree on gathered KV replica append failures --- sdm/execution/_tabular_cache.py | 6 +- sdm/execution/tabular.py | 2 + test/execution/test_bounded_kv_guarded.py | 177 ++++++++++++++++++++++ 3 files changed, 183 insertions(+), 2 deletions(-) create mode 100644 test/execution/test_bounded_kv_guarded.py diff --git a/sdm/execution/_tabular_cache.py b/sdm/execution/_tabular_cache.py index d907f2318e..ad30ed35f9 100644 --- a/sdm/execution/_tabular_cache.py +++ b/sdm/execution/_tabular_cache.py @@ -96,9 +96,11 @@ def materialize_replicas(self, layout: tuple[Any, ...]) -> None: for tensor in pair: gathered = self.exchange.gather(tensor, self.lengths, self.dim) self.exchange.attempt( - lambda gathered=gathered: building_tensors.append(gathered) + lambda gathered=gathered, full=full: ( + building_tensors.append(gathered), + full.append(gathered), + ) ) - full.append(gathered) self.exchange.attempt( lambda key=key, full=full: building.__setitem__( key, KVCacheEntry(*full) diff --git a/sdm/execution/tabular.py b/sdm/execution/tabular.py index a1046cbb72..0fa1a118ac 100644 --- a/sdm/execution/tabular.py +++ b/sdm/execution/tabular.py @@ -865,6 +865,8 @@ def tabular_execution( for a whole generation's gathered ICL K/V. Default zero keeps per-layer gathers. A positive budget is used only when all full layers fit; owned slices remain resident alongside replicas. + Each rank retains O(context rows) replica payload, independent + of group size. Single-rank execution ignores this option. This does not bound transient collective workspace or total allocator/device memory. diff --git a/test/execution/test_bounded_kv_guarded.py b/test/execution/test_bounded_kv_guarded.py new file mode 100644 index 0000000000..b12622d29a --- /dev/null +++ b/test/execution/test_bounded_kv_guarded.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import inspect +import sys +import weakref +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from test.execution.test_bounded_kv_reuse import recipe +from test.execution.test_tabular import fixture + +from sdm.execution import tabular_execution +from sdm.execution._tabular_cache import OwnedCache + + +def worker( + rank: int, world: int, directory: str, append_failure: bool +) -> None: + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=15), + ) + try: + base, x, y, query = fixture("regression") + + def make(budget): + model = copy.deepcopy(base) + run = tabular_execution( + model, + task="regression", + group=dist.group.WORLD, + control_group=dist.group.WORLD, + precision="float16", + gathered_kv_cache_bytes=budget, + ) + return model, run + + def fit(run, rows): + run.fit( + x[:rows] if rank == 0 else None, + y[:rows] if rank == 0 else None, + recipe=recipe(), + ) + + def predict(run): + result = run.predict_batches([query[:3]] if rank == 0 else None) + return result[0].numerical.clone() if rank == 0 else None + + if append_failure: + model, run = make(1 << 20) + fit(run, 8) + cache = run._icl_cache + first = MemoryError("rank-local replica append failure") + lines, start = inspect.getsourcelines( + OwnedCache.materialize_replicas + ) + target = start + next( + index + for index, line in enumerate(lines) + if "full.append(gathered)" in line + ) + filename = inspect.getsourcefile(OwnedCache) + + def trace(frame, event, arg): + # Exercise a real Python-list append allocation boundary, + # including when it lives in the attempted lambda's frame. + if ( + event == "line" + and frame.f_code.co_filename == filename + and frame.f_lineno == target + ): + raise first + return trace + + if rank == 1: + sys.settrace(trace) + try: + with pytest.raises( + (MemoryError, RuntimeError), match="replica append failure" + ) as caught: + predict(run) + finally: + sys.settrace(None) + if rank == 1: + assert caught.value is first + assert run._icl_cache is cache + assert model._execution is run + assert cache._replicas is None + assert len(cache._building_tensors) == 1 + refs = [weakref.ref(tensor) for tensor in cache._building_tensors] + with ( + patch.object( + dist, + "all_reduce", + side_effect=AssertionError("late control"), + ), + patch.object( + dist, + "all_gather_into_tensor", + side_effect=AssertionError("late tensor"), + ), + ): + for action in (lambda: predict(run), run.clear, run.close): + with pytest.raises( + (MemoryError, RuntimeError), + match="replica append failure", + ) as repeated: + action() + assert repeated.value is caught.value + assert all(ref() is not None for ref in refs) + else: + _, probe = make(0) + expected, sizes = {}, {} + for rows in (8, 11, 2): + fit(probe, rows) + sizes[rows] = probe._icl_cache.replica_plan()[1] + expected[rows] = predict(probe) + probe.close() + for budget, contexts in ( + (sizes[8], (8, 11)), + (sizes[8] - 1, (8, 2)), + ): + _, run = make(budget) + previous = None + for rows in contexts: + fit(run, rows) + cache = run._icl_cache + assert cache is not previous + owned_bytes = cache.size() + assert run.predict_batches([] if rank == 0 else None) == ( + [] if rank == 0 else None + ) + assert cache._replicas is None + assert not run._kv_decision_made + actual = predict(run) + if rank == 0: + torch.testing.assert_close( + actual, expected[rows], rtol=0, atol=0 + ) + assert actual.size(-1) == 999 + admitted = sizes[rows] <= budget + assert (cache._replicas is not None) == admitted + assert cache._building_tensors is None + assert cache.size() == owned_bytes + ( + sizes[rows] if admitted else 0 + ) + assert run._kv_decision_made + again = predict(run) + if rank == 0: + assert torch.equal(actual, again) + previous = cache + run.clear() + assert run._icl_cache is None + run.close() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 8]) +def test_budget_boundary_and_context_refit(tmp_path: Path, world: int) -> None: + mp.spawn(worker, args=(world, str(tmp_path), False), nprocs=world) + + +def test_rank_local_replica_append_failure(tmp_path: Path) -> None: + mp.spawn(worker, args=(2, str(tmp_path), True), nprocs=2)