Add routed expert loading for gpt-oss (#52209)
Signed-off-by: Wyett <wyettzeng@gmail.com> Signed-off-by: wyett <wyettzeng@gmail.com>
Wyett committed
Aug 23, 2026 at 02:59 UTC
b26039b09fc97aa00f095a99eda503b7dad594ec
4 files changed
+370
-3
tests/model_executor/model_loader/test_gpt_oss_weight_loading.py
new
+152
@@ -0,0 +1,152 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.gpt_oss import (
9
+ GptOssForCausalLM,
10
+ GptOssModel,
11
+ GptOssRoutedExperts,
12
+)
13
+
14
+
15
+class _TestExperts(GptOssRoutedExperts):
16
+ def __init__(
17
+ self,
18
+ param_name: str,
19
+ shape: tuple[int, ...],
20
+ *,
21
+ tp_rank: int = 0,
22
+ quant_method_name: str = "TestQuantMethod",
23
+ weight_dtype: str = "",
24
+ dtype: torch.dtype = torch.float32,
25
+ ) -> None:
26
+ torch.nn.Module.__init__(self)
27
+ self.moe_config = SimpleNamespace(
28
+ moe_parallel_config=SimpleNamespace(tp_size=2, tp_rank=tp_rank)
29
+ )
30
+ self.quant_method = type(
31
+ quant_method_name, (), {"weight_dtype": weight_dtype}
32
+ )()
33
+
34
+ param = torch.nn.Parameter(torch.zeros(shape, dtype=dtype), requires_grad=False)
35
+ param.weight_loader = self.weight_loader
36
+ self.register_parameter(param_name, param)
37
+
38
+ @staticmethod
39
+ def _map_global_expert_id_to_local_expert_id(expert_id: int) -> int:
40
+ return {3: 0, 7: 1}.get(expert_id, -1)
41
+
42
+
43
+def _load(
44
+ experts: _TestExperts,
45
+ param_name: str,
46
+ loaded_weight: torch.Tensor,
47
+ shard_id: str,
48
+ expert_id: int = 3,
49
+) -> bool:
50
+ param = getattr(experts, param_name)
51
+ return param.weight_loader(
52
+ param,
53
+ loaded_weight,
54
+ weight_name=f"model.layers.0.mlp.experts.routed_experts.{param_name}",
55
+ shard_id=shard_id,
56
+ expert_id=expert_id,
57
+ return_success=True,
58
+ )
59
+
60
+
61
+@pytest.mark.parametrize("base_layer_prefix", ["", "base_layer."])
62
+def test_streamed_expert_resolves_wrapped_parameter(base_layer_prefix: str) -> None:
63
+ experts = _TestExperts("w13_weight", (2, 4, 3))
64
+ checkpoint_name = "model.layers.2.mlp.experts.7.gate_up_proj"
65
+ mapped_name = GptOssForCausalLM.hf_to_vllm_mapper.apply_list([checkpoint_name])[0]
66
+ fused_name = (
67
+ f"model.layers.2.mlp.experts.{base_layer_prefix}routed_experts.w13_weight"
68
+ )
69
+ params_dict = {fused_name: experts.w13_weight}
70
+ loaded_params: set[str] = set()
71
+ loaded_weight = torch.arange(24, dtype=torch.float32).reshape(3, 8)
72
+
73
+ assert GptOssModel._get_streamed_expert_info(mapped_name, params_dict) == (
74
+ 7,
75
+ fused_name,
76
+ "gpt_oss_w13",
77
+ )
78
+ assert GptOssModel._try_load_streamed_expert(
79
+ mapped_name, loaded_weight, params_dict, loaded_params
80
+ )
81
+ assert loaded_params == {fused_name}
82
+ assert torch.equal(experts.w13_weight[1], loaded_weight[:, :4].t())
83
+ assert torch.count_nonzero(experts.w13_weight[0]) == 0
84
+
85
+
86
+def test_streamed_expert_skips_non_local_expert() -> None:
87
+ experts = _TestExperts("w2_bias", (2, 3))
88
+
89
+ assert not _load(experts, "w2_bias", torch.ones(3), "gpt_oss_w2", expert_id=6)
90
+ assert torch.count_nonzero(experts.w2_bias) == 0
91
+
92
+
93
+@pytest.mark.parametrize("param_name", ["w13_weight", "w2_weight"])
94
+def test_streamed_unquantized_expert_tp_slice(param_name: str) -> None:
95
+ if param_name == "w13_weight":
96
+ experts = _TestExperts(param_name, (2, 4, 3), tp_rank=1)
97
+ loaded_weight = torch.arange(24, dtype=torch.float32).reshape(3, 8)
98
+ shard_id = "gpt_oss_w13"
99
+ expected = loaded_weight[:, 4:].t()
100
+ else:
101
+ experts = _TestExperts(param_name, (2, 3, 2), tp_rank=1)
102
+ loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3)
103
+ shard_id = "gpt_oss_w2"
104
+ expected = loaded_weight[2:].t()
105
+
106
+ assert _load(experts, param_name, loaded_weight, shard_id)
107
+ assert torch.equal(getattr(experts, param_name)[0], expected)
108
+
109
+
110
+@pytest.mark.parametrize("tp_rank", [0, 1])
111
+def test_streamed_down_bias_owned_by_tp_rank_zero(tp_rank: int) -> None:
112
+ experts = _TestExperts("w2_bias", (2, 3), tp_rank=tp_rank)
113
+ loaded_weight = torch.arange(1, 4, dtype=torch.float32)
114
+
115
+ assert _load(experts, "w2_bias", loaded_weight, "gpt_oss_w2")
116
+ expected = loaded_weight if tp_rank == 0 else torch.zeros_like(loaded_weight)
117
+ assert torch.equal(experts.w2_bias[0], expected)
118
+
119
+
120
+def test_streamed_gate_up_bias_tp_slice() -> None:
121
+ experts = _TestExperts("w13_bias", (2, 4), tp_rank=1)
122
+ loaded_weight = torch.arange(8, dtype=torch.float32)
123
+
124
+ assert _load(experts, "w13_bias", loaded_weight, "gpt_oss_w13")
125
+ assert torch.equal(experts.w13_bias[0], loaded_weight[4:])
126
+
127
+
128
+@pytest.mark.parametrize("case", ["mxfp4", "nvfp4"])
129
+def test_streamed_quantized_expert_layout(case: str) -> None:
130
+ if case == "mxfp4":
131
+ experts = _TestExperts(
132
+ "w13_weight",
133
+ (2, 4, 4),
134
+ tp_rank=1,
135
+ weight_dtype="gpt_oss_mxfp4",
136
+ dtype=torch.uint8,
137
+ )
138
+ loaded_weight = torch.arange(32, dtype=torch.uint8).reshape(8, 2, 2)
139
+ expected = loaded_weight.reshape(8, 4)[4:]
140
+ else:
141
+ experts = _TestExperts(
142
+ "w13_weight",
143
+ (2, 4, 2),
144
+ tp_rank=1,
145
+ quant_method_name="Nvfp4OnlineMoEMethod",
146
+ dtype=torch.uint8,
147
+ )
148
+ loaded_weight = torch.arange(16, dtype=torch.uint8).reshape(8, 2)
149
+ expected = loaded_weight[4:]
150
+
151
+ assert _load(experts, "w13_weight", loaded_weight, "gpt_oss_w13")
152
+ assert torch.equal(experts.w13_weight[0], expected)
tests/model_executor/model_loader/test_reload.py
+15
@@ -218,6 +218,21 @@ def test_reload_lifecycle():
218
assert tensor.__dict__ == materialized_tensor.__dict__
219
220
221
+def test_restore_layer_replaces_postprocessed_tensor_attribute():
222
+ layer = torch.nn.Linear(2, 3, bias=False)
223
+ info = LayerReloadingInfo(
224
+ restore_metadata=capture_layer_to_meta(layer),
225
+ restore_device=torch.device("cpu"),
226
+ )
227
+ del layer.weight
228
+ layer.weight = torch.empty(3, 2)
229
+
230
+ restore_layer_on_meta(layer, info)
231
+
232
+ assert isinstance(layer.weight, torch.nn.Parameter)
233
+ assert layer.weight.is_meta
234
+
235
+
236
def test_materialize_layer_preserves_non_meta_tensors():
237
"""Ensure that materialize_layer does not overwrite non meta tensors."""
238
layer = torch.nn.Linear(2, 3, bias=True)
vllm/model_executor/model_loader/reload/meta.py
+6
-3
@@ -124,11 +124,14 @@ def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
124
return
125
126
non_persistent = set(layer._non_persistent_buffers_set)
127
- for name in get_layer_tensors(layer):
128
- if name not in SKIP_TENSORS:
127
+ restore_params, restore_buffers = info.restore_metadata
128
+ tensor_names = (
129
+ get_layer_tensors(layer).keys() | restore_params.keys() | restore_buffers.keys()
130
+ )
131
+ for name in tensor_names:
132
+ if name not in SKIP_TENSORS and hasattr(layer, name):
133
delattr(layer, name)
134
131
- restore_params, restore_buffers = info.restore_metadata
135
for name, param in restore_params.items():
136
if name not in SKIP_TENSORS:
137
param = restore_layer_refs(param, layer)
vllm/model_executor/models/gpt_oss.py
+197
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe import (
25
fused_moe_make_expert_params_mapping,
26
)
27
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
28
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
29
from vllm.model_executor.layers.layernorm import RMSNorm
30
from vllm.model_executor.layers.linear import (
31
QKVParallelLinear,
@@ -68,6 +69,15 @@ from .utils import (
69
maybe_prefix,
70
)
71
72
+_GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD = {
73
+ "w13_weight": "gpt_oss_w13",
74
+ "w2_weight": "gpt_oss_w2",
75
+ "w13_bias": "gpt_oss_w13",
76
+ "w2_bias": "gpt_oss_w2",
77
+ "w13_weight_scale": "gpt_oss_w13",
78
+ "w2_weight_scale": "gpt_oss_w2",
79
+}
80
+
81
82
class OAIAttention(nn.Module):
83
# Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves);
@@ -179,6 +189,132 @@ class OAIAttention(nn.Module):
189
return output
190
191
192
+class GptOssRoutedExperts(RoutedExperts):
193
+ """Load one GPT-OSS expert at a time without assembling the global stack."""
194
+
195
+ @staticmethod
196
+ def _narrow_for_rank(
197
+ loaded_weight: torch.Tensor,
198
+ dim: int,
199
+ rank: int,
200
+ size: int,
201
+ ) -> torch.Tensor:
202
+ start = rank * size
203
+ available = loaded_weight.shape[dim] - start
204
+ return loaded_weight.narrow(dim, start, min(size, max(available, 0)))
205
+
206
+ @staticmethod
207
+ def _copy_to_expert(expert_data: torch.Tensor, loaded_weight: torch.Tensor) -> None:
208
+ if expert_data.numel() == loaded_weight.numel() == 1:
209
+ expert_data.copy_(loaded_weight.reshape(()))
210
+ return
211
+ while loaded_weight.ndim > expert_data.ndim and loaded_weight.shape[-1] == 1:
212
+ loaded_weight = loaded_weight.squeeze(-1)
213
+ slices = tuple(slice(0, size) for size in loaded_weight.shape)
214
+ expert_data[slices].copy_(loaded_weight)
215
+
216
+ def _load_expert_bias(
217
+ self,
218
+ expert_data: torch.Tensor,
219
+ loaded_weight: torch.Tensor,
220
+ shard_id: str,
221
+ ) -> None:
222
+ tp_rank = self.moe_config.moe_parallel_config.tp_rank
223
+ if shard_id == "gpt_oss_w13":
224
+ loaded_weight = self._narrow_for_rank(
225
+ loaded_weight, 0, tp_rank, expert_data.shape[0]
226
+ )
227
+ elif tp_rank != 0:
228
+ loaded_weight = torch.zeros_like(loaded_weight)
229
+ self._copy_to_expert(expert_data, loaded_weight)
230
+
231
+ def _load_unquantized_expert(
232
+ self,
233
+ expert_data: torch.Tensor,
234
+ loaded_weight: torch.Tensor,
235
+ shard_id: str,
236
+ ) -> None:
237
+ tp_rank = self.moe_config.moe_parallel_config.tp_rank
238
+ if shard_id == "gpt_oss_w13":
239
+ loaded_weight = self._narrow_for_rank(
240
+ loaded_weight, 1, tp_rank, expert_data.shape[0]
241
+ )
242
+ loaded_weight = loaded_weight.t().contiguous()
243
+ else:
244
+ loaded_weight = self._narrow_for_rank(
245
+ loaded_weight, 0, tp_rank, expert_data.shape[1]
246
+ )
247
+ loaded_weight = loaded_weight.t().contiguous()
248
+ self._copy_to_expert(expert_data, loaded_weight)
249
+
250
+ def _load_packed_expert(
251
+ self,
252
+ expert_data: torch.Tensor,
253
+ loaded_weight: torch.Tensor,
254
+ weight_name: str,
255
+ shard_id: str,
256
+ ) -> None:
257
+ tp_rank = self.moe_config.moe_parallel_config.tp_rank
258
+ is_w13 = shard_id == "gpt_oss_w13"
259
+ is_weight = weight_name.endswith("_weight")
260
+ is_partitioned_scale = weight_name.endswith("_weight_scale")
261
+
262
+ if is_weight:
263
+ loaded_weight = loaded_weight.reshape(loaded_weight.shape[0], -1)
264
+
265
+ if is_w13 and (is_weight or is_partitioned_scale):
266
+ loaded_weight = self._narrow_for_rank(
267
+ loaded_weight, 0, tp_rank, expert_data.shape[0]
268
+ )
269
+ elif not is_w13 and (is_weight or is_partitioned_scale):
270
+ loaded_weight = self._narrow_for_rank(
271
+ loaded_weight, -1, tp_rank, expert_data.shape[-1]
272
+ )
273
+
274
+ self._copy_to_expert(expert_data, loaded_weight)
275
+
276
+ def weight_loader(
277
+ self,
278
+ param: torch.nn.Parameter,
279
+ loaded_weight: torch.Tensor,
280
+ weight_name: str,
281
+ shard_id: str,
282
+ expert_id: int,
283
+ return_success: bool = False,
284
+ ) -> bool | None:
285
+ if shard_id not in ("gpt_oss_w13", "gpt_oss_w2"):
286
+ return super().weight_loader(
287
+ param,
288
+ loaded_weight,
289
+ weight_name,
290
+ shard_id,
291
+ expert_id,
292
+ return_success,
293
+ )
294
+
295
+ expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
296
+ if expert_id == -1:
297
+ return False if return_success else None
298
+
299
+ expert_data = param.data[expert_id]
300
+ weight_dtype = getattr(self.quant_method, "weight_dtype", "")
301
+ quant_method_name = self.quant_method.__class__.__name__
302
+
303
+ if weight_name.endswith("_bias"):
304
+ self._load_expert_bias(expert_data, loaded_weight, shard_id)
305
+ elif weight_dtype in (
306
+ "gpt_oss_mxfp4",
307
+ "mxfp4",
308
+ ) or quant_method_name in (
309
+ "CompressedTensorsW4A4Nvfp4MoEMethod",
310
+ "Nvfp4OnlineMoEMethod",
311
+ ):
312
+ self._load_packed_expert(expert_data, loaded_weight, weight_name, shard_id)
313
+ else:
314
+ self._load_unquantized_expert(expert_data, loaded_weight, shard_id)
315
+ return True if return_success else None
316
+
317
+
318
class MLPBlock(torch.nn.Module):
319
def __init__(
320
self,
@@ -223,6 +359,7 @@ class MLPBlock(torch.nn.Module):
359
has_bias=True,
360
activation="swigluoai",
361
is_sequence_parallel=self.is_sequence_parallel,
362
+ routed_experts_cls=GptOssRoutedExperts,
363
)
364
365
def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -416,6 +553,9 @@ class GptOssModel(nn.Module, EagleModelMixin):
553
if is_pp_missing_parameter(name, self):
554
continue
555
556
+ if self._try_load_streamed_expert(name, weight, params_dict, loaded_params):
557
+ continue
558
+
559
if ".w13_weight_scale" in name:
560
# Handle MLP gate and up projection weights scale
561
if use_ep:
@@ -634,6 +774,7 @@ class GptOssModel(nn.Module, EagleModelMixin):
774
tp_rank_start = tp_rank * per_rank_intermediate_size
775
tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size)
776
expert_params_mapping = self.get_expert_mapping()
777
+ # Streamed per-expert reload is intentionally unsupported for Quark.
778
for name, loaded_weight in weights:
779
if is_pp_missing_parameter(name, self):
780
continue
@@ -1018,6 +1159,9 @@ class GptOssModel(nn.Module, EagleModelMixin):
1159
if is_pp_missing_parameter(name, self):
1160
continue
1161
1162
+ if self._try_load_streamed_expert(name, weight, params_dict, loaded_params):
1163
+ continue
1164
+
1165
if ".w13_weight" in name:
1166
# Handle MLP gate and up projection weights
1167
# Extract gate and up projection parts
@@ -1096,6 +1240,59 @@ class GptOssModel(nn.Module, EagleModelMixin):
1240
loaded_params.add(name)
1241
return loaded_params
1242
1243
+ @staticmethod
1244
+ def _get_streamed_expert_info(
1245
+ name: str,
1246
+ params_dict: dict[str, torch.nn.Parameter],
1247
+ ) -> tuple[int, str, str] | None:
1248
+ """Parse ``...experts.<expert_id>.<fused_param>`` checkpoint keys."""
1249
+ if ".mlp.experts." not in name:
1250
+ return None
1251
+ suffix = name.rsplit(".", 1)[-1]
1252
+ shard_id = _GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD.get(suffix)
1253
+ if shard_id is None:
1254
+ return None
1255
+
1256
+ prefix, expert_suffix = name.split(".mlp.experts.", maxsplit=1)
1257
+ expert_id_str, separator, param_suffix = expert_suffix.partition(".")
1258
+ if not separator or not expert_id_str.isdigit():
1259
+ return None
1260
+ expert_id = int(expert_id_str)
1261
+ for base_layer_prefix in ("", "base_layer."):
1262
+ fused_name = (
1263
+ f"{prefix}.mlp.experts.{base_layer_prefix}routed_experts.{param_suffix}"
1264
+ )
1265
+ if fused_name in params_dict:
1266
+ return expert_id, fused_name, shard_id
1267
+ return None
1268
+
1269
+ @classmethod
1270
+ def _try_load_streamed_expert(
1271
+ cls,
1272
+ name: str,
1273
+ loaded_weight: torch.Tensor,
1274
+ params_dict: dict[str, torch.nn.Parameter],
1275
+ loaded_params: set[str],
1276
+ ) -> bool:
1277
+ expert_info = cls._get_streamed_expert_info(name, params_dict)
1278
+ if expert_info is None:
1279
+ return False
1280
+
1281
+ expert_id, fused_name, shard_id = expert_info
1282
+ param = params_dict[fused_name]
1283
+ weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
1284
+ success = weight_loader(
1285
+ param,
1286
+ loaded_weight,
1287
+ weight_name=fused_name,
1288
+ shard_id=shard_id,
1289
+ expert_id=expert_id,
1290
+ return_success=True,
1291
+ )
1292
+ if success:
1293
+ loaded_params.add(fused_name)
1294
+ return True
1295
+
1296
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1297
stacked_params_mapping = [
1298
# (param_name, shard_name, shard_id)