fix(minimax-h3): keep Qwen3VL vision pos-embed interpolation on the compute device

invokeai/backend/minimax_h3/qwen3vl_vision_device_patch.py

@@ -0,0 +1,81 @@
+"""Keep Qwen3VLVisionModel's pos-embed interpolation on the compute device under partial loading.
+
+``Qwen3VLVisionModel.fast_pos_embed_interpolate`` (transformers 5.5.x,
+``modeling_qwen3_vl.py``) derives its working device from ``self.pos_embed.weight.device`` and
+builds every tensor it makes — index tensor, interpolation weights, the looked-up embeddings —
+on that device. Under InvokeAI's partial loading that is the wrong device to trust:
+``nn.Embedding`` is autocast-wrapped (``CustomEmbedding``), so ``pos_embed.weight`` may
+legitimately reside on the CPU while the model computes on CUDA. The device-autocast wrappers
+cast weights to each *input's* device per-op, which keeps every normal layer on the compute
+device — but this function manufactures its own inputs on the weight's device, so the whole
+interpolation lands on the CPU and ``forward``'s ``hidden_states + pos_embeds`` fails with
+"Expected all tensors to be on the same device, but found at least two devices, cuda:0 and
+cpu!".
+
+Observed in the wild when two simultaneous MiniMax H3 video jobs partial-loaded the 27 GB text
+encoder and a Ref2VA reference image sent a prompt through the vision tower; any VRAM pressure
+that leaves ``visual.pos_embed.weight`` off-device reproduces it. It cannot be triggered on a
+fully loaded model (device autocasting is then disabled and the weight is resident), which is
+why it appears only intermittently, under memory pressure.
+
+The patch is deliberately minimal: ``forward`` records the true compute device from its
+``hidden_states`` input (the tensor ``pos_embeds`` will be added to), and
+``fast_pos_embed_interpolate`` runs UNCHANGED — every internal op stays on the weight's device,
+preserving upstream behavior exactly when the weight is resident — with only its result moved
+to the recorded device. When the weight lives on the CPU the interpolation arithmetic (a few
+thousand rows) runs there and one small tensor crosses the bus; correctness is unaffected
+either way.
+
+Class-level and idempotent, mirroring ``contiguous_attention``: applied by both text-encoder
+load paths (the diffusers-folder and single-file loaders in
+``model_loaders/minimax_h3.py``).
+"""
+
+from functools import wraps
+
+import torch
+
+from invokeai.backend.util.logging import InvokeAILogger
+
+_SENTINEL = "_invokeai_pos_embed_device_patch"
+_DEVICE_ATTR = "_invokeai_vision_input_device"
+
+
+def apply_qwen3vl_vision_pos_embed_device_patch() -> None:
+    """Install the patch on ``Qwen3VLVisionModel`` (idempotent, class-level)."""
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel
+
+    if getattr(Qwen3VLVisionModel, _SENTINEL, False):
+        return
+
+    orig_forward = getattr(Qwen3VLVisionModel, "forward", None)
+    orig_interpolate = getattr(Qwen3VLVisionModel, "fast_pos_embed_interpolate", None)
+    if orig_forward is None or orig_interpolate is None:
+        # A future transformers restructured the class; the residency bug this guards against
+        # may be gone too. Skip rather than crash the model load, but say so.
+        InvokeAILogger.get_logger(__name__).warning(
+            "Qwen3VLVisionModel no longer has forward/fast_pos_embed_interpolate; skipping the "
+            "pos-embed device patch. If partially-loaded H3 text encoders start failing with "
+            "cross-device errors in the vision tower, this patch needs updating."
+        )
+        return
+
+    @wraps(orig_forward)
+    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs):
+        # hidden_states is the tensor pos_embeds will be added to — its device is the one
+        # authoritative answer to "where is this forward computing", independent of where the
+        # partial loader happened to leave any given weight.
+        setattr(self, _DEVICE_ATTR, hidden_states.device)
+        return orig_forward(self, hidden_states, grid_thw, **kwargs)
+
+    @wraps(orig_interpolate)
+    def fast_pos_embed_interpolate(self, grid_thw):
+        pos_embeds = orig_interpolate(self, grid_thw)
+        target = getattr(self, _DEVICE_ATTR, None)
+        if target is not None and pos_embeds.device != target:
+            pos_embeds = pos_embeds.to(target)
+        return pos_embeds
+
+    Qwen3VLVisionModel.forward = forward
+    Qwen3VLVisionModel.fast_pos_embed_interpolate = fast_pos_embed_interpolate
+    setattr(Qwen3VLVisionModel, _SENTINEL, True)

invokeai/backend/model_manager/load/model_loaders/minimax_h3.py

@@ -102,6 +102,14 @@ def _load_model(
 
                 from transformers import AutoConfig, Qwen3VLForConditionalGeneration
 
+                from invokeai.backend.minimax_h3.qwen3vl_vision_device_patch import (
+                    apply_qwen3vl_vision_pos_embed_device_patch,
+                )
+
+                # Keep the vision tower's pos-embed interpolation on the compute device when the
+                # partial loader leaves pos_embed.weight on the CPU; see the patch module.
+                apply_qwen3vl_vision_pos_embed_device_patch()
+
                 te_config = normalize_qwen3vl_rope_config(
                     AutoConfig.from_pretrained(submodel_path, local_files_only=True)
                 )
@@ -318,6 +326,9 @@ def _load_text_encoder_from_singlefile(self, config: AnyModelConfig) -> AnyModel
         from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration
 
         from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear
+        from invokeai.backend.minimax_h3.qwen3vl_vision_device_patch import (
+            apply_qwen3vl_vision_pos_embed_device_patch,
+        )
         from invokeai.backend.minimax_h3.text_conditioning import MINIMAX_H3_TEXT_ENCODER_LAYER
         from invokeai.backend.model_manager.load.model_loaders.minimax_h3_state_dict_utils import (
             convert_minimax_h3_text_encoder_checkpoint,
@@ -367,6 +378,10 @@ def _load_text_encoder_from_singlefile(self, config: AnyModelConfig) -> AnyModel
         config_dict["tie_word_embeddings"] = True
         te_config = normalize_qwen3vl_rope_config(Qwen3VLConfig.from_dict(config_dict))
 
+        # Keep the vision tower's pos-embed interpolation on the compute device when the partial
+        # loader leaves pos_embed.weight on the CPU; see the patch module.
+        apply_qwen3vl_vision_pos_embed_device_patch()
+
         with accelerate.init_empty_weights():
             model = Qwen3VLForConditionalGeneration._from_config(te_config)
 

tests/backend/minimax_h3/test_qwen3vl_vision_device_patch.py

@@ -0,0 +1,106 @@
+"""Tests for the Qwen3VLVisionModel pos-embed device patch.
+
+The bug: ``fast_pos_embed_interpolate`` builds all of its tensors on
+``self.pos_embed.weight.device``. Under partial loading that weight may sit on the CPU while
+the forward computes on CUDA (nn.Embedding is autocast-wrapped), so ``hidden_states +
+pos_embeds`` fails with a cross-device RuntimeError. The patch records the true compute device
+from forward's ``hidden_states`` input and moves only the interpolation's result there.
+
+CI has no GPU, so the cross-device move is exercised against the ``meta`` device; the
+end-to-end CPU forward covers the no-op path (result already on the right device).
+"""
+
+import pytest
+import torch
+
+pytest.importorskip("transformers.models.qwen3_vl")
+
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig  # noqa: E402
+from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel  # noqa: E402
+
+from invokeai.backend.minimax_h3.qwen3vl_vision_device_patch import (  # noqa: E402
+    _DEVICE_ATTR,
+    _SENTINEL,
+    apply_qwen3vl_vision_pos_embed_device_patch,
+)
+
+
+def _tiny_vision_model() -> Qwen3VLVisionModel:
+    config = Qwen3VLVisionConfig(
+        hidden_size=32,
+        num_heads=2,
+        depth=2,
+        intermediate_size=64,
+        num_position_embeddings=16,
+        patch_size=2,
+        temporal_patch_size=1,
+        in_channels=3,
+        out_hidden_size=32,
+        spatial_merge_size=2,
+        deepstack_visual_indexes=[0],
+    )
+    torch.manual_seed(0)
+    return Qwen3VLVisionModel(config)
+
+
+def _inputs() -> tuple[torch.Tensor, torch.Tensor]:
+    t, h, w = 1, 4, 4
+    torch.manual_seed(1)
+    hidden_states = torch.randn(t * h * w, 3 * 1 * 2 * 2)
+    grid_thw = torch.tensor([[t, h, w]])
+    return hidden_states, grid_thw
+
+
+def test_patched_forward_matches_baseline_on_cpu():
+    model = _tiny_vision_model()
+    hidden_states, grid_thw = _inputs()
+    baseline = model(hidden_states, grid_thw).last_hidden_state
+
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    patched = model(hidden_states, grid_thw).last_hidden_state
+
+    assert torch.equal(baseline, patched)
+
+
+def test_forward_records_input_device():
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    model = _tiny_vision_model()
+    hidden_states, grid_thw = _inputs()
+    model(hidden_states, grid_thw)
+    assert getattr(model, _DEVICE_ATTR) == hidden_states.device
+
+
+def test_interpolation_result_moved_to_recorded_device():
+    """The core fix: with the compute device differing from pos_embed.weight's device, the
+    interpolation result must land on the compute device (validated via `meta`, the only
+    second device available on CI)."""
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    model = _tiny_vision_model()
+    _, grid_thw = _inputs()
+
+    setattr(model, _DEVICE_ATTR, torch.device("meta"))
+    pos_embeds = model.fast_pos_embed_interpolate(grid_thw)
+    assert pos_embeds.device.type == "meta"
+
+
+def test_interpolation_without_recorded_device_stays_on_weight_device():
+    """Before any forward has run (attr unset), behavior is upstream's: the result stays on
+    the weight's device rather than being moved to a stale or absent target."""
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    model = _tiny_vision_model()
+    _, grid_thw = _inputs()
+
+    assert not hasattr(model, _DEVICE_ATTR)
+    pos_embeds = model.fast_pos_embed_interpolate(grid_thw)
+    assert pos_embeds.device == model.pos_embed.weight.device
+
+
+def test_apply_is_idempotent():
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    forward_after_first = Qwen3VLVisionModel.forward
+    interpolate_after_first = Qwen3VLVisionModel.fast_pos_embed_interpolate
+    assert getattr(Qwen3VLVisionModel, _SENTINEL) is True
+
+    apply_qwen3vl_vision_pos_embed_device_patch()
+    assert Qwen3VLVisionModel.forward is forward_after_first
+    assert Qwen3VLVisionModel.fast_pos_embed_interpolate is interpolate_after_first