diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index b7ff05638..fd00b4e70 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py @@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import ( MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, MiniMaxH3Qwen3VLConfig, ) +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel @@ -68,7 +69,15 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): @property def device(self) -> torch.device: - return next(self.parameters()).device + """Device this encoder's forward runs on. + + Deliberately not `next(self.parameters()).device`. `--text-encoder-cpu-offload` + loads this component under an FSDP CPU offload policy, which keeps the sharded + parameters on CPU and all-gathers them to the accelerator for the forward. The + parameter device then names the storage side, not the compute side, so inputs + built from it stay on CPU while the forward runs on the accelerator. + """ + return get_local_torch_device() @torch.no_grad() def forward( diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_encoder_device.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_encoder_device.py new file mode 100644 index 000000000..08212dd78 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_encoder_device.py @@ -0,0 +1,54 @@ +import unittest +from unittest import mock + +import torch + +from sglang.multimodal_gen.runtime.models.encoders import minimax_h3_qwen3vl +from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( + MiniMaxH3Qwen3VLEncoder, +) + + +class TestMiniMaxH3EncoderDevice(unittest.TestCase): + """`device` must name the compute side, not the parameter storage side. + + `--text-encoder-cpu-offload` loads this encoder under an FSDP CPU offload + policy: the sharded parameters sit on CPU and are all-gathered to the + accelerator for the forward. Reporting the parameter device there sent + `encode_ids` to build `input_ids`/`attention_mask`/`position_ids` on CPU + while the forward ran on the accelerator, and the rope matmul died with + "Expected all tensors to be on the same device ... mat2 is on cpu". + """ + + def _encoder_with_param_on(self, device: torch.device) -> MiniMaxH3Qwen3VLEncoder: + encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.register_parameter( + "offloaded", torch.nn.Parameter(torch.zeros(1, device=device)) + ) + return encoder + + def test_device_ignores_cpu_offloaded_parameters(self): + encoder = self._encoder_with_param_on(torch.device("cpu")) + compute_device = torch.device("cuda", 3) + + with mock.patch.object( + minimax_h3_qwen3vl, "get_local_torch_device", return_value=compute_device + ): + self.assertEqual(encoder.device, compute_device) + + # The parameter really is on CPU: the property is not just echoing it back. + self.assertEqual(next(encoder.parameters()).device.type, "cpu") + + def test_device_follows_local_device_on_cpu_only_platforms(self): + encoder = self._encoder_with_param_on(torch.device("cpu")) + cpu = torch.device("cpu") + + with mock.patch.object( + minimax_h3_qwen3vl, "get_local_torch_device", return_value=cpu + ): + self.assertEqual(encoder.device, cpu) + + +if __name__ == "__main__": + unittest.main()