feat(minimax-h3): enable 2x RTX 4090/5090 with DLO (#5764)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import (
 from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
     DistributedLayerwiseOffloadHook,
+    PinnedResidentLayerGroup,
 )
+from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
 from vllm_omni.platforms import current_omni_platform
 
@@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch):
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
+    monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
 
 
 class TinyBlock(nn.Module):
@@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook:
         assert shard1.numel() == 2
         assert torch.equal(shard1, torch.tensor([102.0, 0.0]))
 
+    def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+        next_block.weight.mmap_weight_transform_pending = True
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+
+        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
+
+    def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+
+        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
+
+
+class TestPinnedResidentLayerGroup:
+    def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
+        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
+        expected = []
+        for index, block in enumerate(blocks):
+            block.weight.data.fill_(index + 1)
+            block.bias.data.fill_(index + 3)
+            expected.append((block.weight.detach().clone(), block.bias.detach().clone()))
+
+        group = PinnedResidentLayerGroup(
+            blocks,
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+        )
+
+        assert all(block.weight.numel() == 0 for block in blocks)
+
+        group.load()
+        for block, (weight, bias) in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+            assert torch.equal(block.bias, bias)
+
+        # Device-side changes must not overwrite the persistent host master.
+        blocks[0].weight.data.zero_()
+        group.offload()
+        assert all(block.weight.numel() == 0 for block in blocks)
+
+        group.load()
+        for block, (weight, bias) in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+            assert torch.equal(block.bias, bias)
+
+        group.offload()
+
+    def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
+        module = nn.Linear(2, 2, bias=False)
+
+        def transform(tensor):
+            return tensor.flip(0)
+
+        module.weight.mmap_weight_transform = transform
+
+        backend = object.__new__(DistributedLayerwiseOffloadBackend)
+        backend._remember_mmap_param_attrs(module)
+        module.to_empty(device="meta")
+
+        replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
+        backend._attach_mmap_param_attrs("weight", replacement, module.weight)
+
+        assert replacement.mmap_weight_transform is transform
+        assert replacement.mmap_weight_transform_pending is True
+
+    def test_shared_allgather_output_is_narrowed_to_current_block(
+        self, dist_group, patched_offload_runtime, monkeypatch
+    ):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+        # Simulate unified buffers sized for a larger heterogeneous block.
+        hook.gpu_buffers = [
+            {torch.float32: torch.empty(12)},
+            {torch.float32: torch.empty(12)},
+        ]
+        hook.gpu_shard_buffers = [
+            {torch.float32: torch.empty(6)},
+            {torch.float32: torch.empty(6)},
+        ]
+        output_sizes = []
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            output_sizes.append(output.numel())
+            output.zero_()
+            output[: local_shard.numel()].copy_(local_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+
+        hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert output_sizes == [4]
+
+
+class TestPinnedModuleStager:
+    def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime):
+        module = nn.Linear(2, 2)
+        expected_weight = module.weight.detach().clone()
+        expected_bias = module.bias.detach().clone()
+        stager = PinnedModuleStager(
+            module,
+            torch.device("cpu"),
+            pin_memory=False,
+            copy_stream=DummyStream(),
+        )
+
+        stager.load()
+        module.weight.data.zero_()
+        module.bias.data.zero_()
+        stager.offload()
+
+        assert torch.equal(module.weight, expected_weight)
+        assert torch.equal(module.bias, expected_bias)
+        stager.load()
+        assert torch.equal(module.weight, expected_weight)
+        assert torch.equal(module.bias, expected_bias)
+        stager.offload()
+
 
 class _DummyBlock(nn.Module):
     def __init__(self):
@@ -592,6 +753,7 @@ class TestOffloadPlan:
         plan = OffloadPlan(
             block_attrs={"transformer": ("gen_layers",)},
             offload_submodules={"context_encoder": "layers"},
+            resident_dit_paths=frozenset({"transformer"}),
         )
 
         class PipelineWithPlan(nn.Module):
@@ -602,12 +764,14 @@ class TestOffloadPlan:
         assert result is plan
         assert result.block_attrs == {"transformer": ("gen_layers",)}
         assert result.offload_submodules == {"context_encoder": "layers"}
+        assert result.resident_dit_paths == frozenset({"transformer"})
 
     def test_offload_plan_defaults_to_empty(self):
         """OffloadPlan with no arguments should have empty dicts."""
         plan = OffloadPlan()
         assert plan.block_attrs == {}
         assert plan.offload_submodules == {}
+        assert plan.resident_dit_paths == frozenset()
 
     def test_offload_plan_is_frozen(self):
         """OffloadPlan should be immutable (frozen=True)."""
@@ -615,6 +779,50 @@ class TestOffloadPlan:
         with pytest.raises(Exception):
             plan.block_attrs = {"x": ("y",)}  # type: ignore
 
+    def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
+        class TokenRefiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.token_refiner = TokenRefiner()
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"token_refiner": "blocks"},
+                resident_dit_paths=frozenset({"transformer"}),
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                dlo_resident_layers=2,
+            ),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert len(backend._resident_blocks) == 2
+        assert len(backend._all_hook_groups) == 1
+        assert len(backend._all_hook_groups[0]) == 2
+        assert backend.enabled
+
+        backend.disable()
+
 
 class TestMmapValidation:
     """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
@@ -821,7 +1029,7 @@ class TestConfigValidation:
             OffloadConfig.from_od_config(FakeODConfig())
 
     def test_hsdp_without_allgather_allowed(self):
-        """HSDP + DLO + no-AllGather should be allowed (full weights per rank)."""
+        """HSDP + DLO + no-AllGather uses standard-loader rank-local weights."""
         from vllm_omni.diffusion.offloader.base import OffloadConfig
 
         class FakePC:
@@ -843,12 +1051,33 @@ class TestConfigValidation:
             enable_layerwise_offload = False
             enable_distributed_layerwise_offload = True
             dlo_use_allgather = False  # no AllGather → should be allowed
+            dlo_resident_layers = 20
             pin_cpu_memory = True
             parallel_config = FakePC()
             model = "/fake/path"
 
         config = OffloadConfig.from_od_config(FakeODConfig())
         assert config.dp_size == 1  # forced to 1 when no AllGather
+        assert config.dlo_resident_layers == 20
+
+    def test_resident_layers_with_allgather_rejected(self):
+        class FakePC:
+            data_parallel_size = 2
+            use_hsdp = False
+            sequence_parallel_size = 1
+
+        class FakeODConfig:
+            enable_cpu_offload = False
+            enable_layerwise_offload = False
+            enable_distributed_layerwise_offload = True
+            dlo_use_allgather = True
+            dlo_resident_layers = 20
+            pin_cpu_memory = True
+            parallel_config = FakePC()
+            model = "/fake/path"
+
+        with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
+            OffloadConfig.from_od_config(FakeODConfig())
 
     def test_num_inference_steps_none_rejected(self):
         """DP multi-concurrency should reject None num_inference_steps."""