[Bugfix][GLM-5.3-Flash] Run the dense MLP layers on the sequence-parallel shard (#58061)
Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jared Wen committed
Sep 22, 2026 at 21:58 UTC
c9b34fdb2d0c9fdcc20fae6454bf0895a5faf6f5
2 files changed
+115
tests/models/glm5next/test_sequence_parallel.py
new
+114
@@ -0,0 +1,114 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GLM-5.3-Flash sequence-parallel MoE layout.
4
+
5
+With DP > 1, TP > 1 and expert parallelism, ``Glm5NextModel`` shards the token
6
+dimension across the TP group once at the model entry and every layer runs its
7
+MLP on that shard. A module that still does tensor-parallel collectives there
8
+sums the partial results of *different* tokens held by different ranks, so the
9
+dense MLP of the first ``first_k_dense_replace`` layers must hold replicated
10
+weights and run without any collective, like the shared experts already do.
11
+"""
12
+
13
+from types import SimpleNamespace
14
+
15
+import pytest
16
+import torch
17
+from torch import nn
18
+
19
+from vllm.model_executor import parameter
20
+from vllm.model_executor.layers import linear
21
+from vllm.models.glm5next.common import model as glm_model
22
+from vllm.transformers_utils.configs.glm5_next import Glm5NextTextConfig
23
+
24
+HIDDEN, INTERMEDIATE, TP_SIZE = 8, 16, 2
25
+
26
+
27
+class _Attention(nn.Module):
28
+ def __init__(self, *args, **kwargs) -> None:
29
+ super().__init__()
30
+ self.o_proj = SimpleNamespace(reduce_results=True)
31
+
32
+
33
+def _fake_tensor_parallel_world(monkeypatch) -> None:
34
+ """Pretend to be rank 1 of a TP=2 group whose all-reduce must never run."""
35
+
36
+ def _forbidden_all_reduce(*args, **kwargs):
37
+ raise AssertionError("tensor-parallel all-reduce on a sequence-parallel shard")
38
+
39
+ for module in (linear, parameter):
40
+ monkeypatch.setattr(
41
+ module, "get_tensor_model_parallel_world_size", lambda: TP_SIZE
42
+ )
43
+ monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
44
+ monkeypatch.setattr(
45
+ linear, "tensor_model_parallel_all_reduce", _forbidden_all_reduce
46
+ )
47
+
48
+
49
+def _reference_mlp(mlp: glm_model.Glm5NextMLP, x: torch.Tensor) -> torch.Tensor:
50
+ gate_up = x @ mlp.gate_up_proj.weight.T
51
+ gate, up = gate_up.split(INTERMEDIATE, dim=-1)
52
+ return (torch.nn.functional.silu(gate) * up) @ mlp.down_proj.weight.T
53
+
54
+
55
+def _assert_runs_locally_on_shard(mlp: glm_model.Glm5NextMLP) -> None:
56
+ device = "cuda" if torch.cuda.is_available() else "cpu"
57
+ mlp.to(device)
58
+ assert mlp.gate_up_proj.weight.shape == (2 * INTERMEDIATE, HIDDEN)
59
+ assert mlp.down_proj.weight.shape == (HIDDEN, INTERMEDIATE)
60
+ torch.manual_seed(0)
61
+ for param in mlp.parameters():
62
+ param.data.normal_()
63
+ shard = torch.randn(3, HIDDEN, device=device)
64
+ torch.testing.assert_close(mlp(shard), _reference_mlp(mlp, shard))
65
+
66
+
67
+def test_dense_mlp_runs_locally_on_the_sequence_parallel_shard(
68
+ monkeypatch, default_vllm_config
69
+):
70
+ _fake_tensor_parallel_world(monkeypatch)
71
+
72
+ sharded = glm_model.Glm5NextMLP(HIDDEN, INTERMEDIATE, "silu")
73
+ assert sharded.down_proj.tp_size == TP_SIZE
74
+ assert sharded.gate_up_proj.weight.shape == (2 * INTERMEDIATE // TP_SIZE, HIDDEN)
75
+
76
+ _assert_runs_locally_on_shard(
77
+ glm_model.Glm5NextMLP(HIDDEN, INTERMEDIATE, "silu", is_sequence_parallel=True)
78
+ )
79
+
80
+
81
+@pytest.mark.parametrize("use_sequence_parallel_moe", [True, False])
82
+def test_dense_layer_mlp_follows_the_sequence_parallel_layout(
83
+ monkeypatch, default_vllm_config, use_sequence_parallel_moe: bool
84
+):
85
+ _fake_tensor_parallel_world(monkeypatch)
86
+ monkeypatch.setattr(glm_model, "Glm5NextLinearAttention", _Attention)
87
+ monkeypatch.setattr(glm_model, "Glm5NextMLAAttention", _Attention)
88
+
89
+ config = Glm5NextTextConfig(
90
+ hidden_size=HIDDEN,
91
+ intermediate_size=INTERMEDIATE,
92
+ num_hidden_layers=2,
93
+ first_k_dense_replace=1,
94
+ n_routed_experts=4,
95
+ layer_types=["linear_attention", "deepseek_sparse_attention"],
96
+ mhc=False,
97
+ )
98
+ vllm_config = SimpleNamespace(
99
+ cache_config=None,
100
+ quant_config=None,
101
+ parallel_config=SimpleNamespace(
102
+ use_sequence_parallel_moe=use_sequence_parallel_moe
103
+ ),
104
+ kernel_config=SimpleNamespace(enable_jit_warmup=False),
105
+ )
106
+
107
+ layer = glm_model.Glm5NextDecoderLayer(vllm_config, config, layer_idx=0)
108
+
109
+ assert isinstance(layer.mlp, glm_model.Glm5NextMLP)
110
+ assert layer.self_attn.o_proj.reduce_results is not use_sequence_parallel_moe
111
+ if use_sequence_parallel_moe:
112
+ _assert_runs_locally_on_shard(layer.mlp)
113
+ else:
114
+ assert layer.mlp.down_proj.tp_size == TP_SIZE
vllm/models/glm5next/common/model.py
+1
@@ -355,6 +355,7 @@ class Glm5NextDecoderLayer(nn.Module):
355
intermediate_size=config.intermediate_size,
356
hidden_act=config.hidden_act,
357
quant_config=quant_config,
358
+ is_sequence_parallel=self.is_sequence_parallel,
359
prefix=f"{prefix}.mlp",
360
swiglu_limit=config.swiglu_limit,
361
)