Add opt-in AutoEP Python GC policy - #8451
Draft
yh0903 wants to merge 1 commit into
Draft
Conversation
Provide an opt-in engine-managed policy that collects once after initialization, disables automatic cyclic GC during training, restores it on explicit destroy, and keeps forked data-loader workers safe. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>
4 tasks
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Sep 10, 2026
## Summary - Route AutoEP's EP>1 DeepEP path directly from router/utilization accounting into `_deepep_route`, before collective-backend token preparation. - Avoid DeepEP-unused stable argsort, score/expert gathers, `[T*K,H]` routed-input expansion, score preparation, split-count all-to-all, and D2H split materialization. - Share the output finalization tail across DeepEP, standard communication, and EP1 so output shape, shared experts, router logits, and cache clearing retain their existing contracts. This is an opt-in backend cleanup: the default standard communication path and EP1 path are unchanged. Compatibility: - DeepEP continues to bypass fused weighted restore because its combine already restores and reduces token rows. - DeepEP now bypasses split-plan construction entirely, including async split planning. - The existing DeepEP incompatibility with `autoep_non_moe` compile is unchanged. - Folded tensor parallelism and standard communication retain their existing paths. ## Testing Done - [x] Local code review completed - [x] Unit tests added/updated - [x] Integration tests pass - [x] Manual testing performed Correctness: - Repository pre-commit hooks pass for all three changed files. - Rebased H100 targeted suite passed after upstream AutoEP score-correction bias changes. - CPU/mock contracts verify DeepEP does not call argsort, standard score application, or split-plan construction, while standard communication and EP1 still use their existing preparation. - H100 cleanup OFF/ON parity covers output, loss, exact routes, input gradient, routing-score and router-parameter gradients, all expert gradients, optimizer deltas, activation checkpointing on/off, skewed routing, and an empty expert. Performance, fixed-routing Qwen3-30B-A3B, EP16, TP1, BF16, activation checkpointing on, no profiler: - L8 smoke: `115.01 ms -> 109.05 ms`, a `5.18%` improvement. - L48 fresh ABBA + BAAB blocks: `682.86 ms -> 665.42 ms`, a `2.55%` improvement. - Four paired L48 deltas were all positive: `42.50`, `25.32`, `9.56`, and `18.56 ms`. - Paired mean delta was `23.98 ms`, with 95% CI `[1.82, 46.14] ms`. - Maximum loss difference was `0.00682`; p95 did not regress. - Peak allocated/reserved changed by approximately `+2.2 MiB / +22 MiB`, within the no-regression gate. Profiling: - Sparse CUDA-event observer overhead was `-1.69%`, within run-to-run noise. - The early-route path emits no standard split-plan preparation. - One steady-state lean Nsight capture was collected only for explanation; traced wall time is not used as the performance headline. Matched DeepEP V2 context (same-allocation paired blocks; not a merge gate): - Megatron uses NVIDIA/Megatron-LM#5153 head `eb688c4a...` with the same DeepEP `01dc3aaa...`, ElasticBuffer V2, NCCL 2.30.4, 12 SMs, 16 QPs, fixed routing, data, and checkpointing semantics. - Two fresh allocations each discarded one full AutoEP arm and one full Megatron arm before ABBA/BAAB measurement. - All four paired L48 deltas favored AutoEP: `38.62`, `67.21`, `80.91`, and `62.18 ms`. - Paired mean delta was `62.23 ms`, with 95% CI `[34.20, 90.26] ms`; pooled medians were `649.76 ms` for AutoEP and `714.45 ms` for Megatron, an AutoEP speedup of `9.06%`. - Measured-window loss differed by at most `0.00593`. - AutoEP used about `0.63 GiB` more allocated and `2.77 GiB` more reserved memory. - The median result does not represent average elapsed time: across all 80 recorded measured steps, AutoEP averaged `766.67 ms` versus `716.22 ms` for Megatron, so AutoEP was about `7.0%` slower by arithmetic mean. - AutoEP had `32/80` steps over `800 ms` versus `0/80` for Megatron; only `3/80` crossed one second. The repeated slow steps dominate full-run throughput and remain unexplained. Follow-up fixed-routing instrumentation localizes those repeated slow steps: - Two additional AutoEP arms again had fast medians (`633.17` and `638.05 ms`) but slower means (`754.56` and `747.22 ms`), with `16/40` measured steps over `800 ms`. - Slow versus fast median inflation was concentrated in forward (`+281 ms` and `+248 ms` in the two arms); backward and optimizer medians were effectively unchanged. - A detailed follow-up arm attributed essentially all of the slow-forward increase to the DeepEP dispatch call: `+382.07 ms` dispatch versus `+1.12 ms` expert compute, `-3.43 ms` combine, and `+3.94 ms` router. - Cross-rank inspection shows that dispatch is the synchronization surface, not yet the root cause: at every hotspot, `14-15` ranks wait about `233-357 ms`, while one late-arriving rank spends only about `0.55-0.74 ms` in dispatch. The late rank's preceding MoE/combine work is normally only `1.5-3.0 ms`, placing most of the originating delay in the uninstrumented non-MoE forward region between MoE layers; one observed case accumulated the delay in the router call itself. - The synchronized measured-window wallclock was `16893.88 ms`, versus `16894.84 ms` from summing the 20 recorded critical-step times. The step measurements therefore account for the full measured training window; the mean/median reversal is not an omitted gap between steps. - A second detailed run split the inter-MoE region into decoder, attention, and normalization calls. The delayed call site moved between input RMSNorm, self-attention, router, and otherwise uninstrumented Python gaps on different ranks and layers. In each case one rank paused for roughly `260-435 ms`, after which the remaining ranks waited in the next dispatch. This pattern is inconsistent with a specific DeepEP or transformer kernel regression. - A causal run with identical fixed routing and instrumentation but Python cyclic GC disabled after initialization removed the tail completely: mean `933.08 -> 557.31 ms`, p95 `1740.38 -> 581.60 ms`, and steps over `800 ms` `11/20 -> 0/20`. The measured-window mean was `557.24 ms`, matching the recorded `557.31 ms`. This comparison used separate allocations, so a same-allocation paired confirmation is still required before treating the magnitude as final. The same-allocation AutoEP-only confirmation is now complete. A dual-warm `default -> managed -> managed -> default` block used full measured-window timing: - Default automatic GC: `688.57 ms` mean, `535.13 ms` median, `1072.52 ms` median p95, `11/40` steps over `800 ms`, and `4/40` over one second. - `python_gc_policy="disable_during_training"`: `526.42 ms` mean, `522.89 ms` median, `551.12 ms` median p95, and no steps over `800 ms`. - Paired measured-window savings were `166.37` and `157.92 ms`; the paired mean was `162.14 ms`, with 95% CI `[108.45, 215.84] ms`. - Peak allocated/reserved memory was identical, observed routing was identical, and maximum paired loss differences were `0.00266` and `0.01439`. The opt-in engine-managed policy is isolated in Draft PR deepspeedai#8451. It is independent of this PR's DeepEP local-preparation cleanup. Rewriting or overlapping DeepEP dispatch would optimize the waiting point rather than the source of the tail. Natural-routing same-allocation context: - All four paired L48 deltas also favored AutoEP: `67.03`, `55.87`, `80.45`, and `101.26 ms`. - Paired mean delta was `76.15 ms`, with 95% CI `[45.10, 107.21] ms`; pooled medians were `686.49 ms` for AutoEP and `760.23 ms` for Megatron, an AutoEP speedup of `9.70%`. - Measured-window loss differed by at most `0.00970`. - Across all 80 measured steps, AutoEP averaged `793.65 ms` versus `801.08 ms` for Megatron, only a `0.93%` average-time advantage despite the larger median signal. - AutoEP had `28/80` steps over `800 ms` versus `14/80` for Megatron, and `9/80` over one second versus `2/80`; only one AutoEP step exceeded twice its arm median. - A separate tokens-per-expert audit showed that the natural-routing workload is not fully matched. Per-layer sorted expert-load total variation had a `9.84%` median, but indexed expert variation had a `64.72%` median and rank-receive variation had a `46.46%` median. The two frameworks therefore see similar load-shape distributions assigned to different expert/rank identities. - The natural-routing speedup is reported as end-to-end context, not as a pure framework execution gap. The fixed-routing paired comparison remains the controlled cross-framework result. 🤖 Generated with [GitHub Copilot CLI](https://docs.github.com/copilot/github-copilot-cli) --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add an experimental, opt-in AutoEP policy for disabling automatic Python cyclic garbage collection during training:
{ "expert_parallel": { "enabled": true, "python_gc_policy": "disable_during_training" } }The default is
"default"and leaves process GC behavior unchanged.Motivation
Fixed-routing Qwen3-30B-A3B EP16 profiling found repeated single-rank generation-2 GC pauses of
234-327 ms. The other ranks then waited for the paused rank in the next DeepEP dispatch, making communication look slow even though dispatch was only the synchronization point.Across 20 measured steps, 27 generation-2 collections accumulated
7.60 s; all long collections occurred during forward. Generation-0/1 collections had maximum durations of only1.04/2.31 ms.Implementation
engine.collect_python_gc()for explicit collection at application-selected safe boundaries.AutoEP-only performance
One 16xH100 allocation, fixed routing, dual warmup,
default -> managed -> managed -> default, 20 measured L48 steps per arm:166.37 ms,157.92 ms162.14 ms[108.45, 215.84] ms45,106,782,720 / 50,899,976,192bytes.0.00266and0.01439.Also validated in a precision-aligned AutoEP-vs-Megatron E2E benchmark
This policy was re-tested outside the isolated AutoEP-only ABBA above, in a
separate cross-framework paired benchmark (AutoEP vs Megatron, same model/
routing, AutoEP running #8423's DeepEP cleanup plus fp32 gradient
accumulation and a fp32 gradient all-reduce to match Megatron's
grad_reduce_in_fp32=True). That combination is unrelated to this PR's ownscope, but it depends on this policy to produce a usable result:
25% of steps over 800 ms, the same generation-2 GC pause this PR targets.
Median step time favored AutoEP by ~19-20%, but the tail pulled mean step
time to ~18.8% slower than Megatron -- an artifact of this exact pause
pattern, not a real regression.
tail dropped to 1/40 steps over 800 ms (Megatron: 0/40), and median/mean
step time agreed in direction: AutoEP faster by 21.9% (median) / 10.3%
(mean) over 2 paired blocks, 95% CI
[121.0, 184.6] ms, not crossing zero.This is external confirmation that the fix generalizes past the isolated
AutoEP-only benchmark above: without it, a mean-step-time comparison against
another framework is dominated by this pause, not by real per-step cost.
Testing Done
This is independent of the DeepEP local-preparation cleanup in #8423.