diff --git a/examples/README.md b/examples/README.md index 467e72d460..c806a52de1 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`. + +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. 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() 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..6ab48bf3a7 --- /dev/null +++ b/sdm/execution/_collectives.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Caller-owned tensor and control collectives.""" + +from collections.abc import Callable +from typing import Any, TypeVar, cast + +import torch +import torch.distributed as dist +from torch import Tensor + +T = TypeVar("T") + + +class Exchange: + """Keep allocation/compute failures ahead of the next tensor collective.""" + + def __init__( + self, group: dist.ProcessGroup, control_group: dist.ProcessGroup + ) -> None: + self.group, self.control_group = group, control_group + self._failure: BaseException | None = None + self.rank, self.world = ( + dist.get_rank(group), + dist.get_world_size(group), + ) + self.members = tuple(dist.get_process_group_ranks(group)) + if self.members != tuple(dist.get_process_group_ranks(control_group)): + raise ValueError("Tensor and control group membership must match") + if dist.get_backend(control_group) != "gloo": + raise ValueError("The control group must use Gloo") + # Reused only by the owning thread; nested attempt frames complete + # their inner control phase before the outer phase begins. + self._status = torch.zeros(1, dtype=torch.int32, device="cpu") + + def records(self, value: T) -> list[T]: + return self._records(value) + + def _records( + self, value: T, error: BaseException | None = None + ) -> list[T]: + if self._failure is not None: + raise self._failure + values: list[Any] = [None] * self.world + record = (None if error is None else str(error), value) + try: + # Every rank takes the same reduction before deciding whether the + # full envelope is needed. An inner agree(error) may meet a peer's + # records(value), so neither side can skip this phase alone. + self._status.fill_(error is not None or value is not None) + dist.all_reduce( + self._status, + op=dist.ReduceOp.MAX, + group=self.control_group, + ) + if self._status.item() == 0: + return [None] * self.world + dist.all_gather_object(values, record, group=self.control_group) + except BaseException as collective_error: + self._failure = error if error is not None else collective_error + if error is None: + raise + raise error from collective_error + failures = [message for message, _ in values if message is not None] + if failures: + self._failure = ( + error if error is not None else RuntimeError(failures[0]) + ) + raise self._failure + return [item for _, item in values] + + def agree(self, error: BaseException | None) -> None: + # An inner allocation failure can meet a peer's planning record. + # Every control phase consumes the same envelope, then permanently + # stops outer attempt frames from issuing additional collectives. + self._records(None, error) + + def attempt(self, operation: Callable[[], T]) -> T: + error, value = None, None + try: + value = operation() + except BaseException as caught: # noqa: BLE001 + error = caught + self.agree(error) + return cast(T, value) + + def metadata(self, value: T | None) -> T: + values = [value] + dist.broadcast_object_list( + values, src=self.members[0], group=self.control_group + ) + return cast(T, values[0]) + + def broadcast(self, value: Tensor, owner: int = 0) -> Tensor: + if self.world > 1: + dist.broadcast(value, src=self.members[owner], group=self.group) + return value + + def gather( + self, local: Tensor, lengths: tuple[int, ...], dim: int + ) -> Tensor: + if self.world == 1: + return local.clone(memory_format=torch.contiguous_format) + dim %= local.ndim + if ( + dim == 0 + and lengths[0] > 0 + and all(size == lengths[0] for size in lengths) + ): + contiguous, result = self.attempt( + lambda: ( + local.contiguous(), + local.new_empty((sum(lengths), *local.shape[1:])), + ) + ) + dist.all_gather_into_tensor(result, contiguous, group=self.group) + return result + shape = list(local.shape) + shape[dim] = max(lengths) + padded = self.attempt(lambda: local.new_zeros(shape)) + padded.narrow(dim, 0, local.size(dim)).copy_(local) + received = self.attempt( + lambda: [torch.empty_like(padded) for _ in lengths] + ) + dist.all_gather(received, padded, group=self.group) + return self.attempt( + lambda: torch.cat( + [ + part.narrow(dim, 0, size) + for part, size in zip(received, lengths) + ], + dim=dim, + ) + ) + + def exchange( + self, send: Tensor, sent: list[int], received: list[int] + ) -> Tensor: + result = self.attempt(lambda: send.new_empty(sum(received))) + if self.world == 1: + result.copy_(send) + else: + dist.all_to_all_single( + result, + send, + output_split_sizes=received, + input_split_sizes=sent, + group=self.group, + ) + return result diff --git a/sdm/execution/_tabular_cache.py b/sdm/execution/_tabular_cache.py new file mode 100644 index 0000000000..ad30ed35f9 --- /dev/null +++ b/sdm/execution/_tabular_cache.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Owning cache slices, optionally retaining bounded ICL replay replicas.""" + +from collections.abc import Hashable, Iterator +from math import prod +from typing import Any, 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] + 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): + 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) + ): + if self._replicas is not None: + return self._replicas[key] + pair = cast(KVCacheEntry, value) + return KVCacheEntry( + *( + self.exchange.gather( + tensor, self.lengths, self.dim + ).contiguous() + for tensor in pair + ) + ) + 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, full=full: ( + 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(): + if isinstance(value, KVCacheEntry): + 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. + 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..9add05f90d --- /dev/null +++ b/sdm/execution/_tabular_columns.py @@ -0,0 +1,328 @@ +# 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(): + # 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): + 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..86aaabb573 --- /dev/null +++ b/sdm/execution/_tabular_icl.py @@ -0,0 +1,154 @@ +# 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 = 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. + done = torch.cuda.Event() + done.record(torch.cuda.current_stream(result.device)) + pending.append(done) + return result + + # 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) + 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 + * cache_element_size + if len(cache.lengths) > 1 + else compact * cache_element_size + ) + 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 = cache.exchange.attempt(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..ef5df1bf6e --- /dev/null +++ b/sdm/execution/_tabular_rows.py @@ -0,0 +1,406 @@ +# 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, + *, + 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 + 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 * element_size + + self.channels * element_size + ) + 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]] = { + 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, element_size, points, self.total + ), + 65_535, + ) + ow = min( + column.output_block.auto_batch_size_limit( + raw.device, element_size, 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 * 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 + * element_size + ) + 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, + element_size, + 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.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") + + 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 + 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: + 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, + dtype=self.dtype, + ) + 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 + 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) + 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..0fa1a118ac --- /dev/null +++ b/sdm/execution/tabular.py @@ -0,0 +1,898 @@ +# 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 AbstractContextManager, contextmanager, nullcontext +from types import MethodType +from typing import Any, Literal, 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 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, + 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]) + 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.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 + 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._kv_decision_made = 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, + owner_local_frontend, + shard_row_workspace, + precision, + gathered_kv_cache_bytes, + 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, + self.owner_local_frontend, + self.shard_row_workspace, + self.precision, + self.gathered_kv_cache_bytes, + ) + + 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 + ): + 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 + ): + 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" + ) + self._check_precision() + 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() + ), + self._icl_cache.replica_state() + if self._icl_cache is not None + else None, + ) + + @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 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 + ): + raise RuntimeError( + "The borrowed model changed; this execution is stale" + ) + self._check_precision() + 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) + with self._precision_scope(): + 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, + owner_local_frontend=self.owner_local_frontend, + shard_row_workspace=self.shard_row_workspace, + ) + ) + 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() + # 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, + 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)) + 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]: + 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 = 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( + 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 + self._kv_decision_made = False + + 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, + 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. + + 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. + 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. + 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. + 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. + + 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 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. + """ + return TabularExecution( + model, + Task(task), + Exchange(group, control_group), + column_staging_bytes, + attention_workspace_bytes, + fit_head_gather_window, + owner_local_frontend, + shard_row_workspace, + precision, + gathered_kv_cache_bytes, + ) 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..61fc735b12 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,7 +116,12 @@ def forward( *, cache: Cache | None = None, 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) @@ -125,10 +132,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 +154,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 +166,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 +195,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 +222,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 +231,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 +243,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_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) 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) 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 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) 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 new file mode 100644 index 0000000000..7754e8151a --- /dev/null +++ b/test/execution/test_tabular.py @@ -0,0 +1,898 @@ +# 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.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.cache import Cache, KVCacheEntry +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, + 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( + "gloo", + init_method=f"file://{directory}/group", + rank=rank, + world_size=world, + timeout=timedelta(seconds=35), + ) + 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( + TransformerBlock, "auto_batch_size_limit", return_value=2 + ) + if chunked + else nullcontext() + ): + target.fit(*args, **kwargs) + return generator + + 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( + model, + task=task, + 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): + 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, + x if rank == 0 else None, + y if rank == 0 else None, + 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) + 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"]) +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]) +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 + ) + + +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, + ) 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_backpressure.py b/test/execution/test_tabular_backpressure.py new file mode 100644 index 0000000000..6abf42906a --- /dev/null +++ b/test/execution/test_tabular_backpressure.py @@ -0,0 +1,189 @@ +# 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 torch + +import sdm.execution._tabular_icl as icl + + +class Value: + is_cuda = True + device = torch.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: f"stream:{device}", + ), + is_autocast_enabled=lambda device_type: False, + ), + ) + 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] 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) 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)