[ROCm][Perf] Add Fused Shared Expert (FSE) support for GLM-4.5/6/7 (#44313)
Signed-off-by: Olga Miroshnichenko <olga.miroshnichenko@amd.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com> Co-authored-by: Mehdi Ghanimifard <mghanimi@amd.com> Co-authored-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
This commit is contained in:
Olga Miroshnichenko
co-authored by
Mehdi Ghanimifard
Mehdi Ghanimifard
Mehdi Ghanimifard
Mehdi Ghanimifard
parent
c6741b2ad4
commit
c7ca0bccae
@@ -32,6 +32,7 @@ import torch
|
||||
from torch import nn
|
||||
from transformers.models.glm4_moe import Glm4MoeConfig
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
|
||||
from vllm.distributed import (
|
||||
@@ -168,7 +169,16 @@ class Glm4MoE(nn.Module):
|
||||
self.physical_expert_start + self.n_local_physical_experts
|
||||
)
|
||||
|
||||
if config.n_shared_experts is not None:
|
||||
# AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py
|
||||
# pattern (see Glm4MoE / FusedMoE wiring there).
|
||||
self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
|
||||
self.is_fusion_moe_shared_experts_enabled = (
|
||||
rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
)
|
||||
|
||||
if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
|
||||
self.shared_experts = None
|
||||
else:
|
||||
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
||||
self.shared_experts = Glm4MoeMLP(
|
||||
hidden_size=config.hidden_size,
|
||||
@@ -178,8 +188,6 @@ class Glm4MoE(nn.Module):
|
||||
reduce_results=False,
|
||||
prefix=f"{prefix}.shared_experts",
|
||||
)
|
||||
else:
|
||||
self.shared_experts = None
|
||||
|
||||
self.experts = FusedMoE(
|
||||
shared_experts=self.shared_experts,
|
||||
@@ -194,12 +202,18 @@ class Glm4MoE(nn.Module):
|
||||
topk_group=config.topk_group,
|
||||
prefix=f"{prefix}.experts",
|
||||
scoring_func="sigmoid",
|
||||
# aiter applies routed_scaling_factor internally; see deepseek_v2.py.
|
||||
routed_scaling_factor=self.routed_scaling_factor,
|
||||
apply_routed_scale_to_output=True,
|
||||
apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled,
|
||||
e_score_correction_bias=self.gate.e_score_correction_bias,
|
||||
enable_eplb=self.enable_eplb,
|
||||
num_redundant_experts=self.n_redundant_experts,
|
||||
router_logits_dtype=torch.float32,
|
||||
n_shared_experts=(
|
||||
config.n_shared_experts
|
||||
if self.is_fusion_moe_shared_experts_enabled
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
@@ -469,15 +483,25 @@ class Glm4MoeModel(nn.Module):
|
||||
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
|
||||
# Params for weights, fp8 weight scales, fp8 activation scales
|
||||
# (param_name, weight_name, expert_id, shard_id)
|
||||
# FSE widens the mapping by n_shared_experts slots; see deepseek_v2.py.
|
||||
num_experts = self.config.n_routed_experts
|
||||
if (
|
||||
rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
and self.config.n_shared_experts
|
||||
):
|
||||
num_experts += self.config.n_shared_experts
|
||||
return fused_moe_make_expert_params_mapping(
|
||||
self,
|
||||
ckpt_gate_proj_name="gate_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_up_proj_name="up_proj",
|
||||
num_experts=self.config.n_routed_experts,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
rocm_aiter_moe_shared_expert_enabled = (
|
||||
rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
)
|
||||
stacked_params_mapping = [
|
||||
# (param_name, shard_name, shard_id)
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
@@ -494,6 +518,11 @@ class Glm4MoeModel(nn.Module):
|
||||
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
|
||||
if spec_layer is not None:
|
||||
continue
|
||||
|
||||
is_fusion_moe_shared_experts_layer = (
|
||||
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
|
||||
)
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
# Skip non-stacked layers and experts (experts handled below).
|
||||
if weight_name not in name:
|
||||
@@ -506,6 +535,8 @@ class Glm4MoeModel(nn.Module):
|
||||
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
|
||||
if ("mlp.experts." in name) and name not in params_dict:
|
||||
continue
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
continue
|
||||
|
||||
name = name.replace(weight_name, param_name)
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
@@ -527,65 +558,109 @@ class Glm4MoeModel(nn.Module):
|
||||
break
|
||||
else:
|
||||
is_expert_weight = False
|
||||
for mapping in expert_params_mapping:
|
||||
param_name, weight_name, expert_id, shard_id = mapping
|
||||
if weight_name not in name:
|
||||
continue
|
||||
|
||||
# Anyway, this is an expert weight and should not be
|
||||
# attempted to load as other weights later
|
||||
is_expert_weight = True
|
||||
|
||||
# Do not modify `name` since the loop may continue here
|
||||
# Instead, create a new variable
|
||||
name_mapped = name.replace(weight_name, param_name)
|
||||
|
||||
if is_pp_missing_parameter(name_mapped, self):
|
||||
continue
|
||||
|
||||
param = params_dict[name_mapped]
|
||||
# We should ask the weight loader to return success or not
|
||||
# here since otherwise we may skip experts with other
|
||||
# available replicas.
|
||||
weight_loader = typing.cast(
|
||||
Callable[..., bool], param.weight_loader
|
||||
# FSE: split a widened mlp.shared_experts tensor into
|
||||
# n_shared_experts chunks; see deepseek_v2.py for details.
|
||||
num_chunks = 1
|
||||
split_dim = 0
|
||||
chunk_size = 0
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
|
||||
split_dim = (
|
||||
1
|
||||
if ("down_proj.weight" in name and loaded_weight.ndim > 1)
|
||||
else 0
|
||||
)
|
||||
success = weight_loader(
|
||||
param,
|
||||
loaded_weight,
|
||||
name_mapped,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
return_success=True,
|
||||
)
|
||||
if success:
|
||||
name = name_mapped
|
||||
break
|
||||
else:
|
||||
if is_expert_weight:
|
||||
# We've checked that this is an expert weight
|
||||
# However it's not mapped locally to this rank
|
||||
# So we simply skip it
|
||||
continue
|
||||
total = loaded_weight.shape[split_dim]
|
||||
if total % num_chunks != 0:
|
||||
raise ValueError(
|
||||
f"FSE shared-expert weight {name} has dim "
|
||||
f"{total} along axis {split_dim} which is not "
|
||||
f"divisible by n_shared_experts={num_chunks}."
|
||||
)
|
||||
chunk_size = total // num_chunks
|
||||
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
for j in range(num_chunks):
|
||||
chunk_name = name
|
||||
weight_to_load = loaded_weight
|
||||
|
||||
# Remapping the name of FP8 kv-scale.
|
||||
name = maybe_remap_kv_scale_name(name, params_dict)
|
||||
if name is None:
|
||||
continue
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
|
||||
if loaded_weight.ndim == 1:
|
||||
weight_to_load = loaded_weight[chunk_slice]
|
||||
elif split_dim == 0:
|
||||
weight_to_load = loaded_weight[chunk_slice, :]
|
||||
else:
|
||||
weight_to_load = loaded_weight[:, chunk_slice]
|
||||
# Synthesize an expert-style name for expert mapping.
|
||||
chunk_name = name.replace(
|
||||
"mlp.shared_experts",
|
||||
f"mlp.experts.{self.config.n_routed_experts + j}",
|
||||
)
|
||||
|
||||
if is_pp_missing_parameter(name, self):
|
||||
continue
|
||||
for mapping in expert_params_mapping:
|
||||
param_name, weight_name, expert_id, shard_id = mapping
|
||||
if weight_name not in chunk_name:
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
# Anyway, this is an expert weight and should not be
|
||||
# attempted to load as other weights later
|
||||
is_expert_weight = True
|
||||
|
||||
# Do not modify `name` since the loop may continue here
|
||||
# Instead, create a new variable
|
||||
name_mapped = chunk_name.replace(weight_name, param_name)
|
||||
|
||||
if is_pp_missing_parameter(name_mapped, self):
|
||||
continue
|
||||
|
||||
param = params_dict[name_mapped]
|
||||
# We should ask the weight loader to return success
|
||||
# or not here since otherwise we may skip experts
|
||||
# with other available replicas.
|
||||
weight_loader = typing.cast(
|
||||
Callable[..., bool], param.weight_loader
|
||||
)
|
||||
success = weight_loader(
|
||||
param,
|
||||
weight_to_load,
|
||||
name_mapped,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
return_success=True,
|
||||
)
|
||||
if success:
|
||||
if not is_fusion_moe_shared_experts_layer:
|
||||
name = name_mapped
|
||||
else:
|
||||
loaded_params.add(name_mapped)
|
||||
break
|
||||
else:
|
||||
if is_expert_weight:
|
||||
# We've checked that this is an expert weight
|
||||
# However it's not mapped locally to this rank
|
||||
# So we simply skip it
|
||||
continue
|
||||
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
|
||||
# Remapping the name of FP8 kv-scale.
|
||||
name = maybe_remap_kv_scale_name(name, params_dict)
|
||||
if name is None:
|
||||
continue
|
||||
|
||||
if is_pp_missing_parameter(name, self):
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
if name is not None and not is_fusion_moe_shared_experts_layer:
|
||||
loaded_params.add(name)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
@@ -24,12 +24,14 @@
|
||||
"""Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP
|
||||
model compatible with HuggingFace weights."""
|
||||
|
||||
from collections.abc import Iterable
|
||||
import typing
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
MoERunner,
|
||||
@@ -237,6 +239,10 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
|
||||
return self.model.compute_logits(hidden_states, spec_step_idx)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
# FSE weight loading mirrors glm4_moe.py / deepseek_mtp.py.
|
||||
rocm_aiter_moe_shared_expert_enabled = (
|
||||
rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
)
|
||||
stacked_params_mapping = [
|
||||
# (param_name, shard_name, shard_id)
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
@@ -248,12 +254,15 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
|
||||
|
||||
# Params for weights, fp8 weight scales, fp8 activation scales
|
||||
# (param_name, weight_name, expert_id, shard_id)
|
||||
num_experts = self.config.n_routed_experts
|
||||
if rocm_aiter_moe_shared_expert_enabled and self.config.n_shared_experts:
|
||||
num_experts += self.config.n_shared_experts
|
||||
expert_params_mapping = fused_moe_make_expert_params_mapping(
|
||||
self,
|
||||
ckpt_gate_proj_name="gate_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_up_proj_name="up_proj",
|
||||
num_experts=self.config.n_routed_experts,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
@@ -269,6 +278,11 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
|
||||
if spec_layer is None:
|
||||
continue
|
||||
name = self._rewrite_spec_layer_name(spec_layer, name)
|
||||
|
||||
is_fusion_moe_shared_experts_layer = (
|
||||
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
|
||||
)
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
# Skip non-stacked layers and experts (experts handled below).
|
||||
if weight_name not in name:
|
||||
@@ -281,6 +295,8 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
|
||||
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
|
||||
if ("mlp.experts." in name) and name not in params_dict:
|
||||
continue
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
@@ -291,47 +307,105 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
for mapping in expert_params_mapping:
|
||||
param_name, weight_name, expert_id, shard_id = mapping
|
||||
if weight_name not in name:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(
|
||||
param,
|
||||
loaded_weight,
|
||||
name,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
# FSE: split a widened mlp.shared_experts tensor into
|
||||
# n_shared_experts chunks; see deepseek_v2.py for details.
|
||||
num_chunks = 1
|
||||
split_dim = 0
|
||||
chunk_size = 0
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
|
||||
split_dim = (
|
||||
1
|
||||
if ("down_proj.weight" in name and loaded_weight.ndim > 1)
|
||||
else 0
|
||||
)
|
||||
break
|
||||
else:
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
# Some checkpoints include weight scale tensors for the
|
||||
# LM head even when the quantized head isn't built. Skip
|
||||
# them if the model does not expose a matching parameter
|
||||
# to avoid KeyError during load.
|
||||
if name.endswith(".weight_scale") and name not in params_dict:
|
||||
continue
|
||||
total = loaded_weight.shape[split_dim]
|
||||
if total % num_chunks != 0:
|
||||
raise ValueError(
|
||||
f"FSE shared-expert weight {name} has dim "
|
||||
f"{total} along axis {split_dim} which is "
|
||||
f"not divisible by "
|
||||
f"n_shared_experts={num_chunks}."
|
||||
)
|
||||
chunk_size = total // num_chunks
|
||||
|
||||
# According to DeepSeek-V3 Technical Report, MTP modules
|
||||
# shares embedding layer. We only load the first weights.
|
||||
if (
|
||||
spec_layer != self.model.mtp_start_layer_idx
|
||||
and ".layers" not in name
|
||||
):
|
||||
continue
|
||||
for j in range(num_chunks):
|
||||
chunk_name = name
|
||||
weight_to_load = loaded_weight
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
if is_fusion_moe_shared_experts_layer:
|
||||
chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
|
||||
if loaded_weight.ndim == 1:
|
||||
weight_to_load = loaded_weight[chunk_slice]
|
||||
elif split_dim == 0:
|
||||
weight_to_load = loaded_weight[chunk_slice, :]
|
||||
else:
|
||||
weight_to_load = loaded_weight[:, chunk_slice]
|
||||
chunk_name = name.replace(
|
||||
"mlp.shared_experts",
|
||||
f"mlp.experts.{self.config.n_routed_experts + j}",
|
||||
)
|
||||
|
||||
is_expert_weight = False
|
||||
for mapping in expert_params_mapping:
|
||||
param_name, weight_name, expert_id, shard_id = mapping
|
||||
if weight_name not in chunk_name:
|
||||
continue
|
||||
|
||||
is_expert_weight = True
|
||||
name_mapped = chunk_name.replace(weight_name, param_name)
|
||||
|
||||
param = params_dict[name_mapped]
|
||||
# Use return_success so we don't blindly mark
|
||||
# remote-expert replicas as loaded on this rank.
|
||||
weight_loader = typing.cast(
|
||||
Callable[..., bool], param.weight_loader
|
||||
)
|
||||
success = weight_loader(
|
||||
param,
|
||||
weight_to_load,
|
||||
name_mapped,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
return_success=True,
|
||||
)
|
||||
if success:
|
||||
if not is_fusion_moe_shared_experts_layer:
|
||||
name = name_mapped
|
||||
else:
|
||||
loaded_params.add(name_mapped)
|
||||
break
|
||||
else:
|
||||
if is_expert_weight:
|
||||
# Expert weight not local to this rank; skip.
|
||||
continue
|
||||
|
||||
# Skip loading extra bias for GPTQ models.
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
# Some checkpoints include weight scale tensors for
|
||||
# the LM head even when the quantized head isn't
|
||||
# built. Skip them if the model does not expose a
|
||||
# matching parameter to avoid KeyError during load.
|
||||
if name.endswith(".weight_scale") and name not in params_dict:
|
||||
continue
|
||||
|
||||
# According to DeepSeek-V3 Technical Report, MTP
|
||||
# modules share the embedding layer. We only load
|
||||
# the first weights.
|
||||
if (
|
||||
spec_layer != self.model.mtp_start_layer_idx
|
||||
and ".layers" not in name
|
||||
):
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
if not is_fusion_moe_shared_experts_layer:
|
||||
loaded_params.add(name)
|
||||
return loaded_params
|
||||
|
||||
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
|
||||
|
||||
Reference in New Issue
Block a user