DeepSpeed

e9680c7d - Stop the debug name maps from pinning the model they snapshot (#8356)

Stop the debug name maps from pinning the model they snapshot (#8356) 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: ```python 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. --------- Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>