| 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 |