Expert Parallelism for GPT-OSS (#8944)
This commit is contained in:
Cheng Wan
@@ -28,6 +28,7 @@ from sglang.srt.distributed import (
|
||||
get_moe_expert_parallel_rank,
|
||||
get_moe_expert_parallel_world_size,
|
||||
get_moe_tensor_parallel_rank,
|
||||
get_moe_tensor_parallel_world_size,
|
||||
get_pp_group,
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
@@ -96,11 +97,6 @@ class GptOssSparseMoeBlock(nn.Module):
|
||||
self.activation = config.hidden_act
|
||||
self.activation_alpha = getattr(config, "hidden_act_alpha", 1.702)
|
||||
self.swiglu_limit = config.swiglu_limit
|
||||
if self.tp_size > config.num_local_experts:
|
||||
raise ValueError(
|
||||
f"Tensor parallel size {self.tp_size} is greater than "
|
||||
f"the number of experts {config.num_local_experts}."
|
||||
)
|
||||
|
||||
if global_server_args_dict["enable_flashinfer_mxfp4_moe"]:
|
||||
self.topk = None
|
||||
@@ -708,22 +704,26 @@ class GptOssForCausalLM(nn.Module):
|
||||
loaded_params: set[str] = set()
|
||||
mxfp4_block = 32
|
||||
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
moe_tp_rank = get_moe_tensor_parallel_rank()
|
||||
moe_tp_size = get_moe_tensor_parallel_world_size()
|
||||
moe_ep_rank = get_moe_expert_parallel_rank()
|
||||
moe_ep_size = get_moe_expert_parallel_world_size()
|
||||
|
||||
intermediate_size = self.config.intermediate_size
|
||||
intermediate_size_block = intermediate_size // mxfp4_block
|
||||
per_rank_intermediate_size_block = intermediate_size_block // tp_size
|
||||
per_rank_intermediate_size_block = intermediate_size_block // moe_tp_size
|
||||
per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block
|
||||
|
||||
# Calculate common slicing bounds for current rank
|
||||
tp_rank_start = tp_rank * per_rank_intermediate_size
|
||||
tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size)
|
||||
|
||||
# Attention heads per rank
|
||||
heads_per_rank = self.config.num_attention_heads // tp_size
|
||||
head_start = tp_rank * heads_per_rank
|
||||
|
||||
num_experts = self.config.num_local_experts
|
||||
assert self.config.num_local_experts % moe_ep_size == 0
|
||||
moe_num_global_experts = self.config.num_local_experts
|
||||
moe_num_local_experts = self.config.num_local_experts // moe_ep_size
|
||||
moe_tp_rank_start = moe_tp_rank * per_rank_intermediate_size
|
||||
moe_tp_rank_end = min(
|
||||
(moe_tp_rank + 1) * per_rank_intermediate_size, intermediate_size
|
||||
)
|
||||
moe_ep_rank_start = moe_ep_rank * moe_num_local_experts
|
||||
moe_ep_rank_end = (moe_ep_rank + 1) * moe_num_local_experts
|
||||
|
||||
for name, weight in weights:
|
||||
weight = weight.cuda()
|
||||
@@ -735,10 +735,14 @@ class GptOssForCausalLM(nn.Module):
|
||||
# flat weight from (E, 2 * N, block_size, entry_per_block)
|
||||
# to (E, 2 * N, -1), shouldn't trigger copy for contiguous
|
||||
weight = weight.view(
|
||||
num_experts, 2 * intermediate_size, -1
|
||||
moe_num_global_experts, 2 * intermediate_size, -1
|
||||
).contiguous()
|
||||
|
||||
narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...]
|
||||
narrow_weight = weight[
|
||||
moe_ep_rank_start:moe_ep_rank_end,
|
||||
2 * moe_tp_rank_start : 2 * moe_tp_rank_end,
|
||||
...,
|
||||
]
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
@@ -757,9 +761,13 @@ class GptOssForCausalLM(nn.Module):
|
||||
# same flatten here, but since 2 mx4 value are packed in 1
|
||||
# uint8, divide by 2
|
||||
weight = weight.view(
|
||||
num_experts, -1, intermediate_size // 2
|
||||
moe_num_global_experts, -1, intermediate_size // 2
|
||||
).contiguous()
|
||||
narrow_weight = weight[..., tp_rank_start // 2 : tp_rank_end // 2]
|
||||
narrow_weight = weight[
|
||||
moe_ep_rank_start:moe_ep_rank_end,
|
||||
...,
|
||||
moe_tp_rank_start // 2 : moe_tp_rank_end // 2,
|
||||
]
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
@@ -775,7 +783,11 @@ class GptOssForCausalLM(nn.Module):
|
||||
elif "gate_up_proj_scales" in name:
|
||||
# Handle MLP gate and up projection weights scale
|
||||
new_name = name.replace("gate_up_proj_scales", "w13_weight_scale")
|
||||
narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...]
|
||||
narrow_weight = weight[
|
||||
moe_ep_rank_start:moe_ep_rank_end,
|
||||
2 * moe_tp_rank_start : 2 * moe_tp_rank_end,
|
||||
...,
|
||||
]
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
@@ -792,7 +804,9 @@ class GptOssForCausalLM(nn.Module):
|
||||
# Handle MLP down projection weights
|
||||
new_name = name.replace("down_proj_scales", "w2_weight_scale")
|
||||
narrow_weight = weight[
|
||||
..., tp_rank_start // mxfp4_block : tp_rank_end // mxfp4_block
|
||||
moe_ep_rank_start:moe_ep_rank_end,
|
||||
...,
|
||||
moe_tp_rank_start // mxfp4_block : moe_tp_rank_end // mxfp4_block,
|
||||
]
|
||||
|
||||
param = params_dict[new_name]
|
||||
@@ -809,7 +823,10 @@ class GptOssForCausalLM(nn.Module):
|
||||
# Handle MLP gate and up projection biases
|
||||
new_name = name.replace("gate_up_proj_bias", "w13_weight_bias")
|
||||
|
||||
narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end]
|
||||
narrow_weight = weight[
|
||||
moe_ep_rank_start:moe_ep_rank_end,
|
||||
2 * moe_tp_rank_start : 2 * moe_tp_rank_end,
|
||||
]
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
@@ -823,15 +840,20 @@ class GptOssForCausalLM(nn.Module):
|
||||
loaded_params.add(new_name)
|
||||
|
||||
elif "down_proj_bias" in name:
|
||||
if get_moe_tensor_parallel_rank() != 0:
|
||||
weight = torch.zeros_like(weight)
|
||||
narrow_weight = weight[moe_ep_rank_start:moe_ep_rank_end, ...]
|
||||
if moe_tp_rank != 0:
|
||||
narrow_weight = torch.zeros_like(narrow_weight)
|
||||
|
||||
# Handle MLP down projection bias
|
||||
new_name = name.replace("down_proj_bias", "w2_weight_bias")
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(
|
||||
param, weight, weight_name=new_name, shard_id=None, expert_id=None
|
||||
param,
|
||||
narrow_weight,
|
||||
weight_name=new_name,
|
||||
shard_id=None,
|
||||
expert_id=None,
|
||||
)
|
||||
loaded_params.add(new_name)
|
||||
|
||||
@@ -910,27 +932,12 @@ class GptOssForCausalLM(nn.Module):
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
]
|
||||
|
||||
if self.quant_config is not None and (self.quant_config.get_name() == "mxfp4"):
|
||||
expert_params_mapping = (
|
||||
get_moe_impl_class().make_expert_params_mapping_fused_mxfp4(
|
||||
ckpt_gate_up_proj_name="gate_up_proj_blocks",
|
||||
ckpt_down_proj_name="down_proj_blocks",
|
||||
ckpt_gate_up_proj_bias_name="gate_up_proj_bias",
|
||||
ckpt_down_proj_bias_name="down_proj_bias",
|
||||
ckpt_gate_up_proj_scale_name="gate_up_proj_scales",
|
||||
ckpt_down_proj_scale_name="down_proj_scales",
|
||||
)
|
||||
)
|
||||
else:
|
||||
expert_params_mapping = (
|
||||
get_moe_impl_class().make_expert_params_mapping_fused(
|
||||
ckpt_gate_up_proj_name="gate_up_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_gate_up_proj_bias_name="gate_up_proj_bias",
|
||||
ckpt_down_proj_bias_name="down_proj_bias",
|
||||
)
|
||||
)
|
||||
expert_params_mapping = get_moe_impl_class().make_expert_params_mapping_fused(
|
||||
ckpt_gate_up_proj_name="gate_up_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_gate_up_proj_bias_name="gate_up_proj_bias",
|
||||
ckpt_down_proj_bias_name="down_proj_bias",
|
||||
)
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
params_checker = {k: False for k, v in params_dict.items()}
|
||||
|
||||
Reference in New Issue
Block a user