Stop the debug name maps from pinning the model they snapshot

[alanhuangyoo]

Problem 2 of #8353 — the diagnosis there is @pengdurice's, this is the fix for that half. It does not touch problem 1 (the init-time transient in the AutoEP replacement path), which is a separate, larger change.

What retains the model

deepspeed/utils/debug.py keeps two module-level dicts:

module_names = {}
param_names = {}

def debug_extract_module_and_param_names(model):
    module_names = {module: name for name, module in model.named_modules()}
    param_names = {param: name for name, param in model.named_parameters()}

Dict keys are strong references, and these are module-level globals, so they live as long as deepspeed.utils.debug is imported. The call order seals it:

engine.py:402   debug_extract_module_and_param_names(model)   <- snapshots the model
engine.py:965   debug_clear_module_and_param_names()          <- only in destroy()

Anything that swaps a submodule out between those two points cannot release it. setattr(parent, name, replacement) unlinks the old module from the tree, but every one of its parameters is still a live key, so the refcount never reaches zero. Expert-parallel replacement is where #8353 hit it; kernel injection replaces modules the same way.

Replacing all 16 blocks of a 64 MiB model:

master     replaced 64.0 MiB, still resident 64.0 MiB   (16/16 tensors)
this PR    replaced 64.0 MiB, still resident  0.0 MiB   ( 0/16 tensors)

Why not WeakKeyDictionary

It is the obvious fix and it does not work. WeakKeyDictionary stores weakref.ref objects as keys, and weakref.ref.__eq__ forwards to the referents when both are alive. Comparing two live parameters runs Tensor.__eq__, which returns a tensor:

>>> d = weakref.WeakKeyDictionary({p: n for n, p in model.named_parameters()})
>>> d[some_param]
RuntimeError: Boolean value of Tensor with more than one value is ambiguous

So the entries go in fine and every lookup raises.

Keying on id() and dropping the entry from a weakref.finalize keeps exactly the identity semantics the dicts already had — the previous code compared parameters by Tensor.__hash__, which is id-based. deepspeed/utils/pin_memory.py already tracks its allocations this way, in the same package.

The public surface is unchanged: debug_module2name / debug_param2name still do in then [], and still return "unknown" for anything absent.

Test

tests/unit/utils/test_debug_name_maps.py — lookups and the "unknown" fallback, release of a replaced submodule, clear/re-extract, and that a recycled id() is not a stale hit.

On master:

1 failed, 3 passed
tests/unit/utils/test_debug_name_maps.py:65: AssertionError    (test_replaced_submodule_is_released)

With this PR:

4 passed

The other three pass either way — they are there so the behaviour this preserves stays preserved.

tests/unit/utils/   21 passed, 1 skipped, 1 failed
yapf --diff / flake8   clean

The one failure is test_pin_memory_tracker.py::test_checkpoint_emits_info (assert 2 == 1). It fails identically on master with this file reverted, so it is not from this change.

[tohtana]

This is a significant fix, thank you @alanhuangyoo!

The fix overall looks good to me. As test_debug_name_maps.py shows the old copyright, can you update it to DeepSpeed one? (# Copyright (c) DeepSpeed Team.).

Also, can you consider improving these in the tests?

• In test_clear_and_re_extract, _Model() is passed as a temporary object. Since the name maps now hold only weak references, the model can be collected when debug_extract_module_and_param_names(_Model()) returns, removing its entries before the explicit clear call. The empty-map assertions can therefore pass even if debug_clear_module_and_param_names() does nothing. Keep the model in a local variable, assert that both maps are non-empty before clearing, then assert that both are empty while the model is still alive.
• test_a_recycled_id_is_not_a_stale_hit deletes an old parameter and checks that a new parameter resolves to "unknown", but never checks whether their IDs are equal. If the new parameter has a different ID, the lookup returns "unknown" even if the old "ghost" entry was never removed. To claim coverage of ID reuse, the test must confirm that the new parameter actually received the old parameter's ID before checking the lookup. A more deterministic alternative is to assert that the old parameter is collected and its map entry is removed, and rename the test to test_collected_parameter_entry_is_removed.

[alanhuangyoo]

Both points were right, and both tests were vacuous rather than merely weak. Fixed in 2d90ea1, along with the copyright line.

test_clear_and_re_extract. Your reading is exactly what happens. I checked it by neutering the clear and running the original body:

debug_clear_module_and_param_names = lambda: None      # implementation does nothing

original test:   len(module_names) == 0, len(param_names) == 0   -> passes
strengthened:    assert 9 == 0                                    -> fails

So the test asserted nothing about the function it was named for. 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 so it cannot be collected out from under them.

test_a_recycled_id_is_not_a_stale_hit. Agreed, and I took your alternative rather than trying to force id reuse — a test cannot arrange that, so any version of it would be relying on luck to have coverage at all. Renamed to test_collected_parameter_entry_is_removed, asserting the entry count directly:

param_names[doomed.weight] = "ghost"
assert len(param_names) == 1
del doomed; gc.collect()
assert len(param_names) == 0

The "unknown" lookup is kept as a second assertion, but the removal is now what carries the test.

4 passing, yapf and flake8 clean. Thanks — I would rather have this pointed out than ship two tests that pass on a broken implementation.

[alanhuangyoo]

Note on the red check here, since a red X reads as "this PR broke something" and that is not what happened.

The modal-torch-latest / DeepSpeedAI CI job ended with

RuntimeError: run pytest failed with exit code 137

137 is SIGKILL — the sandbox killed pytest mid-suite. There is not a single FAILED or ERROR line anywhere in the log; the run was killed mid-suite, well away from deepspeed/utils/debug.py. collect tests and DCO pass on the same commit.

The same thing hit #8384 on the same day, and three other branches of mine that ran within the same hour (#8362, #8433, #8435) went green, so it is intermittent rather than a property of this tree. I cannot re-run it — that needs write access to the repo. Any maintainer re-running the failed job should be enough; happy to push an empty commit instead if that is easier.

[alanhuangyoo]

@tohtana — re-review request. Both of your points were right, both are fixed in 2d90ea1, and the PR has been sitting on the stale changes-requested since.

test_clear_and_re_extract. Your reading was exactly right: _Model() as a temporary meant the maps could empty themselves before the explicit clear, so the assertions passed whether or not debug_clear_module_and_param_names did anything. The model is now a local, both maps are asserted non-empty before the clear, and the model is kept referenced past the emptiness assertions.

The recycled-id test. Also right — it only held if the new parameter reused the collected one's id, which a test cannot arrange. Renamed to test_collected_parameter_entry_is_removed and it now asserts the removal directly, as you suggested.

Copyright header updated to # Copyright (c) DeepSpeed Team. as well.

The red CI is the 90-minute job timeout rather than a failure; I have merged current master in for #8404's raised limit and it is re-running.

[tohtana]

Thank you for the update, @alanhuangyoo! Looks good to me.