[Bugfix] Keep the Moondream3 MoE all-reduce out of the fused-path try (#54152)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Oxana Korzh committed
Aug 28, 2026 at 16:06 UTC
c274d361031e5bf920af13071794056aa704e690
2 files changed
+205
-6
tests/models/multimodal/generation/test_moondream3_moe.py
new
+182
@@ -0,0 +1,182 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Collective-safety tests for the Moondream3 MoE layer.
4
+
5
+The MoE forward may fall back from the fused kernel to a Python expert loop.
6
+That decision is made per rank from a rank-local exception, so every path has
7
+to issue exactly one all-reduce. A path that issues two (or none) desyncs the
8
+tensor-parallel group and wedges it for the rest of the run.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from tests.utils import ensure_current_vllm_config
15
+from vllm.distributed.parallel_state import (
16
+ destroy_distributed_environment,
17
+ destroy_model_parallel,
18
+ init_distributed_environment,
19
+ initialize_model_parallel,
20
+)
21
+from vllm.model_executor.models import moondream3
22
+from vllm.utils.network_utils import get_open_port
23
+from vllm.utils.system_utils import update_environment_variables
24
+
25
+pytestmark = pytest.mark.cpu_test
26
+
27
+HIDDEN_SIZE = 8
28
+EXPERT_INNER_DIM = 16
29
+NUM_EXPERTS = 4
30
+EXPERTS_PER_TOKEN = 2
31
+NUM_TOKENS = 3
32
+
33
+
34
+class _StubGate(torch.nn.Module):
35
+ """Router stand-in, so the test does not depend on platform-specific
36
+ ``ReplicatedLinear`` weight post-processing."""
37
+
38
+ def __init__(self, hidden_size: int, num_experts: int):
39
+ super().__init__()
40
+ self.linear = torch.nn.Linear(hidden_size, num_experts)
41
+
42
+ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
43
+ return self.linear(x), None
44
+
45
+
46
+class _CountingAllReduce:
47
+ """Stand-in for tensor_model_parallel_all_reduce that counts calls."""
48
+
49
+ def __init__(self):
50
+ self.calls = 0
51
+
52
+ def __call__(self, x: torch.Tensor) -> torch.Tensor:
53
+ self.calls += 1
54
+ return x
55
+
56
+
57
+@pytest.fixture
58
+def tp1_cpu():
59
+ """A single-rank CPU tensor-parallel group, enough to build the layer."""
60
+ update_environment_variables(
61
+ {
62
+ "RANK": "0",
63
+ "LOCAL_RANK": "0",
64
+ "WORLD_SIZE": "1",
65
+ "MASTER_ADDR": "localhost",
66
+ "MASTER_PORT": str(get_open_port()),
67
+ }
68
+ )
69
+ init_distributed_environment(backend="gloo")
70
+ with ensure_current_vllm_config():
71
+ initialize_model_parallel(tensor_model_parallel_size=1)
72
+ yield
73
+ destroy_model_parallel()
74
+ destroy_distributed_environment()
75
+
76
+
77
+@pytest.fixture
78
+def moe(tp1_cpu) -> moondream3.Moondream3TextMoE:
79
+ """A single-rank MoE layer with deterministic weights."""
80
+ layer = moondream3.Moondream3TextMoE(
81
+ hidden_size=HIDDEN_SIZE,
82
+ expert_inner_dim=EXPERT_INNER_DIM,
83
+ num_experts=NUM_EXPERTS,
84
+ experts_per_token=EXPERTS_PER_TOKEN,
85
+ )
86
+ layer.gate = _StubGate(HIDDEN_SIZE, NUM_EXPERTS)
87
+ torch.nn.init.normal_(layer.fc1_weight, std=0.02)
88
+ torch.nn.init.normal_(layer.fc2_weight, std=0.02)
89
+ return layer
90
+
91
+
92
+@pytest.fixture
93
+def all_reduce(monkeypatch: pytest.MonkeyPatch) -> _CountingAllReduce:
94
+ counter = _CountingAllReduce()
95
+ monkeypatch.setattr(moondream3, "tensor_model_parallel_all_reduce", counter)
96
+ return counter
97
+
98
+
99
+def test_expert_loop_all_reduces_once(moe, all_reduce):
100
+ """The Python expert loop issues exactly one all-reduce."""
101
+ moe._use_fused_moe = False
102
+ x = torch.randn(NUM_TOKENS, HIDDEN_SIZE)
103
+
104
+ out = moe(x)
105
+
106
+ assert out.shape == x.shape
107
+ assert all_reduce.calls == 1
108
+
109
+
110
+def test_fused_path_all_reduces_once(moe, all_reduce, monkeypatch):
111
+ """The fused kernel path issues exactly one all-reduce."""
112
+ monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
113
+ monkeypatch.setattr(
114
+ moondream3,
115
+ "fused_experts",
116
+ lambda **kwargs: torch.zeros_like(kwargs["hidden_states"]),
117
+ )
118
+ x = torch.randn(NUM_TOKENS, HIDDEN_SIZE)
119
+
120
+ out = moe(x)
121
+
122
+ assert out.shape == x.shape
123
+ assert all_reduce.calls == 1
124
+
125
+
126
+def test_all_reduce_failure_is_not_retried(moe, monkeypatch):
127
+ """A failing all-reduce must propagate, not trigger a second collective.
128
+
129
+ Regression test. The all-reduce used to sit inside the ``try`` guarding
130
+ the fused kernel, so a ``RuntimeError`` raised by the collective itself
131
+ was caught, and this rank went on to run the fallback loop and all-reduce
132
+ a *second* time -- two collectives on one rank against one on each of its
133
+ peers, which desyncs the tensor-parallel group and wedges it for good.
134
+ """
135
+ monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
136
+ monkeypatch.setattr(
137
+ moondream3,
138
+ "fused_experts",
139
+ lambda **kwargs: torch.zeros_like(kwargs["hidden_states"]),
140
+ )
141
+ calls = []
142
+
143
+ def _failing_all_reduce(x: torch.Tensor) -> torch.Tensor:
144
+ calls.append(x)
145
+ raise RuntimeError("collective failed")
146
+
147
+ monkeypatch.setattr(
148
+ moondream3, "tensor_model_parallel_all_reduce", _failing_all_reduce
149
+ )
150
+ x = torch.randn(NUM_TOKENS, HIDDEN_SIZE)
151
+
152
+ with pytest.raises(RuntimeError, match="collective failed"):
153
+ moe(x)
154
+
155
+ assert len(calls) == 1, (
156
+ "the failed all-reduce was retried on the fallback path; peers issued "
157
+ "one collective and this rank issued two"
158
+ )
159
+
160
+
161
+@pytest.mark.parametrize("exc", [RuntimeError("boom"), NotImplementedError()])
162
+def test_fallback_after_fused_failure_all_reduces_once(
163
+ moe, all_reduce, monkeypatch, exc
164
+):
165
+ """A rank whose fused kernel fails still all-reduces exactly once."""
166
+ monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
167
+
168
+ def _raise(**kwargs):
169
+ raise exc
170
+
171
+ monkeypatch.setattr(moondream3, "fused_experts", _raise)
172
+ x = torch.randn(NUM_TOKENS, HIDDEN_SIZE)
173
+
174
+ out = moe(x)
175
+
176
+ assert out.shape == x.shape
177
+ assert all_reduce.calls == 1
178
+ # The fused path is disabled for subsequent forwards...
179
+ assert moe._use_fused_moe is False
180
+ # ...and those still all-reduce exactly once each.
181
+ moe(x)
182
+ assert all_reduce.calls == 2
vllm/model_executor/models/moondream3.py
+23
-6
@@ -497,7 +497,13 @@ class Moondream3TextMoE(nn.Module):
497
# Softmax over selected experts
498
topk_weights = F.softmax(topk_logits, dim=-1, dtype=torch.float32).to(x.dtype)
499
500
+ out = None
501
if self._use_fused_moe and x.is_cuda:
502
+ # Only the expert computation may fall back. The all-reduce below
503
+ # is deliberately outside this `try`: if it were inside and raised
504
+ # on a single rank, that rank would run the fallback loop and issue
505
+ # a *second* all-reduce while its peers issued one, desyncing the
506
+ # TP group and hanging it for good.
507
try:
508
out = fused_experts(
509
hidden_states=x.contiguous(),
@@ -510,8 +516,6 @@ class Moondream3TextMoE(nn.Module):
516
expert_map=self._expert_map,
517
quant_config=biased_moe_quant_config(self._fused_w1_bias, None),
518
)
513
- out = tensor_model_parallel_all_reduce(out)
514
- return out
519
except (NotImplementedError, RuntimeError) as exc:
520
self._use_fused_moe = False
521
logger.warning_once(
@@ -520,11 +524,27 @@ class Moondream3TextMoE(nn.Module):
524
str(exc),
525
)
526
527
+ if out is None:
528
+ out = self._expert_loop(x, topk_weights, topk_ids)
529
+
530
+ # Combine each rank's partial expert outputs. Reached exactly once per
531
+ # forward, on every rank, whichever expert path ran above.
532
+ return tensor_model_parallel_all_reduce(out)
533
+
534
+ def _expert_loop(
535
+ self,
536
+ x: torch.Tensor,
537
+ topk_weights: torch.Tensor,
538
+ topk_ids: torch.Tensor,
539
+ ) -> torch.Tensor:
540
+ """Fallback for environments where fused kernels are unavailable.
541
+
542
+ Returns this rank's partial output; the caller performs the all-reduce.
543
+ """
544
tp_rank = get_tensor_model_parallel_rank()
545
# Compute local expert range
546
local_expert_start = tp_rank * self.experts_per_rank
547
527
- # Fallback path for environments where fused kernels are unavailable.
548
out = x.new_zeros(x.shape)
549
550
for local_expert_idx in range(self.num_local_experts):
@@ -558,9 +578,6 @@ class Moondream3TextMoE(nn.Module):
578
# Accumulate output
579
out.index_add_(0, token_pos, y)
580
561
- # All-reduce to combine results from all experts across GPUs
562
- out = tensor_model_parallel_all_reduce(out)
563
-
581
return out
582
583