From 148cb943ceba3cd2cb7cd362644669427049949a Mon Sep 17 00:00:00 2001 From: YAN YIXIN Date: Sun, 4 Oct 2026 15:25:51 +0800 Subject: [PATCH] [Bugfix][CI/Build] Fix layerwise offload shutdown and enable it for the LTX-2 distilled T2V accuracy test (#8453) Signed-off-by: nikiwang-125 Co-authored-by: Claude Opus 5.5 --- .../offloader/test_layerwise_backend.py | 18 ++++++++++++++++++ .../ltx/test_ltx_official_similarity.py | 3 ++- .../diffusion/offloader/layerwise_backend.py | 6 ++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 4c965adaa..6411ae246 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -542,6 +542,24 @@ class TestLayerwiseComponentSelection: for name, tensor in pipeline.state_dict().items(): torch.testing.assert_close(tensor, expected[name]) + def test_shutdown_skips_block_restore(self, patched_offload_runtime, monkeypatch): + # Restoring every block copies the whole DiT back onto the device, which + # outlived the executor's shutdown grace period (issue 7625). + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=3) + backend = _layer_backend() + backend.enable(pipeline) + restores = Mock() + for hook in backend._dit_hooks: + monkeypatch.setattr(hook, "restore_next_block", restores) + + backend.shutdown() + + restores.assert_not_called() + assert not backend.enabled + assert not backend._dit_hooks + assert not backend._hooked_dit_blocks + def test_encoder_disable_failure_keeps_host_masters_for_retry( self, patched_offload_runtime, diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index c1b083495..a9f39b77b 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -103,7 +103,7 @@ class LTXAccuracyCase: spatial_upsampler: LTXArtifact | None = None distilled_lora: LTXArtifact | None = None # Official uses block streaming; Omni uses layerwise offload. Both preserve - # bf16 arithmetic while keeping full-resolution two-stage cases on one H100. + # bf16 arithmetic while keeping full-resolution cases on one H100. enable_layerwise_offload: bool = False @@ -182,6 +182,7 @@ TWO_STAGE_CASES = ( gemma_model_id="Lightricks/LTX-2", gemma_model_revision=LTX2_REVISION, gemma_model_env="VLLM_TEST_LTX2_MODEL", + enable_layerwise_offload=True, ), LTXAccuracyCase( name="ltx23_distilled_i2v", diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index e9f8cf5be..660219092 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -521,3 +521,9 @@ class LayerWiseOffloadBackend(OffloadBackend): def disable(self) -> None: self._disable(restore_weights=True) + + def shutdown(self) -> None: + # Restoring every block copies the whole DiT back onto the device, which + # only matters for a later ``enable`` and can outlast the executor's + # shutdown grace period. + self._disable(restore_weights=False)