A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 17:00 UTC
[Performance][Model Loader] Skip non-local expert weights during EP model loading (#37136)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Roy Wang committed
Mar 16, 2026 at 16:33 UTC
821eb80c0d9c3eec0201fda21dbeead83b6ac1fc
4 files changed
+513
-2
tests/model_executor/model_loader/test_ep_weight_filter.py
new
+361
@@ -0,0 +1,361 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for EP weight filtering during model loading."""
4
+
5
+import glob
6
+import tempfile
7
+
8
+import huggingface_hub.constants
9
+import pytest
10
+import torch
11
+
12
+from vllm.model_executor.model_loader.ep_weight_filter import (
13
+ compute_local_expert_ids,
14
+ parse_expert_id,
15
+ should_skip_weight,
16
+)
17
+from vllm.model_executor.model_loader.weight_utils import (
18
+ safetensors_weights_iterator,
19
+)
20
+
21
+# ---------------------------------------------------------------------------
22
+# Unit tests for parse_expert_id
23
+# ---------------------------------------------------------------------------
24
+
25
+
26
+class TestParseExpertId:
27
+ def test_routed_expert(self):
28
+ name = "model.layers.0.mlp.experts.42.gate_proj.weight"
29
+ assert parse_expert_id(name) == 42
30
+
31
+ def test_large_expert_id(self):
32
+ name = "model.layers.60.mlp.experts.383.down_proj.weight"
33
+ assert parse_expert_id(name) == 383
34
+
35
+ def test_shared_expert(self):
36
+ # Shared experts use a different naming convention in most models
37
+ name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
38
+ assert parse_expert_id(name) is None
39
+
40
+ def test_attention_weight(self):
41
+ name = "model.layers.0.self_attn.q_proj.weight"
42
+ assert parse_expert_id(name) is None
43
+
44
+ def test_embedding(self):
45
+ name = "model.embed_tokens.weight"
46
+ assert parse_expert_id(name) is None
47
+
48
+ def test_layernorm(self):
49
+ name = "model.layers.0.input_layernorm.weight"
50
+ assert parse_expert_id(name) is None
51
+
52
+ def test_fused_3d_expert(self):
53
+ # 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
54
+ # They must NOT be filtered — slicing happens later in weight_loader.
55
+ name = "model.layers.0.mlp.experts.gate_proj.weight"
56
+ assert parse_expert_id(name) is None
57
+
58
+ def test_fused_3d_expert_down_proj(self):
59
+ name = "model.layers.10.mlp.experts.down_proj.weight"
60
+ assert parse_expert_id(name) is None
61
+
62
+ def test_expert_scale(self):
63
+ # NVFP4 quantized models have scale tensors for experts
64
+ name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
65
+ assert parse_expert_id(name) == 100
66
+
67
+ def test_expert_zero_id(self):
68
+ name = "model.layers.0.mlp.experts.0.up_proj.weight"
69
+ assert parse_expert_id(name) == 0
70
+
71
+
72
+# ---------------------------------------------------------------------------
73
+# Unit tests for compute_local_expert_ids
74
+# ---------------------------------------------------------------------------
75
+
76
+
77
+class TestComputeLocalExpertIds:
78
+ def test_ep_disabled(self):
79
+ assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
80
+
81
+ def test_even_split(self):
82
+ # 64 experts, EP=8 → 8 per rank
83
+ ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
84
+ assert ids == set(range(0, 8))
85
+
86
+ ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
87
+ assert ids == set(range(56, 64))
88
+
89
+ def test_uneven_split(self):
90
+ # 10 experts, EP=3 → ranks get 4, 3, 3
91
+ ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
92
+ ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
93
+ ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
94
+
95
+ assert len(ids_0) == 4
96
+ assert len(ids_1) == 3
97
+ assert len(ids_2) == 3
98
+ # All experts covered, no overlap
99
+ assert ids_0 | ids_1 | ids_2 == set(range(10))
100
+ assert ids_0.isdisjoint(ids_1)
101
+ assert ids_1.isdisjoint(ids_2)
102
+
103
+ def test_384_experts_ep8(self):
104
+ # Kimi-K2.5 config: 384 experts, EP=8
105
+ for rank in range(8):
106
+ ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
107
+ assert len(ids) == 48
108
+
109
+ # All experts covered
110
+ all_ids = set()
111
+ for rank in range(8):
112
+ ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
113
+ all_ids |= ids
114
+ assert all_ids == set(range(384))
115
+
116
+ def test_384_experts_ep16(self):
117
+ for rank in range(16):
118
+ ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
119
+ assert len(ids) == 24
120
+
121
+ def test_384_experts_ep24(self):
122
+ # 384 / 24 = 16 exactly
123
+ for rank in range(24):
124
+ ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
125
+ assert len(ids) == 16
126
+
127
+ # round_robin placement tests
128
+
129
+ def test_round_robin_basic(self):
130
+ # 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
131
+ rr = "round_robin"
132
+ ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
133
+ ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
134
+ assert ids_0 == {0, 2, 4, 6}
135
+ assert ids_1 == {1, 3, 5, 7}
136
+
137
+ def test_round_robin_full_coverage(self):
138
+ # 384 experts, EP=8: all experts covered, no overlap
139
+ rr = "round_robin"
140
+ all_ids: set[int] = set()
141
+ for rank in range(8):
142
+ ids = compute_local_expert_ids(384, 8, rank, placement=rr)
143
+ assert ids is not None and len(ids) == 48
144
+ assert all_ids.isdisjoint(ids)
145
+ all_ids |= ids
146
+ assert all_ids == set(range(384))
147
+
148
+ def test_round_robin_uneven(self):
149
+ # 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
150
+ rr = "round_robin"
151
+ ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
152
+ ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
153
+ ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
154
+ assert ids_0 == {0, 3, 6, 9}
155
+ assert ids_1 == {1, 4, 7}
156
+ assert ids_2 == {2, 5, 8}
157
+ assert ids_0 | ids_1 | ids_2 == set(range(10))
158
+
159
+
160
+# ---------------------------------------------------------------------------
161
+# Unit tests for should_skip_weight
162
+# ---------------------------------------------------------------------------
163
+
164
+
165
+class TestShouldSkipWeight:
166
+ def setup_method(self):
167
+ # Simulate EP=8, rank=0 → experts 0-47
168
+ self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
169
+
170
+ def test_no_filter(self):
171
+ assert not should_skip_weight("anything", None)
172
+
173
+ def test_dense_not_skipped(self):
174
+ assert not should_skip_weight(
175
+ "model.layers.0.self_attn.q_proj.weight", self.local_ids
176
+ )
177
+
178
+ def test_local_expert_not_skipped(self):
179
+ assert not should_skip_weight(
180
+ "model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
181
+ )
182
+
183
+ def test_remote_expert_skipped(self):
184
+ assert should_skip_weight(
185
+ "model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
186
+ )
187
+
188
+ def test_boundary_expert(self):
189
+ # Expert 47 is local (last one), 48 is not
190
+ assert not should_skip_weight(
191
+ "model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
192
+ )
193
+ assert should_skip_weight(
194
+ "model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
195
+ )
196
+
197
+ def test_shared_expert_not_skipped(self):
198
+ assert not should_skip_weight(
199
+ "model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
200
+ )
201
+
202
+ def test_embedding_not_skipped(self):
203
+ assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
204
+
205
+ def test_fused_3d_expert_not_skipped(self):
206
+ # 3D fused-expert tensors (gpt-oss style) have no numeric id.
207
+ # Must not be skipped — weight_loader handles slicing later.
208
+ assert not should_skip_weight(
209
+ "model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
210
+ )
211
+
212
+
213
+# ---------------------------------------------------------------------------
214
+# Integration test: safetensors_weights_iterator with EP filtering
215
+# ---------------------------------------------------------------------------
216
+
217
+
218
+class TestSafetensorsWeightsIteratorWithEpFilter:
219
+ """Verify that EP filtering produces a strict subset of unfiltered loading
220
+ and that all expected dense + local expert weights are present."""
221
+
222
+ @pytest.fixture(scope="class")
223
+ def gpt2_files(self):
224
+ """Download GPT-2 safetensors to a temp dir (shared across class)."""
225
+ with tempfile.TemporaryDirectory() as tmpdir:
226
+ huggingface_hub.constants.HF_HUB_OFFLINE = False
227
+ from vllm.model_executor.model_loader.weight_utils import (
228
+ download_weights_from_hf,
229
+ )
230
+
231
+ download_weights_from_hf(
232
+ "openai-community/gpt2",
233
+ allow_patterns=["*.safetensors"],
234
+ cache_dir=tmpdir,
235
+ )
236
+ files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
237
+ assert len(files) > 0
238
+ yield files
239
+
240
+ def test_no_filter_returns_all(self, gpt2_files):
241
+ """With local_expert_ids=None, all weights are returned (no MoE)."""
242
+ all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
243
+ filtered_weights = dict(
244
+ safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
245
+ )
246
+ assert set(all_weights.keys()) == set(filtered_weights.keys())
247
+
248
+ def test_empty_filter_skips_experts_only(self, gpt2_files):
249
+ """GPT-2 has no expert weights, so even an empty local_expert_ids
250
+ set should return all weights (all are dense)."""
251
+ all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
252
+ filtered_weights = dict(
253
+ safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
254
+ )
255
+ # GPT-2 has no experts, so nothing should be filtered
256
+ assert set(all_weights.keys()) == set(filtered_weights.keys())
257
+
258
+
259
+class TestEpFilterOnSyntheticMoeWeights:
260
+ """Create synthetic safetensors files with expert-like naming and verify
261
+ that the filter correctly skips non-local experts."""
262
+
263
+ @pytest.fixture
264
+ def synthetic_moe_files(self, tmp_path):
265
+ """Create synthetic safetensors with expert-patterned tensor names."""
266
+ from safetensors.torch import save_file
267
+
268
+ tensors = {}
269
+ # Dense weights
270
+ tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
271
+ tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
272
+ tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
273
+ # Expert weights: 8 experts
274
+ for expert_id in range(8):
275
+ tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
276
+ torch.randn(128, 64)
277
+ )
278
+ tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
279
+ torch.randn(128, 64)
280
+ )
281
+ tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
282
+ torch.randn(64, 128)
283
+ )
284
+ # Shared expert (should never be filtered)
285
+ tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
286
+ 128, 64
287
+ )
288
+
289
+ filepath = str(tmp_path / "model-00001-of-00001.safetensors")
290
+ save_file(tensors, filepath)
291
+ return [filepath], tensors
292
+
293
+ def test_no_filter_returns_all(self, synthetic_moe_files):
294
+ files, expected = synthetic_moe_files
295
+ loaded = dict(safetensors_weights_iterator(files, False))
296
+ assert set(loaded.keys()) == set(expected.keys())
297
+
298
+ def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
299
+ files, expected = synthetic_moe_files
300
+ # EP=2, rank=0 → experts 0-3
301
+ local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
302
+ loaded = dict(
303
+ safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
304
+ )
305
+
306
+ # Should have all dense + shared + experts 0-3 only
307
+ for name in loaded:
308
+ eid = parse_expert_id(name)
309
+ if eid is not None:
310
+ assert eid in local_ids, f"Non-local expert {eid} was loaded"
311
+
312
+ # Check expert count: 4 experts × 3 weights = 12
313
+ expert_names = [n for n in loaded if parse_expert_id(n) is not None]
314
+ assert len(expert_names) == 4 * 3
315
+
316
+ # Check all dense weights present
317
+ assert "model.embed_tokens.weight" in loaded
318
+ assert "model.layers.0.self_attn.q_proj.weight" in loaded
319
+ assert "model.layers.0.input_layernorm.weight" in loaded
320
+ assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
321
+
322
+ def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
323
+ files, expected = synthetic_moe_files
324
+ local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
325
+ loaded = dict(
326
+ safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
327
+ )
328
+
329
+ expert_names = [n for n in loaded if parse_expert_id(n) is not None]
330
+ assert len(expert_names) == 4 * 3
331
+ for name in expert_names:
332
+ assert parse_expert_id(name) in local_ids
333
+
334
+ def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
335
+ files, _ = synthetic_moe_files
336
+ all_expert_names = set()
337
+ for rank in range(8):
338
+ local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
339
+ loaded = dict(
340
+ safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
341
+ )
342
+ expert_names = {n for n in loaded if parse_expert_id(n) is not None}
343
+ # 1 expert × 3 weights
344
+ assert len(expert_names) == 3
345
+ all_expert_names |= expert_names
346
+
347
+ # All 8 experts × 3 weights covered across ranks
348
+ assert len(all_expert_names) == 24
349
+
350
+ def test_tensor_values_match(self, synthetic_moe_files):
351
+ """Filtered tensors have identical values to unfiltered ones."""
352
+ files, _ = synthetic_moe_files
353
+ all_weights = dict(safetensors_weights_iterator(files, False))
354
+
355
+ local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
356
+ filtered = dict(
357
+ safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
358
+ )
359
+
360
+ for name, tensor in filtered.items():
361
+ assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"
vllm/model_executor/model_loader/default_loader.py
+59
@@ -16,6 +16,9 @@ from vllm.config.load import LoadConfig
16
from vllm.logger import init_logger
17
from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
18
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
19
+from vllm.model_executor.model_loader.ep_weight_filter import (
20
+ compute_local_expert_ids,
21
+)
22
from vllm.model_executor.model_loader.weight_utils import (
23
download_safetensors_index_file_from_hf,
24
download_weights_from_hf,
@@ -70,6 +73,7 @@ class DefaultModelLoader(BaseModelLoader):
73
74
def __init__(self, load_config: LoadConfig):
75
super().__init__(load_config)
76
+ self.local_expert_ids: set[int] | None = None
77
78
extra_config = load_config.model_loader_extra_config
79
allowed_keys = {"enable_multithread_load", "num_threads"}
@@ -243,6 +247,7 @@ class DefaultModelLoader(BaseModelLoader):
247
hf_weights_files,
248
self.load_config.use_tqdm_on_load,
249
self.load_config.safetensors_load_strategy,
250
+ local_expert_ids=self.local_expert_ids,
251
)
252
else:
253
if extra_config.get("enable_multithread_load"):
@@ -296,6 +301,58 @@ class DefaultModelLoader(BaseModelLoader):
301
allow_patterns_overrides=None,
302
)
303
304
+ def _init_ep_weight_filter(self, model_config: ModelConfig) -> None:
305
+ """Compute local expert ids for EP weight filtering.
306
+
307
+ When expert parallelism is active, each rank only needs a subset of
308
+ expert weights. By computing the set upfront we can skip non-local
309
+ expert tensors *before* reading them from disk.
310
+ """
311
+ from vllm.config import get_current_vllm_config
312
+
313
+ vllm_config = get_current_vllm_config()
314
+ parallel_config = vllm_config.parallel_config
315
+
316
+ if not (model_config.is_moe and parallel_config.enable_expert_parallel):
317
+ return
318
+
319
+ num_experts = model_config.get_num_experts()
320
+ if num_experts <= 0:
321
+ return
322
+
323
+ # EP size/rank computation mirrors FusedMoEParallelConfig.make():
324
+ # ep_size = dp_size * pcp_size * tp_size (flattened)
325
+ # ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
326
+ from vllm.distributed import (
327
+ get_dp_group,
328
+ get_pcp_group,
329
+ get_tensor_model_parallel_rank,
330
+ )
331
+
332
+ dp_size = parallel_config.data_parallel_size
333
+ tp_size = parallel_config.tensor_parallel_size
334
+ pcp_size = parallel_config.prefill_context_parallel_size
335
+ dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
336
+ tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
337
+ pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
338
+ ep_size = dp_size * pcp_size * tp_size
339
+ ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
340
+
341
+ self.local_expert_ids = compute_local_expert_ids(
342
+ num_experts,
343
+ ep_size,
344
+ ep_rank,
345
+ placement=parallel_config.expert_placement_strategy,
346
+ )
347
+ if self.local_expert_ids is not None:
348
+ logger.info_once(
349
+ "EP weight filter: ep_size=%d, ep_rank=%d, loading %d/%d experts",
350
+ ep_size,
351
+ ep_rank,
352
+ len(self.local_expert_ids),
353
+ num_experts,
354
+ )
355
+
356
@instrument(span_name="Load weights")
357
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
358
if model_config.quantization == "torchao":
@@ -307,6 +364,8 @@ class DefaultModelLoader(BaseModelLoader):
364
):
365
self.load_config.safetensors_load_strategy = "torchao"
366
367
+ self._init_ep_weight_filter(model_config)
368
+
369
weights_to_load = {name for name, _ in model.named_parameters()}
370
loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
371
vllm/model_executor/model_loader/ep_weight_filter.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Filter out non-local expert weights during loading to avoid redundant I/O.
4
+
5
+In DP+EP deployments each rank only needs its own expert shard. Skipping
6
+non-local expert tensors *before* they are read from disk eliminates the
7
+majority of storage I/O for MoE models (experts typically account for
8
+~85-90 % of total weight bytes).
9
+"""
10
+
11
+import regex as re
12
+
13
+# Matches per-expert weight names like ".experts.42.gate_proj.weight".
14
+# Does NOT match 3D fused-expert names like ".experts.gate_proj.weight"
15
+# (no numeric id) — those are intentionally left unfiltered so the full
16
+# tensor is loaded and sliced later by FusedMoE.weight_loader.
17
+_EXPERT_ID_RE = re.compile(r"\.experts\.(\d+)\.")
18
+
19
+
20
+def parse_expert_id(weight_name: str) -> int | None:
21
+ """Return the expert id embedded in *weight_name*, or ``None`` if it is
22
+ not an per-expert weight.
23
+
24
+ Returns ``None`` for dense weights (attention, layernorm, embedding),
25
+ shared experts, and 3D fused-expert tensors where all experts are stored
26
+ in a single tensor without a numeric expert id in the name."""
27
+ m = _EXPERT_ID_RE.search(weight_name)
28
+ return int(m.group(1)) if m else None
29
+
30
+
31
+def compute_local_expert_ids(
32
+ num_experts: int,
33
+ ep_size: int,
34
+ ep_rank: int,
35
+ placement: str = "linear",
36
+) -> set[int] | None:
37
+ """Compute the set of global expert ids owned by *ep_rank*.
38
+
39
+ Returns ``None`` when EP is not active (``ep_size <= 1``), meaning all
40
+ experts are local and no filtering should be performed.
41
+
42
+ The distribution logic mirrors
43
+ :func:`vllm.model_executor.layers.fused_moe.layer.determine_expert_map`.
44
+
45
+ Args:
46
+ placement: ``"linear"`` for contiguous assignment,
47
+ ``"round_robin"`` for interleaved assignment.
48
+ """
49
+ if ep_size <= 1:
50
+ return None
51
+
52
+ if placement == "linear":
53
+ base = num_experts // ep_size
54
+ remainder = num_experts % ep_size
55
+ start = ep_rank * base + min(ep_rank, remainder)
56
+ local_count = base + (1 if ep_rank < remainder else 0)
57
+ return set(range(start, start + local_count))
58
+ elif placement == "round_robin":
59
+ return set(range(ep_rank, num_experts, ep_size))
60
+ else:
61
+ raise ValueError(f"Unknown expert placement strategy: {placement}")
62
+
63
+
64
+def should_skip_weight(
65
+ weight_name: str,
66
+ local_expert_ids: set[int] | None,
67
+) -> bool:
68
+ """Return ``True`` if *weight_name* is an expert weight that does not
69
+ belong to the local rank and should be skipped during loading."""
70
+ if local_expert_ids is None:
71
+ return False
72
+ eid = parse_expert_id(weight_name)
73
+ if eid is None:
74
+ # Not an expert weight (dense / shared-expert / embedding) → keep.
75
+ return False
76
+ return eid not in local_expert_ids
vllm/model_executor/model_loader/weight_utils.py
+17
-2
@@ -35,6 +35,9 @@ from vllm.model_executor.layers.quantization import (
35
QuantizationConfig,
36
get_quantization_config,
37
)
38
+from vllm.model_executor.model_loader.ep_weight_filter import (
39
+ should_skip_weight,
40
+)
41
from vllm.platforms import current_platform
42
from vllm.tracing import instrument
43
from vllm.utils.import_utils import PlaceholderModule
@@ -721,8 +724,14 @@ def safetensors_weights_iterator(
724
hf_weights_files: list[str],
725
use_tqdm_on_load: bool,
726
safetensors_load_strategy: str = "lazy",
727
+ local_expert_ids: set[int] | None = None,
728
) -> Generator[tuple[str, torch.Tensor], None, None]:
725
- """Iterate over the weights in the model safetensor files."""
729
+ """Iterate over the weights in the model safetensor files.
730
+
731
+ When *local_expert_ids* is provided, expert weights not belonging to
732
+ this rank are skipped **before** reading from disk, which drastically
733
+ reduces storage I/O for MoE models under EP.
734
+ """
735
loading_desc = "Loading safetensors checkpoint shards"
736
if safetensors_load_strategy == "eager":
737
loading_desc += " (eager)"
@@ -737,7 +746,9 @@ def safetensors_weights_iterator(
746
if safetensors_load_strategy == "eager":
747
with open(st_file, "rb") as f:
748
state_dict = load(f.read())
740
- yield from state_dict.items()
749
+ for name, param in state_dict.items():
750
+ if not should_skip_weight(name, local_expert_ids):
751
+ yield name, param
752
elif safetensors_load_strategy == "torchao":
753
# we can't load flattened torchao tensor subclasses directly into the model
754
# instead we reconstruct the subclasses here before returning
@@ -753,6 +764,8 @@ def safetensors_weights_iterator(
764
with safe_open(st_file, framework="pt") as f:
765
state_dict = {}
766
for name in f.keys(): # noqa: SIM118
767
+ if should_skip_weight(name, local_expert_ids):
768
+ continue
769
state_dict[name] = f.get_tensor(name)
770
771
# update with leftover tensor data from previous iteration, if any
@@ -769,6 +782,8 @@ def safetensors_weights_iterator(
782
else:
783
with safe_open(st_file, framework="pt") as f:
784
for name in f.keys(): # noqa: SIM118
785
+ if should_skip_weight(name, local_expert_ids):
786
+ continue
787
param = f.get_tensor(name)
788
yield name, param
789