Add opt-in AutoEP Python GC policy (#8451)
## Summary
Add an experimental, opt-in AutoEP policy for disabling automatic Python
cyclic garbage collection during training:
```json
{
"expert_parallel": {
"enabled": true,
"python_gc_policy": "disable_during_training"
}
}
```
The default is `"default"` and does not change Python GC settings.
Opting in affects the entire process, not just one engine. Python
reference counting remains active.
## Motivation
Fixed-routing Qwen3-30B-A3B EP16 profiling found repeated single-rank
generation-2 GC pauses of 234–327 ms. Other ranks then waited for the
paused rank at the next DeepEP dispatch, exposing the host-side pause as
collective wait time.
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
- When the first managed engine has finished initialization, collect
once and disable automatic cyclic GC if it was previously enabled.
- Reference-count the process-wide policy across multiple engines,
preserving a pre-disabled GC state.
- Restore the original GC state when the last managed engine is
explicitly destroyed.
- Restore previously enabled GC in forked workers and use generation
tokens so inherited engines cannot release child-process leases.
- Release the GC lease even if another engine teardown operation raises.
- Expose `engine.collect_python_gc()` for explicit collection at
application-selected safe boundaries.
## AutoEP-only performance
The isolated comparison used one 16-GPU H100 allocation, a 48-layer
model, fixed routing, dual warmup, and the order `default → managed →
managed → default`. Each arm measured 20 steps, giving 40 measured steps
per policy.
| Policy | Full measured-window mean per step | Median of arm medians |
Median of arm p95s | 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 |
The full measured-window mean decreased by 162.14 ms per step, or 23.5%,
primarily through fewer slow steps rather than a comparable reduction in
median step time. Both paired measured-window contrasts favored the
managed policy: 166.37 ms and 157.92 ms per step.
- Peak allocated/reserved GPU 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, within the experiment's 0.02 tolerance.
These are short-run, workload-specific observations from the tested
baseline, not a guaranteed speedup across models or subsequent
implementation changes. This comparison does not establish long-run
training equivalence.
## Testing Done
- Targeted config parsing and validation tests.
- GC manager lifecycle tests covering multiple engines, pre-disabled GC,
explicit collection, fork restoration, stale pre-fork leases, and
reentrant finalizers.
- Changed files pass repository pre-commit hooks.
- The AutoEP-only distributed ABBA experiment checked routing, loss
tolerance, GC state, GPU peak memory, and tail latency.
## Usage and limitations
Disabling automatic cyclic GC defers collection; it does not eliminate
the work or fix live references that retain tensors. Cyclic objects can
accumulate during long runs, including cycles that retain GPU tensors.
Applications should monitor host and device memory and call
`engine.collect_python_gc()` at coordinated, application-selected safe
boundaries, such as after checkpointing. The cost of these explicit
collections must still be included when assessing overall training
throughput.
Call `engine.destroy()` when an engine is no longer needed. Restoration
does not rely on garbage collection of the engine itself. This PR does
not add automatic periodic collection, and the short experiment above
does not establish long-run memory stability.
This policy is independent of the DeepEP local-preparation cleanup in
#8423.
---------
Signed-off-by: yh0903 <helloyu0903@gmail.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>