DeepSpeed
43c53e9b - Add opt-in AutoEP Python GC policy (#8451)

Commit
3 days ago
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>
Author
Parents
Loading