diff --git a/invokeai/backend/minimax_h3/qwen3vl_vision_device_patch.py b/invokeai/backend/minimax_h3/qwen3vl_vision_device_patch.py new file mode 100644 index 00000000000..c86743c1252 --- /dev/null +++ b/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) diff --git a/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py b/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py index 3f5e8a7c343..a8dd404d4d9 100644 --- a/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py +++ b/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) diff --git a/tests/backend/minimax_h3/test_qwen3vl_vision_device_patch.py b/tests/backend/minimax_h3/test_qwen3vl_vision_device_patch.py new file mode 100644 index 00000000000..0127f30878e --- /dev/null +++ b/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