[BUG]Unfreed parameters and unnecessary peak memory in autoEP

[pengdurice]

Describe the bug

Two related problems in the AutoEP module-replacement path under ZeRO-3, both visible as memory that
should not be resident:

1. Init-time transient scaling as 1/autoep_size. _configure_expert_parallel builds each MoE
   layer's local expert tensors as plain nn.Parameters and defers partitioning until every layer
   has been replaced, so the peak carries 2 * P_expert / autoep_size bytes per rank. This term does
   not shrink with world size, so adding GPUs does not help; it OOMs at low autoep_size on large
   MoE models even when steady-state training would fit comfortably.
2. The replaced source modules are never freed. After setattr unlinks an MoE block from the
   module tree its parameters remain strongly referenced by two module-level dicts in
   deepspeed/utils/debug.py, so 2 * P_expert / N of dead weights stays resident for the entire
   run. This one is not AutoEP-specific.

Notation: P_expert is the total parameter count of all routed experts across all layers; P is
total model parameters; N is world size; EP is autoep_size; the leading 2 is bytes per
parameter in bf16.

Phases. Three, and problem 1 lives in the second:

phase	when	EP known?	allocated per rank
A. zero.Init	model construction, before engine.__init__	no	source model sharded over world size -> 2P/N
B. AutoEP replacement	engine.py:331, inside engine.__init__	yes	new GroupedExperts tensors, plain nn.Parameter, unpartitioned -> 2*P_expert/EP
C. re-partition	engine.py:2437 -> parameter_offload.py:258	yes	phase-B tensors sharded over the expert-replica group -> P_expert/N

Phase A is correct and expected: zero.Init has no knowledge of expert parallelism, so it shards
everything over the global data-parallel group.

Root cause of problem 1 is ordering:

• engine.py:331 _configure_expert_parallel(model) replaces every MoE block. Each replacement
  eagerly allocates a full-size GroupedExperts tensor for that rank's local experts
  (auto_ep_layer.py:460-469), and these are plain nn.Parameters, not ZeRO params.
• engine.py:2437 _resolve_zero3_param_placement() -> _convert_to_zero_parameters() is where
  they finally get partitioned. With zero3_init_flag: true the source model already holds ZeRO
  params, so the branch taken is
  zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) at
  parameter_offload.py:258, not the Init(...) fallback at line 264.

The observed peak is therefore:

peak = 2P/N                        source model -- SHOULD be releasable as each layer is
                                    replaced, but is pinned; see problem 2
     + 2 * P_expert / EP           phase-B local experts, not yet partitioned  <-- problem 1
     + 2 * P_expert / n_sparse_layers   one-layer gather during repack (EP-independent)

To Reproduce

1. Take an MoE model;
2. Config:

   {
     "bf16": {"enabled": true, "bf16_master_weights_and_grads": true, "bf16_optimizer_states": true},
     "zero_optimization": {"stage": 3},
     "expert_parallel": {"enabled": true, "autoep_size": 2, "preset_model": "mixtral"},
     "optimizer": {"type": "AdamW", "params": {"lr": 1e-6}},
     "train_micro_batch_size_per_gpu": 1
   }

3. Build the model under deepspeed.zero.Init() (or zero3_init_flag: true via Accelerate) and call
   deepspeed.initialize().

Raising autoep_size is the only workaround, which forces an EP degree chosen by init-time memory
rather than by what is good for throughput.

Expected behavior

Problem 1: the transient should not scale with 1/autoep_size. Partitioning each GroupedExperts
as its layer is replaced -- or constructing it inside a zero.Init context bound to the
expert-replica group -- would leave each rank a partitioned copy rather than its whole EP shard:

peak ~= 2P/N  +  2 * P_expert / N  +  O(2 * P_expert / n_sparse_layers)

Problem 2: once a source MoE block has been replaced its parameters are dead and should be
released, both during replacement (reducing the first term above) and for the remainder of the run.

Why the source model is never freed

after - retained = before in every row above: the source model is still fully resident when
replacement finishes. If its expert weights had been released as each layer was replaced, after
would have been 2*P_dense/N + 2*P_expert/EP -- for EP=2 that is 0.91 + 92.76 = 93.7 GiB, not the
105.30 measured.

The retainer is not AutoEP. deepspeed/utils/debug.py declares two module-level dicts:

# for debug purposes map module and param objects to their fully qualified names
module_names = {}
param_names = {}

def debug_extract_module_and_param_names(model):
    global module_names
    global param_names
    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 because module_names / param_names are module-level
globals they live as long as deepspeed.utils.debug is imported -- i.e. the whole process. The call
order seals it:

engine.py:320   debug_extract_module_and_param_names(model)   <- snapshots the PRE-replacement model
engine.py:331   self._configure_expert_parallel(model)        <- setattr unlinks the old modules

setattr(parent, child_name, replacement) removes the old MoE block from the module tree, but every
one of its parameters is still a live key in param_names, so the refcount never reaches zero. The
only reset is debug_clear_module_and_param_names() at engine.py:946, inside destroy() -- at
teardown, not during init or training.

[alanhuangyoo]

Opened #8356 for problem 2. Your diagnosis was exact — the two module-level dicts in deepspeed/utils/debug.py hold strong keys and are only reset in destroy(), so nothing replaced after engine.py:402 can be released.

Measured on a 64 MiB model with all 16 blocks replaced:

master     replaced 64.0 MiB, still resident 64.0 MiB   (16/16 tensors)
with #8356 replaced 64.0 MiB, still resident  0.0 MiB   ( 0/16 tensors)

One thing worth recording for anyone who tries the obvious fix: weakref.WeakKeyDictionary does not work here. Its keys are weakref.ref objects and ref.__eq__ forwards to the referents, so comparing two live parameters runs Tensor.__eq__ and every lookup raises Boolean value of Tensor with more than one value is ambiguous. Keying on id() with a weakref.finalize keeps the identity semantics the dicts already had — Tensor.__hash__ is id-based, so lookups behaved that way before too.

I left problem 1 alone. The ordering change in _configure_expert_parallel is yours and is a much larger piece of work than this one.