Skip to content

Commit cfa134d

Browse files
authored
[Bugfix/CI] Fixup benchmark_moe.py (#12562)
Fixes `is_marlin` not being passed into `get_default_config` Also allow `--tensor-parallel-size` in addition to `-tp` and `--tp-size` Signed-off-by: Tyler Michael Smith <[email protected]>
1 parent 35b7a05 commit cfa134d

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

benchmarks/kernels/benchmark_moe.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -343,9 +343,13 @@ def benchmark(
343343
op_config = get_moe_configs(num_experts, shard_intermediate_size // 2,
344344
dtype_str)
345345
if op_config is None:
346-
config = get_default_config(num_tokens, num_experts,
347-
shard_intermediate_size, hidden_size,
348-
topk, dtype_str)
346+
config = get_default_config(num_tokens,
347+
num_experts,
348+
shard_intermediate_size,
349+
hidden_size,
350+
topk,
351+
dtype_str,
352+
is_marlin=False)
349353
else:
350354
config = op_config[min(op_config.keys(),
351355
key=lambda x: abs(x - num_tokens))]
@@ -536,7 +540,11 @@ def _distribute(method: str, inputs: List[Any]) -> List[Any]:
536540
parser.add_argument("--model",
537541
type=str,
538542
default="mistralai/Mixtral-8x7B-Instruct-v0.1")
539-
parser.add_argument("--tp-size", "-tp", type=int, default=2)
543+
parser.add_argument("--tp-size",
544+
"-tp",
545+
"--tensor-parallel-size",
546+
type=int,
547+
default=2)
540548
parser.add_argument("--dtype",
541549
type=str,
542550
choices=["auto", "fp8_w8a8", "int8_w8a16"],

0 commit comments

Comments
 (0)