[Bugfix][Diffusion] Fix MiniMax-H3 model-level CPU offload residency (#6072)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

tests/diffusion/offloader/test_sequential_backend.py

@@ -8,7 +8,13 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
+from vllm_omni.diffusion.offloader.sequential_backend import (
+    ModelLevelOffloadBackend,
+    SequentialOffloadHook,
+    apply_sequential_offload,
+    remove_sequential_offload,
+    sequential_offload_component,
+)
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    offloaded: list[nn.Module] = []
+    monkeypatch.setattr(
+        SequentialOffloadHook,
+        "_to_cpu",
+        lambda self, module: offloaded.append(module),
+    )
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=torch.device("cpu"),
+        offload_initial_dits=True,
+    )
+
+    assert offloaded == [dit]
+    remove_sequential_offload([dit, encoder])
+
+
+def test_direct_component_activation_failure_still_offloads(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    component = _create_simple_module()
+    apply_sequential_offload(
+        dit_modules=[_create_simple_module()],
+        encoder_modules=[component],
+        device=torch.device("cpu"),
+    )
+    hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    offloaded: list[nn.Module] = []
+
+    def fail_activation(module: nn.Module) -> None:
+        raise RuntimeError("activation failed")
+
+    monkeypatch.setattr(hook, "pre_forward", fail_activation)
+    monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module))
+
+    with pytest.raises(RuntimeError, match="activation failed"):
+        with sequential_offload_component(component):
+            pass
+
+    assert offloaded == [component]
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""
@@ -85,7 +137,12 @@ class TestMoveParamsPinMemory:
         original_isinstance = isinstance
 
         def fake_isinstance(obj, cls):
-            if cls.__name__ == "DTensor":
+            # torchada's Tensor.to() checks tuples such as
+            # ``(str, torch.device)``.  Keep the DTensor probe compatible with
+            # the normal isinstance contract instead of breaking those calls.
+            if type(cls) is tuple:
+                return any(fake_isinstance(obj, candidate) for candidate in cls)
+            if getattr(cls, "__name__", None) == "DTensor":
                 return True
             return original_isinstance(obj, cls)