From 6fd90412f019ef71ff1031deec92b781eb84be72 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Mon, 7 Sep 2026 02:11:39 -0700 Subject: [PATCH] Add AutoEP Python GC policy 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 --- deepspeed/__init__.py | 2 + deepspeed/module_inject/auto_ep_config.py | 6 + .../module_inject/auto_ep_presets/base.py | 1 + deepspeed/runtime/engine.py | 59 ++++--- deepspeed/runtime/python_gc.py | 75 +++++++++ docs/code-docs/source/autoep.rst | 26 +++ tests/unit/runtime/test_python_gc.py | 159 ++++++++++++++++++ tests/unit/v1/moe/test_autoep_unit.py | 8 + 8 files changed, 316 insertions(+), 20 deletions(-) create mode 100644 deepspeed/runtime/python_gc.py create mode 100644 tests/unit/runtime/test_python_gc.py diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 26bfc7c77df6..31ec83c0c266 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -258,6 +258,8 @@ def initialize( # Restore zero.Init context if necessary zero.partition_parameters.restore_init_context() + engine._configure_python_gc() + return_items = [ engine, engine.optimizer, diff --git a/deepspeed/module_inject/auto_ep_config.py b/deepspeed/module_inject/auto_ep_config.py index 953bb05f6ace..9086e319638e 100644 --- a/deepspeed/module_inject/auto_ep_config.py +++ b/deepspeed/module_inject/auto_ep_config.py @@ -62,6 +62,7 @@ def parse_autoep_config(param_dict: dict) -> AutoEPConfig: config.comm_num_sm = param_dict.get("comm_num_sm", 12) config.comm_qp_margin = param_dict.get("comm_qp_margin", 4) config.comm_max_tokens_per_rank = param_dict.get("comm_max_tokens_per_rank", 0) + config.python_gc_policy = param_dict.get("python_gc_policy", "default") config.num_expert_groups = param_dict.get("num_expert_groups", None) config.num_limited_groups = param_dict.get("num_limited_groups", None) config.score_func = param_dict.get("score_func", "auto") @@ -119,6 +120,11 @@ def validate_autoep_config( if not isinstance(config.validate_folding_routing, bool): raise ValueError("expert_parallel.validate_folding_routing must be a boolean") + valid_python_gc_policies = ("default", "disable_during_training") + if config.python_gc_policy not in valid_python_gc_policies: + raise ValueError(f"python_gc_policy must be one of {valid_python_gc_policies}, " + f"got {config.python_gc_policy!r}") + if not config.enabled: return diff --git a/deepspeed/module_inject/auto_ep_presets/base.py b/deepspeed/module_inject/auto_ep_presets/base.py index 52fda09d2cfc..869c83eda76c 100644 --- a/deepspeed/module_inject/auto_ep_presets/base.py +++ b/deepspeed/module_inject/auto_ep_presets/base.py @@ -115,6 +115,7 @@ class AutoEPConfig: comm_num_sm: int = 12 comm_qp_margin: int = 4 comm_max_tokens_per_rank: int = 0 + python_gc_policy: Literal["default", "disable_during_training"] = "default" num_expert_groups: int | None = None num_limited_groups: int | None = None score_func: Literal["auto", "softmax", "sigmoid"] = "auto" diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..7420abb5e6f8 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -394,6 +394,7 @@ def __init__(self, self.mesh_device = mesh_device self._autoep_folding_spec = None self._autoep_folding_group_handles = None + self._python_gc_generation = None # Flag to indicate that scale() was called before manual backward pass self._manual_backward_expected = False @@ -932,26 +933,44 @@ def __del__(self): logger.debug("DeepSpeedEngine.__del__ cleanup skipped: %s", exc, exc_info=True) def destroy(self): - # DeepEP buffers ask the library not to reclaim them, so they outlive - # the engine unless something releases them here. Only this engine's - # own buffers: another engine in the same process still needs its own. - module = getattr(self, "module", None) - if module is not None: - from deepspeed.module_inject.auto_ep_comm import destroy_exchanges - destroy_exchanges(module) - - self._release_deepcompile_compiled_backward_state() - self._release_deepcompile_dynamo_config() - optimizer = getattr(self, "optimizer", None) - if optimizer is not None and hasattr(optimizer, 'destroy'): - optimizer.destroy() - if self.is_deepcompile_active(): - get_deepcompile_handle().cleanup() - debug_clear_module_and_param_names() - - checkpoint_engine = getattr(self, "checkpoint_engine", None) - if checkpoint_engine is not None and checkpoint_engine.is_decoupled(): - checkpoint_engine.cleanup() + try: + # DeepEP buffers ask the library not to reclaim them, so they outlive + # the engine unless something releases them here. Only this engine's + # own buffers: another engine in the same process still needs its own. + module = getattr(self, "module", None) + if module is not None: + from deepspeed.module_inject.auto_ep_comm import destroy_exchanges + destroy_exchanges(module) + + self._release_deepcompile_compiled_backward_state() + self._release_deepcompile_dynamo_config() + optimizer = getattr(self, "optimizer", None) + if optimizer is not None and hasattr(optimizer, 'destroy'): + optimizer.destroy() + if self.is_deepcompile_active(): + get_deepcompile_handle().cleanup() + debug_clear_module_and_param_names() + + checkpoint_engine = getattr(self, "checkpoint_engine", None) + if checkpoint_engine is not None and checkpoint_engine.is_decoupled(): + checkpoint_engine.cleanup() + finally: + python_gc_generation = getattr(self, "_python_gc_generation", None) + if python_gc_generation is not None: + from deepspeed.runtime.python_gc import python_gc_manager + python_gc_manager.release(python_gc_generation) + self._python_gc_generation = None + + def collect_python_gc(self): + """Run Python cyclic GC at an application-selected safe boundary.""" + from deepspeed.runtime.python_gc import python_gc_manager + return python_gc_manager.collect() + + def _configure_python_gc(self): + autoep_config = self._config.expert_parallel_config + if autoep_config.enabled and autoep_config.python_gc_policy == "disable_during_training": + from deepspeed.runtime.python_gc import python_gc_manager + self._python_gc_generation = python_gc_manager.acquire() def _get_model_parameters(self): if self.autotuning_profile_model_info(): diff --git a/deepspeed/runtime/python_gc.py b/deepspeed/runtime/python_gc.py new file mode 100644 index 000000000000..4aeca412cbec --- /dev/null +++ b/deepspeed/runtime/python_gc.py @@ -0,0 +1,75 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Process-wide Python cyclic garbage collection control.""" + +import gc +import os +from threading import RLock + +from deepspeed.utils import logger + + +class PythonGCManager: + """Coordinate process-wide GC state across multiple DeepSpeed engines.""" + + def __init__(self): + self._lock = RLock() + self._active_engines = 0 + self._restore_enabled = False + self._generation = 0 + if hasattr(os, "register_at_fork"): + os.register_at_fork(before=self._before_fork, + after_in_parent=self._after_fork_parent, + after_in_child=self._after_fork_child) + + def _before_fork(self): + self._lock.acquire() + + def _after_fork_parent(self): + self._lock.release() + + def _after_fork_child(self): + if self._active_engines > 0 and self._restore_enabled and not gc.isenabled(): + gc.enable() + self._active_engines = 0 + self._restore_enabled = False + self._generation += 1 + self._lock = RLock() + + def acquire(self): + with self._lock: + if self._active_engines == 0: + self._restore_enabled = gc.isenabled() + if self._restore_enabled: + collected = gc.collect() + gc.disable() + logger.info("Disabled automatic Python cyclic GC after collecting %d objects", collected) + self._active_engines += 1 + return self._generation + + def release(self, generation): + with self._lock: + if generation != self._generation: + return + if self._active_engines == 0: + return + self._active_engines -= 1 + restore_enabled = self._active_engines == 0 and self._restore_enabled + if self._active_engines == 0: + self._restore_enabled = False + if restore_enabled and not gc.isenabled(): + gc.enable() + logger.info("Restored automatic Python cyclic GC") + + def collect(self): + """Run an explicit collection without changing the automatic-GC policy.""" + return gc.collect() + + @property + def active_engines(self): + with self._lock: + return self._active_engines + + +python_gc_manager = PythonGCManager() diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 38a4f510d260..fa43261f26b4 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -112,6 +112,32 @@ that set nothing keep the existing path unchanged. DeepEP buffer is sized statically and must use the same capacity on every rank. A batch that exceeds it is an error. +**Python cyclic GC policy (experimental):** + +Large Python model graphs can accumulate cyclic objects during training. A +generation-2 collection pauses one rank's Python thread, and the pause can then +be exposed as collective wait time on every expert-parallel rank. AutoEP offers +an opt-in policy that collects once after engine initialization and disables +automatic cyclic collection until the engine is destroyed: + +.. code-block:: json + + { + "expert_parallel": { + "enabled": true, + "autoep_size": 8, + "python_gc_policy": "disable_during_training" + } + } + +The default is ``"default"``, which leaves Python GC unchanged. The policy is +process-wide and reference-counted across DeepSpeed engines. Applications that +create cyclic Python objects during training should call +``engine.collect_python_gc()`` at a safe boundary such as after checkpointing. +Call ``engine.destroy()`` when the engine is no longer needed to restore the +process's original automatic-GC state; restoration does not rely on Python +finalization because disabled cyclic GC cannot reclaim engine reference cycles. + On 16 H100s across two nodes, replaying routing captured from real training, DeepEP reduced payload AllToAll time from roughly 100 ms to 48 ms per step. A full SFT step on Qwen3.5-MoE went from roughly 325 ms to 266 ms, a 1.2x speedup diff --git a/tests/unit/runtime/test_python_gc.py b/tests/unit/runtime/test_python_gc.py new file mode 100644 index 000000000000..eb5627005c1d --- /dev/null +++ b/tests/unit/runtime/test_python_gc.py @@ -0,0 +1,159 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +from deepspeed.runtime.python_gc import PythonGCManager + + +def test_python_gc_manager_restores_enabled_state(monkeypatch): + calls = [] + enabled = True + + def isenabled(): + return enabled + + def collect(): + calls.append("collect") + return 7 + + def disable(): + nonlocal enabled + enabled = False + calls.append("disable") + + def enable(): + nonlocal enabled + enabled = True + calls.append("enable") + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", isenabled) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.collect", collect) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.disable", disable) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.enable", enable) + + manager = PythonGCManager() + generation = manager.acquire() + assert manager.acquire() == generation + assert manager.active_engines == 2 + assert calls == ["collect", "disable"] + + manager.release(generation) + assert manager.active_engines == 1 + assert calls == ["collect", "disable"] + + manager.release(generation) + assert manager.active_engines == 0 + assert calls == ["collect", "disable", "enable"] + + +def test_python_gc_manager_preserves_disabled_state(monkeypatch): + calls = [] + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", lambda: False) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.collect", lambda: calls.append("collect")) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.disable", lambda: calls.append("disable")) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.enable", lambda: calls.append("enable")) + + manager = PythonGCManager() + generation = manager.acquire() + manager.release(generation) + + assert calls == [] + + +def test_python_gc_manager_explicit_collection(monkeypatch): + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.collect", lambda: 11) + assert PythonGCManager().collect() == 11 + + +def test_python_gc_manager_restores_enabled_state_after_fork(monkeypatch): + calls = [] + enabled = False + + def isenabled(): + return enabled + + def enable(): + nonlocal enabled + enabled = True + calls.append("enable") + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", isenabled) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.enable", enable) + + manager = PythonGCManager() + manager._active_engines = 1 + manager._restore_enabled = True + manager._after_fork_child() + + assert calls == ["enable"] + assert manager.active_engines == 0 + + +def test_python_gc_manager_preserves_disabled_state_after_fork(monkeypatch): + calls = [] + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", lambda: False) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.enable", lambda: calls.append("enable")) + + manager = PythonGCManager() + manager._active_engines = 1 + manager._restore_enabled = False + manager._after_fork_child() + + assert calls == [] + assert manager.active_engines == 0 + + +def test_python_gc_manager_ignores_pre_fork_release(monkeypatch): + enabled = True + + def isenabled(): + return enabled + + def disable(): + nonlocal enabled + enabled = False + + def enable(): + nonlocal enabled + enabled = True + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", isenabled) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.collect", lambda: 0) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.disable", disable) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.enable", enable) + + manager = PythonGCManager() + parent_generation = manager.acquire() + manager._after_fork_child() + child_generation = manager.acquire() + + manager.release(parent_generation) + assert manager.active_engines == 1 + assert enabled is False + + manager.release(child_generation) + assert manager.active_engines == 0 + assert enabled is True + + +def test_python_gc_manager_collection_allows_reentrant_release(monkeypatch): + manager = PythonGCManager() + enabled = True + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.isenabled", lambda: enabled) + + def collect(): + manager.release(-1) + return 0 + + def disable(): + nonlocal enabled + enabled = False + + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.collect", collect) + monkeypatch.setattr("deepspeed.runtime.python_gc.gc.disable", disable) + + generation = manager.acquire() + assert manager.active_engines == 1 + manager.release(generation) diff --git a/tests/unit/v1/moe/test_autoep_unit.py b/tests/unit/v1/moe/test_autoep_unit.py index 03dcaa3c4eac..3dfbbb233e6e 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -203,6 +203,7 @@ def test_parse_and_validate_enabled_size_contract(self): assert disabled.enabled is False assert disabled.autoep_size == 1 assert disabled.validate_folding_routing is False + assert disabled.python_gc_policy == "default" assert disabled.load_balance_coeff is None assert disabled._load_balance_coeff_explicit is False @@ -214,12 +215,14 @@ def test_parse_and_validate_enabled_size_contract(self): "score_apply": "pre", "route_scale": 2.0, "validate_folding_routing": True, + "python_gc_policy": "disable_during_training", }) assert config.enabled is True assert config.autoep_size == 4 assert config.preset_model == "mixtral" assert config.validate_folding_routing is True + assert config.python_gc_policy == "disable_during_training" assert config.load_balance_coeff is None assert config._load_balance_coeff_explicit is True assert config.score_apply == "pre" @@ -234,6 +237,11 @@ def test_validate_folding_routing_requires_boolean(self): tp_size=1, sp_size=1) + def test_python_gc_policy_rejects_unknown_value(self): + config = parse_autoep_config({"enabled": True, "python_gc_policy": "aggressive"}) + with pytest.raises(ValueError, match="python_gc_policy must be one of"): + validate_autoep_config(config, world_size=1, pp_size=1, tp_size=1, sp_size=1) + def test_combine_impl_rejects_unknown_value(self): config = parse_autoep_config({"enabled": True, "combine_impl": "triton"}) with pytest.raises(ValueError, match="combine_impl must be one of"):