From a27cde71292815f5113c06c7a0425a463b0f88ef Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 30 Aug 2026 00:57:59 +0800 Subject: [PATCH 1/2] Stop the debug name maps from pinning the model they snapshot module_names and param_names are module-level dicts, and dict keys are strong references. engine.__init__ fills them at engine.py:402 and only destroy() resets them at engine.py:965, so everything in the snapshotted model stays reachable for the life of the process. Anything that swaps a submodule out after that point cannot release it. setattr unlinks the old module from the tree, but its parameters are still live keys, so the refcount never reaches zero. Replacing 16 blocks of a 64 MiB model leaves all 64 MiB resident: before replaced 64.0 MiB, still resident 64.0 MiB (16/16 tensors) after replaced 64.0 MiB, still resident 0.0 MiB (0/16 tensors) weakref.WeakKeyDictionary is not usable here: its keys are weakref.ref objects whose __eq__ forwards to the referents, and comparing two live parameters returns a tensor rather than a bool, so a lookup raises "Boolean value of Tensor with more than one value is ambiguous". Key on id() and drop the entry from a finalizer instead, which keeps the identity semantics the dicts already had. deepspeed/utils/pin_memory.py tracks its allocations the same way. Reported as problem 2 of #8353, with the diagnosis there. This does not touch problem 1, the init-time transient in the AutoEP path. Signed-off-by: alanhuangyoo --- deepspeed/utils/debug.py | 69 +++++++++++++++--- tests/unit/utils/test_debug_name_maps.py | 90 ++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 10 deletions(-) create mode 100644 tests/unit/utils/test_debug_name_maps.py diff --git a/deepspeed/utils/debug.py b/deepspeed/utils/debug.py index f644562deee9..103e0ff6e47c 100644 --- a/deepspeed/utils/debug.py +++ b/deepspeed/utils/debug.py @@ -3,30 +3,79 @@ # DeepSpeed Team +import weakref + import deepspeed.comm as dist # For lazy import with printflock() fcntl = None + +class WeakIdNameMap: + """Maps modules and parameters to their names without keeping them alive. + + These maps are module-level and are only reset in ``destroy()``, so with strong keys + they pin the snapshotted model for the life of the process. Anything that replaces a + submodule afterwards -- expert-parallel replacement, kernel injection -- leaves the + replaced weights resident even though nothing else references them. + + ``weakref.WeakKeyDictionary`` cannot be used: its keys are ``weakref.ref`` objects + whose ``__eq__`` forwards to the referents, and comparing two live parameters yields a + tensor rather than a bool. Keying on ``id()`` and dropping the entry from a finalizer + keeps the identity semantics the previous dicts had. + """ + + def __init__(self): + self._names = {} + self._finalizers = {} + + def __setitem__(self, obj, name): + key = id(obj) + self._names[key] = name + # Replacing an entry whose object is still alive would otherwise leak its finalizer. + finalizer = self._finalizers.pop(key, None) + if finalizer is not None: + finalizer.detach() + self._finalizers[key] = weakref.finalize(obj, self._discard, key) + + def _discard(self, key): + self._names.pop(key, None) + self._finalizers.pop(key, None) + + def __getitem__(self, obj): + return self._names[id(obj)] + + def __contains__(self, obj): + return id(obj) in self._names + + def __len__(self): + return len(self._names) + + def clear(self): + for finalizer in self._finalizers.values(): + finalizer.detach() + self._names.clear() + self._finalizers.clear() + + # for debug purposes map module and param objects to their fully qualified names -module_names = {} -param_names = {} +module_names = WeakIdNameMap() +param_names = WeakIdNameMap() def debug_clear_module_and_param_names(): - global module_names - global param_names - module_names = {} - param_names = {} + module_names.clear() + param_names.clear() def debug_extract_module_and_param_names(model): # extract the fully qualified names as soon as the model is acquired - global module_names - global param_names + debug_clear_module_and_param_names() # XXX: can probably make a map of param2module and vice-versa - module_names = {module: name for name, module in model.named_modules()} - param_names = {param: name for name, param in model.named_parameters()} + for name, module in model.named_modules(): + module_names[module] = name + for name, param in model.named_parameters(): + param_names[param] = name def debug_module2name(module): diff --git a/tests/unit/utils/test_debug_name_maps.py b/tests/unit/utils/test_debug_name_maps.py new file mode 100644 index 000000000000..cf5ba1af6f23 --- /dev/null +++ b/tests/unit/utils/test_debug_name_maps.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""The debug name maps must not keep the model they snapshot alive. + +``debug_extract_module_and_param_names`` runs once during ``engine.__init__`` and the maps +are only reset in ``destroy()``. Anything that swaps a submodule out afterwards -- the +expert-parallel replacement in ``_configure_expert_parallel``, kernel injection -- leaves +the replaced weights reachable from those maps for the rest of the process. +""" + +import gc +import weakref + +import torch.nn as nn + +from deepspeed.utils.debug import ( + debug_clear_module_and_param_names, + debug_extract_module_and_param_names, + debug_module2name, + debug_param2name, + module_names, + param_names, +) + + +class _Block(nn.Module): + + def __init__(self, dim=8): + super().__init__() + self.lin = nn.Linear(dim, dim, bias=False) + + +class _Model(nn.Module): + + def __init__(self, num_blocks=3, dim=8): + super().__init__() + self.blocks = nn.ModuleList([_Block(dim) for _ in range(num_blocks)]) + self.head = nn.Linear(dim, dim, bias=False) + + +def test_names_resolve_and_fall_back(): + model = _Model() + debug_extract_module_and_param_names(model) + + assert debug_param2name(model.blocks[0].lin.weight) == "blocks.0.lin.weight" + assert debug_module2name(model.blocks[0].lin) == "blocks.0.lin" + assert debug_param2name(nn.Linear(2, 2, bias=False).weight) == "unknown" + assert debug_module2name(nn.Identity()) == "unknown" + + +def test_replaced_submodule_is_released(): + model = _Model() + debug_extract_module_and_param_names(model) + + replaced = model.blocks[0] + alive = [weakref.ref(replaced)] + [weakref.ref(p) for p in replaced.parameters()] + entries_before = len(param_names) + + model.blocks[0] = nn.Identity() + del replaced + gc.collect() + + assert all(ref() is None for ref in alive) + assert len(param_names) < entries_before + + +def test_clear_and_re_extract(): + debug_extract_module_and_param_names(_Model()) + debug_clear_module_and_param_names() + + assert len(module_names) == 0 + assert len(param_names) == 0 + + other = _Model(num_blocks=1) + debug_extract_module_and_param_names(other) + + assert debug_param2name(other.head.weight) == "head.weight" + + +def test_a_recycled_id_is_not_a_stale_hit(): + debug_clear_module_and_param_names() + doomed = nn.Linear(4, 4, bias=False) + param_names[doomed.weight] = "ghost" + + del doomed + gc.collect() + + assert debug_param2name(nn.Linear(4, 4, bias=False).weight) == "unknown" From 2d90ea101ddee07f1ebd8ddd1ebe20afa5e692be Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 16:28:41 +0800 Subject: [PATCH 2/2] Make the debug name map tests assert what they claim Both cases could pass against an implementation that does nothing. test_clear_and_re_extract passed the model as a temporary. The maps hold weak references, so it was collectable as soon as extract returned, and the maps could be empty before the clear ran. Neutering debug_clear_module_and_param_names shows the difference: before len(module_names) == 0, len(param_names) == 0 -> passes after assert 9 == 0 -> fails The model is now held in a local, both maps are asserted non-empty before the clear, and the model is kept referenced past the emptiness assertions. test_a_recycled_id_is_not_a_stale_hit built a new parameter and asserted it resolved to "unknown", but never established that it had reused the collected parameter's id. Without that the assertion holds whether or not the stale entry was removed, and id reuse is not something a test can arrange. Renamed to test_collected_parameter_entry_is_removed and asserts the removal directly. Also updates the file's copyright line to the DeepSpeed one. Signed-off-by: alanhuangyoo --- tests/unit/utils/test_debug_name_maps.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/unit/utils/test_debug_name_maps.py b/tests/unit/utils/test_debug_name_maps.py index cf5ba1af6f23..76b716cfa850 100644 --- a/tests/unit/utils/test_debug_name_maps.py +++ b/tests/unit/utils/test_debug_name_maps.py @@ -1,4 +1,4 @@ -# Copyright (c) Microsoft Corporation. +# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team @@ -67,11 +67,20 @@ def test_replaced_submodule_is_released(): def test_clear_and_re_extract(): - debug_extract_module_and_param_names(_Model()) + # The model has to stay alive across the clear. The maps hold weak references, so + # a temporary would be collected when extract returns and the emptiness assertions + # below would pass whether or not the clear did anything. + model = _Model() + debug_extract_module_and_param_names(model) + + assert len(module_names) > 0 + assert len(param_names) > 0 + debug_clear_module_and_param_names() assert len(module_names) == 0 assert len(param_names) == 0 + assert model.head.weight is not None # keeps `model` referenced past the assertions other = _Model(num_blocks=1) debug_extract_module_and_param_names(other) @@ -79,12 +88,21 @@ def test_clear_and_re_extract(): assert debug_param2name(other.head.weight) == "head.weight" -def test_a_recycled_id_is_not_a_stale_hit(): +def test_collected_parameter_entry_is_removed(): + """A collected parameter must leave no entry behind. + + Asserting that a freshly built parameter resolves to "unknown" does not show this: + it only holds if that parameter reused the collected one's id, which is not + something a test can arrange. Assert the removal directly instead. + """ debug_clear_module_and_param_names() doomed = nn.Linear(4, 4, bias=False) param_names[doomed.weight] = "ghost" + assert len(param_names) == 1 + del doomed gc.collect() + assert len(param_names) == 0 assert debug_param2name(nn.Linear(4, 4, bias=False).weight) == "unknown"