566 lines
21 KiB
Python
566 lines
21 KiB
Python
"""
|
||
GQA vs MLA Benchmark 脚本
|
||
对比 MiniMind GQA 与 MLA (Multi-head Latent Attention) 的性能差异：
|
||
1. 模型参数量 & KV-Cache 理论大小
|
||
2. 推理速度 (Prefill / Decode)
|
||
3. KV-Cache 显存占用
|
||
4. 训练吞吐量
|
||
5. 输出质量 (PPL)
|
||
|
||
Usage:
|
||
# 性能测试（随机权重）
|
||
python scripts/benchmark_gqa_vs_mla.py [--device cuda:0] [--hidden_size 768]
|
||
|
||
# 加载 checkpoint 测 PPL（必须提供 checkpoint 路径）
|
||
python scripts/benchmark_gqa_vs_mla.py --device cuda:0 \
|
||
--gqa_ckpt out/pretrain_768.pth \
|
||
--mla_ckpt out/pretrain_768_mla.pth \
|
||
--data_path dataset/pretrain_data.jsonl \
|
||
--tokenizer_path model
|
||
"""
|
||
|
||
import argparse
|
||
import gc
|
||
import math
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import torch
|
||
|
||
sys.path.append(".")
|
||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||
from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM
|
||
|
||
|
||
# ──────────────────────────── helpers ────────────────────────────
|
||
|
||
def count_params(model):
|
||
return sum(p.numel() for p in model.parameters())
|
||
|
||
|
||
def attention_params_per_layer(model):
|
||
"""统计第一个 transformer block 中 self_attn 的参数量"""
|
||
block = model.model.layers[0]
|
||
return sum(p.numel() for p in block.self_attn.parameters())
|
||
|
||
|
||
def fmt_num(n):
|
||
if n >= 1e6:
|
||
return f"{n / 1e6:.2f}M"
|
||
if n >= 1e3:
|
||
return f"{n / 1e3:.2f}K"
|
||
return str(n)
|
||
|
||
|
||
def fmt_mem(bytes_val):
|
||
if bytes_val >= 1 << 30:
|
||
return f"{bytes_val / (1 << 30):.2f} GB"
|
||
if bytes_val >= 1 << 20:
|
||
return f"{bytes_val / (1 << 20):.2f} MB"
|
||
return f"{bytes_val:.0f} B"
|
||
|
||
|
||
def warmup_forward(model, input_ids, device, n=3):
|
||
"""warmup: forward n 次，不计时"""
|
||
model.train(False)
|
||
for _ in range(n):
|
||
with torch.no_grad():
|
||
model(input_ids.to(device))
|
||
torch.cuda.synchronize(device)
|
||
|
||
|
||
def cleanup_gpu(device):
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.reset_peak_memory_stats(device)
|
||
|
||
|
||
# ────────────────────── model factories ──────────────────────────
|
||
|
||
def build_dense_attention(hidden_size, attention_type):
|
||
config = MiniMindConfig(hidden_size=hidden_size, attention_type=attention_type)
|
||
return MiniMindForCausalLM(config)
|
||
|
||
|
||
def build_mla(hidden_size, kv_lora_rank):
|
||
config = MiniMindMLAConfig(hidden_size=hidden_size, kv_lora_rank=kv_lora_rank)
|
||
return MiniMindMLAForCausalLM(config)
|
||
|
||
|
||
# ────────────────────── Benchmark 1: 模型信息 ────────────────────
|
||
|
||
def benchmark_model_info(hidden_size):
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 1: 模型信息对比")
|
||
print("=" * 70)
|
||
|
||
models = {
|
||
"MHA": build_dense_attention(hidden_size, "mha"),
|
||
"GQA": build_dense_attention(hidden_size, "gqa"),
|
||
"MQA": build_dense_attention(hidden_size, "mqa"),
|
||
"MLA(128)": build_mla(hidden_size, 128),
|
||
"MLA(64)": build_mla(hidden_size, 64),
|
||
"MLA(256)": build_mla(hidden_size, 256),
|
||
}
|
||
|
||
# Dense attention: KV-Cache floats/token/layer = 2 * num_kv_heads * head_dim
|
||
gqa_config = MiniMindConfig(hidden_size=hidden_size, attention_type="gqa")
|
||
kv_cache_gqa = 2 * gqa_config.num_key_value_heads * gqa_config.head_dim
|
||
|
||
header = f"{'指标':<28}" + "".join(f"{name:>12}" for name in models)
|
||
print(header)
|
||
print("-" * len(header))
|
||
|
||
total_params = {}
|
||
attn_params = {}
|
||
kv_cache = {}
|
||
|
||
for name, m in models.items():
|
||
total_params[name] = count_params(m)
|
||
attn_params[name] = attention_params_per_layer(m)
|
||
if "MLA" in name:
|
||
rank = int(name.split("(")[1].rstrip(")"))
|
||
kv_cache[name] = rank + m.config.rope_dim # kv_latent + k_rope
|
||
else:
|
||
kv_cache[name] = 2 * m.config.num_key_value_heads * m.config.head_dim
|
||
|
||
row = f"{'总参数量':<28}" + "".join(f"{fmt_num(total_params[n]):>12}" for n in models)
|
||
print(row)
|
||
row = f"{'注意力参数量/层':<28}" + "".join(f"{fmt_num(attn_params[n]):>12}" for n in models)
|
||
print(row)
|
||
row = f"{'KV-Cache (floats/tok/layer)':<28}" + "".join(f"{kv_cache[n]:>12}" for n in models)
|
||
print(row)
|
||
row = f"{'KV-Cache 压缩比':<28}" + "".join(
|
||
f"{kv_cache_gqa / kv_cache[n]:>11.1f}x" for n in models
|
||
)
|
||
print(row)
|
||
|
||
return models, kv_cache, kv_cache_gqa
|
||
|
||
|
||
# ────────────────────── Benchmark 2: 推理速度 ────────────────────
|
||
|
||
def benchmark_inference_speed(models, device):
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 2: 推理速度 (Prefill / Decode)")
|
||
print("=" * 70)
|
||
|
||
prompt_lens = [128, 512, 1024]
|
||
batch_sizes = [1, 4]
|
||
max_new_tokens = 128
|
||
|
||
for bs in batch_sizes:
|
||
print(f"\n--- batch_size = {bs} ---")
|
||
header = (f"{'模型':<16}{'prompt_len':>12}")
|
||
for pl in prompt_lens:
|
||
header += f"{'Prefill(ms)':>14}"
|
||
for pl in prompt_lens:
|
||
header += f"{'Decode(tok/s)':>16}"
|
||
print(header)
|
||
print("-" * len(header))
|
||
|
||
for name, model in models.items():
|
||
model = model.half().to(device).eval()
|
||
row = f"{name:<16}"
|
||
prefill_times = []
|
||
|
||
for pl in prompt_lens:
|
||
input_ids = torch.randint(0, 6400, (bs, pl), device=device)
|
||
warmup_forward(model, input_ids, device, n=2)
|
||
|
||
torch.cuda.synchronize(device)
|
||
t0 = time.perf_counter()
|
||
with torch.no_grad():
|
||
model(input_ids)
|
||
torch.cuda.synchronize(device)
|
||
t1 = time.perf_counter()
|
||
ms = (t1 - t0) * 1000
|
||
prefill_times.append(ms)
|
||
row += f"{ms:>14.2f}"
|
||
|
||
for i, pl in enumerate(prompt_lens):
|
||
input_ids = torch.randint(0, 6400, (bs, pl), device=device)
|
||
attn_mask = torch.ones(bs, pl, device=device)
|
||
warmup_forward(model, input_ids, device, n=2)
|
||
|
||
torch.cuda.synchronize(device)
|
||
t0 = time.perf_counter()
|
||
with torch.no_grad():
|
||
model.generate(
|
||
input_ids=input_ids,
|
||
attention_mask=attn_mask,
|
||
max_new_tokens=max_new_tokens,
|
||
use_cache=True,
|
||
)
|
||
torch.cuda.synchronize(device)
|
||
t1 = time.perf_counter()
|
||
tokens_per_sec = max_new_tokens * bs / (t1 - t0)
|
||
row += f"{tokens_per_sec:>16.1f}"
|
||
|
||
print(row)
|
||
del model
|
||
cleanup_gpu(device)
|
||
|
||
|
||
# ────────────────────── Benchmark 3: KV-Cache 显存 ───────────────
|
||
|
||
def benchmark_kv_cache_memory(models, device):
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 3: KV-Cache 显存占用")
|
||
print("=" * 70)
|
||
|
||
context_lens = [256, 512, 1024, 2048, 4096]
|
||
max_new_tokens = 64
|
||
bs = 1
|
||
|
||
header = f"{'模型':<16}" + "".join(f"{'ctx='+str(c):>12}" for c in context_lens)
|
||
print(header)
|
||
print("-" * len(header))
|
||
|
||
for name, model in models.items():
|
||
model = model.half().to(device).eval()
|
||
row = f"{name:<16}"
|
||
mem_results = []
|
||
|
||
for ctx_len in context_lens:
|
||
cleanup_gpu(device)
|
||
input_ids = torch.randint(0, 6400, (bs, ctx_len), device=device)
|
||
attn_mask = torch.ones(bs, ctx_len, device=device)
|
||
|
||
# 记录 generate 前的显存
|
||
torch.cuda.synchronize(device)
|
||
mem_before = torch.cuda.max_memory_allocated(device)
|
||
|
||
with torch.no_grad():
|
||
model.generate(
|
||
input_ids=input_ids,
|
||
attention_mask=attn_mask,
|
||
max_new_tokens=max_new_tokens,
|
||
use_cache=True,
|
||
)
|
||
|
||
torch.cuda.synchronize(device)
|
||
mem_after = torch.cuda.max_memory_allocated(device)
|
||
delta = mem_after - mem_before
|
||
mem_results.append(delta)
|
||
row += f"{fmt_mem(delta):>12}"
|
||
|
||
print(row)
|
||
|
||
# 打印显存增长趋势
|
||
print(f" {'增长趋势':<14}", end="")
|
||
for i in range(1, len(mem_results)):
|
||
ratio = mem_results[i] / mem_results[i - 1] if mem_results[i - 1] > 0 else 0
|
||
print(f"{'×' + f'{ratio:.2f}':>12}", end="")
|
||
print()
|
||
|
||
del model
|
||
cleanup_gpu(device)
|
||
|
||
|
||
# ────────────────────── Benchmark 4: 训练吞吐量 ─────────────────
|
||
|
||
def benchmark_training_throughput(models, device):
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 4: 训练吞吐量 (forward + backward)")
|
||
print("=" * 70)
|
||
|
||
batch_size = 4
|
||
seq_len = 340
|
||
n_steps = 10
|
||
|
||
header = f"{'模型':<16}{'step时间(ms)':>16}{'steps/s':>12}{'tokens/s':>14}"
|
||
print(header)
|
||
print("-" * len(header))
|
||
|
||
for name, model in models.items():
|
||
model = model.float().to(device).train()
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
||
|
||
# warmup
|
||
for _ in range(3):
|
||
input_ids = torch.randint(0, 6400, (batch_size, seq_len), device=device)
|
||
labels = torch.randint(0, 6400, (batch_size, seq_len), device=device)
|
||
outputs = model(input_ids, labels=labels)
|
||
outputs.loss.backward()
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
|
||
torch.cuda.synchronize(device)
|
||
t0 = time.perf_counter()
|
||
|
||
for _ in range(n_steps):
|
||
input_ids = torch.randint(0, 6400, (batch_size, seq_len), device=device)
|
||
labels = torch.randint(0, 6400, (batch_size, seq_len), device=device)
|
||
outputs = model(input_ids, labels=labels)
|
||
outputs.loss.backward()
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
|
||
torch.cuda.synchronize(device)
|
||
t1 = time.perf_counter()
|
||
|
||
total_time = t1 - t0
|
||
step_ms = total_time / n_steps * 1000
|
||
steps_per_sec = n_steps / total_time
|
||
tokens_per_sec = batch_size * seq_len * n_steps / total_time
|
||
|
||
print(f"{name:<16}{step_ms:>16.2f}{steps_per_sec:>12.2f}{tokens_per_sec:>14.1f}")
|
||
|
||
del model, optimizer
|
||
cleanup_gpu(device)
|
||
|
||
|
||
# ────────────────────── Benchmark 5: 输出质量 (PPL) ──────────────
|
||
|
||
def benchmark_perplexity(args, device, hidden_size):
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 5: 输出质量 (Perplexity)")
|
||
print("=" * 70)
|
||
|
||
from datasets import load_dataset
|
||
from transformers import AutoTokenizer
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||
samples = load_dataset("json", data_files=args.data_path, split="train")
|
||
# 取前 eval_samples 条，截取文本
|
||
max_seq = 512
|
||
all_tokens = []
|
||
for i in range(min(args.eval_samples, len(samples))):
|
||
text = str(samples[i].get("text", ""))
|
||
ids = tokenizer(text, add_special_tokens=False, max_length=max_seq, truncation=True).input_ids
|
||
ids = [tokenizer.bos_token_id] + ids + [tokenizer.eos_token_id]
|
||
all_tokens.append(ids)
|
||
|
||
results = {}
|
||
|
||
for ckpt_name, ckpt_path, config_cls, config_kwargs in [
|
||
("GQA", args.gqa_ckpt, MiniMindConfig, {}),
|
||
("MLA(128)", args.mla_ckpt, MiniMindMLAConfig, {"kv_lora_rank": 128}),
|
||
("MLA(64)", args.mla_ckpt_r64, MiniMindMLAConfig, {"kv_lora_rank": 64}),
|
||
("MLA(256)", args.mla_ckpt_r256, MiniMindMLAConfig, {"kv_lora_rank": 256}),
|
||
]:
|
||
if ckpt_path is None:
|
||
continue
|
||
if not os.path.isfile(ckpt_path):
|
||
print(f" [SKIP] {ckpt_name}: checkpoint not found at {ckpt_path}")
|
||
continue
|
||
|
||
config_kw = {"hidden_size": hidden_size}
|
||
config_kw.update(config_kwargs)
|
||
config = config_cls(**config_kw)
|
||
if isinstance(config, MiniMindMLAConfig):
|
||
model = MiniMindMLAForCausalLM(config)
|
||
else:
|
||
model = MiniMindForCausalLM(config)
|
||
|
||
state_dict = torch.load(ckpt_path, map_location="cpu")
|
||
model.load_state_dict(state_dict, strict=True)
|
||
model = model.half().to(device).eval()
|
||
|
||
total_loss = 0.0
|
||
total_tokens = 0
|
||
batch_input_ids = []
|
||
batch_labels = []
|
||
|
||
for ids in all_tokens:
|
||
pad_len = max_seq - len(ids)
|
||
padded = ids + [tokenizer.pad_token_id] * pad_len
|
||
labels = ids[1:] + [tokenizer.pad_token_id] * (pad_len + 1)
|
||
# 标记 padding 为 -100
|
||
labels = [t if t != tokenizer.pad_token_id else -100 for t in labels]
|
||
batch_input_ids.append(padded)
|
||
batch_labels.append(labels)
|
||
|
||
if len(batch_input_ids) == args.eval_batch_size:
|
||
input_t = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
|
||
label_t = torch.tensor(batch_labels, dtype=torch.long, device=device)
|
||
with torch.no_grad():
|
||
out = model(input_t, labels=label_t)
|
||
# loss 是 cross_entropy 的 mean，乘以有效 token 数得到总 loss
|
||
n_valid = (label_t != -100).sum().item()
|
||
total_loss += out.loss.item() * n_valid
|
||
total_tokens += n_valid
|
||
batch_input_ids, batch_labels = [], []
|
||
|
||
# 处理剩余样本
|
||
if batch_input_ids:
|
||
input_t = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
|
||
label_t = torch.tensor(batch_labels, dtype=torch.long, device=device)
|
||
with torch.no_grad():
|
||
out = model(input_t, labels=label_t)
|
||
n_valid = (label_t != -100).sum().item()
|
||
total_loss += out.loss.item() * n_valid
|
||
total_tokens += n_valid
|
||
|
||
avg_nll = total_loss / total_tokens
|
||
ppl = math.exp(avg_nll)
|
||
results[ckpt_name] = ppl
|
||
print(f" {ckpt_name:<16} Loss: {avg_nll:.4f} PPL: {ppl:.2f} (eval on {total_tokens} tokens)")
|
||
|
||
del model, state_dict
|
||
cleanup_gpu(device)
|
||
|
||
if not results:
|
||
print(" 没有可用的 checkpoint，跳过 PPL 测试")
|
||
print(" 提示：使用 --gqa_ckpt / --mla_ckpt 参数指定权重路径")
|
||
return {}
|
||
|
||
# Markdown PPL 表格
|
||
print()
|
||
cols = " | ".join(results.keys())
|
||
print(f"| Metric | {cols} |")
|
||
sep = "|".join(["------"] * len(results))
|
||
print(f"|--------|{sep}|")
|
||
print(f"| PPL | " + " | ".join(f"{results[k]:.2f}" for k in results) + " |")
|
||
print()
|
||
return results
|
||
|
||
|
||
# ────────────────────── Markdown 汇总 ────────────────────────────
|
||
|
||
def print_markdown_summary(models_info, hidden_size):
|
||
print("\n" + "=" * 70)
|
||
print(" Markdown 汇总（可复制到简历 / README）")
|
||
print("=" * 70)
|
||
print()
|
||
|
||
gqa_config = MiniMindConfig(hidden_size=hidden_size, attention_type="gqa")
|
||
kv_cache_gqa = 2 * gqa_config.num_key_value_heads * gqa_config.head_dim
|
||
|
||
print("### Model Info")
|
||
print()
|
||
print("| Metric | " + " | ".join(models_info.keys()) + " |")
|
||
print("|--------|" + "|".join(["------"] * len(models_info)) + "|")
|
||
|
||
total_params = {}
|
||
attn_params = {}
|
||
kv_cache = {}
|
||
for name, m in models_info.items():
|
||
total_params[name] = count_params(m)
|
||
attn_params[name] = attention_params_per_layer(m)
|
||
if "MLA" in name:
|
||
rank = int(name.split("(")[1].rstrip(")"))
|
||
kv_cache[name] = rank + m.config.rope_dim # kv_latent + k_rope
|
||
else:
|
||
kv_cache[name] = 2 * m.config.num_key_value_heads * m.config.head_dim
|
||
|
||
row = "| Total Params |"
|
||
for n in models_info:
|
||
row += f" {fmt_num(total_params[n])} |"
|
||
print(row)
|
||
|
||
row = "| Attn Params/Layer |"
|
||
for n in models_info:
|
||
row += f" {fmt_num(attn_params[n])} |"
|
||
print(row)
|
||
|
||
row = "| KV-Cache (floats/tok/layer) |"
|
||
for n in models_info:
|
||
row += f" {kv_cache[n]} |"
|
||
print(row)
|
||
|
||
row = "| KV-Cache Compression Ratio |"
|
||
for n in models_info:
|
||
row += f" {kv_cache_gqa / kv_cache[n]:.1f}x |"
|
||
print(row)
|
||
print()
|
||
|
||
print("> hidden_size={}, num_layers={}, num_attention_heads={}, num_kv_heads={}, head_dim={}".format(
|
||
gqa_config.hidden_size, gqa_config.num_hidden_layers,
|
||
gqa_config.num_attention_heads, gqa_config.num_key_value_heads, gqa_config.head_dim
|
||
))
|
||
print()
|
||
|
||
|
||
# ────────────────────── main ─────────────────────────────────────
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="GQA vs MLA Benchmark")
|
||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||
parser.add_argument("--hidden_size", type=int, default=768)
|
||
# Benchmark 5 (PPL) 参数
|
||
parser.add_argument("--gqa_ckpt", type=str, default=None, help="GQA checkpoint 路径 (e.g. out/pretrain_768.pth)")
|
||
parser.add_argument("--mla_ckpt", type=str, default=None, help="MLA checkpoint 路径 (e.g. out/pretrain_768_mla.pth)")
|
||
parser.add_argument("--mla_ckpt_r64", type=str, default=None, help="MLA(64) checkpoint 路径")
|
||
parser.add_argument("--mla_ckpt_r256", type=str, default=None, help="MLA(256) checkpoint 路径")
|
||
parser.add_argument("--data_path", type=str, default="dataset/pretrain_hq.jsonl", help="评估数据集路径")
|
||
parser.add_argument("--tokenizer_path", type=str, default="model", help="tokenizer 路径")
|
||
parser.add_argument("--eval_samples", type=int, default=500, help="评估样本数")
|
||
parser.add_argument("--eval_batch_size", type=int, default=4, help="评估 batch size")
|
||
args = parser.parse_args()
|
||
|
||
device = torch.device(args.device)
|
||
hidden_size = args.hidden_size
|
||
|
||
has_ckpt = any([args.gqa_ckpt, args.mla_ckpt, args.mla_ckpt_r64, args.mla_ckpt_r256])
|
||
|
||
print("=" * 70)
|
||
print(f" MiniMind GQA vs MLA Benchmark")
|
||
print(f" hidden_size={hidden_size} device={device}")
|
||
if has_ckpt:
|
||
print(f" PPL evaluation: ENABLED (data={args.data_path})")
|
||
print("=" * 70)
|
||
|
||
# Benchmark 1: 模型信息（CPU 即可）
|
||
models, kv_cache, kv_cache_gqa = benchmark_model_info(hidden_size)
|
||
|
||
if device.type == "cuda":
|
||
# Benchmark 2: 推理速度
|
||
models_gpu = {
|
||
"MHA": build_dense_attention(hidden_size, "mha"),
|
||
"GQA": build_dense_attention(hidden_size, "gqa"),
|
||
"MQA": build_dense_attention(hidden_size, "mqa"),
|
||
"MLA(128)": build_mla(hidden_size, 128),
|
||
"MLA(64)": build_mla(hidden_size, 64),
|
||
"MLA(256)": build_mla(hidden_size, 256),
|
||
}
|
||
benchmark_inference_speed(models_gpu, device)
|
||
|
||
# Benchmark 3: KV-Cache 显存
|
||
models_gpu = {
|
||
"MHA": build_dense_attention(hidden_size, "mha"),
|
||
"GQA": build_dense_attention(hidden_size, "gqa"),
|
||
"MQA": build_dense_attention(hidden_size, "mqa"),
|
||
"MLA(128)": build_mla(hidden_size, 128),
|
||
"MLA(64)": build_mla(hidden_size, 64),
|
||
"MLA(256)": build_mla(hidden_size, 256),
|
||
}
|
||
benchmark_kv_cache_memory(models_gpu, device)
|
||
|
||
# Benchmark 4: 训练吞吐量
|
||
models_gpu = {
|
||
"MHA": build_dense_attention(hidden_size, "mha"),
|
||
"GQA": build_dense_attention(hidden_size, "gqa"),
|
||
"MQA": build_dense_attention(hidden_size, "mqa"),
|
||
"MLA(128)": build_mla(hidden_size, 128),
|
||
"MLA(64)": build_mla(hidden_size, 64),
|
||
"MLA(256)": build_mla(hidden_size, 256),
|
||
}
|
||
benchmark_training_throughput(models_gpu, device)
|
||
|
||
# Benchmark 5: PPL（需要提供 checkpoint 路径）
|
||
if has_ckpt:
|
||
benchmark_perplexity(args, device, hidden_size)
|
||
else:
|
||
print("\n" + "=" * 70)
|
||
print(" Benchmark 5: 输出质量 (PPL) — SKIPPED")
|
||
print("=" * 70)
|
||
print(" 未指定 checkpoint，跳过 PPL 测试。")
|
||
print(" 使用方法: python scripts/benchmark_gqa_vs_mla.py \\")
|
||
print(" --gqa_ckpt out/pretrain_768.pth \\")
|
||
print(" --mla_ckpt out/pretrain_768_mla.pth \\")
|
||
print(" --data_path dataset/pretrain_hq.jsonl")
|
||
else:
|
||
print("\n[SKIP] Benchmark 2-5 需要 CUDA 设备，当前 device=cpu")
|
||
|
||
# Markdown 汇总（使用 Benchmark 1 的 CPU 模型）
|
||
print_markdown_summary(models, hidden_size)
|
||
|
||
print("Done.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|