[Bugfix][MM] Avoid device sync in FusedInputNorm initialization (#51435)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Aug 8, 2026 at 00:32 UTC
6e18901195fc3726546484fa86254217a747e9fe
2 files changed
+60
-7
tests/models/multimodal/processing/test_qwen2_vl.py
+39
@@ -13,6 +13,45 @@ from ....conftest import ImageTestAssets
13
from ...utils import build_model_context
14
15
16
+@pytest.mark.parametrize(
17
+ ("image_mean", "image_std", "rescale_factor", "is_identity"),
18
+ [
19
+ ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0, True),
20
+ ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255, False),
21
+ ],
22
+)
23
+def test_fused_input_norm_initialization_on_device(
24
+ monkeypatch: pytest.MonkeyPatch,
25
+ image_mean: list[float],
26
+ image_std: list[float],
27
+ rescale_factor: float,
28
+ is_identity: bool,
29
+):
30
+ """Identity detection must not synchronize the default device."""
31
+ original_allclose = torch.allclose
32
+
33
+ def cpu_allclose(input: torch.Tensor, other: torch.Tensor, *args, **kwargs):
34
+ assert input.device.type == "cpu"
35
+ assert other.device.type == "cpu"
36
+ return original_allclose(input, other, *args, **kwargs)
37
+
38
+ monkeypatch.setattr(torch, "allclose", cpu_allclose)
39
+ # Exercise the real accelerator when available. The meta device gives the
40
+ # CPU-only test shard the same non-CPU default-device semantics without
41
+ # requiring a CUDA-enabled PyTorch build.
42
+ default_device = "cuda" if torch.cuda.is_available() else "meta"
43
+ with torch.device(default_device):
44
+ input_norm = FusedInputNorm(image_mean, image_std, rescale_factor)
45
+
46
+ assert input_norm.is_identity is is_identity
47
+ if is_identity:
48
+ assert input_norm.weight is None
49
+ assert input_norm.bias is None
50
+ else:
51
+ assert input_norm.weight.device.type == default_device
52
+ assert input_norm.bias.device.type == default_device
53
+
54
+
55
@pytest.mark.parametrize("model_id", ["Qwen/Qwen2-VL-2B-Instruct"])
56
@pytest.mark.parametrize(
57
("mm_processor_kwargs", "expected_toks_per_img", "expected_pixels_shape"),
vllm/model_executor/models/vision.py
+21
-7
@@ -603,19 +603,33 @@ class FusedInputNorm(nn.Module):
603
604
self.channel = channel
605
606
- image_mean_tensor = torch.tensor(image_mean, dtype=dtype) * (
606
+ # Model construction can set the accelerator as PyTorch's default
607
+ # device. Determine whether the normalization is an identity on CPU
608
+ # so torch.allclose does not introduce a device synchronization while
609
+ # the model is being initialized. The actual buffers below still use
610
+ # the caller's default device.
611
+ image_mean_cpu = torch.tensor(image_mean, dtype=dtype, device="cpu") * (
612
1.0 / rescale_factor
613
)
609
- image_std_tensor = torch.tensor(image_std, dtype=dtype) * (1.0 / rescale_factor)
610
- weight = 1.0 / image_std_tensor
611
- bias = -image_mean_tensor / image_std_tensor
612
-
614
+ image_std_cpu = torch.tensor(image_std, dtype=dtype, device="cpu") * (
615
+ 1.0 / rescale_factor
616
+ )
617
+ weight_cpu = 1.0 / image_std_cpu
618
+ bias_cpu = -image_mean_cpu / image_std_cpu
619
self.is_identity = bool(
614
- torch.allclose(weight, torch.ones_like(weight))
615
- and torch.allclose(bias, torch.zeros_like(bias))
620
+ torch.allclose(weight_cpu, torch.ones_like(weight_cpu))
621
+ and torch.allclose(bias_cpu, torch.zeros_like(bias_cpu))
622
)
623
624
if not self.is_identity:
625
+ image_mean_tensor = torch.tensor(image_mean, dtype=dtype) * (
626
+ 1.0 / rescale_factor
627
+ )
628
+ image_std_tensor = torch.tensor(image_std, dtype=dtype) * (
629
+ 1.0 / rescale_factor
630
+ )
631
+ weight = 1.0 / image_std_tensor
632
+ bias = -image_mean_tensor / image_std_tensor
633
self.register_buffer("weight", weight)
634
self.register_buffer("bias", bias)
635
self.register_buffer("running_mean", torch.zeros_like(image_mean_tensor))