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..76b716cfa850 --- /dev/null +++ b/tests/unit/utils/test_debug_name_maps.py @@ -0,0 +1,108 @@ +# Copyright (c) DeepSpeed Team. +# 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(): + # 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) + + assert debug_param2name(other.head.weight) == "head.weight" + + +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"