Add CPU offloading of row embedding token tables to tabular ICL models

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
 ------
 

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.**",

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
+    <https://docs.pytorch.org/docs/stable/notes/cuda.html#memory-management>`__.
+    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.