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 15, 07:15 UTC
[Model] Optimize nemotron_h implementation (#19249)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Jun 6, 2025 at 18:05 UTC
7661e92ef85e552936195ae4b803e292b9a96776
1 file changed
+16
-8
vllm/model_executor/models/nemotron_h.py
+16
-8
@@ -1,4 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
# Adapted from https://github.com/vllm-project/vllm/blob/94d8ec8d2bcb4ec55e33022b313c7e978edf05e1/vllm/model_executor/models/bamba.py
5
# Copyright 2024 HuggingFace Inc. team. All rights reserved.
@@ -29,7 +30,7 @@ from vllm.distributed.parallel_state import get_pp_group
30
from vllm.forward_context import get_forward_context
31
from vllm.model_executor.layers.activation import ReLUSquaredActivation
32
from vllm.model_executor.layers.layernorm import RMSNorm
32
-from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
33
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
34
QKVParallelLinear,
35
RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -63,19 +64,22 @@ class NemotronHMLP(nn.Module):
64
config: NemotronHConfig,
65
quant_config: Optional[QuantizationConfig] = None,
66
bias: bool = False,
67
+ prefix: str = "",
68
) -> None:
69
super().__init__()
68
- self.up_proj = MergedColumnParallelLinear(
70
+ self.up_proj = ColumnParallelLinear(
71
input_size=config.hidden_size,
70
- output_sizes=[config.intermediate_size],
72
+ output_size=config.intermediate_size,
73
bias=bias,
74
quant_config=quant_config,
75
+ prefix=f"{prefix}.up_proj",
76
)
77
self.down_proj = RowParallelLinear(
78
input_size=config.intermediate_size,
79
output_size=config.hidden_size,
80
bias=bias,
81
quant_config=quant_config,
82
+ prefix=f"{prefix}.down_proj",
83
)
84
self.act_fn = ReLUSquaredActivation()
85
@@ -99,9 +103,12 @@ class NemotronHMLPDecoderLayer(nn.Module):
103
super().__init__()
104
self.config = config
105
102
- self.mixer = NemotronHMLP(config,
103
- quant_config=quant_config,
104
- bias=config.mlp_bias)
106
+ self.mixer = NemotronHMLP(
107
+ config,
108
+ quant_config=quant_config,
109
+ bias=config.mlp_bias,
110
+ prefix=f"{prefix}.mixer",
111
+ )
112
113
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
114
@@ -207,12 +214,14 @@ class NemotronHAttention(nn.Module):
214
self.total_num_kv_heads,
215
bias=False,
216
quant_config=quant_config,
217
+ prefix=f"{prefix}.qkv_proj",
218
)
219
self.o_proj = RowParallelLinear(
220
self.total_num_heads * self.head_dim,
221
config.hidden_size,
222
bias=False,
223
quant_config=quant_config,
224
+ prefix=f"{prefix}.o_proj",
225
)
226
227
self.attn = Attention(
@@ -253,7 +262,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):
262
layer_idx,
263
cache_config,
264
quant_config,
256
- prefix,
265
+ prefix=f"{prefix}.mixer",
266
)
267
268
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -435,7 +444,6 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
444
"k_proj",
445
"v_proj",
446
],
438
- "gate_up_proj": ["up_proj", "down_proj"]
447
}
448
449
# LoRA specific attributes