[Refactor] tuning_fused_moe for MLLM and small refactor (#11224)
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
This commit is contained in:
Xinyuan Tong
co-authored by
Cursor Agent
Cursor Agent
parent
6c1a3f0cb2
commit
82cfcd3bb8
@@ -419,97 +419,25 @@ def get_filename(
|
|||||||
def main(args: argparse.Namespace):
|
def main(args: argparse.Namespace):
|
||||||
print(args)
|
print(args)
|
||||||
|
|
||||||
config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
|
def _calculate_shard_intermediate_size(intermediate_size: int) -> int:
|
||||||
if config.architectures[0] == "DbrxForCausalLM":
|
# In EP mode, use original intermediate_size; otherwise apply TP sharding
|
||||||
E = config.ffn_config.moe_num_experts // args.ep_size
|
return (
|
||||||
topk = config.ffn_config.moe_top_k
|
intermediate_size
|
||||||
intermediate_size = config.ffn_config.ffn_hidden_size
|
if args.ep_size > 1
|
||||||
shard_intermediate_size = (
|
else 2 * intermediate_size // args.tp_size
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] == "JambaForCausalLM":
|
|
||||||
E = config.num_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] in [
|
|
||||||
"Qwen2MoeForCausalLM",
|
|
||||||
"Qwen3MoeForCausalLM",
|
|
||||||
"Qwen3NextForCausalLM",
|
|
||||||
]:
|
|
||||||
E = config.num_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.moe_intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]:
|
|
||||||
E = (config.n_routed_experts // args.ep_size) + (
|
|
||||||
0
|
|
||||||
if args.disable_shared_experts_fusion
|
|
||||||
or config.architectures[0] not in ["DeepseekV3ForCausalLM"]
|
|
||||||
else 1
|
|
||||||
)
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.moe_intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] == "Llama4ForConditionalGeneration":
|
|
||||||
E = config.text_config.num_local_experts // args.ep_size + (
|
|
||||||
0 if args.disable_shared_experts_fusion else 1
|
|
||||||
)
|
|
||||||
topk = config.text_config.num_experts_per_tok
|
|
||||||
intermediate_size = config.text_config.intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] in [
|
|
||||||
"Grok1ForCausalLM",
|
|
||||||
"Grok1ImgGen",
|
|
||||||
"Grok1AForCausalLM",
|
|
||||||
]:
|
|
||||||
E = config.num_local_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.moe_intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] in [
|
|
||||||
"BailingMoEForCausalLM",
|
|
||||||
"BailingMoeForCausalLM",
|
|
||||||
"BailingMoeV2ForCausalLM",
|
|
||||||
]:
|
|
||||||
E = config.num_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.moe_intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
elif config.architectures[0] in ["Glm4MoeForCausalLM"]:
|
|
||||||
E = config.n_routed_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.moe_intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
# Default: Mixtral
|
|
||||||
E = config.num_local_experts // args.ep_size
|
|
||||||
topk = config.num_experts_per_tok
|
|
||||||
intermediate_size = config.intermediate_size
|
|
||||||
shard_intermediate_size = (
|
|
||||||
2 * intermediate_size // (args.tp_size // args.ep_size)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
hidden_size = getattr(config, "hidden_size", None) or config.text_config.hidden_size
|
# Check EP mode constraint: tp_size must be 1 when ep_size > 1
|
||||||
dtype = config.torch_dtype
|
if args.ep_size > 1 and args.tp_size != 1:
|
||||||
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
|
raise ValueError(
|
||||||
use_int8_w8a8 = args.dtype == "int8_w8a8"
|
f"When using Expert Parallelism (ep_size={args.ep_size}), "
|
||||||
use_int8_w8a16 = args.dtype == "int8_w8a16"
|
f"tp_size must be set to 1, but got tp_size={args.tp_size}. "
|
||||||
per_channel_quant = args.per_channel_quant
|
f"Please set --tp-size 1 when using --ep-size > 1."
|
||||||
|
)
|
||||||
|
|
||||||
|
config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
|
||||||
|
|
||||||
|
# Determine block shape for quantization
|
||||||
block_shape = None
|
block_shape = None
|
||||||
if (
|
if (
|
||||||
hasattr(config, "quantization_config")
|
hasattr(config, "quantization_config")
|
||||||
@@ -518,6 +446,82 @@ def main(args: argparse.Namespace):
|
|||||||
block_shape = config.quantization_config["weight_block_size"]
|
block_shape = config.quantization_config["weight_block_size"]
|
||||||
assert len(block_shape) == 2
|
assert len(block_shape) == 2
|
||||||
|
|
||||||
|
architecture = config.architectures[0]
|
||||||
|
# replace config with text_config for encoder-decoder models after getting block_shape and architecture
|
||||||
|
if hasattr(config, "text_config"):
|
||||||
|
config = config.get_text_config()
|
||||||
|
|
||||||
|
if architecture == "DbrxForCausalLM":
|
||||||
|
E = config.ffn_config.moe_num_experts
|
||||||
|
topk = config.ffn_config.moe_top_k
|
||||||
|
intermediate_size = config.ffn_config.ffn_hidden_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture == "JambaForCausalLM":
|
||||||
|
E = config.num_experts
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture in [
|
||||||
|
"Qwen2MoeForCausalLM",
|
||||||
|
"Qwen3MoeForCausalLM",
|
||||||
|
"Qwen3NextForCausalLM",
|
||||||
|
"Qwen3VLMoeForConditionalGeneration",
|
||||||
|
]:
|
||||||
|
E = config.num_experts // args.ep_size
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.moe_intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]:
|
||||||
|
E = (
|
||||||
|
config.n_routed_experts + (0 if args.disable_shared_experts_fusion else 1)
|
||||||
|
if architecture == "DeepseekV3ForCausalLM"
|
||||||
|
else config.n_routed_experts
|
||||||
|
)
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.moe_intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture == "Llama4ForConditionalGeneration":
|
||||||
|
E = config.num_local_experts + (0 if args.disable_shared_experts_fusion else 1)
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture in [
|
||||||
|
"Grok1ForCausalLM",
|
||||||
|
"Grok1ImgGen",
|
||||||
|
"Grok1AForCausalLM",
|
||||||
|
]:
|
||||||
|
E = config.num_local_experts // args.ep_size
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.moe_intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture in [
|
||||||
|
"BailingMoEForCausalLM",
|
||||||
|
"BailingMoeForCausalLM",
|
||||||
|
"BailingMoeV2ForCausalLM",
|
||||||
|
]:
|
||||||
|
E = config.num_experts // args.ep_size
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.moe_intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
elif architecture in ["Glm4MoeForCausalLM"]:
|
||||||
|
E = config.n_routed_experts
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.moe_intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
else:
|
||||||
|
# Default: Mixtral
|
||||||
|
E = config.num_local_experts // args.ep_size
|
||||||
|
topk = config.num_experts_per_tok
|
||||||
|
intermediate_size = config.intermediate_size
|
||||||
|
shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
|
||||||
|
|
||||||
|
hidden_size = config.hidden_size
|
||||||
|
dtype = config.torch_dtype
|
||||||
|
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
|
||||||
|
use_int8_w8a8 = args.dtype == "int8_w8a8"
|
||||||
|
use_int8_w8a16 = args.dtype == "int8_w8a16"
|
||||||
|
per_channel_quant = args.per_channel_quant
|
||||||
|
|
||||||
if args.batch_size is None:
|
if args.batch_size is None:
|
||||||
batch_sizes = [
|
batch_sizes = [
|
||||||
1,
|
1,
|
||||||
|
|||||||
Reference in New Issue
Block a user