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, 08:30 UTC
[Misc] fix olmoe model layer can't laod in tp gt 1 (#18828)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
May 29, 2025 at 01:36 UTC
c68b5c63eba0ecccda8ab34ca1829cb4ced2c430
1 file changed
+25
-4
vllm/model_executor/models/olmoe.py
+25
-4
@@ -13,6 +13,7 @@
13
# limitations under the License.
14
"""Inference-only OLMoE model compatible with HuggingFace weights."""
15
from collections.abc import Iterable
16
+from functools import partial
17
from typing import Any, Optional, Union
18
19
import torch
@@ -22,7 +23,10 @@ from transformers import PretrainedConfig
23
from vllm.attention import Attention
24
from vllm.compilation.decorators import support_torch_compile
25
from vllm.config import CacheConfig, VllmConfig
25
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
26
+from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
27
+ get_tensor_model_parallel_world_size,
28
+ tensor_model_parallel_all_gather)
29
+from vllm.distributed.utils import split_tensor_along_last_dim
30
from vllm.logger import init_logger
31
from vllm.model_executor.layers.fused_moe import FusedMoE
32
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -140,8 +144,11 @@ class OlmoeAttention(nn.Module):
144
bias=False,
145
quant_config=quant_config,
146
)
143
- self.q_norm = RMSNorm(hidden_size, eps=1e-5)
144
- self.k_norm = RMSNorm(hidden_size, eps=1e-5)
147
+ self.tp_size = tp_size
148
+ self.tp_rank = get_tensor_model_parallel_rank()
149
+ self.q_norm = RMSNorm(self.total_num_heads * self.head_dim, eps=1e-5)
150
+ self.k_norm = RMSNorm(self.total_num_kv_heads * self.head_dim,
151
+ eps=1e-5)
152
self.o_proj = RowParallelLinear(
153
self.total_num_heads * self.head_dim,
154
hidden_size,
@@ -165,6 +172,20 @@ class OlmoeAttention(nn.Module):
172
quant_config=quant_config,
173
prefix=f"{prefix}.attn")
174
175
+ def _apply_qk_norm(self, q: torch.Tensor,
176
+ k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
177
+ if self.tp_size > 1:
178
+ q = tensor_model_parallel_all_gather(q.contiguous())
179
+ k = tensor_model_parallel_all_gather(k.contiguous())
180
+ q = self.q_norm(q)
181
+ k = self.k_norm(k)
182
+ if self.tp_size > 1:
183
+ splitter = partial(split_tensor_along_last_dim,
184
+ num_partitions=self.tp_size)
185
+ q = splitter(q)[self.tp_rank]
186
+ k = splitter(k)[self.tp_rank]
187
+ return q, k
188
+
189
def forward(
190
self,
191
positions: torch.Tensor,
@@ -172,7 +193,7 @@ class OlmoeAttention(nn.Module):
193
) -> torch.Tensor:
194
qkv, _ = self.qkv_proj(hidden_states)
195
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
175
- q, k = self.q_norm(q.contiguous()), self.k_norm(k.contiguous())
196
+ q, k = self._apply_qk_norm(q, k)
197
q, k = self.rotary_emb(positions, q, k)
198
attn_output = self.attn(q, k, v)
199
output, _ = self.o_proj(attn_output)