[Bugfix][MAGI-2] Fix layerwise offload OOM by pointing offload block attrs to 'block' (#7523) (#7540)

Signed-off-by: BeatSeat <wendavid552@gmail.com>

tests/diffusion/models/magi2/test_native_preview.py

@@ -90,6 +90,21 @@ def test_hsdp_policy_shards_individual_preview_layers():
     assert not condition("pre_adapter", model.pre_adapter)
 
 
+def test_layerwise_offload_contract() -> None:
+    from vllm_omni.diffusion.offloader import get_blocks_from_dit
+
+    model = Magi2PreviewTransformer(_tiny_config(num_layers=2))
+
+    assert model._layerwise_offload_blocks_attrs == ["block"]
+    assert "block" in dict(model.named_children())
+    assert list(model.block) == list(model.block.layers)
+    assert model.layers is model.block.layers
+
+    attr_names, blocks = get_blocks_from_dit(model)
+    assert attr_names == ["block"]
+    assert blocks == list(model.block.layers)
+
+
 def test_data_proxy_keeps_output_layout_request_scoped() -> None:
     proxy = Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=0))
 

vllm_omni/diffusion/models/magi2/modeling_magi2.py

@@ -670,11 +670,9 @@ def _is_magi2_transformer_layer(_name: str, module: nn.Module) -> bool:
 class Magi2PreviewTransformer(nn.Module):
     """Native preview DiT with the released checkpoint hierarchy."""
 
-    # ``block`` must remain the registered child name because it is part of the
-    # released checkpoint hierarchy.  The ``layers`` property below exposes its
-    # ModuleList through the standard layerwise-offload contract without
-    # changing the registered module names.
-    _layerwise_offload_blocks_attrs = ["layers"]
+    # ``block`` is the registered child module containing the repeated
+    # layers, exposed as the streamable blocks container for layerwise offload.
+    _layerwise_offload_blocks_attrs = ["block"]
     _hsdp_shard_conditions = [_is_magi2_transformer_layer]
     _hsdp_preserve_parameter_dtypes = True
     _EP_SHARDED_SUFFIXES = (
@@ -704,7 +702,7 @@ class Magi2PreviewTransformer(nn.Module):
 
     @property
     def layers(self) -> nn.ModuleList:
-        """Expose Preview layers to the shared offload block discovery API."""
+        """Expose Preview layers for direct access and compatibility."""
 
         return self.block.layers