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/pyproject.toml b/pyproject.toml index b41913d789..013a5075a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -213,6 +213,7 @@ ignorelist = ["id"] [tool.ruff.lint.flake8-tidy-imports] ban-relative-imports = "all" banned-module-level-imports = [ + "cuda", "cudf", "cupy", "graphviz", @@ -254,6 +255,7 @@ banned-module-level-imports = [ [tool.ty.analysis] allowed-unresolved-imports = [ + "cuda.**", "cudf.**", "cupy.**", "pyg_lib.**", diff --git a/sdm/_memory.py b/sdm/_memory.py index ae55f6da48..fa2ee88ce9 100644 --- a/sdm/_memory.py +++ b/sdm/_memory.py @@ -45,6 +45,41 @@ def set_pinned_max_round_threshold( torch._C._accelerator_setAllocatorSettings(",".join(options)) +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: + + - 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, 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(), 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: 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 new file mode 100644 index 0000000000..bbac1a36d5 --- /dev/null +++ b/sdm/_token_buffer.py @@ -0,0 +1,365 @@ +# 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 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 + 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]`` 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 store the buffer: ``None`` stores it on ``device`` + and ``"cpu"`` stores it in pinned host memory, which requires a + CUDA ``device``. + """ + + 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 + *B, R, T, D = self.size + + if offload is None: + self.tensor = torch.empty( + (*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 " + "target classes than the model predicts natively" + ) + 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, + ) + 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"""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)) + 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, :, :].transpose(-2, -3) + 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"""Update rows in place, in chunks of ``N`` rows. + + 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.transpose(-2, -3), slice(0, self.size[-3])) + return + + *B, R, T, D = self.size + 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 * token_bytes, + (chunk.stop - chunk.start) * token_bytes, + T, + ) + for chunk in chunks + ], + compute=lambda slot, k: compute( + slot.view(self.dtype) + .view(*B, T, chunks[k].stop - chunks[k].start, D) + .transpose(-2, -3), + chunks[k], + ), + load=load, + store=store, + ) + + def columns( + self, + 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, 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, :, :], tokens) + return + + *B, R, _, D = self.size + 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) + ] + 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, R, D + ), + chunks[k], + ), + load=load, + ) + + 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 + + @torch.compiler.disable + 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, + 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. + + 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``. + 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 + 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, 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, + 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, load=embed is None) + + 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), + ) + ) 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 0f7a0f5b08..e54fd50cc2 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,26 @@ 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. + ``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. + :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: 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 +207,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 +246,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 +278,23 @@ 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. ``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. + :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) if context_only_columns and not self.supports_context_only_columns: raise ValueError( f"{self.__class__.__name__!r} does not support context-only " @@ -301,56 +339,61 @@ 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: - # Keep class labels on GPU for compiled prediction. - classes = batch_cache.pop("classes") - 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: + # Keep class labels on GPU for compiled prediction. + classes = batch_cache.pop("classes") + 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: - batch_cache["classes"] = classes - torch.cuda.current_stream(x.device).synchronize() + finally: + batch_cache["classes"] = classes + torch.cuda.current_stream(x.device).synchronize() - cache[i] = batch_cache + cache[i] = batch_cache + finally: + _release_offload_memory(offload) self._cache = cache.freeze() @@ -492,10 +535,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 @@ -521,6 +569,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` 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. Returns: @@ -533,6 +585,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], @@ -772,6 +840,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 aa3e2d6bc9..c465396530 100644 --- a/sdm/models/kumo/tabular/cell_embedding.py +++ b/sdm/models/kumo/tabular/cell_embedding.py @@ -37,44 +37,55 @@ def __init__( dtype=dtype, ) - def forward( # type: ignore + def impute_mean( self, x: Tensor, # [..., R, C] - categorical_mask: Tensor, # [..., C] *, train_size: int | Tensor | None = 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: # [..., 1, C] if cache is not None and cache.is_replaying: - impute_mean = cast(Tensor, cache["cell_impute_mean"]) - else: - assert train_size is not None - if isinstance(train_size, Tensor): - context = torch.arange(x.size(-2), device=x.device) - context = context < train_size[..., None] - observed = x.masked_fill(~context[..., None], torch.nan) - else: - observed = x[..., :train_size, :] - impute_mean = observed.nanmean(dim=-2, keepdim=True) - impute_mean = impute_mean.masked_fill(impute_mean.isnan(), 0.0) + return cast(Tensor, cache["cell_impute_mean"]) + assert train_size is not None + if isinstance(train_size, Tensor): + context = torch.arange(x.size(-2), device=x.device) + context = context < train_size[..., None] + observed = x.masked_fill(~context[..., None], torch.nan) + else: + observed = x[..., :train_size, :] + impute_mean = observed.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 - x = torch.where(missing, impute_mean, x) + def forward( # type: ignore + self, + x: Tensor, # [..., R, C] + categorical_mask: Tensor, # [..., C] + *, + train_size: int | Tensor | 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, 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 + ) 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 @@ -84,9 +95,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/model.py b/sdm/models/kumo/tabular/model.py index b28cf4f55e..b7bbb7e2bf 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 862c2af601..7a666f6ec4 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 from sdm.cache import Cache, KVCacheEntry from sdm.models.kumo.tabular.block import KumoTabularTransformerBlock from sdm.models.kumo.tabular.cell_embedding import CellEmbedding @@ -115,6 +116,7 @@ def forward( cache: Cache | None = None, num_columns: Tensor | None = None, train_size: Tensor | None = None, + offload: Literal["cpu"] | None = None, ) -> Tensor: # [..., R, K * D] *B, R, C = x.size() @@ -127,61 +129,173 @@ 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 train_size is not None: + context = torch.arange(R_train, device=y.device) + context = context < train_size[..., None] + y_emb = y_emb.masked_fill(~context[..., None], 0) + if torch.is_grad_enabled() or train_size is not None: - 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 if train_size is None else train_size, - num_columns=num_columns, + num_train_rows=R_train, + train_size=train_size, 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 if train_size is None else train_size, + impute_mean=impute_mean, num_columns=num_columns, - cache=cache, batch_size_limit="auto", - out=buffer[..., K:, :], + out=tokens[..., K:, :], ) + 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) - 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) - if train_size is not None: - context = torch.arange(R_train, device=y.device) - context = context < train_size[..., None] - y_emb = y_emb.masked_fill(~context[..., None, None], 0) - x[..., :R_train, :, :] += y_emb.to(x.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] + key_value: KVCacheEntry | None = None + for i in range(L + 1): + summary: tuple[Tensor, Tensor] | None = None - if buffer is not None: - x = x.transpose(-2, -3) # [..., C, R, D] + def update( + tokens: Tensor, # [..., N, K + C, D] + rows: slice, + i: int = i, + key_value: KVCacheEntry | None = key_value, + ) -> None: + 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", + ) + + 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) + + def _forward_autograd( + self, + x: Tensor, # [..., R, C] + y_emb: Tensor | None, # [..., R_train, D] + categorical_mask: Tensor, # [..., C] + *, + num_train_rows: int, + train_size: Tensor | None, # [...] + cache: Cache | None, + num_columns: Tensor | None, + row_key_lengths: Tensor | None, + ) -> Tensor: # [..., R, K * D] + *B, R, _ = x.size() + R_train = num_train_rows + K = self.readout_token.size(-2) + D = self.channels + + x = self.cell_embedding( + x=x, + categorical_mask=categorical_mask, + train_size=R_train if train_size is None else train_size, + 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( @@ -194,39 +308,22 @@ def forward( else train_size.unsqueeze(-1), 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/cell_embedding.py b/sdm/models/tabfm/cell_embedding.py index babe20c5d4..3f8c036f81 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 chunk_index in range((R + rows_per_chunk - 1) // rows_per_chunk): start = chunk_index * rows_per_chunk self._forward( 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..5a8e9d8ce2 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, 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,131 @@ 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( + + 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(cells: Tensor, columns: slice) -> None: # [..., N, R, D] + self.cell_embedding( x, categorical_mask, + columns=slice(columns.start - K, columns.stop - K), batch_size_limit="auto", - out=buffer[..., K:, :], + out=cells.transpose(-2, -3), ) + if y_emb is not None: + cells[..., :R_train, :] += y_emb.unsqueeze(-3).to(cells.dtype) - 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: + 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, + 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( + tokens: Tensor, # [..., N, K + C, D] + 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: + 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, + ) + # 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) + + 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 +296,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 +304,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 +326,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..49da62d1d0 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, 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,126 @@ 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(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, out=buffer[0, ..., K:, :]) - buffer[1:, ..., K:, :] = buffer[:1, ..., K:, :] - x = buffer[..., K:, :] + self.lin(group, out=cells[0]) + cells[1:] = cells[:1] else: - x = self.lin(x, out=buffer[..., K:, :]) # [..., R, C, D] + self.lin(group, out=cells) + if y_emb is None: + return + rows = slice(R_train) if train_mask is None else train_mask + cells[..., rows, :] += y_emb.unsqueeze(-3).to(cells.dtype) + + with buffer: + 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, + embed=embed, + ) + y_emb = None - 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 +282,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 +296,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/sdm/nn/attention.py b/sdm/nn/attention.py index f58aaca488..58fbbc4e88 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 @@ -213,6 +214,91 @@ 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. + # 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 + ) + 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): r"""Multi-head attention layer with grouped-query attention support. @@ -384,7 +470,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 @@ -443,27 +561,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): @@ -659,26 +797,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, @@ -714,81 +835,51 @@ 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) - + key: Tensor | None = None + value: Tensor | None = None + for index in _batch_chunks(batch_shape, batch_size_limit): 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, @@ -831,6 +922,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( + 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:]) + 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) @@ -839,13 +1040,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, @@ -875,6 +1075,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, @@ -903,8 +1121,7 @@ def _chunk( tensor: Tensor, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> Tensor: ... @@ -913,8 +1130,7 @@ def _chunk( tensor: KVCacheEntry, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> KVCacheEntry: ... @@ -923,8 +1139,7 @@ def _chunk( tensor: None, batch_shape: torch.Size, trailing_dims: int, - start: int, - end: int, + index: tuple[int | slice, ...], ) -> None: ... @@ -932,8 +1147,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: @@ -941,18 +1155,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..e55f180fb6 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, strict=True): + 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)) diff --git a/sdm/nn/set_transformer.py b/sdm/nn/set_transformer.py index 9247d48af2..b6b31252b1 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 = 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 + + 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 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..0f6936f446 --- /dev/null +++ b/test/models/test_offload.py @@ -0,0 +1,265 @@ +# 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"]) +@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) + 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(), + torch.autocast(device.type, dtype=dtype, enabled=dtype is not None), + ): + 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, None] = embedding(x, y, offload=offload) + + # Chunked inducing-point summaries round differently in half precision. + 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 +@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") + 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]) + + +@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)) + 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", + ) diff --git a/test/nn/test_set_transformer.py b/test/nn/test_set_transformer.py index fee0217e87..01988ad68d 100644 --- a/test/nn/test_set_transformer.py +++ b/test/nn/test_set_transformer.py @@ -62,6 +62,60 @@ def test_induced_transformer_block( torch.testing.assert_close(out_seqused, out_cross) +@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( + channels=channels, + num_inducing_points=4, + inducing_block=TransformerBlock( + channels=channels, + num_query_heads=2, + num_key_value_heads=num_key_value_heads, + 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