Stop the debug name maps from pinning the model they snapshot

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

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"