From 95781c31e082c0182e0ab1f037ed03b1a740bd13 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:33:32 +0800 Subject: [PATCH] [Bugfix][Diffusion] Skip AllGather DLO weight restore on worker shutdown (#8103) Signed-off-by: specture724 Co-authored-by: Claude Opus 5.5 Co-authored-by: Gao Han --- .../test_distributed_layerwise_backend.py | 29 +++++++++++++++++++ tests/diffusion/test_worker_wrapper_base.py | 4 +-- vllm_omni/diffusion/offloader/base.py | 7 +++++ .../distributed_layerwise_backend.py | 12 ++++++-- .../diffusion/worker/diffusion_worker.py | 4 +-- 5 files changed, 49 insertions(+), 7 deletions(-) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index fed60be1f..73e45124e 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -55,6 +55,7 @@ from vllm_omni.diffusion.offloader.offload_plan import ( ) from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform @@ -2866,6 +2867,34 @@ class TestDistributedComponentSelection: with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"): backend.enable(nn.Module()) + def test_worker_shutdown_skips_allgather_restore(self, patched_offload_runtime, mocker): + # Restoring AllGather blocks rebuilds the full DiT on every rank's host, + # which outlived the executor's shutdown grace period (issue 8081). + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"dit", "text_encoder"}), + dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"}, + ), + torch.device("cpu"), + ) + allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1)) + rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1)) + backend._all_hook_groups = [[allgather_hook], [rank_local_hook]] + backend.enabled = True + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env") + worker = DiffusionWorker.__new__(DiffusionWorker) + worker.model_runner = SimpleNamespace(offload_backend=backend, kv_transfer_manager=None) + + worker.shutdown() + + allgather_hook.restore_next_block_to_cpu.assert_not_called() + rank_local_hook.restore_next_block_to_cpu.assert_called_once_with() + assert not backend.enabled + assert not backend._all_hook_groups + def test_disable_runs_collectives_before_best_effort_local_cleanup( self, patched_offload_runtime, diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index d49f3304b..aeedd25b7 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -290,7 +290,7 @@ class TestWorkerWrapperBaseDelegation: def test_worker_shutdown_releases_model_before_cumem_pools(self, mocker: MockerFixture, native_kv): events: list[str] = [] offload_backend = mocker.Mock() - offload_backend.disable.side_effect = lambda: events.append("offload") + offload_backend.shutdown.side_effect = lambda: events.append("offload") kv_manager = mocker.Mock() kv_manager.close.side_effect = lambda: events.append("kv") worker = DiffusionWorker.__new__(DiffusionWorker) @@ -389,7 +389,7 @@ class TestWorkerWrapperBaseDelegation: mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector") offload_backend = mocker.Mock() if cleanup_fails: - offload_backend.disable.side_effect = RuntimeError("offload cleanup failed") + offload_backend.shutdown.side_effect = RuntimeError("offload cleanup failed") worker = DiffusionWorker.__new__(DiffusionWorker) weights = torch.ones(1) weights_ref = weakref.ref(weights) diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index bb047a7eb..2438cedcf 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -244,5 +244,12 @@ class OffloadBackend(ABC): """ raise NotImplementedError + def shutdown(self) -> None: + """Release offload resources at process exit. + + Backends may skip work that only matters for a later ``enable``. + """ + self.disable() + def is_enabled(self) -> bool: return self.enabled diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e79d1f197..f359a4385 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1889,10 +1889,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): skipped_allgather = bool(allgather_hooks) and not restore_allgather_weights if skipped_allgather: # Startup rollback cannot safely enter a collective that a failed - # peer may never reach. Those blocks cannot be reconstructed, so - # make accidental reuse explicit instead of accepting zero weights. + # peer may never reach, and shutdown has no later user. Those blocks + # cannot be reconstructed, so make accidental reuse explicit + # instead of accepting zero weights. self._poisoned_reason = ( - "Distributed layerwise offload startup skipped AllGather weight restoration; " + "Distributed layerwise offload skipped AllGather weight restoration; " "recreate the backend and reload the pipeline before retrying" ) @@ -1981,6 +1982,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def disable(self) -> None: self._disable(restore_allgather_weights=True) + def shutdown(self) -> None: + # Restoring AllGather blocks rebuilds the full DiT in every rank's host + # memory, which can outlast the executor's shutdown grace period. + self._disable(restore_allgather_weights=False) + @staticmethod def _allocate_shared_buffers( hooks: list[DistributedLayerwiseOffloadHook], diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 3f4a39ac8..f43e66906 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -1048,9 +1048,9 @@ class DiffusionWorker: try: offload_backend = getattr(model_runner, "offload_backend", None) if offload_backend is not None: - offload_backend.disable() + offload_backend.shutdown() except Exception: - logger.exception("Failed to disable diffusion offload backend during shutdown") + logger.exception("Failed to shut down diffusion offload backend during shutdown") try: if mgr is not None: mgr.close()