Skip to content

Add opt-in AutoEP Python GC policy - #8451

Draft
yh0903 wants to merge 1 commit into
deepspeedai:masterfrom
yh0903:yh0903/autoep-python-gc-policy
Draft

Add opt-in AutoEP Python GC policy#8451
yh0903 wants to merge 1 commit into
deepspeedai:masterfrom
yh0903:yh0903/autoep-python-gc-policy

Conversation

@yh0903

@yh0903 yh0903 commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

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 only 1.04/2.31 ms.

Implementation

  • Collect once after the concrete DeepSpeed engine has initialized, then disable automatic cyclic GC.
  • Reference-count the process-wide policy across multiple engines.
  • Restore the original GC state when the last managed engine is explicitly destroyed.
  • Restore GC in forked data-loader workers and use generation tokens so inherited engines cannot release child-process leases.
  • Guarantee lease release even if another engine teardown operation raises.
  • Expose 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:

Policy Full measured-window mean Median Median p95 Steps >800 ms Steps >1 s
Default 688.57 ms 535.13 ms 1072.52 ms 11/40 4/40
Managed GC 526.42 ms 522.89 ms 551.12 ms 0/40 0/40
  • Paired measured-window savings: 166.37 ms, 157.92 ms
  • Paired mean: 162.14 ms
  • 95% CI: [108.45, 215.84] ms
  • Both paired deltas favored the managed policy.
  • Peak allocated/reserved memory was unchanged: 45,106,782,720 / 50,899,976,192 bytes.
  • Observed routes were identical across all arms.
  • Maximum paired measured-window loss differences were 0.00266 and 0.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 own
scope, but it depends on this policy to produce a usable result:

  • Without this policy (precision alignment alone): both AutoEP arms had
    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.
  • With this policy enabled, on the same precision-aligned setup: AutoEP
    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

  • Targeted config parsing and validation tests.
  • GC manager lifecycle tests for multiple engines, pre-disabled GC, explicit collection, fork restoration, stale pre-fork leases, and reentrant finalizers.
  • Changed files pass repository pre-commit hooks.
  • AutoEP-only distributed ABBA benchmark passed routing, loss, GC-state, memory, and tail gates.

This is independent of the DeepEP local-preparation cleanup in #8423.

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>
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant