From cb31cb2a68081a4f1ff7f426accdf647323450e0 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Thu, 15 Oct 2026 22:14:28 +0000 Subject: [PATCH 1/7] Preserve fitted cache pinning across model serialization --- sdm/models/base.py | 27 +++++++++++++ test/models/test_base.py | 86 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+) diff --git a/sdm/models/base.py b/sdm/models/base.py index b533dfd24e..1a3c93880d 100644 --- a/sdm/models/base.py +++ b/sdm/models/base.py @@ -528,11 +528,38 @@ def __getstate__(self) -> dict[str, object]: stream.synchronize() state = super().__getstate__() state.pop("_transfer_streams", None) + # Tensor serialization does not preserve pinned CPU allocations. + state["_pinned_cache_batches"] = ( + tuple( + index + for index in range(cast(int, self._cache["num_batches"])) + if cast(Cache, self._cache[index]).is_pinned() + ) + if self._cache is not None + else () + ) return state def __setstate__(self, state: dict[str, object]) -> None: + pinned_batches = cast( + tuple[int, ...], state.pop("_pinned_cache_batches", ()) + ) super().__setstate__(state) self._transfer_streams = {} + if not pinned_batches or not torch.cuda.is_available(): + return + + cache = cast(Cache, self._cache) + items: dict[Hashable, object] | None = None + for index in pinned_batches: + batch = cast(Cache, cache[index]) + if not batch.is_pinned(): + if items is None: + set_pinned_max_round_threshold(mb=1) + items = dict(cache.items()) + items[index] = batch.pin_memory() + if items is not None: + self._cache = Cache(items).freeze() def __repr__(self) -> str: device = next(self.parameters()).device diff --git a/test/models/test_base.py b/test/models/test_base.py index 91f1784969..2cbe807bd2 100644 --- a/test/models/test_base.py +++ b/test/models/test_base.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import copy +import io +import pickle from dataclasses import dataclass from typing import Any, ClassVar, cast @@ -19,6 +22,7 @@ from sdm.models import ICLModel from sdm.models.callback import Callback from sdm.processing import InvertibleMixin, Processor +from sdm.testing import withCUDA @dataclass @@ -128,6 +132,88 @@ class _UnsupportedRecordingModel(_RecordingModel): supports_related_tables = False +@withCUDA +@pytest.mark.parametrize("serialization", ["pickle", "torch", "deepcopy"]) +def test_fitted_model_serialization_preserves_pinning( + device: torch.device, + serialization: str, +) -> None: + model = _ClassFrequencyModel() + x = torch.randn(12, 2, device=device) + y = TableTensor( + categorical=CategoricalTensor( + code=torch.arange(12, device=device).remainder(3).view(-1, 1), + categories=(torch.tensor([10, 20, 30], device=device),), + ) + ) + model.fit(x, y, num_estimators=2) + expected = model.predict(x).numerical + cache = cast(Cache, model._cache) + assert all( + cast(Cache, cache[index]).is_pinned() == (device.type == "cuda") + for index in range(cast(int, cache["num_batches"])) + ) + assert copy.copy(model)._cache is cache + + if serialization == "pickle": + restored = pickle.loads(pickle.dumps(model)) + elif serialization == "deepcopy": + restored = copy.deepcopy(model) + else: + stream = io.BytesIO() + torch.save(model, stream) + stream.seek(0) + restored = torch.load(stream, weights_only=False) + + restored_cache = cast(Cache, restored._cache) + assert all( + cast(Cache, restored_cache[index]).is_pinned() + == (device.type == "cuda") + for index in range(cast(int, restored_cache["num_batches"])) + ) + torch.testing.assert_close( + restored.predict(x).numerical, expected, rtol=0, atol=0 + ) + restored.clear() + with pytest.raises(RuntimeError, match="not yet fitted"): + restored.predict(x) + restored.fit(x, y, num_estimators=2) + torch.testing.assert_close( + restored.predict(x).numerical, expected, rtol=0, atol=0 + ) + + +@withCUDA +def test_fitted_model_serialization_without_cuda_pinning( + device: torch.device, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _ClassFrequencyModel() + x = torch.randn(12, 2, device=device) + y = TableTensor( + categorical=CategoricalTensor( + code=torch.arange(12, device=device).remainder(3).view(-1, 1), + categories=(torch.tensor([10, 20, 30], device=device),), + ) + ) + model.fit(x, y, num_estimators=2) + expected = model.predict(x).numerical.cpu() + stream = io.BytesIO() + torch.save(model, stream) + stream.seek(0) + with monkeypatch.context() as patch: + patch.setattr(torch.cuda, "is_available", lambda: False) + restored = torch.load(stream, map_location="cpu", weights_only=False) + torch.testing.assert_close( + restored.predict(x.cpu()).numerical, expected, rtol=0, atol=0 + ) + cache = cast(Cache, restored._cache) + assert all( + not cast(Cache, cache[index]).is_pinned() + for index in range(cast(int, cache["num_batches"])) + ) + + class MyCallback(Callback): """Callback used by callback lifecycle tests.""" From 284950bcd8ea52617192791bde0e0a75d80e6226 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Thu, 15 Oct 2026 22:19:55 +0000 Subject: [PATCH 2/7] Document fitted cache pinning restoration and load cost --- docs/source/icl.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/source/icl.md b/docs/source/icl.md index c8174aee44..fa4d95a593 100644 --- a/docs/source/icl.md +++ b/docs/source/icl.md @@ -73,6 +73,11 @@ Use one-shot {py:meth}`~sdm.models.ICLModel.forward` calls for one-time calls wh Unlike {py:meth}`~sdm.models.ICLModel.forward`, {py:meth}`~sdm.models.ICLModel.predict` does not support gradient-based fine-tuning and raises if the model is in train mode. +CUDA ensembles store fitted caches in pinned CPU memory so transfers can overlap prediction. +Saving and loading a whole fitted model, or making a deep copy, restores this pinning when CUDA is available. +Restoring pinning copies the cached CPU tensors during loading and can temporarily increase host-memory use. +Previously saved models without pinning metadata retain their existing loading behavior. + ## Model Concepts Structured data foundation models are not bound to a specific task type. From e0c3bfac6fdb4a7c88047bacab72db0ba2ae081f Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Thu, 15 Oct 2026 22:28:28 +0000 Subject: [PATCH 3/7] Cover fitted model loading onto a requested device --- test/models/test_base.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/test/models/test_base.py b/test/models/test_base.py index 2cbe807bd2..0fe8ac8910 100644 --- a/test/models/test_base.py +++ b/test/models/test_base.py @@ -133,7 +133,9 @@ class _UnsupportedRecordingModel(_RecordingModel): @withCUDA -@pytest.mark.parametrize("serialization", ["pickle", "torch", "deepcopy"]) +@pytest.mark.parametrize( + "serialization", ["pickle", "torch", "torch_device", "deepcopy"] +) def test_fitted_model_serialization_preserves_pinning( device: torch.device, serialization: str, @@ -163,14 +165,24 @@ def test_fitted_model_serialization_preserves_pinning( stream = io.BytesIO() torch.save(model, stream) stream.seek(0) - restored = torch.load(stream, weights_only=False) + restored = torch.load( + stream, + map_location=device if serialization == "torch_device" else None, + weights_only=False, + ) restored_cache = cast(Cache, restored._cache) assert all( cast(Cache, restored_cache[index]).is_pinned() - == (device.type == "cuda") + == (device.type == "cuda" and serialization != "torch_device") for index in range(cast(int, restored_cache["num_batches"])) ) + if serialization == "torch_device": + assert all( + tensor.device == x.device + for index in range(cast(int, restored_cache["num_batches"])) + for tensor in cast(Cache, restored_cache[index])._tensors() + ) torch.testing.assert_close( restored.predict(x).numerical, expected, rtol=0, atol=0 ) From 0e964ee793e9e58239c81e708062f51d243cba72 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Thu, 15 Oct 2026 22:37:51 +0000 Subject: [PATCH 4/7] Respect explicit device placement when restoring pinned caches --- docs/source/icl.md | 1 + sdm/models/base.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/source/icl.md b/docs/source/icl.md index fa4d95a593..ac55e18479 100644 --- a/docs/source/icl.md +++ b/docs/source/icl.md @@ -76,6 +76,7 @@ Unlike {py:meth}`~sdm.models.ICLModel.forward`, {py:meth}`~sdm.models.ICLModel.p CUDA ensembles store fitted caches in pinned CPU memory so transfers can overlap prediction. Saving and loading a whole fitted model, or making a deep copy, restores this pinning when CUDA is available. Restoring pinning copies the cached CPU tensors during loading and can temporarily increase host-memory use. +Tensors explicitly mapped to another device retain that placement. Previously saved models without pinning metadata retain their existing loading behavior. ## Model Concepts diff --git a/sdm/models/base.py b/sdm/models/base.py index 1a3c93880d..3664db35ab 100644 --- a/sdm/models/base.py +++ b/sdm/models/base.py @@ -553,11 +553,18 @@ def __setstate__(self, state: dict[str, object]) -> None: items: dict[Hashable, object] | None = None for index in pinned_batches: batch = cast(Cache, cache[index]) - if not batch.is_pinned(): + if any( + tensor.is_cpu and not tensor.is_pinned() + for tensor in batch._tensors() + ): if items is None: set_pinned_max_round_threshold(mb=1) items = dict(cache.items()) - items[index] = batch.pin_memory() + items[index] = batch._apply_tensor( + lambda tensor: ( + tensor.pin_memory() if tensor.is_cpu else tensor + ) + ) if items is not None: self._cache = Cache(items).freeze() From ffcf34ee17bfa0d6b54b5529a2b2353aa3d8c35b Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Thu, 15 Oct 2026 23:53:46 +0000 Subject: [PATCH 5/7] Remove cache-pinning details from ICL guide --- docs/source/icl.md | 6 ------ 1 file changed, 6 deletions(-) diff --git a/docs/source/icl.md b/docs/source/icl.md index ac55e18479..c8174aee44 100644 --- a/docs/source/icl.md +++ b/docs/source/icl.md @@ -73,12 +73,6 @@ Use one-shot {py:meth}`~sdm.models.ICLModel.forward` calls for one-time calls wh Unlike {py:meth}`~sdm.models.ICLModel.forward`, {py:meth}`~sdm.models.ICLModel.predict` does not support gradient-based fine-tuning and raises if the model is in train mode. -CUDA ensembles store fitted caches in pinned CPU memory so transfers can overlap prediction. -Saving and loading a whole fitted model, or making a deep copy, restores this pinning when CUDA is available. -Restoring pinning copies the cached CPU tensors during loading and can temporarily increase host-memory use. -Tensors explicitly mapped to another device retain that placement. -Previously saved models without pinning metadata retain their existing loading behavior. - ## Model Concepts Structured data foundation models are not bound to a specific task type. From 88e8ab0f32d1284d3e55b3dfb0d0b79ffdd0bcd3 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 22 Oct 2026 16:05:12 -0700 Subject: [PATCH 6/7] refactor: Move the new state hooks to `Cache` #1152 (#1226) --- sdm/cache.py | 36 ++++++++++++++++++++++++++++++++++- sdm/models/base/icl.py | 43 ------------------------------------------ 2 files changed, 35 insertions(+), 44 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 08ab744c70..f644afe5f1 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -10,10 +10,12 @@ MutableMapping, ) from enum import StrEnum -from typing import NamedTuple, Self +from typing import NamedTuple, Self, cast +import torch from torch import Tensor +from sdm._memory import set_pinned_max_round_threshold from sdm.tensor.mixin import DeviceMixin @@ -126,6 +128,38 @@ def __len__(self) -> int: def __repr__(self) -> str: return repr(self._items) + def __getstate__(self) -> dict[str, object]: + state = self.__dict__.copy() + # Stable serialization metadata: positions in the tensor traversal. + state["_pinned_tensors"] = tuple( + index + for index, tensor in enumerate(self._tensors()) + if tensor.is_pinned() + ) + return state + + def __setstate__(self, state: dict[str, object]) -> None: + pinned = set(cast(tuple[int, ...], state.pop("_pinned_tensors", ()))) + self.__dict__.update(state) + if not pinned or not torch.cuda.is_available(): + return + + replacements: dict[int, Tensor] = {} + for index, tensor in enumerate(self._tensors()): + if ( + index in pinned + and tensor.is_cpu + and not tensor.is_pinned() + and id(tensor) not in replacements + ): + if not replacements: + set_pinned_max_round_threshold(mb=1) + replacements[id(tensor)] = tensor.pin_memory() + if replacements: + self._items = self._apply_tensor( + lambda tensor: replacements.get(id(tensor), tensor) + )._items + def _tensors(self) -> Iterator[Tensor]: def _iter_tensors(value: object) -> Iterator[Tensor]: if isinstance(value, Tensor): diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index f3f2107781..00b1dfdce1 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -787,49 +787,6 @@ def _validate_query( "same schema" ) - # Python Builtins ######################################################### - - def __getstate__(self) -> dict[str, object]: - state = super().__getstate__() - # Tensor serialization does not preserve pinned CPU allocations. - state["_pinned_cache_batches"] = ( - tuple( - index - for index in range(cast(int, self._cache["num_batches"])) - if cast(Cache, self._cache[index]).is_pinned() - ) - if self._cache is not None - else () - ) - return state - - def __setstate__(self, state: dict[str, object]) -> None: - pinned_batches = cast( - tuple[int, ...], state.pop("_pinned_cache_batches", ()) - ) - super().__setstate__(state) - if not pinned_batches or not torch.cuda.is_available(): - return - - cache = cast(Cache, self._cache) - items: dict[Hashable, object] | None = None - for index in pinned_batches: - batch = cast(Cache, cache[index]) - if any( - tensor.is_cpu and not tensor.is_pinned() - for tensor in batch._tensors() - ): - if items is None: - set_pinned_max_round_threshold(mb=1) - items = dict(cache.items()) - items[index] = batch._apply_tensor( - lambda tensor: ( - tensor.pin_memory() if tensor.is_cpu else tensor - ) - ) - if items is not None: - self._cache = Cache(items).freeze() - def _stack(tables: Sequence[TableTensor]) -> TableTensor: ref = tables[0] From 3d420a7339780f2c1f65d033233e1d9889234089 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 22 Oct 2026 17:07:23 -0700 Subject: [PATCH 7/7] nits --- sdm/cache.py | 7 ++- test/models/base/test_icl.py | 98 ------------------------------------ test/test_cache.py | 75 ++++++++++++++++++++++++++- 3 files changed, 79 insertions(+), 101 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index f644afe5f1..59c6c0db06 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -131,7 +131,7 @@ def __repr__(self) -> str: def __getstate__(self) -> dict[str, object]: state = self.__dict__.copy() # Stable serialization metadata: positions in the tensor traversal. - state["_pinned_tensors"] = tuple( + state["_pinned_tensor_indices"] = tuple( index for index, tensor in enumerate(self._tensors()) if tensor.is_pinned() @@ -139,7 +139,9 @@ def __getstate__(self) -> dict[str, object]: return state def __setstate__(self, state: dict[str, object]) -> None: - pinned = set(cast(tuple[int, ...], state.pop("_pinned_tensors", ()))) + pinned = set( + cast(tuple[int, ...], state.pop("_pinned_tensor_indices", ())) + ) self.__dict__.update(state) if not pinned or not torch.cuda.is_available(): return @@ -155,6 +157,7 @@ def __setstate__(self, state: dict[str, object]) -> None: if not replacements: set_pinned_max_round_threshold(mb=1) replacements[id(tensor)] = tensor.pin_memory() + if replacements: self._items = self._apply_tensor( lambda tensor: replacements.get(id(tensor), tensor) diff --git a/test/models/base/test_icl.py b/test/models/base/test_icl.py index 74731db406..9ec4d576d1 100644 --- a/test/models/base/test_icl.py +++ b/test/models/base/test_icl.py @@ -1,9 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import copy -import io -import pickle from dataclasses import dataclass from typing import Any, ClassVar, cast @@ -21,7 +18,6 @@ from sdm.cache import Cache from sdm.models import Callback, ICLModel from sdm.processing import InvertibleMixin, Processor -from sdm.testing import withCUDA @dataclass @@ -131,100 +127,6 @@ class _UnsupportedRecordingModel(_RecordingModel): supports_related_tables = False -@withCUDA -@pytest.mark.parametrize( - "serialization", ["pickle", "torch", "torch_device", "deepcopy"] -) -def test_fitted_model_serialization_preserves_pinning( - device: torch.device, - serialization: str, -) -> None: - model = _ClassFrequencyModel() - x = torch.randn(12, 2, device=device) - y = TableTensor( - categorical=CategoricalTensor( - code=torch.arange(12, device=device).remainder(3).view(-1, 1), - categories=(torch.tensor([10, 20, 30], device=device),), - ) - ) - model.fit(x, y, num_estimators=2) - expected = model.predict(x).numerical - cache = cast(Cache, model._cache) - assert all( - cast(Cache, cache[index]).is_pinned() == (device.type == "cuda") - for index in range(cast(int, cache["num_batches"])) - ) - assert copy.copy(model)._cache is cache - - if serialization == "pickle": - restored = pickle.loads(pickle.dumps(model)) - elif serialization == "deepcopy": - restored = copy.deepcopy(model) - else: - stream = io.BytesIO() - torch.save(model, stream) - stream.seek(0) - restored = torch.load( - stream, - map_location=device if serialization == "torch_device" else None, - weights_only=False, - ) - - restored_cache = cast(Cache, restored._cache) - assert all( - cast(Cache, restored_cache[index]).is_pinned() - == (device.type == "cuda" and serialization != "torch_device") - for index in range(cast(int, restored_cache["num_batches"])) - ) - if serialization == "torch_device": - assert all( - tensor.device == x.device - for index in range(cast(int, restored_cache["num_batches"])) - for tensor in cast(Cache, restored_cache[index])._tensors() - ) - torch.testing.assert_close( - restored.predict(x).numerical, expected, rtol=0, atol=0 - ) - restored.clear() - with pytest.raises(RuntimeError, match="not yet fitted"): - restored.predict(x) - restored.fit(x, y, num_estimators=2) - torch.testing.assert_close( - restored.predict(x).numerical, expected, rtol=0, atol=0 - ) - - -@withCUDA -def test_fitted_model_serialization_without_cuda_pinning( - device: torch.device, - monkeypatch: pytest.MonkeyPatch, -) -> None: - model = _ClassFrequencyModel() - x = torch.randn(12, 2, device=device) - y = TableTensor( - categorical=CategoricalTensor( - code=torch.arange(12, device=device).remainder(3).view(-1, 1), - categories=(torch.tensor([10, 20, 30], device=device),), - ) - ) - model.fit(x, y, num_estimators=2) - expected = model.predict(x).numerical.cpu() - stream = io.BytesIO() - torch.save(model, stream) - stream.seek(0) - with monkeypatch.context() as patch: - patch.setattr(torch.cuda, "is_available", lambda: False) - restored = torch.load(stream, map_location="cpu", weights_only=False) - torch.testing.assert_close( - restored.predict(x.cpu()).numerical, expected, rtol=0, atol=0 - ) - cache = cast(Cache, restored._cache) - assert all( - not cast(Cache, cache[index]).is_pinned() - for index in range(cast(int, cache["num_batches"])) - ) - - class MyCallback(Callback): """Callback used by callback lifecycle tests.""" diff --git a/test/test_cache.py b/test/test_cache.py index 36865b1376..d05102f090 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -1,12 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import io from typing import cast +import pytest import torch from sdm.cache import Cache, KVCacheEntry -from sdm.testing import withCUDA +from sdm.testing import onlyCUDA, withCUDA @withCUDA @@ -45,3 +47,74 @@ def test_cache_size() -> None: ) assert cache.size() == 3 * 4 + 2 * 8 + 5 * 1 + 4 * 2 + + +@withCUDA +@pytest.mark.parametrize("map_to_device", [False, True]) +def test_cache_serialization( + device: torch.device, + map_to_device: bool, +) -> None: + nested = Cache( + entry=KVCacheEntry( + key=torch.arange(6).view(2, 3), value=torch.ones(2, 3) + ) + ) + if device.type == "cuda": + nested = nested.pin_memory() + cache = Cache(foo="foo", tensor=torch.arange(3), nested=nested).freeze() + + stream = io.BytesIO() + torch.save(cache, stream) + stream.seek(0) + restored = torch.load( + stream, + map_location=device if map_to_device else None, + weights_only=False, + ) + + assert restored is not cache + assert restored["foo"] == "foo" + assert restored.is_replaying + restored_nested = cast(Cache, restored["nested"]) + assert restored_nested.is_replaying + assert restored_nested.is_pinned() == ( + device.type == "cuda" and not map_to_device + ) + assert not cast(torch.Tensor, restored["tensor"]).is_pinned() + + load_device = device if map_to_device else torch.device("cpu") + expected_entry = cast(KVCacheEntry, nested["entry"]) + restored_entry = cast(KVCacheEntry, restored_nested["entry"]) + for actual, expected in ( + ( + cast(torch.Tensor, restored["tensor"]), + cast(torch.Tensor, cache["tensor"]), + ), + (restored_entry.key, expected_entry.key), + (restored_entry.value, expected_entry.value), + ): + assert actual.device == load_device + torch.testing.assert_close( + actual, expected.to(load_device), rtol=0, atol=0 + ) + + +@onlyCUDA +def test_cache_serialization_without_cuda( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache = Cache(tensor=torch.arange(3).pin_memory()) + stream = io.BytesIO() + torch.save(cache, stream) + stream.seek(0) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + restored = torch.load(stream, map_location="cpu", weights_only=False) + + assert restored.is_recording + assert restored.is_cpu + assert not restored.is_pinned() + torch.testing.assert_close( + restored["tensor"], cache["tensor"], rtol=0, atol=0 + )