[Bugfix] Fix layer-wise offload incompatibility with HSDP (#2021)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

tests/diffusion/offloader/test_layerwise_backend.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LayerwiseOffloadHook."""
+
+import gc
+import os
+import socket
+from contextlib import contextmanager
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
+
+import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class DummyEvent:
+    def record(self, _stream) -> None:
+        return None
+
+
+@contextmanager
+def dummy_stream(_stream):
+    yield None
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return int(s.getsockname()[1])
+
+
+def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+
+
+def _cleanup_distributed() -> None:
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
+
+    gc.collect()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+
+@pytest.fixture(scope="module")
+def dist_group():
+    master_port = _find_free_port()
+    _set_dist_env(rank=0, world_size=1, master_port=master_port)
+
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        _cleanup_distributed()
+
+
+@pytest.fixture
+def patched_offload_runtime(mocker):
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+
+
+class TinyBlock(nn.Module):
+    def __init__(self, values: torch.Tensor):
+        super().__init__()
+        mesh = DeviceMesh("cpu", [0])
+        dtensor = DTensor.from_local(values, mesh, [Replicate()])
+        self.weight = nn.Parameter(dtensor)
+
+
+def _make_values(start: float) -> torch.Tensor:
+    return torch.arange(start, start + 4, dtype=torch.float32)
+
+
+class TestLayerwiseOffloadHook:
+    def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
+        current_block = TinyBlock(_make_values(1.0))
+        next_block = TinyBlock(_make_values(10.0))
+
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+
+        hook.initialize_hook(current_block)
+
+        assert isinstance(next_block.weight, DTensor)
+        assert next_block.weight.to_local().is_meta
+        assert next_block.weight.to_local().shape == torch.Size([4])
+        assert hook.dtype_metadata[next_block.weight.dtype][0]["shape"] == torch.Size([4])
+
+        hook.prefetch_layer(non_blocking=False)
+        assert isinstance(next_block.weight, DTensor)
+        assert torch.equal(next_block.weight.to_local(), _make_values(10.0))
+        assert next_block.weight.to_local().shape == torch.Size([4])
+
+        hook.offload_layer()
+        assert isinstance(current_block.weight, DTensor)
+        assert current_block.weight.to_local().is_meta
+        assert current_block.weight.to_local().shape == torch.Size([4])
+        assert not hook.is_materialized