@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1e15aaef56

Fork 0

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, 22:15 UTC

[Bugfix][Quantization] Fix FP8 + EP (#13784)

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Tyler Michael Smith committed Feb 24, 2025 at 21:54 UTC 1e15aaef562de78b4684e88ed56bc18b3860e724

6 files changed +22 -22

vllm/model_executor/layers/fused_moe/layer.py

+15 -15

@@ -260,7 +260,7 @@ class FusedMoE(torch.nn.Module):

260

261 def __init__(

262 self,

263 - num_experts: int,

263 + num_experts: int, # Global number of experts

264 top_k: int,

265 hidden_size: int,

266 intermediate_size: int,

@@ -291,7 +291,8 @@ class FusedMoE(torch.nn.Module):

291 else:

292 self.ep_size = 1

293 self.top_k = top_k

294 - self.num_experts = num_experts # Global number of experts

294 + self.global_num_experts = num_experts

295 + self.local_num_experts = self.global_num_experts // self.ep_size

296 assert intermediate_size % self.tp_size == 0

297 self.intermediate_size_per_partition = intermediate_size // self.tp_size

298 self.reduce_results = reduce_results

@@ -308,27 +309,29 @@ class FusedMoE(torch.nn.Module):

309

310 if self.ep_size > 1:

311 # Create a tensor of size num_experts filled with -1

311 - self.expert_map = torch.full((self.num_experts, ),

312 + self.expert_map = torch.full((self.global_num_experts, ),

313 -1,

314 dtype=torch.int32)

315 # Create a expert map for the local experts

315 - local_num_experts = num_experts // self.ep_size

316 ep_rank = get_tensor_model_parallel_rank()

317 if ep_rank < (self.ep_size - 1):

318 # Each non-last rank gets local_num_experts experts.

319 - self.expert_map[ep_rank * local_num_experts:

320 - (ep_rank + 1) * local_num_experts] = \

321 - torch.arange(0, local_num_experts, dtype=torch.int32)

319 + self.expert_map[ep_rank * self.local_num_experts:

320 + (ep_rank + 1) * self.local_num_experts] = \

321 + torch.arange(0, self.local_num_experts, dtype=torch.int32)

322 else:

323 # All remaining experts are assigned to the last rank.

324 - local_num_experts = num_experts - ep_rank * local_num_experts

325 - self.expert_map[-local_num_experts:] = \

326 - torch.arange(0, local_num_experts, dtype=torch.int32)

324 + self.local_num_experts = (self.global_num_experts -

325 + ep_rank * self.local_num_experts)

326 + self.expert_map[-self.local_num_experts:] = \

327 + torch.arange(0, self.local_num_experts, dtype=torch.int32)

328

329 if self.scoring_func != "softmax" and not self.use_grouped_topk:

330 raise ValueError("Only softmax scoring function is supported for "

331 "non-grouped topk.")

332

333 + # Note: get_quant_method will look at the layer's local_num_experts

334 + # for heuristic purposes, so it must be initialized first.

335 if quant_config is None:

336 self.quant_method: Optional[QuantizeMethodBase] = (

337 UnquantizedFusedMoEMethod())

@@ -336,11 +339,8 @@ class FusedMoE(torch.nn.Module):

339 self.quant_method = quant_config.get_quant_method(self, prefix)

340 assert self.quant_method is not None

341

339 - local_num_experts = torch.sum(self.expert_map != -1) \

340 - if self.expert_map is not None else num_experts

341 -

342 moe_quant_params = {

343 - "num_experts": local_num_experts,

343 + "num_experts": self.local_num_experts,

344 "hidden_size": hidden_size,

345 "intermediate_size_per_partition":

346 self.intermediate_size_per_partition,

@@ -647,7 +647,7 @@ class FusedMoE(torch.nn.Module):

647 top_k=self.top_k,

648 renormalize=self.renormalize,

649 use_grouped_topk=self.use_grouped_topk,

650 - global_num_experts=self.num_experts,

650 + global_num_experts=self.global_num_experts,

651 expert_map=self.expert_map,

652 topk_group=self.topk_group,

653 num_expert_group=self.num_expert_group,

vllm/model_executor/layers/quantization/awq_marlin.py

+1 -1

@@ -136,7 +136,7 @@ class AWQMarlinConfig(QuantizationConfig):

136 self.full_config).get_quant_method(layer, prefix)

137 return AWQMarlinLinearMethod(self)

138 elif isinstance(layer, FusedMoE):

139 - if layer.num_experts > 32:

139 + if layer.local_num_experts > 32:

140 # For MoEs with many experts the moe_wna16 kernel is faster

141 return MoeWNA16Config.from_config(

142 self.full_config).get_quant_method(layer, prefix)

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+1 -1

@@ -190,7 +190,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

190 assert layer.w13_weight_scale is not None

191 shard_size = layer.intermediate_size_per_partition

192 max_w13_scales = layer.w13_weight_scale.max(dim=1).values

193 - for expert_id in range(layer.num_experts):

193 + for expert_id in range(layer.local_num_experts):

194 start = 0

195 for shard_id in range(2):

196 dq_weight = per_tensor_dequantize(

vllm/model_executor/layers/quantization/fp8.py

+3 -3

@@ -573,11 +573,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

573 # Re-initialize w13_scale because we directly quantize

574 # merged w13 weights and generate a single scaling factor.

575 layer.w13_weight_scale = torch.nn.Parameter(torch.ones(

576 - layer.num_experts,

576 + layer.local_num_experts,

577 dtype=torch.float32,

578 device=w13_weight.device),

579 requires_grad=False)

580 - for expert in range(layer.num_experts):

580 + for expert in range(layer.local_num_experts):

581 w13_weight[expert, :, :], layer.w13_weight_scale[

582 expert] = ops.scaled_fp8_quant(

583 layer.w13_weight.data[expert, :, :])

@@ -644,7 +644,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

644 assert layer.w13_weight_scale is not None

645 shard_size = layer.intermediate_size_per_partition

646 max_w13_scales = layer.w13_weight_scale.max(dim=1).values

647 - for expert_id in range(layer.num_experts):

647 + for expert_id in range(layer.local_num_experts):

648 start = 0

649 for shard_id in range(2):

650 dq_weight = per_tensor_dequantize(

vllm/model_executor/layers/quantization/gptq_marlin.py

+1 -1

@@ -153,7 +153,7 @@ class GPTQMarlinConfig(QuantizationConfig):

153 def get_quant_method(self, layer: torch.nn.Module,

154 prefix: str) -> Optional["QuantizeMethodBase"]:

155 if isinstance(layer, FusedMoE):

156 - if layer.num_experts > 32:

156 + if layer.local_num_experts > 32:

157 # For MoEs with many experts the moe_wna16 kernel is faster

158 return MoeWNA16Config.from_config(

159 self.full_config).get_quant_method(layer, prefix)

vllm/model_executor/layers/quantization/quark/quark_moe.py

+1 -1

@@ -174,7 +174,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

174 assert layer.w13_weight_scale is not None

175 shard_size = layer.intermediate_size_per_partition

176 max_w13_scales = layer.w13_weight_scale.max(dim=1).values

177 - for expert_id in range(layer.num_experts):

177 + for expert_id in range(layer.local_num_experts):

178 start = 0

179 for shard_id in range(2):

180 dq_weight = per_tensor_dequantize(