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 14, 05:30 UTC
[Model] Add MoE support for NemotronH (#25863)
Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com>
tomeras91 committed
Oct 23, 2025 at 13:27 UTC
61089465a6101790635ed96c26df3e9a57d8d2c9
7 files changed
+413
-39
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -823,6 +823,8 @@ class FusedMoEConfig:
823
824
has_bias: bool = False
825
826
+ is_act_and_mul: bool = True
827
+
828
def __post_init__(self):
829
if self.dp_size > 1:
830
logger.debug_once(
vllm/model_executor/layers/fused_moe/fused_moe.py
+3
-1
@@ -1647,6 +1647,7 @@ def fused_experts(
1647
1648
SILU_NO_MUL: str = activation_without_mul("silu")
1649
GELU_NO_MUL: str = activation_without_mul("gelu")
1650
+RELU2_NO_MUL: str = activation_without_mul("relu2")
1651
1652
1653
def _get_config_quant_dtype(
@@ -1914,7 +1915,8 @@ def fused_experts_impl(
1915
intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))
1916
elif activation == GELU_NO_MUL:
1917
intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))
1917
-
1918
+ elif activation == RELU2_NO_MUL:
1919
+ intermediate_cache2 = torch.square(F.relu(intermediate_cache1.view(-1, N)))
1920
else:
1921
raise ValueError(f"Unsupported FusedMoe activation: {activation}.")
1922
vllm/model_executor/layers/fused_moe/layer.py
+30
-5
@@ -411,11 +411,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
411
params_dtype: torch.dtype,
412
**extra_weight_attrs,
413
):
414
+ if self.moe.is_act_and_mul:
415
+ w13_up_dim = 2 * intermediate_size_per_partition
416
+ else:
417
+ w13_up_dim = intermediate_size_per_partition
418
# Fused gate_up_proj (column parallel)
419
w13_weight = torch.nn.Parameter(
420
torch.empty(
421
num_experts,
418
- 2 * intermediate_size_per_partition,
422
+ w13_up_dim,
423
hidden_size,
424
dtype=params_dtype,
425
),
@@ -425,9 +429,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
429
set_weight_attrs(w13_weight, extra_weight_attrs)
430
if self.moe.has_bias:
431
w13_bias = torch.nn.Parameter(
428
- torch.zeros(
429
- num_experts, 2 * intermediate_size_per_partition, dtype=params_dtype
430
- ),
432
+ torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
433
requires_grad=False,
434
)
435
layer.register_parameter("w13_bias", w13_bias)
@@ -1073,6 +1075,7 @@ class FusedMoE(CustomOp):
1075
e_score_correction_bias: torch.Tensor | None = None,
1076
apply_router_weight_on_input: bool = False,
1077
activation: str = "silu",
1078
+ is_act_and_mul: bool = True,
1079
enable_eplb: bool = False,
1080
num_redundant_experts: int = 0,
1081
has_bias: bool = False,
@@ -1263,6 +1266,7 @@ class FusedMoE(CustomOp):
1266
in_dtype=moe_in_dtype,
1267
max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
1268
has_bias=has_bias,
1269
+ is_act_and_mul=is_act_and_mul,
1270
)
1271
self.moe_config = moe
1272
self.moe_quant_config: FusedMoEQuantConfig | None = None
@@ -1283,6 +1287,24 @@ class FusedMoE(CustomOp):
1287
assert isinstance(quant_method, FusedMoEMethodBase)
1288
self.quant_method = quant_method
1289
1290
+ if not self.moe_config.is_act_and_mul:
1291
+ # Avoid circular import
1292
+ from vllm.model_executor.layers.quantization.modelopt import (
1293
+ ModelOptFp8MoEMethod,
1294
+ )
1295
+
1296
+ if not isinstance(
1297
+ quant_method, (UnquantizedFusedMoEMethod, ModelOptFp8MoEMethod)
1298
+ ):
1299
+ raise NotImplementedError(
1300
+ "is_act_and_mul=False is supported only for unquantized "
1301
+ "and ModelOpt FP8 moe for now"
1302
+ )
1303
+ if not current_platform.is_cuda():
1304
+ raise NotImplementedError(
1305
+ "is_act_and_mul=False is supported only for CUDA for now"
1306
+ )
1307
+
1308
if self.enable_eplb:
1309
from vllm.model_executor.layers.quantization.fp8 import Fp8MoEMethod
1310
@@ -1531,7 +1553,10 @@ class FusedMoE(CustomOp):
1553
):
1554
# Index the loaded weight for tp sharding.
1555
# gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim
1534
- shard_size = expert_data.shape[shard_dim] // 2
1556
+ if self.moe_config.is_act_and_mul:
1557
+ shard_size = expert_data.shape[shard_dim] // 2
1558
+ else:
1559
+ shard_size = expert_data.shape[shard_dim]
1560
if not load_full:
1561
loaded_weight = loaded_weight.narrow(
1562
shard_dim, shard_size * tp_rank, shard_size
vllm/model_executor/layers/quantization/modelopt.py
+26
-5
@@ -354,7 +354,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
354
355
self.cutlass_fp8_supported = cutlass_fp8_supported()
356
self.flashinfer_moe_backend: FlashinferMoeBackend | None = None
357
- if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
357
+ if (
358
+ envs.VLLM_USE_FLASHINFER_MOE_FP8
359
+ and has_flashinfer_moe()
360
+ and self.moe.is_act_and_mul
361
+ ):
362
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
363
logger.info_once(
364
f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
@@ -405,10 +409,15 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
409
)
410
weight_loader = extra_weight_attrs.get("weight_loader")
411
412
+ if self.moe.is_act_and_mul:
413
+ w13_up_dim = 2 * intermediate_size_per_partition
414
+ else:
415
+ w13_up_dim = intermediate_size_per_partition
416
+
417
w13_weight = ModelWeightParameter(
418
data=torch.empty(
419
num_experts,
411
- 2 * intermediate_size_per_partition,
420
+ w13_up_dim,
421
hidden_size,
422
dtype=weight_dtype,
423
),
@@ -433,11 +442,16 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
442
443
if self.quant_config.is_checkpoint_fp8_serialized:
444
# WEIGHT SCALES - Per-tensor scaling for ModelOpts
436
- # Allocate 2 scales for w1 and w3 respectively.
445
+ # For gated MoE, allocate 2 scales for w1 and w3 respectively.
446
# They will be combined to a single scale after weight loading.
447
+ # For non-gated MoE, allocate 1 scale for w13.
448
+ if self.moe.is_act_and_mul:
449
+ w13_weight_scale_shape = (num_experts, 2)
450
+ else:
451
+ w13_weight_scale_shape = (num_experts, 1)
452
w13_weight_scale = PerTensorScaleParameter(
453
data=torch.full(
440
- (num_experts, 2),
454
+ w13_weight_scale_shape,
455
1.0,
456
dtype=torch.float32,
457
),
@@ -485,7 +499,14 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
499
# Fp8 moe kernel needs single weight scale for w13 per expert.
500
# We take the max of the w1 and w3 scales
501
# then dequant and requant each expert.
488
- if layer.w13_weight_scale.dim() == 2:
502
+ if (
503
+ layer.w13_weight_scale.dim() == 2
504
+ and layer.w13_weight_scale.shape[1] == 2
505
+ ):
506
+ assert self.moe.is_act_and_mul, (
507
+ "w13_weight_scale should have 2 elements per expert "
508
+ "only for gated MoE"
509
+ )
510
# Get the maximum scale across w1 and w3 for each expert
511
max_w13_scales = layer.w13_weight_scale.max(dim=1).values
512
vllm/model_executor/models/interfaces.py
+3
-1
@@ -673,7 +673,9 @@ class MixtureOfExperts(Protocol):
673
674
675
def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
676
- return isinstance(model, MixtureOfExperts)
676
+ return (
677
+ isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
678
+ )
679
680
681
@runtime_checkable
vllm/model_executor/models/nemotron_h.py
+329
-27
@@ -18,7 +18,8 @@
18
# limitations under the License.
19
"""Inference-only NemotronH model."""
20
21
-from collections.abc import Iterable
21
+import typing
22
+from collections.abc import Callable, Iterable
23
24
import torch
25
from torch import nn
@@ -26,13 +27,18 @@ from torch import nn
27
from vllm.attention.layer import Attention
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, ModelConfig, VllmConfig
29
-from vllm.distributed import get_tensor_model_parallel_world_size
30
+from vllm.config.parallel import ParallelConfig
31
+from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
32
+from vllm.distributed.communication_op import tensor_model_parallel_all_gather
33
from vllm.distributed.parallel_state import get_pp_group
34
from vllm.model_executor.layers.activation import ReLUSquaredActivation
35
+from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE
36
+from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
37
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import (
39
ColumnParallelLinear,
40
QKVParallelLinear,
41
+ ReplicatedLinear,
42
RowParallelLinear,
43
)
44
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -54,6 +60,7 @@ from vllm.model_executor.model_loader.weight_utils import (
60
from vllm.model_executor.models.interfaces import (
61
HasInnerState,
62
IsHybrid,
63
+ MixtureOfExperts,
64
SupportsLoRA,
65
SupportsPP,
66
SupportsQuant,
@@ -61,9 +68,11 @@ from vllm.model_executor.models.interfaces import (
68
from vllm.model_executor.models.utils import (
69
AutoWeightsLoader,
70
WeightsMapper,
71
+ is_pp_missing_parameter,
72
make_empty_intermediate_tensors_factory,
73
make_layers,
74
maybe_prefix,
75
+ sequence_parallel_chunk,
76
)
77
from vllm.sequence import IntermediateTensors
78
from vllm.transformers_utils.configs import NemotronHConfig
@@ -73,28 +82,21 @@ class NemotronHMLP(nn.Module):
82
def __init__(
83
self,
84
config: NemotronHConfig,
76
- layer_idx: int,
85
+ intermediate_size: int,
86
quant_config: QuantizationConfig | None = None,
87
bias: bool = False,
88
+ reduce_results: bool = True,
89
+ is_sequence_parallel: bool = False,
90
prefix: str = "",
91
) -> None:
92
super().__init__()
93
83
- hybrid_override_pattern = config.hybrid_override_pattern
84
- mlp_index = hybrid_override_pattern[: layer_idx + 1].count("-") - 1
85
- if isinstance(config.intermediate_size, list):
86
- if len(config.intermediate_size) == 1:
87
- intermediate_size = config.intermediate_size[0]
88
- else:
89
- intermediate_size = config.intermediate_size[mlp_index]
90
- else:
91
- intermediate_size = config.intermediate_size
92
-
94
self.up_proj = ColumnParallelLinear(
95
input_size=config.hidden_size,
96
output_size=intermediate_size,
97
bias=bias,
98
quant_config=quant_config,
99
+ disable_tp=is_sequence_parallel,
100
prefix=f"{prefix}.up_proj",
101
)
102
self.down_proj = RowParallelLinear(
@@ -102,6 +104,8 @@ class NemotronHMLP(nn.Module):
104
output_size=config.hidden_size,
105
bias=bias,
106
quant_config=quant_config,
107
+ reduce_results=reduce_results,
108
+ disable_tp=is_sequence_parallel,
109
prefix=f"{prefix}.down_proj",
110
)
111
self.act_fn = ReLUSquaredActivation()
@@ -113,6 +117,130 @@ class NemotronHMLP(nn.Module):
117
return x
118
119
120
+class NemotronHMoE(nn.Module):
121
+ def __init__(
122
+ self,
123
+ config: NemotronHConfig,
124
+ quant_config: QuantizationConfig | None = None,
125
+ parallel_config: ParallelConfig | None = None,
126
+ prefix: str = "",
127
+ ):
128
+ super().__init__()
129
+ self.tp_size = get_tensor_model_parallel_world_size()
130
+ self.routed_scaling_factor = config.routed_scaling_factor
131
+
132
+ self.ep_group = get_ep_group().device_group
133
+ self.ep_rank = self.ep_group.rank()
134
+ self.ep_size = self.ep_group.size()
135
+ self.n_routed_experts: int = config.n_routed_experts
136
+ self.n_shared_experts: int = config.n_shared_experts
137
+
138
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
139
+
140
+ self.gate = ReplicatedLinear(
141
+ config.hidden_size,
142
+ config.n_routed_experts,
143
+ bias=False,
144
+ params_dtype=torch.float32,
145
+ quant_config=None,
146
+ prefix=f"{prefix}.gate",
147
+ )
148
+
149
+ self.gate.e_score_correction_bias = nn.Parameter(
150
+ torch.empty(config.n_routed_experts, dtype=torch.float32)
151
+ )
152
+ # Load balancing settings.
153
+ self.enable_eplb = parallel_config.enable_eplb
154
+
155
+ self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts # noqa: E501
156
+ self.n_logical_experts = self.n_routed_experts
157
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
158
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
159
+
160
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
161
+ self.physical_expert_end = (
162
+ self.physical_expert_start + self.n_local_physical_experts
163
+ )
164
+
165
+ if config.n_shared_experts is None or config.n_shared_experts == 0:
166
+ self.shared_experts = None
167
+ else:
168
+ intermediate_size = (
169
+ config.moe_shared_expert_intermediate_size * config.n_shared_experts
170
+ )
171
+
172
+ self.shared_experts = NemotronHMLP(
173
+ config=config,
174
+ intermediate_size=intermediate_size,
175
+ quant_config=quant_config,
176
+ reduce_results=False,
177
+ is_sequence_parallel=self.is_sequence_parallel,
178
+ prefix=f"{prefix}.shared_experts",
179
+ )
180
+
181
+ self.experts = SharedFusedMoE(
182
+ shared_experts=self.shared_experts,
183
+ num_experts=config.n_routed_experts,
184
+ top_k=config.num_experts_per_tok,
185
+ hidden_size=config.hidden_size,
186
+ intermediate_size=config.moe_intermediate_size,
187
+ reduce_results=False,
188
+ renormalize=config.norm_topk_prob,
189
+ quant_config=quant_config,
190
+ use_grouped_topk=True,
191
+ num_expert_group=config.n_group,
192
+ topk_group=config.topk_group,
193
+ prefix=f"{prefix}.experts",
194
+ scoring_func="sigmoid",
195
+ e_score_correction_bias=self.gate.e_score_correction_bias,
196
+ activation=activation_without_mul(config.mlp_hidden_act),
197
+ is_act_and_mul=False, # non-gated MoE
198
+ enable_eplb=self.enable_eplb,
199
+ num_redundant_experts=self.n_redundant_experts,
200
+ is_sequence_parallel=self.is_sequence_parallel,
201
+ )
202
+
203
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
204
+ num_tokens, hidden_dim = hidden_states.shape
205
+ hidden_states = hidden_states.view(-1, hidden_dim)
206
+
207
+ if self.is_sequence_parallel:
208
+ hidden_states = sequence_parallel_chunk(hidden_states)
209
+
210
+ # router_logits: (num_tokens, n_experts)
211
+ router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
212
+
213
+ fused_moe_out = self.experts(
214
+ hidden_states=hidden_states, router_logits=router_logits
215
+ )
216
+
217
+ shared_output, final_hidden_states = fused_moe_out
218
+
219
+ # Fix FP16 overflow
220
+ # See DeepseekV2DecoderLayer for more details.
221
+ if hidden_states.dtype != torch.float16:
222
+ final_hidden_states *= self.routed_scaling_factor
223
+ elif self.shared_experts is not None:
224
+ assert shared_output is not None
225
+ shared_output *= 1.0 / self.routed_scaling_factor
226
+
227
+ if self.shared_experts is not None:
228
+ assert shared_output is not None
229
+ final_hidden_states += shared_output
230
+
231
+ if self.is_sequence_parallel:
232
+ final_hidden_states = tensor_model_parallel_all_gather(
233
+ final_hidden_states, 0
234
+ )
235
+ final_hidden_states = final_hidden_states[:num_tokens]
236
+ elif self.tp_size > 1:
237
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
238
+ final_hidden_states
239
+ )
240
+
241
+ return final_hidden_states.view(num_tokens, hidden_dim)
242
+
243
+
244
class NemotronHMLPDecoderLayer(nn.Module):
245
def __init__(
246
self,
@@ -121,20 +249,70 @@ class NemotronHMLPDecoderLayer(nn.Module):
249
model_config: ModelConfig | None = None,
250
cache_config: CacheConfig | None = None,
251
quant_config: QuantizationConfig | None = None,
252
+ parallel_config: ParallelConfig | None = None,
253
prefix: str = "",
254
) -> None:
255
super().__init__()
256
self.config = config
257
258
+ hybrid_override_pattern = config.hybrid_override_pattern
259
+ mlp_index = hybrid_override_pattern[: layer_idx + 1].count("-") - 1
260
+ if isinstance(config.intermediate_size, list):
261
+ if len(config.intermediate_size) == 1:
262
+ intermediate_size = config.intermediate_size[0]
263
+ else:
264
+ intermediate_size = config.intermediate_size[mlp_index]
265
+ else:
266
+ intermediate_size = config.intermediate_size
267
+
268
self.mixer = NemotronHMLP(
269
config,
270
+ intermediate_size=intermediate_size,
271
quant_config=quant_config,
272
bias=config.mlp_bias,
273
prefix=f"{prefix}.mixer",
134
- layer_idx=layer_idx,
274
)
275
137
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
276
+ self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
277
+
278
+ def forward(
279
+ self,
280
+ hidden_states: torch.Tensor,
281
+ residual: torch.Tensor | None,
282
+ **kwargs,
283
+ ):
284
+ if residual is None:
285
+ residual = hidden_states
286
+ hidden_states = self.norm(hidden_states)
287
+ else:
288
+ hidden_states, residual = self.norm(hidden_states, residual)
289
+
290
+ hidden_states = self.mixer(hidden_states)
291
+ return hidden_states, residual
292
+
293
+
294
+class NemotronHMoEDecoderLayer(nn.Module):
295
+ def __init__(
296
+ self,
297
+ config: NemotronHConfig,
298
+ layer_idx: int,
299
+ model_config: ModelConfig | None = None,
300
+ cache_config: CacheConfig | None = None,
301
+ quant_config: QuantizationConfig | None = None,
302
+ parallel_config: ParallelConfig | None = None,
303
+ prefix: str = "",
304
+ ) -> None:
305
+ super().__init__()
306
+ self.config = config
307
+
308
+ self.mixer = NemotronHMoE(
309
+ config,
310
+ quant_config=quant_config,
311
+ parallel_config=parallel_config,
312
+ prefix=f"{prefix}.mixer",
313
+ )
314
+
315
+ self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
316
317
def forward(
318
self,
@@ -160,6 +338,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
338
model_config: ModelConfig | None = None,
339
cache_config: CacheConfig | None = None,
340
quant_config: QuantizationConfig | None = None,
341
+ parallel_config: ParallelConfig | None = None,
342
prefix: str = "",
343
) -> None:
344
super().__init__()
@@ -174,7 +353,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
353
n_groups=config.n_groups,
354
num_heads=config.mamba_num_heads,
355
head_dim=config.mamba_head_dim,
177
- rms_norm_eps=config.rms_norm_eps,
356
+ rms_norm_eps=config.layer_norm_epsilon,
357
activation=config.mamba_hidden_act,
358
model_config=model_config,
359
cache_config=cache_config,
@@ -182,7 +361,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
361
prefix=f"{prefix}.mixer",
362
)
363
185
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
364
+ self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
365
366
def forward(
367
self,
@@ -281,6 +460,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):
460
model_config: ModelConfig | None = None,
461
cache_config: CacheConfig | None = None,
462
quant_config: QuantizationConfig | None = None,
463
+ parallel_config: ParallelConfig | None = None,
464
prefix: str = "",
465
) -> None:
466
super().__init__()
@@ -294,7 +474,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):
474
prefix=f"{prefix}.mixer",
475
)
476
297
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
477
+ self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
478
479
def forward(
480
self,
@@ -317,6 +497,7 @@ ALL_DECODER_LAYER_TYPES = {
497
"M": NemotronHMambaDecoderLayer,
498
"-": NemotronHMLPDecoderLayer,
499
"*": NemotronHAttentionDecoderLayer,
500
+ "E": NemotronHMoEDecoderLayer,
501
}
502
503
@@ -329,6 +510,7 @@ class NemotronHModel(nn.Module):
510
model_config = vllm_config.model_config
511
cache_config = vllm_config.cache_config
512
quant_config = vllm_config.quant_config
513
+ parallel_config = vllm_config.parallel_config
514
lora_config = vllm_config.lora_config
515
516
self.config = config
@@ -346,17 +528,20 @@ class NemotronHModel(nn.Module):
528
org_num_embeddings=config.vocab_size,
529
)
530
531
+ self.has_moe = "E" in config.hybrid_override_pattern
532
+
533
def get_layer(prefix: str):
534
layer_idx = int(prefix.rsplit(".", 1)[1])
535
layer_class = ALL_DECODER_LAYER_TYPES[
536
config.hybrid_override_pattern[layer_idx]
537
]
538
return layer_class(
355
- config,
356
- layer_idx,
357
- model_config,
358
- cache_config,
539
+ config=config,
540
+ layer_idx=layer_idx,
541
+ model_config=model_config,
542
+ cache_config=cache_config,
543
quant_config=quant_config,
544
+ parallel_config=parallel_config,
545
prefix=prefix,
546
)
547
@@ -367,7 +552,7 @@ class NemotronHModel(nn.Module):
552
["hidden_states", "residual"], config.hidden_size
553
)
554
370
- self.norm_f = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
555
+ self.norm_f = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
556
557
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
558
return self.embed_tokens(input_ids)
@@ -413,6 +598,22 @@ class NemotronHModel(nn.Module):
598
("qkv_proj", "v_proj", "v"),
599
]
600
601
+ if self.has_moe:
602
+ # (param_name, weight_name, expert_id, shard_id)
603
+ expert_params_mapping = FusedMoE.make_expert_params_mapping(
604
+ # - FusedMoe.w1 (aka gate_proj) should be up_proj since that's
605
+ # what the activation is applied to
606
+ # - FusedMoe.w3 (aka up_proj) should be ignored since we're
607
+ # using non-gated MoE
608
+ ckpt_gate_proj_name="up_proj",
609
+ ckpt_down_proj_name="down_proj",
610
+ ckpt_up_proj_name="",
611
+ num_experts=self.config.n_routed_experts,
612
+ num_redundant_experts=getattr(self, "num_redundant_experts", 0),
613
+ )
614
+ else:
615
+ expert_params_mapping = []
616
+
617
params_dict = dict(self.named_parameters())
618
loaded_params: set[str] = set()
619
for name, loaded_weight in weights:
@@ -438,16 +639,62 @@ class NemotronHModel(nn.Module):
639
640
# load other params
641
else:
441
- param = params_dict[name]
442
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
443
- weight_loader(param, loaded_weight)
642
+ is_expert_weight = False
643
+ for mapping in expert_params_mapping:
644
+ param_name, weight_name, expert_id, shard_id = mapping
645
+ if weight_name not in name:
646
+ continue
647
+
648
+ # Anyway, this is an expert weight and should not be
649
+ # attempted to load as other weights later
650
+ is_expert_weight = True
651
+
652
+ # Do not modify `name` since the loop may continue here
653
+ # Instead, create a new variable
654
+ name_mapped = name.replace(weight_name, param_name)
655
+
656
+ if is_pp_missing_parameter(name_mapped, self):
657
+ continue
658
+ param = params_dict[name_mapped]
659
+ # We should ask the weight loader to return success or not
660
+ # here since otherwise we may skip experts with other
661
+ # available replicas.
662
+ weight_loader = typing.cast(
663
+ Callable[..., bool], param.weight_loader
664
+ )
665
+ success = weight_loader(
666
+ param,
667
+ loaded_weight,
668
+ name_mapped,
669
+ shard_id=shard_id,
670
+ expert_id=expert_id,
671
+ return_success=True,
672
+ )
673
+ if success:
674
+ name = name_mapped
675
+ break
676
+ else:
677
+ if is_expert_weight:
678
+ continue
679
+
680
+ param = params_dict[name]
681
+ weight_loader = getattr(
682
+ param, "weight_loader", default_weight_loader
683
+ )
684
+ weight_loader(param, loaded_weight)
685
686
loaded_params.add(name)
687
return loaded_params
688
689
690
class NemotronHForCausalLM(
450
- nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid, SupportsQuant
691
+ nn.Module,
692
+ HasInnerState,
693
+ SupportsLoRA,
694
+ SupportsPP,
695
+ IsHybrid,
696
+ SupportsQuant,
697
+ MixtureOfExperts,
698
):
699
hf_to_vllm_mapper = WeightsMapper(
700
orig_to_new_prefix={"backbone": "model"},
@@ -545,6 +792,61 @@ class NemotronHForCausalLM(
792
793
self.make_empty_intmd_tensors = self.model.make_empty_intmd_tensors
794
795
+ # Set MoE hyperparameters
796
+ if self.model.has_moe:
797
+ self.expert_weights = []
798
+ self.num_expert_groups = config.n_group
799
+
800
+ self.moe_layers: list[SharedFusedMoE] = []
801
+ example_moe = None
802
+ for layer in self.model.layers:
803
+ if isinstance(layer, NemotronHMoEDecoderLayer):
804
+ # Pick last one layer since the first ones
805
+ # may be dense layers.
806
+ example_moe = layer.mixer
807
+ self.moe_layers.append(layer.mixer.experts)
808
+
809
+ self.num_moe_layers = len(self.moe_layers)
810
+ self.num_logical_experts = example_moe.n_logical_experts
811
+ self.num_physical_experts = example_moe.n_physical_experts
812
+ self.num_local_physical_experts = example_moe.n_local_physical_experts # noqa: E501
813
+ self.num_routed_experts = example_moe.n_routed_experts
814
+ self.num_shared_experts = example_moe.n_shared_experts
815
+ self.num_redundant_experts = example_moe.n_redundant_experts
816
+
817
+ def set_eplb_state(
818
+ self,
819
+ expert_load_view: torch.Tensor,
820
+ logical_to_physical_map: torch.Tensor,
821
+ logical_replica_count: torch.Tensor,
822
+ ) -> None:
823
+ for layer_idx, layer in enumerate(self.moe_layers):
824
+ # Register the expert weights.
825
+ self.expert_weights.append(layer.get_expert_weights())
826
+ layer.set_eplb_state(
827
+ moe_layer_idx=layer_idx,
828
+ expert_load_view=expert_load_view,
829
+ logical_to_physical_map=logical_to_physical_map,
830
+ logical_replica_count=logical_replica_count,
831
+ )
832
+
833
+ def update_physical_experts_metadata(
834
+ self,
835
+ num_physical_experts: int,
836
+ num_local_physical_experts: int,
837
+ ) -> None:
838
+ assert self.num_local_physical_experts == num_local_physical_experts
839
+ self.num_physical_experts = num_physical_experts
840
+ self.num_local_physical_experts = num_local_physical_experts
841
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
842
+ for layer in self.model.layers:
843
+ if isinstance(layer, NemotronHMoEDecoderLayer):
844
+ moe = layer.mixer
845
+ moe.n_local_physical_experts = num_local_physical_experts
846
+ moe.n_physical_experts = num_physical_experts
847
+ moe.n_redundant_experts = self.num_redundant_experts
848
+ moe.experts.update_expert_map()
849
+
850
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
851
return self.model.get_input_embeddings(input_ids)
852
vllm/transformers_utils/configs/nemotron_h.py
+20
@@ -185,6 +185,15 @@ class NemotronHConfig(PretrainedConfig):
185
mamba_proj_bias=False,
186
mamba_chunk_size=256,
187
rescale_prenorm_residual=True,
188
+ n_routed_experts=8,
189
+ n_shared_experts=1,
190
+ moe_intermediate_size=7688,
191
+ moe_shared_expert_intermediate_size=7688,
192
+ num_experts_per_tok=2,
193
+ routed_scaling_factor=1.0,
194
+ n_group=1,
195
+ topk_group=1,
196
+ norm_topk_prob=True,
197
**kwargs,
198
):
199
self.vocab_size = vocab_size
@@ -241,6 +250,15 @@ class NemotronHConfig(PretrainedConfig):
250
self.mamba_proj_bias = mamba_proj_bias
251
self.chunk_size = mamba_chunk_size
252
self.rescale_prenorm_residual = rescale_prenorm_residual
253
+ self.n_routed_experts = n_routed_experts
254
+ self.n_shared_experts = n_shared_experts
255
+ self.moe_intermediate_size = moe_intermediate_size
256
+ self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size # noqa: E501
257
+ self.num_experts_per_tok = num_experts_per_tok
258
+ self.routed_scaling_factor = routed_scaling_factor
259
+ self.n_group = n_group
260
+ self.topk_group = topk_group
261
+ self.norm_topk_prob = norm_topk_prob
262
263
super().__init__(
264
pad_token_id=pad_token_id,
@@ -258,5 +276,7 @@ class NemotronHConfig(PretrainedConfig):
276
else "attention"
277
if self.hybrid_override_pattern[i] == "*"
278
else "mlp"
279
+ if self.hybrid_override_pattern[i] == "-"
280
+ else "moe"
281
for i in range(self.num_hidden_layers)
282
]