DeepSpeed

b16a3fab - Release AutoEP router-logit caches on checkpoint replay exits (#8581)

Release AutoEP router-logit caches on checkpoint replay exits (#8581) ## Summary Fix AutoEP router-logit cache lifetime during activation checkpointing without changing model return contracts. - Remove the module-owned router-logit cache and its forward hook. Hugging Face model-level recorders already capture the gate output and do not need a second projection. - For MoE blocks that return logits directly, compute them locally in the existing shared output tail. Preserve raw and post-score return values and their gradients. Non-reentrant checkpoint replay can stop before `_finalize_output`. Previously, a tensor left in `_cached_router_logits` could keep the replay autograd graph alive between training steps. Keeping returned logits local removes this module-owned reference rather than adding cleanup paths for early-stop and exceptions. The standard and DeepEP paths already share `_finalize_output`. ## Testing Done - [x] Local code review completed. - [x] Repository pre-commit hooks passed for all four changed files. - [x] Focused CPU router semantics, real Hugging Face Mixtral integration, and DeepEP output-tail cases: **19 passed, 6 CUDA cases skipped, 2 subtests passed** on the local CPU run. - [x] Complete communication backend test module: **37 passed, 4 subtests passed**. - [x] Mutation check rejected the original implementation with exactly four expected regression failures and seven passing cases. - [x] Actual **NVIDIA H100 80GB HBM3** validation of the cache-free revision: **18 passed, 0 failed, 0 skipped**, including all six CUDA combinations of checkpoint off / non-reentrant / reentrant and returning / not returning router logits. Runtime: PyTorch 2.10.0.2+cu130, CUDA 13.0, Transformers 5.12.0. The two-step training regressions compare outputs, loss, input gradients, every named parameter gradient, and optimizer updates against an uncheckpointed reference. Weak references verify that gate tensors are released after each step and after an expert exception, without relying on a private cache attribute. Additional cases check single-capture behavior and raw/softmax/sigmoid returned values and gradients against an independent gate reference. --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>