From 41f9599729143d428d0050123ff55542f550b613 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 29 Oct 2026 11:48:10 -0700 Subject: [PATCH 01/12] initial offloading implementation --- pyproject.toml | 2 + sdm/_token_buffer.py | 351 ++++++++++++++++++++++ sdm/models/base/icl.py | 174 +++++++---- sdm/models/kumo/tabular/cell_embedding.py | 33 +- sdm/models/kumo/tabular/model.py | 13 +- sdm/models/kumo/tabular/row_embedding.py | 185 ++++++++---- sdm/models/tabfm/model.py | 21 +- sdm/models/tabfm/row_embedding.py | 168 ++++++++--- sdm/models/tabiclv2/model.py | 27 +- sdm/models/tabiclv2/row_embedding.py | 196 ++++++++---- test/models/tabiclv2/test_model.py | 1 + test/models/test_offload.py | 213 +++++++++++++ 12 files changed, 1147 insertions(+), 237 deletions(-) create mode 100644 sdm/_token_buffer.py create mode 100644 test/models/test_offload.py diff --git a/pyproject.toml b/pyproject.toml index b0cf8eeb29..386910bec4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -204,6 +204,7 @@ ignorelist = ["id"] [tool.ruff.lint.flake8-tidy-imports] ban-relative-imports = "all" banned-module-level-imports = [ + "cuda", "cudf", "cupy", "graphviz", @@ -245,6 +246,7 @@ banned-module-level-imports = [ [tool.ty.analysis] allowed-unresolved-imports = [ + "cuda.**", "cudf.**", "cupy.**", "pyg_lib.**", diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py new file mode 100644 index 0000000000..6e67e47bc7 --- /dev/null +++ b/sdm/_token_buffer.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from collections.abc import Callable, Sequence +from types import TracebackType +from typing import Any, Literal, Self, cast + +import torch +from torch import Tensor + +from sdm._memory import chunk_memory_limit, set_pinned_max_round_threshold +from sdm.cache import Cache, KVCacheEntry + +# Lets one chunk load, another compute, and a third store concurrently. +_NUM_SLOTS = 3 + + +class TokenBuffer: + r"""A ``[..., R, T, D]`` buffer of ``T`` tokens per row for row embeddings. + + Row embeddings update the buffer in passes: :meth:`rows` updates rows + independently and :meth:`columns` updates token columns independently. + + Without offloading, the buffer is a device tensor and each pass computes + on the whole buffer at once. With ``offload="cpu"``, the buffer lives in + pinned host memory, and each pass streams chunks of it through a few + device slots. Host-to-device copies, compute and device-to-host copies run + on separate streams, so copies of neighboring chunks overlap with compute. + + Args: + size: The buffer size ``[..., R, T, D]``. Offloading requires leading + dimensions of size one, *i.e.*, a single estimator and no class + decomposition. + dtype: The buffer dtype. + device: The compute device. + offload: Where to keep the buffer: ``None`` keeps it on ``device`` and + ``"cpu"`` keeps it in pinned host memory. + """ + + def __init__( + self, + size: Sequence[int], + *, + dtype: torch.dtype, + device: torch.device, + offload: Literal["cpu"] | None = None, + ) -> None: + self.size = torch.Size(size) + self.dtype = dtype + self.device = device + self.offload = offload + + if offload is None: + self.tensor = torch.empty(self.size, dtype=dtype, device=device) + return + + if math.prod(self.size[:-3]) != 1: + raise ValueError( + "'offload' requires 'estimator_batch_size=1' and no more " + "target classes than the model predicts natively (got " + f"buffer size {self.size})" + ) + R, T, D = self.size[-3:] + self._pitch = T * D * dtype.itemsize + set_pinned_max_round_threshold(mb=1) + self._host = torch.empty( + self.size.numel(), dtype=dtype, pin_memory=True + ) + slot_bytes = max( + chunk_memory_limit(device) // 2, + self._pitch, + R * D * dtype.itemsize, + ) + self._slots = [ + torch.empty(slot_bytes, dtype=torch.uint8, device=device) + for _ in range(_NUM_SLOTS) + ] + self._h2d = torch.cuda.Stream(device) + self._d2h = torch.cuda.Stream(device) + self._free = [torch.cuda.Event() for _ in self._slots] + + @classmethod + def from_tensor(cls, tensor: Tensor) -> Self: + r"""Wraps a device tensor as a buffer without offloading.""" + buffer = cls.__new__(cls) + buffer.size = tensor.size() + buffer.dtype = tensor.dtype + buffer.device = tensor.device + buffer.offload = None + buffer.tensor = tensor + return buffer + + def __enter__(self) -> Self: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + if self.offload is None: + return + # The pinned allocator does not track these copies, so the host table + # may only be freed or reused once they complete. + self._h2d.synchronize() + self._d2h.synchronize() + + def readout(self, num_tokens: int) -> Tensor: + r"""Device tensor for the first ``num_tokens`` tokens of every row. + + Without offloading, this is a view of the buffer. Otherwise, it is a + separate ``[..., R, num_tokens, D]`` tensor for the last row pass to + write into. + """ + if self.offload is None: + return self.tensor[..., :num_tokens, :] + return torch.empty( + (*self.size[:-2], num_tokens, self.size[-1]), + dtype=self.dtype, + device=self.device, + ) + + def rows( + self, + compute: Callable[[Tensor, slice], object], + *, + load: bool = True, + store: bool = True, + ) -> None: + r"""Updates rows in place. + + Args: + compute: Called as ``compute(tokens, rows)`` with ``tokens`` of + shape ``[..., N, T, D]`` holding the buffer rows ``rows``. + load: Whether ``compute`` reads ``tokens``. If ``False``, + ``compute`` must fill them. + store: Whether the updated rows are needed afterwards. + """ + if self.offload is None: + compute(self.tensor, slice(0, self.size[-3])) + return + + *B, R, T, D = self.size + step = max(1, len(self._slots[0]) // self._pitch) + chunks = [ + slice(start, min(start + step, R)) for start in range(0, R, step) + ] + self._stream( + chunks=[ + ( + chunk.start * self._pitch, + self._pitch, + chunk.stop - chunk.start, + ) + for chunk in chunks + ], + compute=lambda slot, k: compute( + slot.view(self.dtype).view( + *B, chunks[k].stop - chunks[k].start, T, D + ), + chunks[k], + ), + load=load, + store=store, + ) + + def columns( + self, + compute: Callable[[Tensor], object], + tokens: slice, + ) -> None: + r"""Updates the token columns ``tokens`` in place. + + Args: + compute: Called on chunks of consecutive token columns, in order, + with ``cells`` of shape ``[..., N, R, D]``. + tokens: The token columns to update. + """ + if self.offload is None: + compute(self.tensor[..., tokens, :].transpose(-2, -3)) + return + + *B, R, _, D = self.size + token_bytes = D * self.dtype.itemsize + step = max(1, len(self._slots[0]) // (R * token_bytes)) + chunks = [ + slice(start, min(start + step, tokens.stop)) + for start in range(tokens.start, tokens.stop, step) + ] + self._stream( + chunks=[ + ( + chunk.start * token_bytes, + (chunk.stop - chunk.start) * token_bytes, + R, + ) + for chunk in chunks + ], + compute=lambda slot, k: compute( + slot.view(self.dtype) + .view(*B, R, chunks[k].stop - chunks[k].start, D) + .transpose(-2, -3) + ), + ) + + def _stream( + self, + chunks: list[tuple[int, int, int]], + compute: Callable[[Tensor, int], object], + *, + load: bool = True, + store: bool = True, + ) -> None: + # A chunk holds `height` host rows of `width` bytes at byte `offset`, + # stored densely in its slot. + from cuda.bindings import runtime + + def copy( + dst: int, dpitch: int, src: int, spitch: int, chunk, kind, stream + ) -> None: + _, width, height = chunk + (error,) = runtime.cudaMemcpy2DAsync( + dst, + dpitch, + src, + spitch, + width, + height, + kind, + stream.cuda_stream, + ) + if error != runtime.cudaError_t.cudaSuccess: + raise RuntimeError(f"Host offloading copy failed ({error})") + + host = self._host.data_ptr() + compute_stream = torch.cuda.current_stream(self.device) + loaded = [torch.cuda.Event() for _ in self._slots] + # Stores of the previous pass must land before this pass loads. + self._h2d.wait_stream(self._d2h) + + def issue_load(k: int) -> None: + i = k % len(self._slots) + self._h2d.wait_event(self._free[i]) + copy( + dst=self._slots[i].data_ptr(), + dpitch=chunks[k][1], + src=host + chunks[k][0], + spitch=self._pitch, + chunk=chunks[k], + kind=runtime.cudaMemcpyKind.cudaMemcpyHostToDevice, + stream=self._h2d, + ) + loaded[i].record(self._h2d) + + if load and chunks: + issue_load(0) + for k, (offset, width, height) in enumerate(chunks): + i = k % len(self._slots) + if load and k + 1 < len(chunks): + issue_load(k + 1) + compute_stream.wait_event(loaded[i] if load else self._free[i]) + compute(self._slots[i][: width * height], k) + if not store: + self._free[i].record(compute_stream) + continue + self._d2h.wait_stream(compute_stream) + copy( + dst=host + offset, + dpitch=self._pitch, + src=self._slots[i].data_ptr(), + spitch=width, + chunk=chunks[k], + kind=runtime.cudaMemcpyKind.cudaMemcpyDeviceToHost, + stream=self._d2h, + ) + self._free[i].record(self._d2h) + + +def run_column_layers( + buffer: TokenBuffer, + layers: Sequence[tuple[str, torch.nn.Module]], + *, + tokens: slice, + key_rows: Sequence[Any], + cache: Cache | None, + epilogue: Callable[[Tensor], object] | None = None, +) -> None: + r"""Runs consecutive column layers over ``buffer`` in a single pass. + + Records or replays each layer's key/value projections under its cache key. + + Args: + buffer: The token buffer. + layers: Pairs of cache entry name and column layer, in order. + tokens: The token columns to update. + key_rows: Per layer, the row index of context cells to attend to. + cache: The cache to record into or replay from, or ``None``. + epilogue: Called on each chunk of cells after the last layer. + """ + recording = cache is not None and cache.is_recording + replaying = cache is not None and cache.is_replaying + # Replayed entries span all token columns, so they need a single chunk. + assert not replaying or buffer.offload is None + entries: dict[str, list[KVCacheEntry]] = {key: [] for key, _ in layers} + + def compute(cells: Tensor) -> None: # [..., N, R, D] + for (key, layer), rows in zip(layers, key_rows, strict=True): + result = layer( + query=cells, + key_value=cast(KVCacheEntry, cast(Cache, cache)[key]) + if replaying + else cells[..., rows, :], + return_key_value=recording, + batch_size_limit="auto", + out=cells, + ) + if recording: + entries[key].append(result[1]) + del result + if epilogue is not None: + epilogue(cells) + + buffer.columns(compute, tokens) + + if recording: + assert cache is not None + dim = len(buffer.size) - 3 + for key, parts in entries.items(): + cache[key] = ( + parts[0] + if len(parts) == 1 + else KVCacheEntry( + key=torch.cat([part.key for part in parts], dim=dim), + value=torch.cat([part.value for part in parts], dim=dim), + ) + ) + + +def add_labels( + cells: Tensor, # [..., N, C, D] + y_emb: Tensor, # [..., R_train, D] + rows: slice, +) -> None: + r"""Adds label embeddings to the context rows among ``rows``.""" + stop = min(rows.stop, y_emb.size(-2)) + if stop > rows.start: + cells[..., : stop - rows.start, :, :] += y_emb[ + ..., rows.start : stop, None, : + ].to(cells.dtype) diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 78734c8892..5dc441a525 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -4,7 +4,7 @@ import abc import copy from collections.abc import Hashable, Iterable, Mapping, Sequence -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -57,6 +57,9 @@ class ICLModel(Model): #: Whether this model supports additional related context. supports_related_tables: ClassVar[bool] + #: Whether this model supports offloading intermediate activations. + supports_offload: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) @@ -100,6 +103,7 @@ def forward( estimator_batch_size: int | None = 1, callbacks: Sequence[Callback] | None = None, generator: torch.Generator | None = None, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> TableTensor: # Recipe-defined output shape. r"""The in-context learning forward pass. @@ -129,12 +133,24 @@ def forward( callbacks: Callbacks applied in sequence to this model call. generator: Pseudorandom number generator used for sampling during pre-processing and model execution. + offload: Where to store the largest intermediate activations, + such as the per-cell embeddings of row embedding layers. + Offloading requires inputs on a CUDA device. ``None`` stores + them on the device of the inputs. ``"cpu"`` stores them in + pinned host memory and streams them through the device in + chunks. This lowers peak device memory at the cost of + additional host-device transfers, which overlap with compute + to limit the slowdown. Offloading also requires + :attr:`supports_offload`, disabled gradients, + ``estimator_batch_size=1``, and no more target classes than + the model predicts natively. kwargs: Additional keyword arguments passed to the model. Returns: The processed prediction after applying ``recipe.output`` to the stacked estimator outputs with shape ``[E, ..., R_query, *]``. """ + self._validate_offload(offload, x_query.device) callbacks = () if callbacks is None else callbacks requires_grad = self.training requires_grad |= any(callback.requires_grad for callback in callbacks) @@ -189,14 +205,18 @@ def forward( related_tables=related_query_tables, ) - outs = self._forward_members( - contexts=contexts, - queries=queries, - estimator_batch_size=estimator_batch_size, - callbacks=callbacks, - generator=generator, - **kwargs, - ) + try: + outs = self._forward_members( + contexts=contexts, + queries=queries, + estimator_batch_size=estimator_batch_size, + callbacks=callbacks, + generator=generator, + offload=offload, + **kwargs, + ) + finally: + _release_offload_memory(offload) # Regression: invert target before stacking estimator outputs. if contexts[0].y.numerical.size(-1) > 0: @@ -224,6 +244,7 @@ def fit( estimator_batch_size: int | None = 1, callbacks: Sequence[Callback] | None = None, generator: torch.Generator | None = None, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> None: r"""Fit and cache in-context examples. @@ -255,8 +276,19 @@ def fit( callbacks: Callbacks applied in sequence to this model call. generator: Pseudorandom number generator used for sampling during pre-processing and model execution. + offload: Where to store the largest intermediate activations + while fitting, such as the per-cell embeddings of row + embedding layers. Offloading requires inputs on a CUDA + device. ``None`` stores them on the device of the inputs. + ``"cpu"`` stores them in pinned host memory and streams them + through the device in chunks. This lowers peak device memory + at the cost of additional host-device transfers, which + overlap with compute to limit the slowdown. Offloading also + requires :attr:`supports_offload`, ``estimator_batch_size=1``, + and no more target classes than the model predicts natively. kwargs: Additional keyword arguments passed to the model. """ + self._validate_offload(offload, x.device) if context_only_columns and not self.supports_context_only_columns: raise ValueError( f"{self.__class__.__name__!r} does not support context-only " @@ -301,53 +333,58 @@ def fit( kwargs=kwargs, num_batches=len(batches), ) - for i, batch in enumerate(batches): - with inference_mode("no_grad"): - context = _stack_context(contexts[batch]) - categorical_mask = _categorical_mask(contexts[batch]) - batch_cache = Cache( - x_schemas=tuple( - context.x.schema for context in contexts[batch] - ), - y_schema=context.y.schema, - related_tables_schema=context.related_tables.schema - if context.related_tables is not None - else None, - classes=( - context.y.categorical.categories[0] - if context.y.categorical.size(-1) > 0 - else None - ), - class_values=class_values[batch], - categorical_mask=categorical_mask, - ) - self._forward( - x_context=context.x, - y_context=context.y, - x_query=None, - x_context_only=context.x_context_only, - related_context_tables=context.related_tables, - related_query_tables=None, - cache=batch_cache, - generator=generator, - categorical_mask=categorical_mask, - **kwargs, - ) + try: + for i, batch in enumerate(batches): + with inference_mode("no_grad"): + context = _stack_context(contexts[batch]) + categorical_mask = _categorical_mask(contexts[batch]) + batch_cache = Cache( + x_schemas=tuple( + context.x.schema for context in contexts[batch] + ), + y_schema=context.y.schema, + related_tables_schema=context.related_tables.schema + if context.related_tables is not None + else None, + classes=( + context.y.categorical.categories[0] + if context.y.categorical.size(-1) > 0 + else None + ), + class_values=class_values[batch], + categorical_mask=categorical_mask, + ) + self._forward( + x_context=context.x, + y_context=context.y, + x_query=None, + x_context_only=context.x_context_only, + related_context_tables=context.related_tables, + related_query_tables=None, + cache=batch_cache, + generator=generator, + categorical_mask=categorical_mask, + offload=offload, + **kwargs, + ) - if x.is_cuda and len(contexts) > 1: - try: # Copy to pinned CPU memory: - batch_cache = batch_cache._apply_tensor( - lambda tensor: torch.ops.aten._to_copy.default( - tensor, - device="cpu", - pin_memory=True, - non_blocking=True, # Required for `pin_memory`. + if x.is_cuda and len(contexts) > 1: + try: # Copy to pinned CPU memory: + batch_cache = batch_cache._apply_tensor( + lambda tensor: torch.ops.aten._to_copy.default( + tensor, + device="cpu", + pin_memory=True, + # Required for `pin_memory`: + non_blocking=True, + ) ) - ) - finally: - torch.cuda.current_stream(x.device).synchronize() + finally: + torch.cuda.current_stream(x.device).synchronize() - cache[i] = batch_cache + cache[i] = batch_cache + finally: + _release_offload_memory(offload) self._cache = cache.freeze() @@ -485,10 +522,15 @@ def _forward( generator: torch.Generator | None, *, categorical_mask: Tensor, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> TableTensor: # [..., R_query, *] r"""Run the model on preprocessed tables of one estimator batch. + :meth:`forward` calls this without a cache, :meth:`fit` records the + context into a cache, and :meth:`predict` replays that cache with + new queries. + Tables carry shape ``[E, ..., R, D]`` when ``E > 1`` estimators run together and ``[..., R, D]`` otherwise. Column names and category vocabularies are those of the first estimator; per-estimator column @@ -514,6 +556,10 @@ def _forward( numerical feature columns that were categorical before preprocessing. Passed on recording, replaying and uncached calls alike. + offload: Where to store the largest intermediate activations, as + in :meth:`forward`. Only passed to models with + :attr:`supports_offload`. Always ``None`` when replaying a + cache, *i.e.*, in :meth:`predict`. kwargs: Additional keyword arguments passed by the caller. Returns: @@ -526,6 +572,22 @@ def _forward( # Helpers ################################################################# + def _validate_offload( + self, + offload: Literal["cpu"] | None, + device: torch.device, + ) -> None: + if offload is None: + return + if not self.supports_offload: + raise ValueError( + f"{self.__class__.__name__!r} does not support 'offload'" + ) + if device.type != "cuda": + raise ValueError( + f"'offload' requires inputs on a CUDA device (got {device})" + ) + def _forward_members( self, contexts: Sequence[MemberContext], @@ -765,6 +827,12 @@ def _validate_query( ) +def _release_offload_memory(offload: Literal["cpu"] | None) -> None: + # The pinned allocator otherwise keeps offloaded buffers cached for reuse. + if offload == "cpu": + torch._C._host_emptyCache() + + def _stack(tables: Sequence[TableTensor]) -> TableTensor: ref = tables[0] if len(tables) == 1: diff --git a/sdm/models/kumo/tabular/cell_embedding.py b/sdm/models/kumo/tabular/cell_embedding.py index 7529820b50..1fca97b458 100644 --- a/sdm/models/kumo/tabular/cell_embedding.py +++ b/sdm/models/kumo/tabular/cell_embedding.py @@ -37,29 +37,40 @@ def __init__( dtype=dtype, ) + def impute_mean( + self, + x: Tensor, # [..., R, C] + *, + train_size: int | None = None, + cache: Cache | None = None, + ) -> Tensor: # [..., 1, C] + if cache is not None and cache.is_replaying: + return cast(Tensor, cache["cell_impute_mean"]) + + assert train_size is not None + impute_mean = x[..., :train_size, :].nanmean(dim=-2, keepdim=True) + impute_mean = impute_mean.masked_fill(impute_mean.isnan(), 0.0) + if cache is not None and cache.is_recording: + cache["cell_impute_mean"] = impute_mean + return impute_mean + def forward( # type: ignore self, x: Tensor, # [..., R, C] categorical_mask: Tensor, # [..., C] *, train_size: int | None = None, + impute_mean: Tensor | None = None, # [..., 1, C] num_columns: Tensor | None = None, # [...] cache: Cache | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, ) -> Tensor: # [..., R, C, D] missing = x.isnan() - - if cache is not None and cache.is_replaying: - impute_mean = cast(Tensor, cache["cell_impute_mean"]) - else: - assert train_size is not None - impute_mean = x[..., :train_size, :].nanmean(dim=-2, keepdim=True) - impute_mean = impute_mean.masked_fill(impute_mean.isnan(), 0.0) - - if cache is not None and cache.is_recording: - cache["cell_impute_mean"] = impute_mean - + if impute_mean is None: + impute_mean = self.impute_mean( + x, train_size=train_size, cache=cache + ) x = torch.where(missing, impute_mean, x) return super().forward( diff --git a/sdm/models/kumo/tabular/model.py b/sdm/models/kumo/tabular/model.py index b946f963e3..1fc5dfa810 100644 --- a/sdm/models/kumo/tabular/model.py +++ b/sdm/models/kumo/tabular/model.py @@ -112,6 +112,7 @@ class KumoTabular(ICLModel): supports_multi_target: ClassVar[bool] = False supports_context_only_columns: ClassVar[bool] = False supports_related_tables: ClassVar[bool] = False + supports_offload: ClassVar[bool] = True def __init__( self, @@ -180,6 +181,7 @@ def _forward( generator: torch.Generator | None, *, categorical_mask: Tensor, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> TableTensor: # [..., R_query, num_classes or 999] @@ -214,6 +216,7 @@ def _forward( y=y, categorical_mask=categorical_mask, cache=cache, + offload=offload, ) return TableTensor( columns={ @@ -231,6 +234,7 @@ def _forward( cache=cache, generator=generator, categorical_mask=categorical_mask, + offload=offload, ) if cache is not None and cache.is_replaying: columns = cast(tuple[str, ...], cache["class_columns"]) @@ -304,8 +308,15 @@ def forward( categorical_mask: Tensor, # [..., C] *, cache: Cache | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R_test, num_classes or num_quantiles] - x = self.row_embedding(x, y, categorical_mask, cache=cache) + x = self.row_embedding( + x, + y, + categorical_mask, + cache=cache, + offload=offload, + ) x = self.row_project(x) return self.icl_block(x=x, y=y, cache=cache) diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index 6714c5ce5e..2da8762d70 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -3,12 +3,13 @@ # ruff: noqa: D101, D102 -from typing import Any, cast +from typing import Any, Literal, cast import torch from torch import Tensor from torch.nn import Embedding, Linear, ModuleList, Parameter +from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.kumo.tabular.block import KumoTabularTransformerBlock from sdm.models.kumo.tabular.cell_embedding import CellEmbedding @@ -114,6 +115,7 @@ def forward( *, cache: Cache | None = None, num_columns: Tensor | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R, K * D] *B, R, C = x.size() @@ -126,57 +128,129 @@ def forward( else (K + num_columns).unsqueeze(-1).to(torch.int32) ) - buffer: Tensor | None = None + y_emb: Tensor | None = None # [..., R_train, D] + if y.numel() > 0: + if self.y_emb is not None: + y_emb = self.y_emb(y) + else: + assert self.y_lin is not None + y_emb = self.y_lin(y.unsqueeze(-1)) + if torch.is_grad_enabled(): - x = self.cell_embedding( + if offload is not None: + raise RuntimeError( + "'offload' is only supported when gradients are disabled" + ) + return self._forward_autograd( x=x, + y_emb=y_emb, categorical_mask=categorical_mask, train_size=R_train, - num_columns=num_columns, cache=cache, - ) # [..., R, C, D] - else: - buffer = torch.empty( - (*B, R, K + C, D), - device=x.device, - dtype=torch.get_autocast_dtype(x.device.type) - if torch.is_autocast_enabled(x.device.type) - else x.dtype, + num_columns=num_columns, + row_key_lengths=row_key_lengths, ) - buffer[..., :K, :] = self.readout_token.to(buffer.dtype) - x = self.cell_embedding( - x=x, + + impute_mean = self.cell_embedding.impute_mean( + x, + train_size=R_train, + cache=cache, + ) + buffer = TokenBuffer( + (*B, R, K + C, D), + dtype=torch.get_autocast_dtype(x.device.type) + if torch.is_autocast_enabled(x.device.type) + else x.dtype, + device=x.device, + offload=offload, + ) + + def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] + tokens[..., :K, :] = self.readout_token.to(tokens.dtype) + self.cell_embedding( + x=x[..., rows, :], categorical_mask=categorical_mask, - train_size=R_train, + impute_mean=impute_mean, num_columns=num_columns, - cache=cache, batch_size_limit="auto", - out=buffer[..., K:, :], + out=tokens[..., K:, :], ) + if y_emb is not None: + add_labels(tokens[..., K:, :], y_emb, rows) - if y.numel() > 0: - if self.y_emb is not None: - y_emb = self.y_emb(y).unsqueeze(-2) # [..., R_train, 1, D] - else: - assert self.y_lin is not None - y_emb = self.y_lin(y.unsqueeze(-1)).unsqueeze(-2) - x[..., :R_train, :, :] += y_emb.to(x.dtype) + with buffer: + buffer.rows(embed, load=False) + y_emb = None - if buffer is not None: - x = x.transpose(-2, -3) # [..., C, R, D] + readout = buffer.readout(K) # [..., R, K, D] + for i, (col_block, row_block) in enumerate( + zip(self.col_blocks, self.row_blocks) + ): + run_column_layers( + buffer, + [(f"row_embedding.col_block{i}", col_block)], + tokens=slice(K, K + C), + key_rows=[slice(R_train)], + cache=cache, + ) + + last = i == len(self.row_blocks) - 1 + + def update( + tokens: Tensor, # [..., N, K + C, D] + rows: slice, + row_block: torch.nn.Module = row_block, + last: bool = last, + ) -> None: + row_block( + query=tokens[..., :K, :] if last else tokens, + key_value=tokens, + seqused_key_value=row_key_lengths, + batch_size_limit="auto", + out=readout[..., rows, :, :] if last else tokens, + ) + + buffer.rows(update, store=not last) + + return self.norm(readout).flatten(-2) + + def _forward_autograd( + self, + x: Tensor, # [..., R, C] + y_emb: Tensor | None, # [..., R_train, D] + categorical_mask: Tensor, # [..., C] + *, + train_size: int, + cache: Cache | None, + num_columns: Tensor | None, + row_key_lengths: Tensor | None, + ) -> Tensor: # [..., R, K * D] + *B, R, _ = x.size() + R_train = train_size + K = self.readout_token.size(-2) + D = self.channels + + x = self.cell_embedding( + x=x, + categorical_mask=categorical_mask, + train_size=R_train, + num_columns=num_columns, + cache=cache, + ) # [..., R, C, D] + if y_emb is not None: + x[..., :R_train, :, :] += y_emb.unsqueeze(-2).to(x.dtype) for i, (col_block, row_block) in enumerate( zip(self.col_blocks, self.row_blocks) ): - if buffer is None: - if i > 0: - readout_token, x = x.split([K, x.size(-2) - K], dim=-2) - readout_token = readout_token.clone() - else: - readout_token = self.readout_token - readout_token = readout_token.view(*(1,) * len(B), 1, K, D) - readout_token = readout_token.expand(*B, R, K, D) - x = x.transpose(-2, -3) # [..., C, R, D] + if i > 0: + readout_token, x = x.split([K, x.size(-2) - K], dim=-2) + readout_token = readout_token.clone() + else: + readout_token = self.readout_token + readout_token = readout_token.view(*(1,) * len(B), 1, K, D) + readout_token = readout_token.expand(*B, R, K, D) + x = x.transpose(-2, -3) # [..., C, R, D] key = f"row_embedding.col_block{i}" result = col_block( @@ -186,39 +260,22 @@ def forward( else x[..., :R_train, :], return_key_value=cache is not None and cache.is_recording, batch_size_limit="auto", - out=None if buffer is None else x, ) - if cache is not None and cache.is_recording: x, cache[key] = result else: x = result del result - if buffer is None: - x = torch.cat( - [readout_token.to(x.dtype), x.transpose(-2, -3)], - dim=-2, - ) - x = row_block( - query=x[..., :K, :] - if i == len(self.row_blocks) - 1 - else x, - key_value=x, - seqused_key_value=row_key_lengths, - batch_size_limit="auto", - ) - else: - buffer = row_block( - query=buffer[..., :K, :] - if i == len(self.row_blocks) - 1 - else buffer, - key_value=buffer, - seqused_key_value=row_key_lengths, - batch_size_limit="auto", - out=buffer[..., :K, :] - if i == len(self.row_blocks) - 1 - else buffer, - ) + x = torch.cat( + [readout_token.to(x.dtype), x.transpose(-2, -3)], + dim=-2, + ) + x = row_block( + query=x[..., :K, :] if i == len(self.row_blocks) - 1 else x, + key_value=x, + seqused_key_value=row_key_lengths, + batch_size_limit="auto", + ) - return self.norm(buffer if buffer is not None else x).flatten(-2) + return self.norm(x).flatten(-2) diff --git a/sdm/models/tabfm/model.py b/sdm/models/tabfm/model.py index 3919981287..866c71e0e1 100644 --- a/sdm/models/tabfm/model.py +++ b/sdm/models/tabfm/model.py @@ -20,7 +20,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -82,6 +82,7 @@ class TabFM(ICLModel): supports_multi_target: ClassVar[bool] = False supports_context_only_columns: ClassVar[bool] = False supports_related_tables: ClassVar[bool] = False + supports_offload: ClassVar[bool] = True def __init__( self, @@ -158,6 +159,7 @@ def _forward( generator: torch.Generator | None, *, categorical_mask: Tensor, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> TableTensor: # [..., R_query, num_classes or 1] @@ -193,7 +195,13 @@ def _forward( categorical_mask = categorical_mask.expand(*x.size()[:-2], -1) task = Task.classification if classes is not None else Task.regression - out = self.models[task](x, y, categorical_mask, cache=cache) + out = self.models[task]( + x, + y, + categorical_mask, + cache=cache, + offload=offload, + ) if classes is None: return TableTensor( @@ -257,6 +265,13 @@ def forward( categorical_mask: Tensor, # [..., C] *, cache: Cache | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R_test, 1 or num_classes] - x = self.row_embedding(x, y, categorical_mask, cache=cache) + x = self.row_embedding( + x, + y, + categorical_mask, + cache=cache, + offload=offload, + ) return self.icl_block(x, y, cache=cache) diff --git a/sdm/models/tabfm/row_embedding.py b/sdm/models/tabfm/row_embedding.py index d49a7d22d4..3294c6abbf 100644 --- a/sdm/models/tabfm/row_embedding.py +++ b/sdm/models/tabfm/row_embedding.py @@ -17,7 +17,7 @@ # ruff: noqa: D101, D102 -from typing import Any, cast +from typing import Any, Literal, cast import torch from torch import Tensor @@ -31,6 +31,7 @@ Sequential, ) +from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.tabfm.block import TabFMTransformerBlock from sdm.models.tabfm.cell_embedding import CellEmbedding @@ -142,6 +143,7 @@ def forward( categorical_mask: Tensor, # [..., C] *, cache: Cache | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R, K * D] *B, R, C = x.size() @@ -149,32 +151,123 @@ def forward( K = self.readout_token.size(-2) D = self.channels - buffer: Tensor | None = None + y_emb: Tensor | None = None # [..., R_train, D] + if y.numel() > 0: + if self.y_emb is not None: + y_emb = self.y_emb(y) + else: + assert self.y_mlp is not None + y_emb = self.y_mlp(y.unsqueeze(-1)) + if torch.is_grad_enabled(): - x = self.cell_embedding(x, categorical_mask) # [..., R, C, D] - else: - buffer = torch.empty( - (*B, R, K + C, D), - device=x.device, - dtype=torch.get_autocast_dtype(x.device.type) - if torch.is_autocast_enabled(x.device.type) - else x.dtype, + if offload is not None: + raise RuntimeError( + "'offload' is only supported when gradients are disabled" + ) + return self._forward_autograd( + x=x, + y_emb=y_emb, + categorical_mask=categorical_mask, + train_size=R_train, + cache=cache, ) - buffer[..., :K, :] = self.readout_token.to(buffer.dtype) - x = self.cell_embedding( - x, + + buffer = TokenBuffer( + (*B, R, K + C, D), + dtype=torch.get_autocast_dtype(x.device.type) + if torch.is_autocast_enabled(x.device.type) + else x.dtype, + device=x.device, + offload=offload, + ) + + def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] + tokens[..., :K, :] = self.readout_token.to(tokens.dtype) + self.cell_embedding( + x[..., rows, :], categorical_mask, batch_size_limit="auto", - out=buffer[..., K:, :], + out=tokens[..., K:, :], ) + if y_emb is not None: + add_labels(tokens[..., K:, :], y_emb, rows) - if y.numel() > 0: - if self.y_emb is not None: - y_emb = self.y_emb(y).unsqueeze(-2) # [..., R_train, 1, D] - else: - assert self.y_mlp is not None - y_emb = self.y_mlp(y.unsqueeze(-1)).unsqueeze(-2) - x[..., :R_train, :, :] += y_emb.to(x.dtype) + with buffer: + buffer.rows(embed, load=False) + y_emb = None + + readout = buffer.readout(K) # [..., R, K, D] + for i, (col_layers, col_proj, row_layers, row_norm) in enumerate( + zip( + self.col_blocks, + self.col_projections, + self.row_blocks, + self.row_norms, + ) + ): + # Later repeats also attend over readout tokens column-wise. + run_column_layers( + buffer, + [ + (f"row_embedding.col_layer{i}.{j}", col_layer) + for j, col_layer in enumerate(col_layers) + ], + tokens=slice(K if i == 0 else 0, K + C), + key_rows=[slice(R_train)] * len(col_layers), + cache=cache, + epilogue=lambda cells, col_proj=col_proj: cells.copy_( + col_proj(cells) + ), + ) + + last = i == len(self.row_blocks) - 1 + + def update( + tokens: Tensor, # [..., N, K + C, D] + rows: slice, + row_layers: ModuleList = row_layers, + row_norm: torch.nn.Module = row_norm, + last: bool = last, + ) -> None: + x = tokens + for j, row_layer in enumerate(row_layers): + if last and j == len(row_layers) - 1: + x = row_layer( + query=x[..., :K, :], + key_value=x, + batch_size_limit="auto", + out=readout[..., rows, :, :], + ) + else: + x = row_layer( + query=x, + key_value=x, + batch_size_limit="auto", + out=x, + ) + x.copy_(row_norm(x)) + + buffer.rows(update, store=not last) + + return readout.flatten(-2) + + def _forward_autograd( + self, + x: Tensor, # [..., R, C] + y_emb: Tensor | None, # [..., R_train, D] + categorical_mask: Tensor, # [..., C] + *, + train_size: int, + cache: Cache | None, + ) -> Tensor: # [..., R, K * D] + *B, R, _ = x.size() + R_train = train_size + K = self.readout_token.size(-2) + D = self.channels + + x = self.cell_embedding(x, categorical_mask) # [..., R, C, D] + if y_emb is not None: + x[..., :R_train, :, :] += y_emb.unsqueeze(-2).to(x.dtype) for i, (col_layers, col_proj, row_layers, row_norm) in enumerate( zip( @@ -195,7 +288,6 @@ def forward( else x[..., :R_train, :], return_key_value=cache is not None and cache.is_recording, batch_size_limit="auto", - out=None if buffer is None else x, ) # [..., C, R, D] if cache is not None and cache.is_recording: @@ -204,21 +296,17 @@ def forward( x = result del result - if buffer is None: - x = col_proj(x.transpose(-2, -3)) # [..., R, C, D] - if i == 0: # Prepend readout tokens before row-wise attention. - x = torch.cat( - [ - self.readout_token.to(x.dtype) - .view(*(1,) * len(B), 1, K, D) - .expand(*B, R, K, D), - x, # [..., R, C, D] - ], - dim=-2, - ) # [..., R, K + C, D] - else: - x.copy_(col_proj(x)) # NOTE Double peak memory. - x = buffer # [..., R, K + C, D] + x = col_proj(x.transpose(-2, -3)) # [..., R, C, D] + if i == 0: # Prepend readout tokens before row-wise attention. + x = torch.cat( + [ + self.readout_token.to(x.dtype) + .view(*(1,) * len(B), 1, K, D) + .expand(*B, R, K, D), + x, # [..., R, C, D] + ], + dim=-2, + ) # [..., R, K + C, D] # Row-wise attention (B * R as the batch axis). for j, row_layer in enumerate(row_layers): @@ -230,12 +318,8 @@ def forward( query=query, # [..., R, K + C, D] or [..., R, K, D] key_value=x, # [..., R, K + C, D] batch_size_limit="auto", - out=None if buffer is None else query, ) # [..., R, K + C, D] or [..., R, K, D] - if buffer is None: - x = row_norm(x) - else: - x.copy_(row_norm(x)) # NOTE Double peak memory. + x = row_norm(x) return x.flatten(-2) diff --git a/sdm/models/tabiclv2/model.py b/sdm/models/tabiclv2/model.py index 30bdc52784..f72f5c99a2 100644 --- a/sdm/models/tabiclv2/model.py +++ b/sdm/models/tabiclv2/model.py @@ -9,7 +9,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -127,6 +127,7 @@ class TabICLv2(ICLModel): supports_multi_target: ClassVar[bool] = False supports_context_only_columns: ClassVar[bool] = False supports_related_tables: ClassVar[bool] = False + supports_offload: ClassVar[bool] = True def __init__( self, @@ -188,6 +189,8 @@ def _forward( related_query_tables: RelatedTables[TableTensor] | None, cache: Cache | None, generator: torch.Generator | None, + *, + offload: Literal["cpu"] | None = None, **kwargs: Any, ) -> TableTensor: # [..., R_query, num_classes or 999] @@ -216,7 +219,12 @@ def _forward( ) if classes is None: - out = self.models[Task.regression](x, y, cache=cache) + out = self.models[Task.regression]( + x, + y, + cache=cache, + offload=offload, + ) return TableTensor( columns={ Stype.numerical: [f"q{i:03d}" for i in range(1, 1000)] @@ -225,7 +233,11 @@ def _forward( ) out = self.models[Task.classification]( - x, y, cache=cache, num_classes=len(classes) + x, + y, + cache=cache, + num_classes=len(classes), + offload=offload, ) return TableTensor( columns={Stype.numerical: [str(i) for i in classes.tolist()]}, @@ -282,11 +294,18 @@ def forward( *, cache: Cache | None = None, num_classes: int | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R_test, out_channels or num_classes] if not y.is_floating_point(): assert num_classes is not None - x = self.row_embedding(x, y, num_classes=num_classes, cache=cache) + x = self.row_embedding( + x, + y, + num_classes=num_classes, + cache=cache, + offload=offload, + ) return self.icl_block( x=x, y=y, diff --git a/sdm/models/tabiclv2/row_embedding.py b/sdm/models/tabiclv2/row_embedding.py index caf0ab3a4b..15f49cf18d 100644 --- a/sdm/models/tabiclv2/row_embedding.py +++ b/sdm/models/tabiclv2/row_embedding.py @@ -7,12 +7,13 @@ # ruff: noqa: D101, D102 import math -from typing import Any, cast +from typing import Any, Literal, cast import torch from torch import Tensor from torch.nn import Embedding, LayerNorm, ModuleList, Parameter +from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.tabiclv2.block import TabICLv2TransformerBlock from sdm.nn import InducedTransformerBlock, RotaryEmbedding @@ -104,18 +105,41 @@ def forward( num_classes: int | None = None, cache: Cache | None = None, generator: torch.Generator | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R, K * D] + if offload is not None and train_mask is not None: + raise ValueError("'offload' does not support 'train_mask'") + *B, R, C = x.size() R_train = y.size(-1) G, D = self.lin.in_features, self.lin.out_features K = self.readout_token.size(-2) - train_mask: Any = slice(R_train) if train_mask is None else train_mask # Feature grouping: gather G columns into each token. shift = 2 ** torch.arange(G, device=x.device) index = torch.arange(C, device=x.device) index = (index.view(C, 1) + shift.view(1, G)) % C # [C, G] - x = x[..., index] # [..., R, C, G] + + # Context rows each column layer attends to: + key_rows: list[Any] = [] + if cache is None or not cache.is_replaying: + rows = slice(R_train) if train_mask is None else train_mask + for _ in self.col_layers: + if max_keys is None or R_train <= max_keys: + key_rows.append(rows) + continue + perm = torch.randperm( + R_train, + device=x.device, + generator=generator, + )[:max_keys] + key_rows.append( + perm + if train_mask is None + else train_mask.nonzero().squeeze(-1)[perm] + ) + else: + key_rows = [None] * len(self.col_layers) num_digits = 1 if ( @@ -128,63 +152,130 @@ def forward( if y.numel() > 0: y = _mixed_radix_digits(y, bases) # [F, ..., R_train] - buffer: Tensor | None = None + y_emb: Tensor | None = None # [..., R_train, D] + if y.numel() > 0: + if self.y_emb is not None: + y_emb = self.y_emb(y) + else: + assert self.y_lin is not None + y_emb = self.y_lin(y.unsqueeze(-1)) + if torch.is_grad_enabled(): - x = self.lin(x) # [..., R, C, D] - if num_digits > 1: - if y.numel() > 0: - x = x.repeat(num_digits, *(1,) * x.dim()) - else: - x = x.expand(num_digits, *x.size()) - else: - buffer = torch.empty( - ((num_digits,) if num_digits > 1 else ()) - + (*x.size()[:-3], R, K + C, D), - device=x.device, - dtype=torch.get_autocast_dtype(x.device.type) - if torch.is_autocast_enabled(x.device.type) - else x.dtype, + if offload is not None: + raise RuntimeError( + "'offload' is only supported when gradients are disabled" + ) + return self._forward_autograd( + x=x[..., index], + y_emb=y_emb, + train_mask=slice(R_train) + if train_mask is None + else train_mask, + key_rows=key_rows, + num_digits=num_digits, + cache=cache, ) + + buffer = TokenBuffer( + ((num_digits,) if num_digits > 1 else ()) + (*B, R, K + C, D), + dtype=torch.get_autocast_dtype(x.device.type) + if torch.is_autocast_enabled(x.device.type) + else x.dtype, + device=x.device, + offload=offload, + ) + + def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] + cells = tokens[..., K:, :] if num_digits > 1: - self.lin(x, out=buffer[0, ..., K:, :]) - buffer[1:, ..., K:, :] = buffer[:1, ..., K:, :] - x = buffer[..., K:, :] + self.lin(x[..., rows, :][..., index], out=cells[0]) + cells[1:] = cells[:1] + else: + self.lin(x[..., rows, :][..., index], out=cells) + if y_emb is None: + return + if train_mask is None: + add_labels(cells, y_emb, rows) else: - x = self.lin(x, out=buffer[..., K:, :]) # [..., R, C, D] + cells[..., train_mask, :, :] += y_emb.unsqueeze(-2).to( + cells.dtype + ) + + with buffer: + buffer.rows(embed, load=False) + y_emb = None + + run_column_layers( + buffer, + [ + (f"row_embedding.col_layer{i}", col_layer) + for i, col_layer in enumerate(self.col_layers) + ], + tokens=slice(K, K + C), + key_rows=key_rows, + cache=cache, + ) - if y.numel() > 0: - if self.y_emb is not None: - y_emb = self.y_emb(y).unsqueeze(-2) # [..., R_train, 1, D] + if num_digits > 1: # Average over mixed-radix digits. + buffer = TokenBuffer.from_tensor(buffer.tensor.mean(dim=0)) + + with buffer: + readout = buffer.readout(K) # [..., R, K, D] + + def update( + tokens: Tensor, # [..., N, K + C, D] + rows: slice, + ) -> None: + tokens[..., :K, :] = self.readout_token.to(tokens.dtype) + for i, row_layer in enumerate(self.row_layers): + last = i == len(self.row_layers) - 1 + row_layer( + query=tokens[..., :K, :] if last else tokens, + key_value=tokens, # [..., N, K + C, D] + batch_size_limit="auto", + out=readout[..., rows, :, :] if last else tokens, + ) + + buffer.rows(update, store=False) + + return self.norm(readout).view(*B, R, K * D) # [..., R, K * D] + + def _forward_autograd( + self, + x: Tensor, # [..., R, C, G] + y_emb: Tensor | None, # [..., R_train, D] + *, + train_mask: Any, + key_rows: list[Any], + num_digits: int, + cache: Cache | None, + ) -> Tensor: # [..., R, K * D] + *B, R, _, _ = x.size() + D = self.lin.out_features + K = self.readout_token.size(-2) + + x = self.lin(x) # [..., R, C, D] + if num_digits > 1: + if y_emb is not None: + x = x.repeat(num_digits, *(1,) * x.dim()) else: - assert self.y_lin is not None - y_emb = self.y_lin(y.unsqueeze(-1)).unsqueeze(-2) + x = x.expand(num_digits, *x.size()) - x[..., train_mask, :, :] += y_emb.to(x.dtype) + if y_emb is not None: + x[..., train_mask, :, :] += y_emb.unsqueeze(-2).to(x.dtype) # Column-wise induced set attention (B * C as the batch axis). x = x.transpose(-2, -3) # [..., C, R, D] - for i, col_layer in enumerate(self.col_layers): + for i, (col_layer, rows) in enumerate(zip(self.col_layers, key_rows)): key = f"row_embedding.col_layer{i}" - if cache is not None and cache.is_replaying: - key_value = cast(KVCacheEntry, cache[key]) - else: - key_value = x[..., train_mask, :] - if max_keys is not None and key_value.size(-2) > max_keys: - index = torch.randperm( - key_value.size(-2), - device=key_value.device, - generator=generator, - )[:max_keys] - key_value = key_value[..., index, :] - result = col_layer( query=x, # [..., C, R, D] - key_value=key_value, # [..., C, R_train, D] + key_value=cast(KVCacheEntry, cache[key]) + if cache is not None and cache.is_replaying + else x[..., rows, :], # [..., C, R_train, D] return_key_value=cache is not None and cache.is_recording, batch_size_limit="auto", - out=None if torch.is_grad_enabled() else x, ) # [..., C, R, D] - del key_value if cache is not None and cache.is_recording: x, cache[key] = result @@ -195,21 +286,13 @@ def forward( x = x.transpose(-2, -3) # [..., R, C, D] if num_digits > 1: # Average over mixed-radix digits. - if buffer is not None: - buffer = buffer.mean(dim=0) # [..., R, K + C, D] - else: - x = x.mean(dim=0) # [..., R, C, D] + x = x.mean(dim=0) # [..., R, C, D] # Prepend readout tokens before row-wise attention. readout_token = self.readout_token.to(x.dtype) readout_token = readout_token.view(*(1,) * len(B), 1, K, D) readout_token = readout_token.expand(*B, R, K, D) - - if buffer is not None: - buffer[..., :, :K, :] = readout_token - x = buffer - else: - x = torch.cat([readout_token, x], dim=-2) # [..., R, K + C, D] + x = torch.cat([readout_token, x], dim=-2) # [..., R, K + C, D] # Row-wise attention (B * R as the batch axis). for i, row_layer in enumerate(self.row_layers): @@ -217,11 +300,6 @@ def forward( query=x[..., :K, :] if i == len(self.row_layers) - 1 else x, key_value=x, # [..., R, K + C, D] batch_size_limit="auto", - out=None - if torch.is_grad_enabled() - else x[..., :K, :] - if i == len(self.row_layers) - 1 - else x, ) # [..., R, K + C, D] or [..., R, K, D] return self.norm(x).view(*B, R, K * D) # [..., R, K * D] diff --git a/test/models/tabiclv2/test_model.py b/test/models/tabiclv2/test_model.py index e33cd6dcf1..f0695260d2 100644 --- a/test/models/tabiclv2/test_model.py +++ b/test/models/tabiclv2/test_model.py @@ -211,6 +211,7 @@ def forward( *, num_classes: int | None = None, cache: object | None = None, + offload: str | None = None, ) -> torch.Tensor: return x diff --git a/test/models/test_offload.py b/test/models/test_offload.py new file mode 100644 index 0000000000..935c94cd26 --- /dev/null +++ b/test/models/test_offload.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Callable +from functools import partial +from typing import Any + +import pytest +import torch +from torch import Tensor + +import sdm.processing as sp +from sdm import CategoricalTensor, Stype, TableTensor +from sdm.cache import Cache +from sdm.models import ICLModel, KumoRelational, KumoTabular, TabFM, TabICLv2 +from sdm.models.kumo.tabular.row_embedding import ( + RowEmbedding as KumoTabularRowEmbedding, +) +from sdm.models.tabfm.row_embedding import RowEmbedding as TabFMRowEmbedding +from sdm.models.tabiclv2.row_embedding import ( + RowEmbedding as TabICLv2RowEmbedding, +) +from sdm.testing import onlyCUDA + + +def _randomize(module: torch.nn.Module) -> torch.nn.Module: + # Zero-initialized residual branches would hide the encoder layers. + for parameter in module.parameters(): + if not parameter.any(): + torch.nn.init.normal_(parameter, std=0.02) + return module + + +def _kumo_tabular(device: torch.device) -> Callable[..., Tensor]: + module = _randomize( + KumoTabularRowEmbedding( + num_classes=3, + channels=16, + num_layers=2, + num_heads=2, + group_size=3, + num_frequencies=8, + num_inducing_points=4, + num_readout_tokens=2, + device=device, + ) + ) + return lambda x, y, **kwargs: module( + x, + y, + torch.zeros(x.size(-1), dtype=torch.bool, device=device), + **kwargs, + ) + + +def _tabfm(device: torch.device) -> Callable[..., Tensor]: + module = _randomize( + TabFMRowEmbedding( + num_classes=3, + channels=16, + num_layers=2, + num_repeats=2, + num_col_heads=2, + num_row_heads=2, + group_size=3, + num_frequencies=8, + num_inducing_points=4, + num_readout_tokens=2, + device=device, + ) + ) + return lambda x, y, **kwargs: module( + x, + y, + torch.zeros(x.size(-1), dtype=torch.bool, device=device), + **kwargs, + ) + + +def _tabiclv2(device: torch.device) -> Callable[..., Tensor]: + module = _randomize( + TabICLv2RowEmbedding( + num_classes=3, + channels=16, + num_layers=2, + num_heads=2, + group_size=3, + num_inducing_points=4, + num_readout_tokens=2, + norm_bias=True, + device=device, + ) + ) + return lambda x, y, **kwargs: module(x, y, num_classes=3, **kwargs) + + +@onlyCUDA +@pytest.mark.parametrize("build", [_kumo_tabular, _tabfm, _tabiclv2]) +@pytest.mark.parametrize("chunk_memory_fraction", ["0.05", "1e-9"]) +def test_row_embedding_offload( + build: Callable[[torch.device], Callable[..., Tensor]], + chunk_memory_fraction: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SDM_CHUNK_MEMORY_FRACTION", chunk_memory_fraction) + device = torch.device("cuda") + embedding = build(device) + x = torch.randn(1, 40, 5, device=device) + y = torch.randint(0, 3, (1, 30), device=device) + + outs: dict[Any, Tensor] = {} + with torch.no_grad(): + for offload in (None, "cpu"): + cache = Cache() + outs[offload, "fit"] = embedding( + x[:, :30], y, cache=cache, offload=offload + ) + outs[offload, "predict"] = embedding( + x[:, 30:], y[:, :0], cache=cache.freeze() + ) + outs[offload] = embedding(x, y, offload=offload) + + for key in ("fit", "predict"): + torch.testing.assert_close(outs["cpu", key], outs[None, key]) + torch.testing.assert_close(outs["cpu"], outs[None]) + + +@onlyCUDA +def test_row_embedding_offload_subsamples_keys() -> None: + device = torch.device("cuda") + embedding = _tabiclv2(device) + x = torch.randn(1, 40, 5, device=device) + y = torch.randint(0, 3, (1, 30), device=device) + + with torch.no_grad(): + outs = [ + embedding( + x, + y, + max_keys=10, + generator=torch.Generator(device).manual_seed(0), + offload=offload, + ) + for offload in (None, "cpu") + ] + + torch.testing.assert_close(outs[1], outs[0]) + + +def test_offload_requires_cuda_inputs() -> None: + model = KumoTabular(task="classification", size="small", pretrained=False) + x = TableTensor.from_tensor(torch.randn(8, 5)) + y = TableTensor( + columns={Stype.categorical: ("target",)}, + categorical=CategoricalTensor( + code=torch.randint(0, 3, (6, 1)), + categories=(torch.arange(3),), + ), + ) + + with pytest.raises(ValueError, match="CUDA"): + model.fit(x[:6], y, offload="cpu") + with pytest.raises(ValueError, match="CUDA"): + model(x[:6], y, x[6:], offload="cpu") + + +def test_offload_requires_model_support() -> None: + model = KumoRelational(pretrained=False) + x = TableTensor.from_tensor(torch.randn(6, 5)) + y = TableTensor.from_tensor(torch.randn(6, 1)) + + with pytest.raises(ValueError, match="does not support 'offload'"): + model.fit(x, y, offload="cpu") + + +@onlyCUDA +@pytest.mark.parametrize( + "model_cls", + [partial(KumoTabular, size="small"), TabFM, TabICLv2], +) +def test_fit_predict_offload(model_cls: Callable[..., ICLModel]) -> None: + device = torch.device("cuda") + model = model_cls(task="classification", pretrained=False, device=device) + _randomize(model) + x = TableTensor.from_tensor(torch.randn(48, 5, device=device)) + y = TableTensor( + columns={Stype.categorical: ("target",)}, + categorical=CategoricalTensor( + code=torch.randint(0, 3, (32, 1), device=device), + categories=(torch.arange(3, device=device),), + ), + ) + recipe = sp.Recipe( + features=[sp.ToNumerical()], + output=[sp.AverageEstimators()], + ) + + outs = [] + for offload in (None, "cpu"): + model.fit(x[:32], y, recipe=recipe, num_estimators=2, offload=offload) + outs.append(model.predict(x[32:])) + + torch.testing.assert_close(outs[1].numerical, outs[0].numerical) + + with pytest.raises(ValueError, match="estimator_batch_size"): + model.fit( + x[:32], + y, + recipe=recipe, + num_estimators=2, + estimator_batch_size=None, + offload="cpu", + ) From e20505c2d746e5d3736803bb0ea5b59ea061fb59 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 30 Oct 2026 00:12:27 -0700 Subject: [PATCH 02/12] store token buffers column-major and align attention chunks to batch dims --- sdm/_token_buffer.py | 58 ++++++++------- sdm/models/tabfm/row_embedding.py | 3 + sdm/nn/attention.py | 117 +++++++++++------------------- sdm/nn/linear.py | 7 ++ 4 files changed, 86 insertions(+), 99 deletions(-) diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 6e67e47bc7..9f377c8eee 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -17,11 +17,14 @@ class TokenBuffer: - r"""A ``[..., R, T, D]`` buffer of ``T`` tokens per row for row embeddings. + r"""A buffer of ``T`` tokens per row for row embeddings. Row embeddings update the buffer in passes: :meth:`rows` updates rows independently and :meth:`columns` updates token columns independently. + The buffer is stored column-major as ``[..., T, R, D]``. Column passes + see slices of it, and row passes see ``[..., R, T, D]`` views. + Without offloading, the buffer is a device tensor and each pass computes on the whole buffer at once. With ``offload="cpu"``, the buffer lives in pinned host memory, and each pass streams chunks of it through a few @@ -50,27 +53,29 @@ def __init__( self.dtype = dtype self.device = device self.offload = offload + *B, R, T, D = self.size if offload is None: - self.tensor = torch.empty(self.size, dtype=dtype, device=device) + self.tensor = torch.empty( + (*B, T, R, D), dtype=dtype, device=device + ) return - if math.prod(self.size[:-3]) != 1: + if math.prod(B) != 1: raise ValueError( "'offload' requires 'estimator_batch_size=1' and no more " "target classes than the model predicts natively (got " f"buffer size {self.size})" ) - R, T, D = self.size[-3:] - self._pitch = T * D * dtype.itemsize + self._pitch = R * D * dtype.itemsize set_pinned_max_round_threshold(mb=1) self._host = torch.empty( self.size.numel(), dtype=dtype, pin_memory=True ) slot_bytes = max( chunk_memory_limit(device) // 2, + T * D * dtype.itemsize, self._pitch, - R * D * dtype.itemsize, ) self._slots = [ torch.empty(slot_bytes, dtype=torch.uint8, device=device) @@ -82,9 +87,10 @@ def __init__( @classmethod def from_tensor(cls, tensor: Tensor) -> Self: - r"""Wraps a device tensor as a buffer without offloading.""" + r"""Wraps a column-major ``[..., T, R, D]`` device tensor.""" buffer = cls.__new__(cls) - buffer.size = tensor.size() + *B, T, R, D = tensor.size() + buffer.size = torch.Size((*B, R, T, D)) buffer.dtype = tensor.dtype buffer.device = tensor.device buffer.offload = None @@ -115,7 +121,7 @@ def readout(self, num_tokens: int) -> Tensor: write into. """ if self.offload is None: - return self.tensor[..., :num_tokens, :] + return self.tensor[..., :num_tokens, :, :].transpose(-2, -3) return torch.empty( (*self.size[:-2], num_tokens, self.size[-1]), dtype=self.dtype, @@ -139,27 +145,28 @@ def rows( store: Whether the updated rows are needed afterwards. """ if self.offload is None: - compute(self.tensor, slice(0, self.size[-3])) + compute(self.tensor.transpose(-2, -3), slice(0, self.size[-3])) return *B, R, T, D = self.size - step = max(1, len(self._slots[0]) // self._pitch) + token_bytes = D * self.dtype.itemsize + step = max(1, len(self._slots[0]) // (T * token_bytes)) chunks = [ slice(start, min(start + step, R)) for start in range(0, R, step) ] self._stream( chunks=[ ( - chunk.start * self._pitch, - self._pitch, - chunk.stop - chunk.start, + chunk.start * token_bytes, + (chunk.stop - chunk.start) * token_bytes, + T, ) for chunk in chunks ], compute=lambda slot, k: compute( - slot.view(self.dtype).view( - *B, chunks[k].stop - chunks[k].start, T, D - ), + slot.view(self.dtype) + .view(*B, T, chunks[k].stop - chunks[k].start, D) + .transpose(-2, -3), chunks[k], ), load=load, @@ -179,12 +186,11 @@ def columns( tokens: The token columns to update. """ if self.offload is None: - compute(self.tensor[..., tokens, :].transpose(-2, -3)) + compute(self.tensor[..., tokens, :, :]) return *B, R, _, D = self.size - token_bytes = D * self.dtype.itemsize - step = max(1, len(self._slots[0]) // (R * token_bytes)) + step = max(1, len(self._slots[0]) // self._pitch) chunks = [ slice(start, min(start + step, tokens.stop)) for start in range(tokens.start, tokens.stop, step) @@ -192,16 +198,16 @@ def columns( self._stream( chunks=[ ( - chunk.start * token_bytes, - (chunk.stop - chunk.start) * token_bytes, - R, + chunk.start * self._pitch, + self._pitch, + chunk.stop - chunk.start, ) for chunk in chunks ], compute=lambda slot, k: compute( - slot.view(self.dtype) - .view(*B, R, chunks[k].stop - chunks[k].start, D) - .transpose(-2, -3) + slot.view(self.dtype).view( + *B, chunks[k].stop - chunks[k].start, R, D + ) ), ) diff --git a/sdm/models/tabfm/row_embedding.py b/sdm/models/tabfm/row_embedding.py index 3294c6abbf..da2e013ce0 100644 --- a/sdm/models/tabfm/row_embedding.py +++ b/sdm/models/tabfm/row_embedding.py @@ -245,6 +245,9 @@ def update( batch_size_limit="auto", out=x, ) + # Normalize in column-major storage order, since + # strided inputs make RMSNorm copy. + x = x.transpose(-2, -3) x.copy_(row_norm(x)) buffer.rows(update, store=not last) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index fa4721ff7e..c45b3cf526 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -3,6 +3,7 @@ """Attention modules for structured tensor models.""" +import itertools import math from typing import Any, Literal, overload @@ -692,81 +693,64 @@ def forward( out=out, ) - flat_out: Tensor | None = None - flat_key: Tensor | None = None - flat_value: Tensor | None = None - - if out is not None and (out.dim() <= 3 or out.is_contiguous()): - flat_out = out.view(batch_size, *query.size()[-2:]) - - for start in range(0, batch_size, batch_size_limit): - end = min(start + batch_size_limit, batch_size) + # Slice batch dimension `dim` and take later batch dimensions whole, + # so that every chunk is a view of up to `batch_size_limit` elements. + dim, inner = len(batch_shape) - 1, 1 + while inner * batch_shape[dim] <= batch_size_limit: + inner *= batch_shape[dim] + dim -= 1 + step = batch_size_limit // inner + chunks = [ + (*index, slice(start, min(start + step, batch_shape[dim]))) + for index in itertools.product(*map(range, batch_shape[:dim])) + for start in range(0, batch_shape[dim], step) + ] + key: Tensor | None = None + value: Tensor | None = None + for index in chunks: result = self._forward( - query=_chunk(query, batch_shape, 2, start, end), - key_value=_chunk(key_value, batch_shape, 2, start, end), + query=_chunk(query, batch_shape, 2, index), + key_value=_chunk(key_value, batch_shape, 2, index), seqused_key_value=_chunk( - seqused_key_value, - batch_shape, - trailing_dims=0, - start=start, - end=end, + seqused_key_value, batch_shape, 0, index ), - attn_mask=_chunk(attn_mask, batch_shape, 3, start, end), - query_pos=_chunk(query_pos, batch_shape, 1, start, end), - key_pos=_chunk(key_pos, batch_shape, 1, start, end), + attn_mask=_chunk(attn_mask, batch_shape, 3, index), + query_pos=_chunk(query_pos, batch_shape, 1, index), + key_pos=_chunk(key_pos, batch_shape, 1, index), is_causal=is_causal, return_key_value=return_key_value, - out=flat_out[start:end] if flat_out is not None else None, + out=None if out is None else out[index], ) if isinstance(result, tuple): chunk, chunk_kv = result - - if flat_key is None: - flat_key = chunk_kv.key.new_empty( - batch_size, *chunk_kv.key.size()[-3:] + if key is None or value is None: + key = chunk_kv.key.new_empty( + batch_shape + chunk_kv.key.size()[-3:] ) - if flat_value is None: - flat_value = chunk_kv.value.new_empty( - batch_size, *chunk_kv.value.size()[-3:] + value = chunk_kv.value.new_empty( + batch_shape + chunk_kv.value.size()[-3:] ) - - flat_key[start:end] = chunk_kv.key - flat_value[start:end] = chunk_kv.value - + key[index] = chunk_kv.key + value[index] = chunk_kv.value del chunk_kv - else: chunk = result - if flat_out is None and out is not None: - flat_index = torch.arange(start, end, device=out.device) - batch_indices: list[Tensor] = [] - for size in reversed(batch_shape): - batch_indices.append(flat_index % size) - flat_index = flat_index // size - out[tuple(reversed(batch_indices))] = chunk.to(out.dtype) - elif flat_out is None: - flat_out = chunk.new_empty((batch_size, *query.size()[-2:])) - flat_out[start:end] = chunk + if out is None: + out = chunk.new_empty(batch_shape + query.size()[-2:]) + out[index] = chunk del chunk, result - if out is None: - assert flat_out is not None - out = flat_out.view(*batch_shape, *query.size()[-2:]) - + assert out is not None if not return_key_value: return out - assert flat_key is not None - assert flat_value is not None - - return out, KVCacheEntry( - key=flat_key.view(*batch_shape, *flat_key.size()[-3:]), - value=flat_value.view(*batch_shape, *flat_value.size()[-3:]), - ) + assert key is not None + assert value is not None + return out, KVCacheEntry(key=key, value=value) def _forward( self, @@ -885,8 +869,7 @@ def _chunk( tensor: Tensor, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> Tensor: ... @@ -895,8 +878,7 @@ def _chunk( tensor: KVCacheEntry, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> KVCacheEntry: ... @@ -905,8 +887,7 @@ def _chunk( tensor: None, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> None: ... @@ -914,8 +895,7 @@ def _chunk( tensor: Tensor | KVCacheEntry | None, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> Tensor | KVCacheEntry | None: if tensor is None: @@ -923,18 +903,9 @@ def _chunk( if isinstance(tensor, KVCacheEntry): return KVCacheEntry( - key=_chunk(tensor.key, batch_shape, 3, start, end), - value=_chunk(tensor.value, batch_shape, 3, start, end), + key=_chunk(tensor.key, batch_shape, 3, index), + value=_chunk(tensor.value, batch_shape, 3, index), ) trailing_shape = tensor.size()[-trailing_dims:] if trailing_dims else () - tensor = tensor.expand(batch_shape + trailing_shape) - if len(batch_shape) == 1: - return tensor[start:end] - - flat_index = torch.arange(start, end, device=tensor.device) - batch_indices: list[Tensor] = [] - for size in reversed(batch_shape): - batch_indices.append(flat_index % size) - flat_index = flat_index // size - return tensor[tuple(reversed(batch_indices))] + return tensor.expand(batch_shape + trailing_shape)[index] diff --git a/sdm/nn/linear.py b/sdm/nn/linear.py index cce1e1a2d1..72e75584c3 100644 --- a/sdm/nn/linear.py +++ b/sdm/nn/linear.py @@ -31,6 +31,13 @@ def forward( out.copy_(torch.matmul(input, weight)) elif input.dim() == 2: torch.matmul(input, weight, out=out) + elif any( + out.stride(i) != out.size(i + 1) * out.stride(i + 1) + for i in range(out.dim() - 3) + ): + for input_i, out_i in zip(input, out): + self.forward(input_i, out=out_i) + return out else: batch_size = math.prod(input.shape[:-2]) input = input.view(batch_size, input.size(-2), input.size(-1)) From aec83b5062b2118c99d8067ef39e82f7d737e28d Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 30 Oct 2026 00:29:22 -0700 Subject: [PATCH 03/12] update documentation --- sdm/_token_buffer.py | 30 ++++++++++++++++-------------- sdm/models/base/icl.py | 32 +++++++++++++++----------------- 2 files changed, 31 insertions(+), 31 deletions(-) diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 9f377c8eee..186c56e6a7 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -32,13 +32,15 @@ class TokenBuffer: on separate streams, so copies of neighboring chunks overlap with compute. Args: - size: The buffer size ``[..., R, T, D]``. Offloading requires leading - dimensions of size one, *i.e.*, a single estimator and no class - decomposition. + size: The buffer size ``[..., R, T, D]`` of ``R`` rows with ``T`` + tokens of ``D`` channels each. Offloading requires leading + dimensions of size one, *i.e.*, a single estimator and no more + target classes than the model predicts natively. dtype: The buffer dtype. device: The compute device. - offload: Where to keep the buffer: ``None`` keeps it on ``device`` and - ``"cpu"`` keeps it in pinned host memory. + offload: Where to store the buffer: ``None`` stores it on ``device`` + and ``"cpu"`` stores it in pinned host memory, which requires a + CUDA ``device``. """ def __init__( @@ -64,8 +66,7 @@ def __init__( if math.prod(B) != 1: raise ValueError( "'offload' requires 'estimator_batch_size=1' and no more " - "target classes than the model predicts natively (got " - f"buffer size {self.size})" + "target classes than the model predicts natively" ) self._pitch = R * D * dtype.itemsize set_pinned_max_round_threshold(mb=1) @@ -87,7 +88,7 @@ def __init__( @classmethod def from_tensor(cls, tensor: Tensor) -> Self: - r"""Wraps a column-major ``[..., T, R, D]`` device tensor.""" + r"""Wrap a column-major ``[..., T, R, D]`` device tensor.""" buffer = cls.__new__(cls) *B, T, R, D = tensor.size() buffer.size = torch.Size((*B, R, T, D)) @@ -135,7 +136,7 @@ def rows( load: bool = True, store: bool = True, ) -> None: - r"""Updates rows in place. + r"""Update rows in place, in chunks of ``N`` rows. Args: compute: Called as ``compute(tokens, rows)`` with ``tokens`` of @@ -178,11 +179,12 @@ def columns( compute: Callable[[Tensor], object], tokens: slice, ) -> None: - r"""Updates the token columns ``tokens`` in place. + r"""Update the token columns ``tokens`` in place, in chunks. Args: - compute: Called on chunks of consecutive token columns, in order, - with ``cells`` of shape ``[..., N, R, D]``. + compute: Called as ``compute(cells)`` on chunks of ``N`` + consecutive token columns, in order, with ``cells`` of shape + ``[..., N, R, D]``. tokens: The token columns to update. """ if self.offload is None: @@ -293,7 +295,7 @@ def run_column_layers( cache: Cache | None, epilogue: Callable[[Tensor], object] | None = None, ) -> None: - r"""Runs consecutive column layers over ``buffer`` in a single pass. + r"""Run consecutive column layers over ``buffer`` in a single pass. Records or replays each layer's key/value projections under its cache key. @@ -349,7 +351,7 @@ def add_labels( y_emb: Tensor, # [..., R_train, D] rows: slice, ) -> None: - r"""Adds label embeddings to the context rows among ``rows``.""" + r"""Add label embeddings to the context rows among ``rows``.""" stop = min(rows.stop, y_emb.size(-2)) if stop > rows.start: cells[..., : stop - rows.start, :, :] += y_emb[ diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 5dc441a525..ab10dbfba6 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -135,15 +135,13 @@ def forward( pre-processing and model execution. offload: Where to store the largest intermediate activations, such as the per-cell embeddings of row embedding layers. - Offloading requires inputs on a CUDA device. ``None`` stores - them on the device of the inputs. ``"cpu"`` stores them in - pinned host memory and streams them through the device in - chunks. This lowers peak device memory at the cost of - additional host-device transfers, which overlap with compute - to limit the slowdown. Offloading also requires - :attr:`supports_offload`, disabled gradients, - ``estimator_batch_size=1``, and no more target classes than - the model predicts natively. + ``None`` stores them on the device of the inputs. ``"cpu"`` + stores them in pinned host memory and streams them through + the device in chunks. This lowers peak device memory at the + cost of host-device transfers, which overlap with compute. + Offloading requires :attr:`supports_offload`, inputs on a + CUDA device, disabled gradients, ``estimator_batch_size=1``, + and no more target classes than the model predicts natively. kwargs: Additional keyword arguments passed to the model. Returns: @@ -278,14 +276,14 @@ def fit( pre-processing and model execution. offload: Where to store the largest intermediate activations while fitting, such as the per-cell embeddings of row - embedding layers. Offloading requires inputs on a CUDA - device. ``None`` stores them on the device of the inputs. - ``"cpu"`` stores them in pinned host memory and streams them - through the device in chunks. This lowers peak device memory - at the cost of additional host-device transfers, which - overlap with compute to limit the slowdown. Offloading also - requires :attr:`supports_offload`, ``estimator_batch_size=1``, - and no more target classes than the model predicts natively. + embedding layers. ``None`` stores them on the device of the + inputs. ``"cpu"`` stores them in pinned host memory and + streams them through the device in chunks. This lowers peak + device memory at the cost of host-device transfers, which + overlap with compute. :meth:`predict` does not offload. + Offloading requires :attr:`supports_offload`, inputs on a + CUDA device, ``estimator_batch_size=1``, and no more target + classes than the model predicts natively. kwargs: Additional keyword arguments passed to the model. """ self._validate_offload(offload, x.device) From e406fc1ceaefe3d51bc10be41617bcc77066d170 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 30 Oct 2026 11:28:22 -0700 Subject: [PATCH 04/12] use multiple threads to allocate pinned memory --- sdm/_memory.py | 30 ++++++++++++++++++++++++++++++ sdm/_token_buffer.py | 11 ++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/sdm/_memory.py b/sdm/_memory.py index f9688ce10f..0460b0cbb2 100644 --- a/sdm/_memory.py +++ b/sdm/_memory.py @@ -35,6 +35,36 @@ def set_pinned_max_round_threshold( torch._C._accelerator_setAllocatorSettings(",".join(options)) +def set_pinned_register_threads( + num_threads: int, + force: bool = False, +) -> None: + r"""Set the process-wide number of threads that pin new allocations. + + Registering the pages of a pinned allocation from several threads pins + large allocations several times faster than a single thread. + + Args: + num_threads: The number of threads, a power of two of at most + ``128``. + force: Whether to replace an existing setting. + """ + if torch.__version__ < "2.14": + return + keys = ("pinned_use_cuda_host_register", "pinned_num_register_threads") + settings = torch._C._accelerator_getAllocatorSettings() + if not force and any(key in settings for key in keys): + return + options = [ + option + for option in settings.split(",") + if option and option.split(":")[0].strip() not in keys + ] + options.append(f"{keys[0]}:True") + options.append(f"{keys[1]}:{num_threads}") + torch._C._accelerator_setAllocatorSettings(",".join(options)) + + def chunk_memory_limit(device: torch.device) -> int: r"""Bytes one chunk of a chunked operation may occupy on a CUDA device. diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 186c56e6a7..06f55e60a2 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -9,7 +9,11 @@ import torch from torch import Tensor -from sdm._memory import chunk_memory_limit, set_pinned_max_round_threshold +from sdm._memory import ( + chunk_memory_limit, + set_pinned_max_round_threshold, + set_pinned_register_threads, +) from sdm.cache import Cache, KVCacheEntry # Lets one chunk load, another compute, and a third store concurrently. @@ -70,6 +74,11 @@ def __init__( ) self._pitch = R * D * dtype.itemsize set_pinned_max_round_threshold(mb=1) + num_threads = min(torch.get_num_threads(), 16) + # PyTorch requires the register thread count to be a power of two. + set_pinned_register_threads( + num_threads=1 << (num_threads.bit_length() - 1) + ) self._host = torch.empty( self.size.numel(), dtype=dtype, pin_memory=True ) From 5d983b98a67a30938ad0f6381499f892a4882ba1 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Tue, 3 Nov 2026 17:46:17 -0800 Subject: [PATCH 05/12] embed cells in the first column pass instead of a separate row pass --- sdm/_token_buffer.py | 40 +++++++++++------------ sdm/models/kumo/tabular/cell_embedding.py | 12 ++++--- sdm/models/kumo/tabular/row_embedding.py | 23 ++++++++----- sdm/models/tabfm/cell_embedding.py | 22 +++++++------ sdm/models/tabfm/row_embedding.py | 23 ++++++++----- sdm/models/tabiclv2/row_embedding.py | 24 ++++++-------- 6 files changed, 76 insertions(+), 68 deletions(-) diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 06f55e60a2..493cf079b6 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -185,19 +185,23 @@ def rows( def columns( self, - compute: Callable[[Tensor], object], + compute: Callable[[Tensor, slice], object], tokens: slice, + *, + load: bool = True, ) -> None: r"""Update the token columns ``tokens`` in place, in chunks. Args: - compute: Called as ``compute(cells)`` on chunks of ``N`` - consecutive token columns, in order, with ``cells`` of shape - ``[..., N, R, D]``. + compute: Called as ``compute(cells, columns)`` on chunks of ``N`` + consecutive token columns ``columns``, in order, with + ``cells`` of shape ``[..., N, R, D]``. tokens: The token columns to update. + load: Whether ``compute`` reads ``cells``. If ``False``, + ``compute`` must fill them. """ if self.offload is None: - compute(self.tensor[..., tokens, :, :]) + compute(self.tensor[..., tokens, :, :], tokens) return *B, R, _, D = self.size @@ -218,8 +222,10 @@ def columns( compute=lambda slot, k: compute( slot.view(self.dtype).view( *B, chunks[k].stop - chunks[k].start, R, D - ) + ), + chunks[k], ), + load=load, ) def _stream( @@ -302,6 +308,7 @@ def run_column_layers( tokens: slice, key_rows: Sequence[Any], cache: Cache | None, + embed: Callable[[Tensor, slice], object] | None = None, epilogue: Callable[[Tensor], object] | None = None, ) -> None: r"""Run consecutive column layers over ``buffer`` in a single pass. @@ -314,6 +321,8 @@ def run_column_layers( tokens: The token columns to update. key_rows: Per layer, the row index of context cells to attend to. cache: The cache to record into or replay from, or ``None``. + embed: If given, called as ``embed(cells, columns)`` to fill each + chunk of cells before the first layer, instead of loading them. epilogue: Called on each chunk of cells after the last layer. """ recording = cache is not None and cache.is_recording @@ -322,7 +331,9 @@ def run_column_layers( assert not replaying or buffer.offload is None entries: dict[str, list[KVCacheEntry]] = {key: [] for key, _ in layers} - def compute(cells: Tensor) -> None: # [..., N, R, D] + def compute(cells: Tensor, columns: slice) -> None: # [..., N, R, D] + if embed is not None: + embed(cells, columns) for (key, layer), rows in zip(layers, key_rows, strict=True): result = layer( query=cells, @@ -339,7 +350,7 @@ def compute(cells: Tensor) -> None: # [..., N, R, D] if epilogue is not None: epilogue(cells) - buffer.columns(compute, tokens) + buffer.columns(compute, tokens, load=embed is None) if recording: assert cache is not None @@ -353,16 +364,3 @@ def compute(cells: Tensor) -> None: # [..., N, R, D] value=torch.cat([part.value for part in parts], dim=dim), ) ) - - -def add_labels( - cells: Tensor, # [..., N, C, D] - y_emb: Tensor, # [..., R_train, D] - rows: slice, -) -> None: - r"""Add label embeddings to the context rows among ``rows``.""" - stop = min(rows.stop, y_emb.size(-2)) - if stop > rows.start: - cells[..., : stop - rows.start, :, :] += y_emb[ - ..., rows.start : stop, None, : - ].to(cells.dtype) diff --git a/sdm/models/kumo/tabular/cell_embedding.py b/sdm/models/kumo/tabular/cell_embedding.py index 1fca97b458..c218ed3df0 100644 --- a/sdm/models/kumo/tabular/cell_embedding.py +++ b/sdm/models/kumo/tabular/cell_embedding.py @@ -61,25 +61,25 @@ def forward( # type: ignore *, train_size: int | None = None, impute_mean: Tensor | None = None, # [..., 1, C] + columns: slice = slice(None), num_columns: Tensor | None = None, # [...] cache: Cache | None = None, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, - ) -> Tensor: # [..., R, C, D] - missing = x.isnan() + ) -> Tensor: # [..., R, N, D] for the N columns in `columns` if impute_mean is None: impute_mean = self.impute_mean( x, train_size=train_size, cache=cache ) - x = torch.where(missing, impute_mean, x) return super().forward( x, categorical_mask, + columns=columns, num_columns=num_columns, batch_size_limit=batch_size_limit, out=out, - missing=missing, + impute_mean=impute_mean.expand_as(x), ) def _forward( # type: ignore @@ -89,9 +89,11 @@ def _forward( # type: ignore weight: Tensor, # [..., 1, C, G, D, 2F] bias: Tensor, # [..., 1, C, D] *, - missing: Tensor, # [..., R, C, G] + impute_mean: Tensor, # [..., R, C, G] out: Tensor | None = None, ) -> Tensor: + missing = x.isnan() + x = torch.where(missing, impute_mean, x) out = super()._forward(x, freq, weight, bias, out=out) out += self.nan_lin(missing.to(out.dtype)) diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index 2da8762d70..f32550a697 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -9,7 +9,7 @@ from torch import Tensor from torch.nn import Embedding, Linear, ModuleList, Parameter -from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers +from sdm._token_buffer import TokenBuffer, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.kumo.tabular.block import KumoTabularTransformerBlock from sdm.models.kumo.tabular.cell_embedding import CellEmbedding @@ -165,23 +165,20 @@ def forward( offload=offload, ) - def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] - tokens[..., :K, :] = self.readout_token.to(tokens.dtype) + def embed(cells: Tensor, columns: slice) -> None: # [..., N, R, D] self.cell_embedding( - x=x[..., rows, :], + x=x, categorical_mask=categorical_mask, impute_mean=impute_mean, + columns=slice(columns.start - K, columns.stop - K), num_columns=num_columns, batch_size_limit="auto", - out=tokens[..., K:, :], + out=cells.transpose(-2, -3), ) if y_emb is not None: - add_labels(tokens[..., K:, :], y_emb, rows) + cells[..., :R_train, :] += y_emb.unsqueeze(-3).to(cells.dtype) with buffer: - buffer.rows(embed, load=False) - y_emb = None - readout = buffer.readout(K) # [..., R, K, D] for i, (col_block, row_block) in enumerate( zip(self.col_blocks, self.row_blocks) @@ -192,16 +189,24 @@ def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] tokens=slice(K, K + C), key_rows=[slice(R_train)], cache=cache, + embed=embed if i == 0 else None, ) + y_emb = None + first = i == 0 last = i == len(self.row_blocks) - 1 def update( tokens: Tensor, # [..., N, K + C, D] rows: slice, row_block: torch.nn.Module = row_block, + first: bool = first, last: bool = last, ) -> None: + if first: + tokens[..., :K, :] = self.readout_token.to( + tokens.dtype + ) row_block( query=tokens[..., :K, :] if last else tokens, key_value=tokens, diff --git a/sdm/models/tabfm/cell_embedding.py b/sdm/models/tabfm/cell_embedding.py index eeb700a87c..59bdc6f08b 100644 --- a/sdm/models/tabfm/cell_embedding.py +++ b/sdm/models/tabfm/cell_embedding.py @@ -57,27 +57,29 @@ def forward( x: Tensor, # [..., R, C], categorical_mask: Tensor, # [..., C], *, + columns: slice = slice(None), num_columns: Tensor | None = None, # [...] batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, **kwargs: Tensor, # [..., R, C], - ) -> Tensor: # [..., R, C, D] + ) -> Tensor: # [..., R, N, D] for the N columns in `columns` if out is not None and torch.is_grad_enabled(): raise RuntimeError( "'out' is only supported when gradients are disabled" ) *B, R, C = x.size() - if out is not None and out.size() != (*B, R, C, self.channels): + index = torch.arange(C, device=x.device)[columns] + N = index.size(0) + if out is not None and out.size() != (*B, R, N, self.channels): raise ValueError("'out' has an incorrect shape") # Feature grouping: - index = torch.arange(C, device=x.device) shift = 2 ** torch.arange(self.group_size, device=x.device) - 1 - index = index.view(C, 1) + shift.view(1, self.group_size) + index = index.view(N, 1) + shift.view(1, self.group_size) if num_columns is None: index = index % C - x = x[..., index] # [..., R, C, G] + x = x[..., index] # [..., R, N, G] kwargs = {key: value[..., index] for key, value in kwargs.items()} categorical_mask = categorical_mask[..., None, index] else: @@ -85,16 +87,16 @@ def forward( index = index % num_columns[..., None, None] index = index.flatten(-2) row_index = index.unsqueeze(-2).expand(*B, R, -1) - x = x.gather(-1, row_index).unflatten(-1, (C, self.group_size)) + x = x.gather(-1, row_index).unflatten(-1, (N, self.group_size)) kwargs = { key: value.gather(-1, row_index).unflatten( - -1, (C, self.group_size) + -1, (N, self.group_size) ) for key, value in kwargs.items() } categorical_mask = categorical_mask.expand(*B, C).gather(-1, index) categorical_mask = categorical_mask.unflatten( - -1, (C, self.group_size) + -1, (N, self.group_size) ) categorical_mask = categorical_mask.unsqueeze(-3) @@ -164,13 +166,13 @@ def forward( if out is None: out = x.new_empty( - (*B, R, C, self.channels), + (*B, R, N, self.channels), dtype=torch.get_autocast_dtype(x.device.type) if torch.is_autocast_enabled(x.device.type) else x.dtype, ) - rows_per_chunk = max(1, batch_size_limit // max(math.prod(B) * C, 1)) + rows_per_chunk = max(1, batch_size_limit // max(math.prod(B) * N, 1)) for start in range(0, R, rows_per_chunk): self._forward( x=x[..., start : start + rows_per_chunk, :, :], diff --git a/sdm/models/tabfm/row_embedding.py b/sdm/models/tabfm/row_embedding.py index da2e013ce0..5a8e9d8ce2 100644 --- a/sdm/models/tabfm/row_embedding.py +++ b/sdm/models/tabfm/row_embedding.py @@ -31,7 +31,7 @@ Sequential, ) -from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers +from sdm._token_buffer import TokenBuffer, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.tabfm.block import TabFMTransformerBlock from sdm.models.tabfm.cell_embedding import CellEmbedding @@ -181,21 +181,18 @@ def forward( offload=offload, ) - def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] - tokens[..., :K, :] = self.readout_token.to(tokens.dtype) + def embed(cells: Tensor, columns: slice) -> None: # [..., N, R, D] self.cell_embedding( - x[..., rows, :], + x, categorical_mask, + columns=slice(columns.start - K, columns.stop - K), batch_size_limit="auto", - out=tokens[..., K:, :], + out=cells.transpose(-2, -3), ) if y_emb is not None: - add_labels(tokens[..., K:, :], y_emb, rows) + cells[..., :R_train, :] += y_emb.unsqueeze(-3).to(cells.dtype) with buffer: - buffer.rows(embed, load=False) - y_emb = None - readout = buffer.readout(K) # [..., R, K, D] for i, (col_layers, col_proj, row_layers, row_norm) in enumerate( zip( @@ -215,11 +212,14 @@ def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] tokens=slice(K if i == 0 else 0, K + C), key_rows=[slice(R_train)] * len(col_layers), cache=cache, + embed=embed if i == 0 else None, epilogue=lambda cells, col_proj=col_proj: cells.copy_( col_proj(cells) ), ) + y_emb = None + first = i == 0 last = i == len(self.row_blocks) - 1 def update( @@ -227,8 +227,13 @@ def update( rows: slice, row_layers: ModuleList = row_layers, row_norm: torch.nn.Module = row_norm, + first: bool = first, last: bool = last, ) -> None: + if first: + tokens[..., :K, :] = self.readout_token.to( + tokens.dtype + ) x = tokens for j, row_layer in enumerate(row_layers): if last and j == len(row_layers) - 1: diff --git a/sdm/models/tabiclv2/row_embedding.py b/sdm/models/tabiclv2/row_embedding.py index 15f49cf18d..49da62d1d0 100644 --- a/sdm/models/tabiclv2/row_embedding.py +++ b/sdm/models/tabiclv2/row_embedding.py @@ -13,7 +13,7 @@ from torch import Tensor from torch.nn import Embedding, LayerNorm, ModuleList, Parameter -from sdm._token_buffer import TokenBuffer, add_labels, run_column_layers +from sdm._token_buffer import TokenBuffer, run_column_layers from sdm.cache import Cache, KVCacheEntry from sdm.models.tabiclv2.block import TabICLv2TransformerBlock from sdm.nn import InducedTransformerBlock, RotaryEmbedding @@ -185,26 +185,20 @@ def forward( offload=offload, ) - def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] - cells = tokens[..., K:, :] + def embed(cells: Tensor, columns: slice) -> None: # [..., N, R, D] + group = index[columns.start - K : columns.stop - K] # [N, G] + group = x[..., group].transpose(-2, -3).contiguous() if num_digits > 1: - self.lin(x[..., rows, :][..., index], out=cells[0]) + self.lin(group, out=cells[0]) cells[1:] = cells[:1] else: - self.lin(x[..., rows, :][..., index], out=cells) + self.lin(group, out=cells) if y_emb is None: return - if train_mask is None: - add_labels(cells, y_emb, rows) - else: - cells[..., train_mask, :, :] += y_emb.unsqueeze(-2).to( - cells.dtype - ) + rows = slice(R_train) if train_mask is None else train_mask + cells[..., rows, :] += y_emb.unsqueeze(-3).to(cells.dtype) with buffer: - buffer.rows(embed, load=False) - y_emb = None - run_column_layers( buffer, [ @@ -214,7 +208,9 @@ def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] tokens=slice(K, K + C), key_rows=key_rows, cache=cache, + embed=embed, ) + y_emb = None if num_digits > 1: # Average over mixed-radix digits. buffer = TokenBuffer.from_tensor(buffer.tensor.mean(dim=0)) From b41c65d16bff7f8d231c9dd96a927b8128cf60ab Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 4 Nov 2026 15:15:44 -0800 Subject: [PATCH 06/12] fuse KumoTabular column blocks into row passes via chunked inducing attention --- sdm/models/kumo/tabular/row_embedding.py | 105 ++++--- sdm/nn/attention.py | 336 +++++++++++++++++++---- sdm/nn/set_transformer.py | 63 +++++ 3 files changed, 410 insertions(+), 94 deletions(-) diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index f32550a697..209655b7c2 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -9,7 +9,7 @@ from torch import Tensor from torch.nn import Embedding, Linear, ModuleList, Parameter -from sdm._token_buffer import TokenBuffer, run_column_layers +from sdm._token_buffer import TokenBuffer from sdm.cache import Cache, KVCacheEntry from sdm.models.kumo.tabular.block import KumoTabularTransformerBlock from sdm.models.kumo.tabular.cell_embedding import CellEmbedding @@ -165,57 +165,90 @@ def forward( offload=offload, ) - def embed(cells: Tensor, columns: slice) -> None: # [..., N, R, D] + def embed(tokens: Tensor, rows: slice) -> None: # [..., N, K + C, D] + tokens[..., :K, :] = self.readout_token.to(tokens.dtype) self.cell_embedding( - x=x, + x=x[..., rows, :], categorical_mask=categorical_mask, impute_mean=impute_mean, - columns=slice(columns.start - K, columns.stop - K), num_columns=num_columns, batch_size_limit="auto", - out=cells.transpose(-2, -3), + out=tokens[..., K:, :], ) - if y_emb is not None: - cells[..., :R_train, :] += y_emb.unsqueeze(-3).to(cells.dtype) + stop = min(rows.stop, R_train) + if y_emb is not None and stop > rows.start: + tokens[..., : stop - rows.start, K:, :] += y_emb[ + ..., rows.start : stop, None, : + ].to(tokens.dtype) + # Column block `i` lets every cell attend to `H_i`, its column's + # inducing points after attending to the column's context cells. Only + # `H_i` needs whole columns, and since attention over the context rows + # can be accumulated chunk by chunk (`summarize` merges chunks by + # log-sum-exp), it can be built during a row pass. So instead of + # alternating column and row passes (`2 * L` passes), pass `i` + # streams the buffer once and, for each row chunk: + # 1. completes column block `i - 1` with `H_{i-1}` (pass 0 embeds + # the cells instead), + # 2. runs row block `i - 1` (none in pass 0), and + # 3. accumulates the chunk's context rows into `H_i` (none in the + # last pass). + # `H_i` is complete when pass `i` ends, ready for pass `i + 1`, which + # gives `L + 1` passes in total. + L = len(self.row_blocks) + col_blocks = cast(list[InducedTransformerBlock], self.col_blocks) + replaying = cache is not None and cache.is_replaying with buffer: readout = buffer.readout(K) # [..., R, K, D] - for i, (col_block, row_block) in enumerate( - zip(self.col_blocks, self.row_blocks) - ): - run_column_layers( - buffer, - [(f"row_embedding.col_block{i}", col_block)], - tokens=slice(K, K + C), - key_rows=[slice(R_train)], - cache=cache, - embed=embed if i == 0 else None, - ) - y_emb = None - - first = i == 0 - last = i == len(self.row_blocks) - 1 + key_value: KVCacheEntry | None = None + for i in range(L + 1): + summary: tuple[Tensor, Tensor] | None = None def update( tokens: Tensor, # [..., N, K + C, D] rows: slice, - row_block: torch.nn.Module = row_block, - first: bool = first, - last: bool = last, + i: int = i, + key_value: KVCacheEntry | None = key_value, ) -> None: - if first: - tokens[..., :K, :] = self.readout_token.to( - tokens.dtype + nonlocal summary + if i == 0: + embed(tokens, rows) + else: + cells = tokens[..., K:, :].transpose(-2, -3) + col_blocks[i - 1]( + query=cells, # [..., C, N, D] + key_value=key_value, + batch_size_limit="auto", + out=cells, + ) + self.row_blocks[i - 1]( + query=tokens[..., :K, :] if i == L else tokens, + key_value=tokens, + seqused_key_value=row_key_lengths, + batch_size_limit="auto", + out=readout[..., rows, :, :] if i == L else tokens, + ) + stop = min(rows.stop, R_train) - rows.start + if i < L and not replaying and stop > 0: + summary = col_blocks[i].summarize( + tokens[..., :stop, K:, :].transpose(-2, -3), + key_len=R_train, + summary=summary, + batch_size_limit="auto", ) - row_block( - query=tokens[..., :K, :] if last else tokens, - key_value=tokens, - seqused_key_value=row_key_lengths, - batch_size_limit="auto", - out=readout[..., rows, :, :] if last else tokens, - ) - buffer.rows(update, store=not last) + buffer.rows(update, load=i > 0, store=i < L) + y_emb = None + if i < L: + key = f"row_embedding.col_block{i}" + if replaying: + assert cache is not None + key_value = cast(KVCacheEntry, cache[key]) + else: + assert summary is not None + key_value = col_blocks[i].induced_key_value(summary) + if cache is not None and cache.is_recording: + cache[key] = key_value return self.norm(readout).flatten(-2) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index c45b3cf526..3b1c489ac4 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -192,6 +192,77 @@ def forward( return out.view(batch_shape + out.size()[-3:]) # [..., Q, Hq, C] + def partial( + self, + query: Tensor, # [..., Q, Hq, C] + key: Tensor, # [..., KV, Hkv, C] + value: Tensor, # [..., KV, Hkv, C] + *, + key_len: int, + ) -> tuple[Tensor, Tensor]: # [..., Q, Hq, C], [..., Q, Hq] + r"""Attention over one chunk of a longer key/value sequence. + + Returns the attention output over the chunk together with the + log-sum-exp of its attention logits, which lets outputs over + consecutive chunks be merged into the output over all of them. + + Args: + query: The query tensor with shape ``[..., Q, Hq, C]``. + key: The key chunk with shape ``[..., KV, Hkv, C]``. + value: The value chunk with shape ``[..., KV, Hkv, C]``. + key_len: The length of the whole key/value sequence, used for + query scaling. + """ + if self.query_scaling is not None: + query = self.query_scaling(query, key_len=key_len) + if torch.is_autocast_enabled(query.device.type): + dtype = torch.get_autocast_dtype(query.device.type) + query, key, value = query.to(dtype), key.to(dtype), value.to(dtype) + + batch_shape = torch.broadcast_shapes( + query.size()[:-3], key.size()[:-3], value.size()[:-3] + ) + query, key, value = ( + x.expand(batch_shape + x.size()[-3:]) + .reshape(-1, *x.size()[-3:]) + .transpose(-3, -2) # [B, H, S, C] + for x in (query, key, value) + ) + out, lse = _attention_with_lse(query, key, value, scale=self.scale) + out = out.transpose(-3, -2) # [B, Q, Hq, C] + lse = lse.transpose(-2, -1) # [B, Q, Hq] + return ( + out.view(batch_shape + out.size()[-3:]), + lse.view(batch_shape + lse.size()[-2:]), + ) + + +def _attention_with_lse( + query: Tensor, # [B, Hq, Q, C] + key: Tensor, # [B, Hkv, KV, C] + value: Tensor, # [B, Hkv, KV, C] + scale: float | None, +) -> tuple[Tensor, Tensor]: # [B, Hq, Q, C], [B, Hq, Q] + # `F.scaled_dot_product_attention` does not return the log-sum-exp. + if not query.is_cuda: + out, lse = torch.ops.aten._scaled_dot_product_flash_attention_for_cpu( + query, key, value, scale=scale + ) + return out, lse + if ( + query.dtype in (torch.float16, torch.bfloat16) + and query.size(-1) % 8 == 0 + and query.size(-1) <= 256 + ): + out, lse, *_ = torch.ops.aten._scaled_dot_product_flash_attention( + query, key, value, scale=scale + ) + return out, lse + out, lse, *_ = torch.ops.aten._scaled_dot_product_efficient_attention( + query, key, value, None, True, scale=scale + ) + return out, lse[..., : query.size(-2)] + class Attention(torch.nn.Module): r"""Multi-head attention layer with grouped-query attention support. @@ -363,7 +434,39 @@ def forward( :class:`~sdm.cache.KVCacheEntry`. """ gate = self.gate(query) if self.gate is not None else None + query, key, value = self._project( + query=query, + key_value=key_value, + query_pos=query_pos, + key_pos=key_pos, + ) + + out = self.sdpa( + query=query, # [..., Q, Hq, C // Hq] + key=key, # [..., KV, Hkv, C // Hq] + value=value, # [..., KV, Hkv, C // Hq] + seqused_key_value=seqused_key_value, # [...] + attn_mask=attn_mask, # [..., Hq, Q, KV] + is_causal=is_causal, + ) # [..., Q, Hq, C // Hq] + + out = self._output(out, gate) # [..., Q, C] + if return_key_value: + # CUDA autocast runs RMSNorm in fp32, so cast explicitly before + # caching: https://github.com/pytorch/pytorch/blob/v2.13.0/aten/src/ATen/autocast_mode.h#L875 + return out, KVCacheEntry( + key=key.to(value.dtype).contiguous(), + value=value.contiguous(), + ) + return out + def _project( + self, + query: Tensor, # [..., Q, C] + key_value: Tensor | KVCacheEntry | None, # [..., KV, C] + query_pos: Tensor | None, + key_pos: Tensor | None, + ) -> tuple[Tensor, Tensor, Tensor]: # [..., S, H, C // H] if key_value is None and key_pos is None: key_pos = query_pos @@ -422,27 +525,47 @@ def forward( else: key = self.key_transform(key, pos=key_pos) - out = self.sdpa( - query=query, # [..., Q, Hq, C // Hq] - key=key, # [..., KV, Hkv, C // Hq] - value=value, # [..., KV, Hkv, C // Hq] - seqused_key_value=seqused_key_value, # [...] - attn_mask=attn_mask, # [..., Hq, Q, KV] - is_causal=is_causal, - ) # [..., Q, Hq, C // Hq] + return query, key, value - out = out.flatten(-2, -1) # [..., Q, C] + def partial( + self, + query: Tensor, # [..., Q, C] + key_value: Tensor, # [..., KV, C] + *, + key_len: int, + ) -> tuple[Tensor, Tensor]: # [..., Q, Hq, C // Hq], [..., Q, Hq] + r""":meta private:""" # noqa: D415 + query, key, value = self._project( + query=query, + key_value=key_value, + query_pos=None, + key_pos=None, + ) + return self.sdpa.partial( + query=query, + key=key, + value=value, + key_len=key_len, + ) + + def finish( + self, + query: Tensor, # [..., Q, C] + out: Tensor, # [..., Q, Hq, C // Hq] + ) -> Tensor: # [..., Q, C] + r""":meta private:""" # noqa: D415 + gate = self.gate(query) if self.gate is not None else None + return self._output(out.to(query.dtype), gate) + + def _output( + self, + out: Tensor, # [..., Q, Hq, C // Hq] + gate: Tensor | None, # [..., Q, C] + ) -> Tensor: # [..., Q, C] + out = out.flatten(-2, -1) if gate is not None: out = out * gate - out = self.out_lin(out) # [..., Q, C] - if return_key_value: - # CUDA autocast runs RMSNorm in fp32, so cast explicitly before - # caching: https://github.com/pytorch/pytorch/blob/v2.13.0/aten/src/ATen/autocast_mode.h#L875 - return out, KVCacheEntry( - key=key.to(value.dtype).contiguous(), - value=value.contiguous(), - ) - return out + return self.out_lin(out) class TransformerBlock(torch.nn.Module): @@ -638,26 +761,9 @@ def forward( out=out, ) - if batch_size_limit == "auto": - key_value_length: int | None = None - if isinstance(key_value, Tensor): - key_value_length = key_value.size(-2) - elif isinstance(key_value, KVCacheEntry): - key_value_length = key_value.key.size(-3) - - batch_size_limit = self.auto_batch_size_limit( - device=query.device, - element_size=torch.get_autocast_dtype( - query.device.type - ).itemsize - if torch.is_autocast_enabled(query.device.type) - else query.element_size(), - query_length=query.size(-2), - key_value_length=key_value_length, - ) - - batch_size_limit = min(batch_size_limit or 65_535, 65_535) - + batch_size_limit = self._batch_size_limit( + batch_size_limit, query=query, key_value=key_value + ) batch_shape = _batch_shape( query=query, key_value=key_value, @@ -693,22 +799,9 @@ def forward( out=out, ) - # Slice batch dimension `dim` and take later batch dimensions whole, - # so that every chunk is a view of up to `batch_size_limit` elements. - dim, inner = len(batch_shape) - 1, 1 - while inner * batch_shape[dim] <= batch_size_limit: - inner *= batch_shape[dim] - dim -= 1 - step = batch_size_limit // inner - chunks = [ - (*index, slice(start, min(start + step, batch_shape[dim]))) - for index in itertools.product(*map(range, batch_shape[:dim])) - for start in range(0, batch_shape[dim], step) - ] - key: Tensor | None = None value: Tensor | None = None - for index in chunks: + for index in _batch_chunks(batch_shape, batch_size_limit): result = self._forward( query=_chunk(query, batch_shape, 2, index), key_value=_chunk(key_value, batch_shape, 2, index), @@ -793,6 +886,116 @@ def _forward( else: attn_out = result + out = self._residual(query, attn_out, out=out) + return (out, kv) if return_key_value else out + + def partial( + self, + query: Tensor, # [..., Q, C] + key_value: Tensor, # [..., KV, C] + *, + key_len: int, + batch_size_limit: int | Literal["auto"] | None = None, + ) -> tuple[Tensor, Tensor]: # [..., Q, Hq, C // Hq], [..., Q, Hq] + r"""Attention over one chunk of a longer key/value sequence. + + Returns the attention output over the chunk, before the output + projection, together with the log-sum-exp of its attention logits. + :meth:`finish` completes the block from outputs merged over all + chunks. + + Args: + query: The query tensor with shape ``[..., Q, C]``. + key_value: The key/value chunk with shape ``[..., KV, C]``. + key_len: The length of the whole key/value sequence. + batch_size_limit: Maximum number of batch elements processed at + once. + """ + batch_size_limit = self._batch_size_limit( + batch_size_limit, query=query, key_value=key_value + ) + batch_shape = torch.broadcast_shapes( + query.size()[:-2], key_value.size()[:-2] + ) + + def compute(q: Tensor, kv: Tensor) -> tuple[Tensor, Tensor]: + return self.attn.partial( + query=q if self.query_norm is None else self.query_norm(q), + key_value=kv + if self.key_value_norm is None + else self.key_value_norm(kv), + key_len=key_len, + ) + + if math.prod(batch_shape) <= batch_size_limit: + return compute(query, key_value) + + out: Tensor | None = None + lse: Tensor | None = None + for index in _batch_chunks(batch_shape, batch_size_limit): + chunk_out, chunk_lse = compute( + _chunk(query, batch_shape, 2, index), + _chunk(key_value, batch_shape, 2, index), + ) + if out is None or lse is None: + out = chunk_out.new_empty(batch_shape + chunk_out.size()[-3:]) + lse = chunk_lse.new_empty(batch_shape + chunk_lse.size()[-2:]) + out[index] = chunk_out + lse[index] = chunk_lse + assert out is not None + assert lse is not None + return out, lse + + def finish( + self, + query: Tensor, # [..., Q, C] + out: Tensor, # [..., Q, Hq, C // Hq] + ) -> Tensor: # [..., Q, C] + r"""Complete the block from attention outputs merged over chunks. + + Args: + query: The query tensor passed to :meth:`partial`. + out: The attention output over the whole key/value sequence. + """ + attn_out = self.attn.finish( + query=query if self.query_norm is None else self.query_norm(query), + out=out, + ) + return self._residual(query, attn_out) + + def _batch_size_limit( + self, + batch_size_limit: int | Literal["auto"] | None, + *, + query: Tensor, # [..., Q, C] + key_value: Tensor | KVCacheEntry | None, + ) -> int: + if batch_size_limit == "auto": + key_value_length: int | None = None + if isinstance(key_value, Tensor): + key_value_length = key_value.size(-2) + elif isinstance(key_value, KVCacheEntry): + key_value_length = key_value.key.size(-3) + + batch_size_limit = self.auto_batch_size_limit( + device=query.device, + element_size=torch.get_autocast_dtype( + query.device.type + ).itemsize + if torch.is_autocast_enabled(query.device.type) + else query.element_size(), + query_length=query.size(-2), + key_value_length=key_value_length, + ) + + return min(batch_size_limit or 65_535, 65_535) + + def _residual( + self, + query: Tensor, # [..., Q, C] + attn_out: Tensor, # [..., Q, C] + out: Tensor | None = None, + ) -> Tensor: # [..., Q, C] if self.post_attn_norm is not None: attn_out = self.post_attn_norm(attn_out) @@ -805,13 +1008,12 @@ def _forward( if self.mlp is not None: tmp += self.mlp(tmp) out.copy_(tmp) - else: - tmp = torch.add(attn_out, query, out=out) - if self.mlp is not None: - tmp = torch.add(tmp, self.mlp(tmp), out=out) - out = tmp + return out - return (out, kv) if return_key_value else out + tmp = torch.add(attn_out, query, out=out) + if self.mlp is not None: + tmp = torch.add(tmp, self.mlp(tmp), out=out) + return tmp def peak_bytes_per_example( self, @@ -841,6 +1043,24 @@ def auto_batch_size_limit( return min(max(limit, 1), 65_535) +def _batch_chunks( + batch_shape: torch.Size, + batch_size_limit: int, +) -> list[tuple[int | slice, ...]]: + # Slice batch dimension `dim` and take later batch dimensions whole, + # so that every chunk is a view of up to `batch_size_limit` elements. + dim, inner = len(batch_shape) - 1, 1 + while inner * batch_shape[dim] <= batch_size_limit: + inner *= batch_shape[dim] + dim -= 1 + step = batch_size_limit // inner + return [ + (*index, slice(start, min(start + step, batch_shape[dim]))) + for index in itertools.product(*map(range, batch_shape[:dim])) + for start in range(0, batch_shape[dim], step) + ] + + def _batch_shape( query: Tensor, key_value: Tensor | KVCacheEntry | None, diff --git a/sdm/nn/set_transformer.py b/sdm/nn/set_transformer.py index 9247d48af2..f62525ded5 100644 --- a/sdm/nn/set_transformer.py +++ b/sdm/nn/set_transformer.py @@ -160,3 +160,66 @@ def forward( batch_size_limit=batch_size_limit, out=out, ) # [..., Q, C] + + def summarize( + self, + key_value: Tensor, # [..., KV, C] + *, + key_len: int, + summary: tuple[Tensor, Tensor] | None = None, + batch_size_limit: int | Literal["auto"] | None = None, + ) -> tuple[Tensor, Tensor]: + r"""Accumulate the inducing points' attention over a key/value chunk. + + Calling this on consecutive chunks of a key/value sequence and passing + the result to :meth:`induced_key_value` gives the same key/value + projections of :math:`H` as attending to the whole sequence at once, + up to rounding. + + Args: + key_value: The key/value chunk with shape ``[..., KV, C]``. + key_len: The length of the whole key/value sequence. + summary: The summary of the preceding chunks, if any. + batch_size_limit: Maximum number of batch elements processed at + once. + + Returns: + The summary of the chunks so far. + """ + out, lse = self.inducing_block.partial( + query=self.inducing_points, # [M, C] + key_value=key_value, # [..., KV, C] + key_len=key_len, + batch_size_limit=batch_size_limit, + ) # [..., M, H, C // H], [..., M, H] + if summary is None: + return out, lse + summary_out, summary_lse = summary + total = torch.logaddexp(summary_lse, lse) # [..., M, H] + out = out.float() * (lse - total).exp()[..., None] + out += summary_out.float() * (summary_lse - total).exp()[..., None] + return out, total + + def induced_key_value( + self, + summary: tuple[Tensor, Tensor], + ) -> KVCacheEntry: # [..., M, H, C // H] + r"""The key/value projections of :math:`H` from a summary. + + The result can be passed as ``key_value`` to :meth:`forward`. + + Args: + summary: The summary of a whole key/value sequence returned by + :meth:`summarize`. + """ + h = self.inducing_block.finish( + query=self.inducing_points, # [M, C] + out=summary[0], # [..., M, H, C // H] + ) # [..., M, C] + # An empty query only projects the key/value. + _, key_value = self.output_block( + query=h[..., :0, :], + key_value=h, + return_key_value=True, + ) + return key_value From 90507ebb2b1712b93bc9e473dcbeb9abbff13df1 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 4 Nov 2026 22:30:43 -0800 Subject: [PATCH 07/12] pin estimator caches with multiple threads when fitting several estimators --- sdm/_memory.py | 7 +++++-- sdm/_token_buffer.py | 6 +----- sdm/models/base/icl.py | 6 +++++- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/sdm/_memory.py b/sdm/_memory.py index 0460b0cbb2..c0985051f6 100644 --- a/sdm/_memory.py +++ b/sdm/_memory.py @@ -36,7 +36,7 @@ def set_pinned_max_round_threshold( def set_pinned_register_threads( - num_threads: int, + num_threads: int | None = None, force: bool = False, ) -> None: r"""Set the process-wide number of threads that pin new allocations. @@ -46,11 +46,14 @@ def set_pinned_register_threads( Args: num_threads: The number of threads, a power of two of at most - ``128``. + ``128``. ``None`` uses :func:`torch.get_num_threads`, rounded + down to a power of two and capped at ``16``. force: Whether to replace an existing setting. """ if torch.__version__ < "2.14": return + if num_threads is None: + num_threads = 1 << (min(torch.get_num_threads(), 16).bit_length() - 1) keys = ("pinned_use_cuda_host_register", "pinned_num_register_threads") settings = torch._C._accelerator_getAllocatorSettings() if not force and any(key in settings for key in keys): diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 493cf079b6..593d4c2079 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -74,11 +74,7 @@ def __init__( ) self._pitch = R * D * dtype.itemsize set_pinned_max_round_threshold(mb=1) - num_threads = min(torch.get_num_threads(), 16) - # PyTorch requires the register thread count to be a power of two. - set_pinned_register_threads( - num_threads=1 << (num_threads.bit_length() - 1) - ) + set_pinned_register_threads() self._host = torch.empty( self.size.numel(), dtype=dtype, pin_memory=True ) diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index ab10dbfba6..5b91aeffb7 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -19,7 +19,10 @@ TaskLike, ) from sdm._inference import inference_mode -from sdm._memory import set_pinned_max_round_threshold +from sdm._memory import ( + set_pinned_max_round_threshold, + set_pinned_register_threads, +) from sdm.cache import Cache from sdm.models.base import Callback, Model from sdm.processing.execution import ( @@ -318,6 +321,7 @@ def fit( ] if x.is_cuda and len(contexts) > 1: set_pinned_max_round_threshold(mb=1) + set_pinned_register_threads() class_values = _class_values(contexts, estimator_batch_size) batches = _batch_slices( From 3152301da9d867f799a6befc207749cbdef74ecd Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 5 Nov 2026 10:28:31 -0800 Subject: [PATCH 08/12] Support torch.compile with offloaded TokenBuffer --- sdm/_token_buffer.py | 8 ++++++++ test/models/test_offload.py | 25 +++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index 593d4c2079..b12307e7f4 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -66,7 +66,14 @@ def __init__( (*B, T, R, D), dtype=dtype, device=device ) return + self._init_offload() + # Pinned allocation and the copy streams need real CUDA objects, which + # `torch.compile` cannot trace. + @torch.compiler.disable + def _init_offload(self) -> None: + *B, R, T, D = self.size + dtype, device = self.dtype, self.device if math.prod(B) != 1: raise ValueError( "'offload' requires 'estimator_batch_size=1' and no more " @@ -236,6 +243,7 @@ def _stream( # stored densely in its slot. from cuda.bindings import runtime + @torch.compiler.disable def copy( dst: int, dpitch: int, src: int, spitch: int, chunk, kind, stream ) -> None: diff --git a/test/models/test_offload.py b/test/models/test_offload.py index 935c94cd26..df6a27ac26 100644 --- a/test/models/test_offload.py +++ b/test/models/test_offload.py @@ -125,6 +125,31 @@ def test_row_embedding_offload( torch.testing.assert_close(outs["cpu"], outs[None]) +@onlyCUDA +@pytest.mark.parametrize("build", [_kumo_tabular, _tabfm, _tabiclv2]) +@pytest.mark.usefixtures("compilation") +def test_row_embedding_offload_compile( + build: Callable[[torch.device], Callable[..., Tensor]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SDM_CHUNK_MEMORY_FRACTION", "1e-9") + device = torch.device("cuda") + embedding = build(device) + x = torch.randn(1, 40, 5, device=device) + y = torch.randint(0, 3, (1, 30), device=device) + + def fit_predict(x: Tensor, y: Tensor) -> tuple[Tensor, Tensor]: + cache = Cache() + fit = embedding(x[:, :30], y, cache=cache, offload="cpu") + return fit, embedding(x[:, 30:], y[:, :0], cache=cache.freeze()) + + with torch.no_grad(): + expected = fit_predict(x, y) + out = torch.compile(fit_predict, backend="eager")(x, y) + + torch.testing.assert_close(out, expected) + + @onlyCUDA def test_row_embedding_offload_subsamples_keys() -> None: device = torch.device("cuda") From 52cf941057ac916338e3b40654476b60510ba7b2 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 5 Nov 2026 15:40:09 -0800 Subject: [PATCH 09/12] address automated review comments --- docs/source/api/models.rst | 9 ++++++ sdm/_memory.py | 50 ++++++++++++++++---------------- sdm/_token_buffer.py | 7 +---- sdm/models/__init__.py | 9 +++++- sdm/models/base/__init__.py | 2 ++ sdm/models/base/icl.py | 18 +++++++----- sdm/nn/attention.py | 5 ++-- sdm/nn/linear.py | 2 +- sdm/nn/set_transformer.py | 2 +- test/models/test_offload.py | 13 +++++++-- test/nn/test_set_transformer.py | 51 +++++++++++++++++++++++++++++++++ 11 files changed, 123 insertions(+), 45 deletions(-) diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 421874aa5d..14f289530e 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -73,6 +73,15 @@ Model API ~sdm.models.ICLModel ~sdm.models.Callback +Runtime Configuration +--------------------- + +.. autosummary:: + :toctree: generated + :nosignatures: + + ~sdm.models.set_pinned_register_threads + Models ------ diff --git a/sdm/_memory.py b/sdm/_memory.py index c2473b2fda..fa2ee88ce9 100644 --- a/sdm/_memory.py +++ b/sdm/_memory.py @@ -45,37 +45,39 @@ def set_pinned_max_round_threshold( torch._C._accelerator_setAllocatorSettings(",".join(options)) -def set_pinned_register_threads( - num_threads: int | None = None, - force: bool = False, -) -> None: - r"""Set the process-wide number of threads that pin new allocations. +def set_pinned_register_threads(num_threads: int | None = None) -> None: + r"""Set the number of threads that pin new host allocations. + + By default, PyTorch pins host memory from a single thread. Pinning from + several threads can reduce the allocation time of large pinned buffers, + such as the activations of ``offload="cpu"``. The setting applies to the + whole process. Equivalent alternatives are: - Registering the pages of a pinned allocation from several threads pins - large allocations several times faster than a single thread. + - the ``PYTORCH_ALLOC_CONF`` environment variable, which must be set + before the process starts, *e.g.*, + ``PYTORCH_ALLOC_CONF=pinned_use_cuda_host_register:True,pinned_num_register_threads:8``. + - the private PyTorch call + ``torch._C._accelerator_setAllocatorSettings("pinned_use_cuda_host_register:True,pinned_num_register_threads:8")``. + + See the ``pinned_num_register_threads`` option in the + `PyTorch CUDA memory management notes + `__. + Requires PyTorch 2.14 or newer and does nothing otherwise. Args: - num_threads: The number of threads, a power of two of at most - ``128``. ``None`` uses :func:`torch.get_num_threads`, rounded - down to a power of two and capped at ``16``. - force: Whether to replace an existing setting. + num_threads: The number of threads, which PyTorch requires to be a + power of two of at most ``128``. ``None`` uses ``8``, the value + PyTorch recommends, or :func:`torch.get_num_threads` rounded down + to a power of two if smaller. """ if torch.__version__ < "2.14": return if num_threads is None: - num_threads = 1 << (min(torch.get_num_threads(), 16).bit_length() - 1) - keys = ("pinned_use_cuda_host_register", "pinned_num_register_threads") - settings = torch._C._accelerator_getAllocatorSettings() - if not force and any(key in settings for key in keys): - return - options = [ - option - for option in settings.split(",") - if option and option.split(":")[0].strip() not in keys - ] - options.append(f"{keys[0]}:True") - options.append(f"{keys[1]}:{num_threads}") - torch._C._accelerator_setAllocatorSettings(",".join(options)) + num_threads = 1 << (min(torch.get_num_threads(), 8).bit_length() - 1) + torch._C._accelerator_setAllocatorSettings( + "pinned_use_cuda_host_register:True," + f"pinned_num_register_threads:{num_threads}" + ) def chunk_memory_limit(device: torch.device) -> int: diff --git a/sdm/_token_buffer.py b/sdm/_token_buffer.py index b12307e7f4..bbac1a36d5 100644 --- a/sdm/_token_buffer.py +++ b/sdm/_token_buffer.py @@ -9,11 +9,7 @@ import torch from torch import Tensor -from sdm._memory import ( - chunk_memory_limit, - set_pinned_max_round_threshold, - set_pinned_register_threads, -) +from sdm._memory import chunk_memory_limit, set_pinned_max_round_threshold from sdm.cache import Cache, KVCacheEntry # Lets one chunk load, another compute, and a third store concurrently. @@ -81,7 +77,6 @@ def _init_offload(self) -> None: ) self._pitch = R * D * dtype.itemsize set_pinned_max_round_threshold(mb=1) - set_pinned_register_threads() self._host = torch.empty( self.size.numel(), dtype=dtype, pin_memory=True ) diff --git a/sdm/models/__init__.py b/sdm/models/__init__.py index 3791affdb5..fad9529584 100644 --- a/sdm/models/__init__.py +++ b/sdm/models/__init__.py @@ -11,7 +11,13 @@ "sdm.models.timesfm3", } -from sdm.models.base import Callback, Model, ICLModel, ECOC +from sdm.models.base import ( + Callback, + Model, + ICLModel, + ECOC, + set_pinned_register_threads, +) from sdm.models.tabiclv2 import TabICLv2 from sdm.models.tabfm import TabFM from sdm.models.kumo import ( @@ -29,6 +35,7 @@ "Model", "ICLModel", "ECOC", + "set_pinned_register_threads", "TabICLv2", "TabFM", "KumoTabular", diff --git a/sdm/models/base/__init__.py b/sdm/models/base/__init__.py index 5bce05a58a..23e4139f42 100644 --- a/sdm/models/base/__init__.py +++ b/sdm/models/base/__init__.py @@ -7,10 +7,12 @@ from sdm.models.base.model import Model from sdm.models.base.icl import ICLModel from sdm.models.base.ecoc import ECOC +from sdm._memory import set_pinned_register_threads __all__ = [ "Callback", "Model", "ICLModel", "ECOC", + "set_pinned_register_threads", ] diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 5aa4aa806b..e54fd50cc2 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -19,10 +19,7 @@ TaskLike, ) from sdm._inference import inference_mode -from sdm._memory import ( - set_pinned_max_round_threshold, - set_pinned_register_threads, -) +from sdm._memory import set_pinned_max_round_threshold from sdm.cache import Cache from sdm.models.base import Callback, Model from sdm.processing.execution import ( @@ -145,6 +142,10 @@ def forward( Offloading requires :attr:`supports_offload`, inputs on a CUDA device, disabled gradients, ``estimator_batch_size=1``, and no more target classes than the model predicts natively. + :func:`~sdm.models.set_pinned_register_threads` speeds up + pinning large activations. Offloading also requires the + ``cuda-bindings`` package, which PyTorch's CUDA wheels + install on Linux with Python 3.14 or older. kwargs: Additional keyword arguments passed to the model. Returns: @@ -287,6 +288,10 @@ def fit( Offloading requires :attr:`supports_offload`, inputs on a CUDA device, ``estimator_batch_size=1``, and no more target classes than the model predicts natively. + :func:`~sdm.models.set_pinned_register_threads` speeds up + pinning large activations. Offloading also requires the + ``cuda-bindings`` package, which PyTorch's CUDA wheels + install on Linux with Python 3.14 or older. kwargs: Additional keyword arguments passed to the model. """ self._validate_offload(offload, x.device) @@ -321,7 +326,6 @@ def fit( ] if x.is_cuda and len(contexts) > 1: set_pinned_max_round_threshold(mb=1) - set_pinned_register_threads() class_values = _class_values(contexts, estimator_batch_size) batches = _batch_slices( @@ -566,8 +570,8 @@ def _forward( preprocessing. Passed on recording, replaying and uncached calls alike. offload: Where to store the largest intermediate activations, as - in :meth:`forward`. Only passed to models with - :attr:`supports_offload`. Always ``None`` when replaying a + in :meth:`forward` and :meth:`fit`. Always ``None`` for + models without :attr:`supports_offload` and when replaying a cache, *i.e.*, in :meth:`predict`. kwargs: Additional keyword arguments passed by the caller. diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index 67d874f16d..55eba53ba3 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -275,6 +275,7 @@ def _attention_with_lse( query.dtype in (torch.float16, torch.bfloat16) and query.size(-1) % 8 == 0 and query.size(-1) <= 256 + and torch.cuda.get_device_capability(query.device)[0] >= 8 ): out, lse, *_ = torch.ops.aten._scaled_dot_product_flash_attention( query, key, value, scale=scale @@ -956,8 +957,8 @@ def compute(q: Tensor, kv: Tensor) -> tuple[Tensor, Tensor]: lse: Tensor | None = None for index in _batch_chunks(batch_shape, batch_size_limit): chunk_out, chunk_lse = compute( - _chunk(query, batch_shape, 2, index), - _chunk(key_value, batch_shape, 2, index), + q=_chunk(query, batch_shape, 2, index), + kv=_chunk(key_value, batch_shape, 2, index), ) if out is None or lse is None: out = chunk_out.new_empty(batch_shape + chunk_out.size()[-3:]) diff --git a/sdm/nn/linear.py b/sdm/nn/linear.py index 72e75584c3..e55f180fb6 100644 --- a/sdm/nn/linear.py +++ b/sdm/nn/linear.py @@ -35,7 +35,7 @@ def forward( out.stride(i) != out.size(i + 1) * out.stride(i + 1) for i in range(out.dim() - 3) ): - for input_i, out_i in zip(input, out): + for input_i, out_i in zip(input, out, strict=True): self.forward(input_i, out=out_i) return out else: diff --git a/sdm/nn/set_transformer.py b/sdm/nn/set_transformer.py index f62525ded5..b6b31252b1 100644 --- a/sdm/nn/set_transformer.py +++ b/sdm/nn/set_transformer.py @@ -195,7 +195,7 @@ def summarize( if summary is None: return out, lse summary_out, summary_lse = summary - total = torch.logaddexp(summary_lse, lse) # [..., M, H] + total = summary_lse.logaddexp(lse) # [..., M, H] out = out.float() * (lse - total).exp()[..., None] out += summary_out.float() * (summary_lse - total).exp()[..., None] return out, total diff --git a/test/models/test_offload.py b/test/models/test_offload.py index df6a27ac26..8369e43095 100644 --- a/test/models/test_offload.py +++ b/test/models/test_offload.py @@ -97,9 +97,11 @@ def _tabiclv2(device: torch.device) -> Callable[..., Tensor]: @onlyCUDA @pytest.mark.parametrize("build", [_kumo_tabular, _tabfm, _tabiclv2]) @pytest.mark.parametrize("chunk_memory_fraction", ["0.05", "1e-9"]) +@pytest.mark.parametrize("dtype", [None, torch.float16, torch.bfloat16]) def test_row_embedding_offload( build: Callable[[torch.device], Callable[..., Tensor]], chunk_memory_fraction: str, + dtype: torch.dtype | None, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("SDM_CHUNK_MEMORY_FRACTION", chunk_memory_fraction) @@ -109,7 +111,10 @@ def test_row_embedding_offload( y = torch.randint(0, 3, (1, 30), device=device) outs: dict[Any, Tensor] = {} - with torch.no_grad(): + with ( + torch.no_grad(), + torch.autocast(device.type, dtype=dtype, enabled=dtype is not None), + ): for offload in (None, "cpu"): cache = Cache() outs[offload, "fit"] = embedding( @@ -120,9 +125,11 @@ def test_row_embedding_offload( ) outs[offload] = embedding(x, y, offload=offload) + # Chunked inducing-point summaries round differently in half precision. + tol = {} if dtype is None else {"atol": 5e-2, "rtol": 5e-2} for key in ("fit", "predict"): - torch.testing.assert_close(outs["cpu", key], outs[None, key]) - torch.testing.assert_close(outs["cpu"], outs[None]) + torch.testing.assert_close(outs["cpu", key], outs[None, key], **tol) + torch.testing.assert_close(outs["cpu"], outs[None], **tol) @onlyCUDA diff --git a/test/nn/test_set_transformer.py b/test/nn/test_set_transformer.py index fee0217e87..fd2221e3bc 100644 --- a/test/nn/test_set_transformer.py +++ b/test/nn/test_set_transformer.py @@ -62,6 +62,57 @@ def test_induced_transformer_block( torch.testing.assert_close(out_seqused, out_cross) +@withCUDA +@pytest.mark.parametrize("batch_size_limit", [None, 2, 4]) +def test_induced_transformer_block_summarize( + device: torch.device, + batch_size_limit: int | None, +) -> None: + channels = 8 + module = InducedTransformerBlock( + channels=channels, + num_inducing_points=4, + inducing_block=TransformerBlock( + channels=channels, + num_query_heads=2, + mlp=torch.nn.Identity(), + device=device, + ), + output_block=TransformerBlock( + channels=channels, + num_query_heads=2, + mlp=torch.nn.Identity(), + device=device, + ), + device=device, + ) + # Zero-initialized output projections would hide the attention. + for parameter in module.parameters(): + if not parameter.any(): + torch.nn.init.normal_(parameter) + key_value = torch.randn(2, 3, 7, channels, device=device) + + with torch.no_grad(): + _, expected = module( + query=key_value[..., :0, :], + key_value=key_value, + return_key_value=True, + ) + summary = None + for chunk in key_value.split([3, 4], dim=-2): + summary = module.summarize( + chunk, + key_len=key_value.size(-2), + summary=summary, + batch_size_limit=batch_size_limit, + ) + assert summary is not None + out = module.induced_key_value(summary) + + torch.testing.assert_close(out.key, expected.key) + torch.testing.assert_close(out.value, expected.value) + + def test_induced_transformer_block_kv_cache() -> None: batch_size = 2 set_size = 6 From 14683acc46382e881c441c275fe786324982989b Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 5 Nov 2026 16:37:35 -0800 Subject: [PATCH 10/12] fix precommit ty error --- test/models/test_offload.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/test/models/test_offload.py b/test/models/test_offload.py index 8369e43095..40abb2b3d6 100644 --- a/test/models/test_offload.py +++ b/test/models/test_offload.py @@ -123,13 +123,17 @@ def test_row_embedding_offload( outs[offload, "predict"] = embedding( x[:, 30:], y[:, :0], cache=cache.freeze() ) - outs[offload] = embedding(x, y, offload=offload) + outs[offload, None] = embedding(x, y, offload=offload) # Chunked inducing-point summaries round differently in half precision. - tol = {} if dtype is None else {"atol": 5e-2, "rtol": 5e-2} - for key in ("fit", "predict"): - torch.testing.assert_close(outs["cpu", key], outs[None, key], **tol) - torch.testing.assert_close(outs["cpu"], outs[None], **tol) + tol = None if dtype is None else 5e-2 + for key in ("fit", "predict", None): + torch.testing.assert_close( + outs["cpu", key], + outs[None, key], + atol=tol, + rtol=tol, + ) @onlyCUDA From 5215194f5d93ecbf208a9ad180068b925eb10d65 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 5 Nov 2026 17:18:22 -0800 Subject: [PATCH 11/12] address new coderabiit review comments --- sdm/nn/attention.py | 45 +++++++++++++++++++++------------ test/nn/test_set_transformer.py | 3 +++ 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index 55eba53ba3..58fbbc4e88 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -266,25 +266,38 @@ def _attention_with_lse( scale: float | None, ) -> tuple[Tensor, Tensor]: # [B, Hq, Q, C], [B, Hq, Q] # `F.scaled_dot_product_attention` does not return the log-sum-exp. - if not query.is_cuda: + # Query heads that share a key/value head attend as one longer query, + # since kernels of older PyTorch versions don't group key/value heads. + B, H, Q, C = query.size() + query = query.reshape(B, key.size(1), -1, C) # [B, Hkv, Hq // Hkv * Q, C] + if query.device.type == "cpu": out, lse = torch.ops.aten._scaled_dot_product_flash_attention_for_cpu( query, key, value, scale=scale ) - return out, lse - if ( - query.dtype in (torch.float16, torch.bfloat16) - and query.size(-1) % 8 == 0 - and query.size(-1) <= 256 - and torch.cuda.get_device_capability(query.device)[0] >= 8 - ): - out, lse, *_ = torch.ops.aten._scaled_dot_product_flash_attention( - query, key, value, scale=scale - ) - return out, lse - out, lse, *_ = torch.ops.aten._scaled_dot_product_efficient_attention( - query, key, value, None, True, scale=scale - ) - return out, lse[..., : query.size(-2)] + elif query.is_cuda: + if ( + query.dtype in (torch.float16, torch.bfloat16) + and query.size(-1) % 8 == 0 + and query.size(-1) <= 256 + and torch.cuda.get_device_capability(query.device)[0] >= 8 + ): + out, lse, *_ = torch.ops.aten._scaled_dot_product_flash_attention( + query, key, value, scale=scale + ) + else: + out, lse, *_ = ( + torch.ops.aten._scaled_dot_product_efficient_attention( + query, key, value, None, True, scale=scale + ) + ) + lse = lse[..., : query.size(-2)] + else: + # Other devices, e.g., MPS, have no kernel returning the log-sum-exp. + scale = C**-0.5 if scale is None else scale + logits = (query @ key.mT).float() * scale + lse = logits.logsumexp(dim=-1) # [B, Hkv, Hq // Hkv * Q] + out = (logits - lse[..., None]).exp().to(value.dtype) @ value + return out.reshape(B, H, Q, C), lse.reshape(B, H, Q) class Attention(torch.nn.Module): diff --git a/test/nn/test_set_transformer.py b/test/nn/test_set_transformer.py index fd2221e3bc..01988ad68d 100644 --- a/test/nn/test_set_transformer.py +++ b/test/nn/test_set_transformer.py @@ -64,9 +64,11 @@ def test_induced_transformer_block( @withCUDA @pytest.mark.parametrize("batch_size_limit", [None, 2, 4]) +@pytest.mark.parametrize("num_key_value_heads", [2, 1]) def test_induced_transformer_block_summarize( device: torch.device, batch_size_limit: int | None, + num_key_value_heads: int, ) -> None: channels = 8 module = InducedTransformerBlock( @@ -75,6 +77,7 @@ def test_induced_transformer_block_summarize( inducing_block=TransformerBlock( channels=channels, num_query_heads=2, + num_key_value_heads=num_key_value_heads, mlp=torch.nn.Identity(), device=device, ), From 2807ccc2c534d024f6a5b788d53ffc3c323bc8fc Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 5 Nov 2026 17:30:02 -0800 Subject: [PATCH 12/12] codecov fix --- test/models/test_offload.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/models/test_offload.py b/test/models/test_offload.py index 40abb2b3d6..0f6936f446 100644 --- a/test/models/test_offload.py +++ b/test/models/test_offload.py @@ -183,6 +183,22 @@ def test_row_embedding_offload_subsamples_keys() -> None: torch.testing.assert_close(outs[1], outs[0]) +@pytest.mark.parametrize("build", [_kumo_tabular, _tabfm, _tabiclv2]) +def test_row_embedding_grad( + build: Callable[[torch.device], Callable[..., Tensor]], +) -> None: + embedding = build(torch.device("cpu")) + x = torch.randn(1, 40, 5) + y = torch.randint(0, 3, (1, 30)) + + with torch.no_grad(): + expected = embedding(x, y) + torch.testing.assert_close(embedding(x, y), expected) + + with pytest.raises(RuntimeError, match="gradients are disabled"): + embedding(x, y, offload="cpu") + + def test_offload_requires_cuda_inputs() -> None: model = KumoTabular(task="classification", size="small", pretrained=False) x = TableTensor.from_tensor(torch.randn(8, 5))