Skip to content

GitLab

Commit dc6f80a5 authored by zhuwenwen's avatar zhuwenwen
Browse files

update tp_size

parent a441a5d9
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 4 deletions
benchmarks/kernels/benchmark_moe.py
View file @ dc6f80a5
......@@ -718,12 +718,12 @@ def main(args: argparse.Namespace):
intermediate_size = config.intermediate_size
enable_ep = bool(args.enable_expert_parallel)
if enable_ep:
ensure_divisibility(E, args.tp_size, "Number of experts")
E = E // args.tp_size
ensure_divisibility(E, tp_size, "Number of experts")
E = E // tp_size
shard_intermediate_size = 2 * intermediate_size
else:
ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
shard_intermediate_size = 2 * intermediate_size // args.tp_size
ensure_divisibility(intermediate_size, tp_size, "intermediate_size")
shard_intermediate_size = 2 * intermediate_size // tp_size
hidden_size = config.hidden_size
dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment