From 293cf7ec4f7c6431b7186908ae00391778d626fc Mon Sep 17 00:00:00 2001 From: sunteng Date: Mon, 6 Jul 2026 08:33:54 +0800 Subject: [PATCH] feat(qwen3_moe): adapt & optimize Qwen3-30B-A3B-Thinking MoE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - MoE experts 前向：stacked weights + grouped_gemm + swiglu + TP allreduce （相对朴素 per-expert 循环，bs=1 输出吞吐 +42%） - P0：可选 host group counts 跳过 grouped_gemm 内部 D2H 同步（graph guard） - 正确性脚本：prefill 首步/前缀 token 对齐 HF 两阶段校验 （MetaX TP=2 + NVIDIA TP=4 均 4/4 PASS）+ CORRECTNESS_README - bench/qwen3_moe：无依赖并发吞吐基准（serve + load client + summarize） Signed-off-by: sunteng --- bench/qwen3_moe/README.md | 111 +++++++++++ bench/qwen3_moe/load_client.py | 176 ++++++++++++++++++ bench/qwen3_moe/run_bench.sh | 69 +++++++ bench/qwen3_moe/serve.sh | 54 ++++++ bench/qwen3_moe/summarize.py | 105 +++++++++++ csrc/models/qwen3_moe/qwen3_moe_experts.cpp | 166 ++++++++++++++--- csrc/models/qwen3_moe/qwen3_moe_experts.hpp | 37 +++- test/models/qwen3_moe/CORRECTNESS_README.md | 59 ++++++ test/models/qwen3_moe/check_prefill_logits.py | 119 ++++++++++++ test/models/qwen3_moe/dump_hf_reference.py | 99 ++++++++++ 10 files changed, 955 insertions(+), 40 deletions(-) create mode 100644 bench/qwen3_moe/README.md create mode 100755 bench/qwen3_moe/load_client.py create mode 100755 bench/qwen3_moe/run_bench.sh create mode 100755 bench/qwen3_moe/serve.sh create mode 100755 bench/qwen3_moe/summarize.py create mode 100644 test/models/qwen3_moe/CORRECTNESS_README.md create mode 100644 test/models/qwen3_moe/check_prefill_logits.py create mode 100644 test/models/qwen3_moe/dump_hf_reference.py diff --git a/bench/qwen3_moe/README.md b/bench/qwen3_moe/README.md new file mode 100644 index 000000000..861b20d29 --- /dev/null +++ b/bench/qwen3_moe/README.md @@ -0,0 +1,111 @@ +# Qwen3-MoE 并发吞吐基准（MetaX） + +T2-1-3 的服务吞吐基准，对齐赛题「测试风格参考长文本优化赛题（T2-1-2）」的要求： +起 InfiniLM 服务 + `vllm bench serve` 压测 **concurrency × input-len × output-len** 矩阵， +产出 **base vs this** 的输出吞吐 / 总吞吐对比（赛题「性能」维度的打分依据）。 + +平台：沐曦 MetaX C500，TP=2。 + +## 文件 + +| 文件 | 作用 | +|---|---| +| `serve.sh` | 起 OpenAI 兼容服务（Qwen3-30B-A3B-Thinking，metax，TP=2，paged-attn + flash-attn + ignore-eos）| +| `run_bench.sh` | 客户端 `vllm bench serve` 遍历矩阵（**需装 vllm**）| +| `load_client.py` | **零依赖**并发压测客户端（纯标准库，无需 vllm/aiohttp）——vllm 装不上时用这个 | +| `summarize.py` | 解析结果文件成表格；给两个目录则输出 base-vs-this 的 Δ% | + +> 两个客户端产出**同一格式**的结果文件（`bench_results//bsX_inY_outZ.txt`）， +> summarize.py 都能解析。**MetaX 容器一般没有 vllm**（且与 maca torch 易冲突），推荐用 `load_client.py`。 + +## 服务依赖 + +服务端只需轻量依赖： + +```bash +pip install fastapi uvicorn +``` + +## 用 load_client.py（推荐，无 vllm） + +```bash +# 终端 A：起服务（见下） +# 终端 B： +python3 bench/qwen3_moe/load_client.py --tag this # 或 --tag base +python3 bench/qwen3_moe/summarize.py bench_results/this +``` + +矩阵/并发/模型名均可用 `--batch-sizes/--input-lens/--output-lens/--model/--base-url` 覆盖。 +`load_client.py` 用非流式请求测**输出吞吐 / 总吞吐**（赛题打分项）；不测 TTFT/TPOT（需流式，表里显示 `-`）。 +输入长度用 CJK 填充字近似（~1 token/字）；要精确可加 `--tokenizer <模型路径>`（需 transformers）。 + +## 默认压测矩阵 + +- concurrency：`1 8 32` +- input-len：`32 256 2048` +- output-len：`256 1024` + +（均可用环境变量覆盖，见各脚本头部注释。） + +## base vs this 怎么跑 + +「this」= T2-1-3 分支（grouped_gemm 批化 MoE）。「base」= 改造前的**朴素逐 token×逐 expert** MoE +（`main` 分支 / grouped_gemm 之前的实现）。两者用**同一套客户端脚本**分别压测，只是换部署的 `.so`。 + +### 1) 跑 this（当前 T2-1-3） + +```bash +# 终端 A：起服务（保持前台） +cd /data/InfiniLM +MODEL=/data/huggingface_home/Qwen3-30B-A3B-Thinking-2507 \ + bash bench/qwen3_moe/serve.sh +# 等打印 “load weights over” + uvicorn 起监听 + +# 终端 B：等服务 ready 后压测 +cd /data/InfiniLM +curl -s http://127.0.0.1:8102/health && echo " <- server ready" +TAG=this bash bench/qwen3_moe/run_bench.sh +``` + +### 2) 跑 base（朴素 MoE 基线） + +```bash +# 切到基线实现并重建 + 部署 .so（示例，按你的基线 commit 调整） +cd /data/InfiniLM && git stash && git checkout main +cd /data/InfiniCore && xmake install +cd /data/InfiniLM && xmake build _infinilm +cp build/linux/x86_64/release/_infinilm.cpython-310-x86_64-linux-gnu.so python/infinilm/lib/ + +# 起服务 + 压测（同上，改 TAG=base） +MODEL=/data/huggingface_home/Qwen3-30B-A3B-Thinking-2507 bash bench/qwen3_moe/serve.sh # 终端 A +TAG=base bash bench/qwen3_moe/run_bench.sh # 终端 B +``` + +> 注：base 与 this 必须用**完全相同的客户端参数**（矩阵、seed 策略、ignore-eos）才可比。 +> 脚本已固定这些；两次只改 `TAG` 和部署的 `.so`。 + +### 3) 汇总对比 + +```bash +python3 bench/qwen3_moe/summarize.py bench_results/base bench_results/this +``` + +输出每个 (bs,in,out) 的 `out_tok/s` 与 `total_tok/s` 的 base→this 及 Δ%。 +只看单侧：`python3 bench/qwen3_moe/summarize.py bench_results/this`。 + +## 显存 / 参数提示 + +- `NUM_BLOCKS` 默认 1024（≈ 26 万 token 的 KV，覆盖 concurrency 32 × (2048+1024) 有余）。 + **OOM 就调小**（`NUM_BLOCKS=512`），**报 KV 不足就调大**。 +- `MAX_CACHE`（--max-cache-len）需 ≥ 最大 input+output（默认 4096 覆盖 2048+1024）。 +- **注意力后端 `ATTN` 默认 `paged-attn`**：MetaX build **未编 flash-attn**（用 `--attn flash-attn` + 会报 "FlashAttention is not enabled in this build" 并崩 forward）。若 paged-attn 在你的 build 也不可用， + 退回 static cache：`PAGED=0 ATTN=default bash bench/qwen3_moe/serve.sh`（并发批处理能力会下降）。 +- `--enable-graph` 默认**关**：MoE 前向含数据相关的 host dispatch + syncStream，graph capture 可能不兼容。 + 想试开：`GRAPH=1 bash bench/qwen3_moe/serve.sh`，若起不来或结果异常就关掉。 +- 采样固定 `top-k 1`（greedy）+ `ignore-eos`，保证每请求生成满 output-len、结果可比。 + +## 关注指标 + +赛题打分看 **Output token throughput (tok/s)** 与 **Total Token throughput (tok/s)**； +小规模（bs=1）作为「不回退」门槛。summarize.py 已抽取这两项 + req/s + TTFT + TPOT。 diff --git a/bench/qwen3_moe/load_client.py b/bench/qwen3_moe/load_client.py new file mode 100755 index 000000000..3020b3369 --- /dev/null +++ b/bench/qwen3_moe/load_client.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Dependency-free concurrent load client for the InfiniLM OpenAI server. + +A stdlib-only alternative to `vllm bench serve` (no vllm / aiohttp needed), for +environments where vllm is not installed (e.g. MetaX containers). Sweeps the +concurrency × input-len × output-len matrix and writes one result file per +config in the exact label format `summarize.py` parses. + +Because the server runs with `--ignore-eos`, each request generates exactly +`output_len` tokens, so output-token throughput is measured against the nominal +count (also reconciled with the response `usage` field when present). + +Usage: + python3 load_client.py --tag this + python3 load_client.py --tag base --base-url http://127.0.0.1:8102 \ + --batch-sizes 1,8,32 --input-lens 32,256,2048 --output-lens 256,1024 + +Input length is approximated by repeating a CJK filler char (~1 token each); +this is fine for relative base-vs-this comparison. Pass --tokenizer to count +prompt tokens exactly if `transformers` is available. +""" +import argparse +import json +import os +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor, as_completed + +FILLER = "的" # ~1 Qwen token each; used to build an approximately input_len prompt. + + +def build_prompt(input_len, tokenizer=None): + """Return a user message string of approximately `input_len` tokens.""" + if tokenizer is not None: + # Grow the filler until the tokenizer reports >= input_len tokens. + s = FILLER * input_len + while len(tokenizer(s)["input_ids"]) < input_len: + s += FILLER * 8 + return s + return FILLER * input_len + + +def one_request(base_url, model, prompt, output_len, timeout): + """Send one non-streaming chat request; return (ok, latency_s, completion_tokens).""" + payload = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": output_len, + "ignore_eos": True, + "stream": False, + "temperature": 1.0, + }).encode("utf-8") + req = urllib.request.Request( + base_url.rstrip("/") + "/v1/chat/completions", + data=payload, headers={"Content-Type": "application/json"}, method="POST") + t0 = time.perf_counter() + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + body = json.loads(resp.read().decode("utf-8")) + except Exception as e: # noqa: BLE001 - report and keep going + return False, time.perf_counter() - t0, 0, str(e) + lat = time.perf_counter() - t0 + # Prefer server-reported usage; fall back to nominal output_len. + comp = output_len + usage = body.get("usage") or {} + if isinstance(usage.get("completion_tokens"), int) and usage["completion_tokens"] > 0: + comp = usage["completion_tokens"] + return True, lat, comp, None + + +def run_config(base_url, model, bs, il, ol, repeat, timeout, tokenizer): + num_prompts = bs * repeat + prompt = build_prompt(il, tokenizer) + results = [] + err_sample = None + + t_start = time.perf_counter() + with ThreadPoolExecutor(max_workers=bs) as pool: + futs = [pool.submit(one_request, base_url, model, prompt, ol, timeout) + for _ in range(num_prompts)] + for f in as_completed(futs): + ok, lat, comp, err = f.result() + results.append((ok, lat, comp)) + if not ok and err_sample is None: + err_sample = err + wall = time.perf_counter() - t_start + + ok_n = sum(1 for ok, _, _ in results if ok) + out_tokens = sum(c for ok, _, c in results if ok) + in_tokens = ok_n * il # approximate (see build_prompt) + lat_ok = [lat for ok, lat, _ in results if ok] + mean_lat_ms = 1000.0 * sum(lat_ok) / len(lat_ok) if lat_ok else 0.0 + + metrics = { + "successful": ok_n, + "num_prompts": num_prompts, + "duration_s": wall, + "req_s": ok_n / wall if wall > 0 else 0.0, + "out_tok_s": out_tokens / wall if wall > 0 else 0.0, + "total_tok_s": (in_tokens + out_tokens) / wall if wall > 0 else 0.0, + "mean_latency_ms": mean_lat_ms, + "err_sample": err_sample, + } + return metrics + + +def write_result(path, bs, il, ol, m): + # Labels chosen to match summarize.py's regexes. + lines = [ + "============ Serving Benchmark Result ============", + f"config: bs={bs} in={il} out={ol}", + f"Successful requests: {m['successful']}/{m['num_prompts']}", + f"Benchmark duration (s): {m['duration_s']:.2f}", + f"Request throughput (req/s): {m['req_s']:.4f}", + f"Output token throughput (tok/s): {m['out_tok_s']:.2f}", + f"Total Token throughput (tok/s): {m['total_tok_s']:.2f}", + f"Mean E2E latency (ms): {m['mean_latency_ms']:.2f}", + ] + if m["err_sample"]: + lines.append(f"error_sample: {m['err_sample']}") + open(path, "w", encoding="utf-8").write("\n".join(lines) + "\n") + + +def parse_list(s): + return [int(x) for x in s.replace(",", " ").split()] + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--base-url", default="http://127.0.0.1:8102") + ap.add_argument("--model", default="Qwen3-30B-A3B-Thinking-2507") + ap.add_argument("--tag", default="this", help="result subdir under --outdir-root") + ap.add_argument("--outdir-root", default="bench_results") + ap.add_argument("--batch-sizes", type=parse_list, default=[1, 8, 32]) + ap.add_argument("--input-lens", type=parse_list, default=[32, 256, 2048]) + ap.add_argument("--output-lens", type=parse_list, default=[256, 1024]) + ap.add_argument("--repeat", type=int, default=3, + help="requests per config = concurrency * repeat") + ap.add_argument("--timeout", type=float, default=1200.0) + ap.add_argument("--tokenizer", default=None, + help="optional HF tokenizer path for exact prompt token counts") + args = ap.parse_args() + + tok = None + if args.tokenizer: + try: + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + except Exception as e: # noqa: BLE001 + print(f"(tokenizer load failed, using approximate prompt: {e})") + + outdir = os.path.join(args.outdir_root, args.tag) + os.makedirs(outdir, exist_ok=True) + print(f">>> tag={args.tag} url={args.base_url} model={args.model} -> {outdir}") + print(f">>> concurrency={args.batch_sizes} input={args.input_lens} output={args.output_lens}") + + for bs in args.batch_sizes: + for il in args.input_lens: + for ol in args.output_lens: + path = os.path.join(outdir, f"bs{bs}_in{il}_out{ol}.txt") + print(f">>> bs={bs} in={il} out={ol} ...", end="", flush=True) + m = run_config(args.base_url, args.model, bs, il, ol, + args.repeat, args.timeout, tok) + write_result(path, bs, il, ol, m) + if m["successful"] == 0: + print(f" FAIL ({m['err_sample']})") + else: + print(f" ok out={m['out_tok_s']:.1f} tok/s total={m['total_tok_s']:.1f} tok/s" + f" ({m['successful']}/{m['num_prompts']})") + + print(f">>> done. 汇总: python3 {os.path.join(os.path.dirname(__file__), 'summarize.py')} {outdir}") + + +if __name__ == "__main__": + main() diff --git a/bench/qwen3_moe/run_bench.sh b/bench/qwen3_moe/run_bench.sh new file mode 100755 index 000000000..6a278764f --- /dev/null +++ b/bench/qwen3_moe/run_bench.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# 客户端并发吞吐压测（vllm bench serve），对齐 T2-1-2 的服务测法。 +# 遍历 concurrency × input-len × output-len 矩阵，每组结果落一个文件。 +# +# 前置：serve.sh 已在同机 PORT 上把服务起好。 +# +# 用法： +# TAG=base ./run_bench.sh # 跑基线版（如 main 的 naive MoE），结果存 bench_results/base +# TAG=this ./run_bench.sh # 跑本方案（T2-1-3 grouped_gemm），结果存 bench_results/this +# +# 环境变量： +# TAG 结果子目录名（base / this），用于 base-vs-this 对比 +# PORT 服务端口（默认 8102） +# MODEL_NAME 请求里的 model 字段（服务单模型，一般被忽略；默认见下） +# TOKENIZER 分词器路径，用于准确统计 token 数 +# BATCH_SIZES concurrency 列表（默认 "1 8 32"） +# INPUT_LENS 输入长度列表（默认 "32 256 2048"） +# OUTPUT_LENS 输出长度列表（默认 "256 1024"） +# REPEAT 每组的请求数 = concurrency × REPEAT（默认 3） +# --------------------------------------------------------------------------- +set -euo pipefail + +TAG="${TAG:-this}" +PORT="${PORT:-8102}" +MODEL_NAME="${MODEL_NAME:-Qwen3-30B-A3B-Thinking-2507}" +TOKENIZER="${TOKENIZER:-/data/huggingface_home/Qwen3-30B-A3B-Thinking-2507}" +REPEAT="${REPEAT:-3}" + +read -r -a BATCH_SIZES <<< "${BATCH_SIZES:-1 8 32}" +read -r -a INPUT_LENS <<< "${INPUT_LENS:-32 256 2048}" +read -r -a OUTPUT_LENS <<< "${OUTPUT_LENS:-256 1024}" + +OUTDIR="${OUTDIR:-bench_results/${TAG}}" +mkdir -p "${OUTDIR}" +unset http_proxy https_proxy all_proxy ALL_PROXY 2>/dev/null || true + +echo ">>> TAG=${TAG} port=${PORT} -> ${OUTDIR}" +echo ">>> concurrency=[${BATCH_SIZES[*]}] input=[${INPUT_LENS[*]}] output=[${OUTPUT_LENS[*]}]" + +for bs in "${BATCH_SIZES[@]}"; do + for il in "${INPUT_LENS[@]}"; do + for ol in "${OUTPUT_LENS[@]}"; do + sleep 1 + seed=$(date +%s) + num_prompts=$(( bs * REPEAT )) + out="${OUTDIR}/bs${bs}_in${il}_out${ol}.txt" + echo ">>> bs=${bs} in=${il} out=${ol} num_prompts=${num_prompts} -> ${out}" + + # 用 --extra-body 把 max_tokens/ignore_eos 塞进 chat 请求，保证每请求生成满 output-len， + # 吞吐才可比（与 T2-1-2 一致）。--request-rate inf = 爆发模式。 + vllm bench serve \ + --backend openai-chat \ + --model "${MODEL_NAME}" \ + --tokenizer "${TOKENIZER}" \ + --endpoint /v1/chat/completions \ + --port "${PORT}" \ + --request-rate inf \ + --seed "${seed}" \ + --num-prompts "${num_prompts}" \ + --max-concurrency "${bs}" \ + --random-input-len "${il}" \ + --extra-body "{\"max_tokens\": ${ol}, \"ignore_eos\": true}" \ + > "${out}" 2>&1 || { echo "!! 该组失败，见 ${out}"; continue; } + done + done +done + +echo ">>> done. 汇总： python3 $(dirname "$0")/summarize.py ${OUTDIR}" diff --git a/bench/qwen3_moe/serve.sh b/bench/qwen3_moe/serve.sh new file mode 100755 index 000000000..0b4adb634 --- /dev/null +++ b/bench/qwen3_moe/serve.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# 启动 InfiniLM OpenAI 兼容服务，用于 Qwen3-30B-A3B-Thinking-2507 的并发吞吐基准。 +# 目标平台：沐曦 MetaX C500，TP=2（两张 64GB 卡）。 +# +# 用法： +# MODEL=/data/huggingface_home/Qwen3-30B-A3B-Thinking-2507 ./serve.sh +# MAX_CON=32 NUM_BLOCKS=1024 ./serve.sh # 覆盖默认值 +# GRAPH=1 ./serve.sh # 试开 graph（MoE 含 host dispatch，可能不兼容） +# +# 环境变量（均可覆盖）： +# MODEL 模型路径 +# PORT 服务端口（默认 8102，与客户端脚本一致） +# TP 张量并行度（默认 2） +# MAX_CON --max-batch-size，必须 >= 你要压测的最大 concurrency（默认 32） +# NUM_BLOCKS KV cache 分块数（默认 1024；OOM 就调小，不够就调大，见 README） +# MAX_NEW 单请求最大生成 token（默认 4096，需 >= 压测的最大 output-len） +# MAX_CACHE --max-cache-len（默认 4096，需 >= 最大 input+output） +# ATTN 注意力后端（默认 paged-attn；MetaX 未编 flash-attn，不要用 flash-attn） +# PAGED 1=开 paged KV cache（默认，并发批处理需要）；0=退回 static cache +# --------------------------------------------------------------------------- +set -euo pipefail + +MODEL="${MODEL:-/data/huggingface_home/Qwen3-30B-A3B-Thinking-2507}" +DEVICE="${DEVICE:-metax}" +PORT="${PORT:-8102}" +TP="${TP:-2}" +MAX_CON="${MAX_CON:-32}" +NUM_BLOCKS="${NUM_BLOCKS:-1024}" +MAX_NEW="${MAX_NEW:-4096}" +MAX_CACHE="${MAX_CACHE:-4096}" +ATTN="${ATTN:-paged-attn}" + +REPO="$(cd "$(dirname "$0")/../.." && pwd)" +export PYTHONPATH="${REPO}/python:${PYTHONPATH:-}" + +OPT_FLAGS=() +[[ "${GRAPH:-0}" == "1" ]] && OPT_FLAGS+=(--enable-graph) +[[ "${PAGED:-1}" == "1" ]] && OPT_FLAGS+=(--enable-paged-attn --num-blocks "${NUM_BLOCKS}") + +echo ">>> serving ${MODEL}" +echo ">>> device=${DEVICE} tp=${TP} port=${PORT} max_batch=${MAX_CON} attn=${ATTN} paged=${PAGED:-1} num_blocks=${NUM_BLOCKS} graph=${GRAPH:-0}" + +exec python3 "${REPO}/python/infinilm/server/inference_server.py" \ + --model "${MODEL}" \ + --device "${DEVICE}" --tp "${TP}" \ + --port "${PORT}" \ + --max-batch-size "${MAX_CON}" \ + --max-new-tokens "${MAX_NEW}" \ + --max-cache-len "${MAX_CACHE}" \ + --temperature 1.0 --top-p 0.8 --top-k 1 \ + --attn "${ATTN}" \ + --ignore-eos \ + "${OPT_FLAGS[@]}" diff --git a/bench/qwen3_moe/summarize.py b/bench/qwen3_moe/summarize.py new file mode 100755 index 000000000..ffaa42be0 --- /dev/null +++ b/bench/qwen3_moe/summarize.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Summarize `vllm bench serve` result files into a throughput table. + +Parses the `bsX_inY_outZ.txt` files produced by run_bench.sh and prints one row +per (concurrency, input, output) config. Pass one result dir to just tabulate, +or two dirs (base first, then this) to get a base-vs-this comparison with the +output/total throughput deltas that the 赛题 scores on. + +Usage: + python3 summarize.py bench_results/this + python3 summarize.py bench_results/base bench_results/this +""" +import os +import re +import sys +import glob + +# Label -> canonical metric. vllm 各版本标签略有差异，用宽松正则匹配数值。 +_METRICS = { + "req_s": r"Request throughput \(req/s\):\s*([\d.]+)", + "out_tok_s": r"Output token throughput \(tok/s\):\s*([\d.]+)", + "total_tok_s": r"Total Token throughput \(tok/s\):\s*([\d.]+)", + "ttft_ms": r"Mean TTFT \(ms\):\s*([\d.]+)", + "tpot_ms": r"Mean TPOT \(ms\):\s*([\d.]+)", +} +_FNAME = re.compile(r"bs(\d+)_in(\d+)_out(\d+)\.txt$") + + +def parse_file(path): + text = open(path, encoding="utf-8", errors="replace").read() + row = {} + for key, pat in _METRICS.items(): + m = re.search(pat, text) + row[key] = float(m.group(1)) if m else None + return row + + +def load_dir(d): + """Return {(bs, in, out): metrics} for one result dir.""" + out = {} + for path in sorted(glob.glob(os.path.join(d, "bs*_in*_out*.txt"))): + m = _FNAME.search(os.path.basename(path)) + if not m: + continue + cfg = tuple(int(x) for x in m.groups()) + out[cfg] = parse_file(path) + return out + + +def fmt(v, width=10): + return ("%.2f" % v).rjust(width) if isinstance(v, float) else "-".rjust(width) + + +def print_single(d): + data = load_dir(d) + if not data: + print(f"(无结果文件: {d})") + return + print(f"\n=== {d} ===") + hdr = ["bs", "in", "out", "out_tok/s", "total_tok/s", "req/s", "TTFT_ms", "TPOT_ms"] + print(" ".join(h.rjust(10) for h in hdr)) + for cfg in sorted(data): + r = data[cfg] + cols = [str(cfg[0]), str(cfg[1]), str(cfg[2]), + r["out_tok_s"], r["total_tok_s"], r["req_s"], r["ttft_ms"], r["tpot_ms"]] + print(" ".join(c.rjust(10) if isinstance(c, str) else fmt(c) for c in cols)) + + +def print_compare(base_dir, this_dir): + base, this = load_dir(base_dir), load_dir(this_dir) + keys = sorted(set(base) | set(this)) + if not keys: + print("(两目录都无结果)") + return + print(f"\n=== base={base_dir} vs this={this_dir} ===") + print("列: out_tok/s (base -> this, Δ%) | total_tok/s (base -> this, Δ%)") + hdr = ["bs", "in", "out", "out_base", "out_this", "out_Δ%", "tot_base", "tot_this", "tot_Δ%"] + print(" ".join(h.rjust(10) for h in hdr)) + for cfg in keys: + b, t = base.get(cfg, {}), this.get(cfg, {}) + + def delta(bv, tv): + if isinstance(bv, float) and isinstance(tv, float) and bv > 0: + return (tv - bv) / bv * 100.0 + return None + + cols = [str(cfg[0]), str(cfg[1]), str(cfg[2]), + b.get("out_tok_s"), t.get("out_tok_s"), delta(b.get("out_tok_s"), t.get("out_tok_s")), + b.get("total_tok_s"), t.get("total_tok_s"), delta(b.get("total_tok_s"), t.get("total_tok_s"))] + print(" ".join(c.rjust(10) if isinstance(c, str) else fmt(c) for c in cols)) + + +def main(): + args = sys.argv[1:] + if len(args) == 1: + print_single(args[0]) + elif len(args) == 2: + print_compare(args[0], args[1]) + else: + print(__doc__) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp index d08b3bf77..b8221672e 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp +++ b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp @@ -1,21 +1,61 @@ #include "qwen3_moe_experts.hpp" +#include "../../global_state/global_state.hpp" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" -#include +#include +#include +#include namespace infinilm::models::qwen3_moe { Qwen3MoeExperts::Qwen3MoeExperts(std::shared_ptr model_config, const infinicore::Device &device) { + dtype_ = model_config->get_dtype(); + device_ = device; + num_experts_ = model_config->get("num_experts"); num_experts_per_tok_ = model_config->get("num_experts_per_tok"); + hidden_size_ = model_config->get("hidden_size"); + size_t moe_intermediate_size = model_config->get("moe_intermediate_size"); ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tp_rank_ = rank_info.tp_rank; + tp_size_ = rank_info.tp_size; + communicator_ = rank_info.comm; + ASSERT(moe_intermediate_size % size_t(tp_size_) == 0); + moe_intermediate_size_per_rank_ = moe_intermediate_size / size_t(tp_size_); + + gate_proj_stacked_ = infinicore::nn::Parameter( + {num_experts_, moe_intermediate_size_per_rank_, hidden_size_}, + dtype_, device_); + up_proj_stacked_ = infinicore::nn::Parameter( + {num_experts_, moe_intermediate_size_per_rank_, hidden_size_}, + dtype_, device_); + down_proj_stacked_ = infinicore::nn::Parameter( + {num_experts_, hidden_size_, moe_intermediate_size_per_rank_}, + dtype_, device_); + for (size_t i = 0; i < num_experts_; ++i) { - experts_.push_back(this->register_module(std::to_string(i), model_config, device)); + const std::string suffix = std::to_string(i) + "."; + + auto gate_slab = gate_proj_stacked_->narrow({{0, i, 1}})->squeeze(0); + this->register_parameter(suffix + "gate_proj.weight", + infinicore::nn::Parameter(gate_slab, /*tp_dim=*/0, tp_rank_, tp_size_)); + + auto up_slab = up_proj_stacked_->narrow({{0, i, 1}})->squeeze(0); + this->register_parameter(suffix + "up_proj.weight", + infinicore::nn::Parameter(up_slab, /*tp_dim=*/0, tp_rank_, tp_size_)); + + auto down_slab = down_proj_stacked_->narrow({{0, i, 1}})->squeeze(0); + this->register_parameter(suffix + "down_proj.weight", + infinicore::nn::Parameter(down_slab, /*tp_dim=*/1, tp_rank_, tp_size_)); } } @@ -23,42 +63,106 @@ infinicore::Tensor Qwen3MoeExperts::forward(const infinicore::Tensor &hidden_sta const infinicore::Tensor &top_k_index, const infinicore::Tensor &top_k_weights) const { ASSERT(hidden_states->ndim() == 2); + ASSERT(top_k_index->ndim() == 2); + ASSERT(top_k_weights->ndim() == 2); - auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU); - auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); - - int *top_k_index_ptr = reinterpret_cast(top_k_index_cpu->data()); - float *top_k_weights_ptr = reinterpret_cast(top_k_weights_cpu->data()); - - size_t ntoken = hidden_states->shape()[0]; - int index; - float score; - - auto final_hidden_states = infinicore::Tensor::empty(hidden_states->shape(), hidden_states->dtype(), hidden_states->device()); - for (size_t itok = 0; itok < ntoken; ++itok) { - auto hidden_states_i = hidden_states->narrow({{0, itok, 1}}); - const size_t route_row = itok * num_experts_per_tok_; + const size_t ntoken = hidden_states->shape()[0]; + const size_t k = top_k_index->shape()[1]; + const size_t total_assignments = ntoken * k; - infinicore::Tensor final_hidden_states_i; - for (size_t k = 0; k < num_experts_per_tok_; ++k) { - index = top_k_index_ptr[route_row + k]; - score = top_k_weights_ptr[route_row + k]; - - ASSERT(index >= 0 && static_cast(index) < num_experts_); - - experts_[index]->set_alpha(score); - auto expert_out = experts_[index]->forward(hidden_states_i); + // ---- Step 1: device->host sync of routing tables. + infinicore::context::syncStream(); + auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); + auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU); + const int32_t *idx_host = reinterpret_cast(top_k_index_cpu->data()); + const float *w_host = reinterpret_cast(top_k_weights_cpu->data()); + + // ---- Step 2: CPU-side bucketing. + std::vector counts(num_experts_, 0); + for (size_t a = 0; a < total_assignments; ++a) { + int32_t e = idx_host[a]; + ASSERT(e >= 0 && size_t(e) < num_experts_); + ++counts[e]; + } + std::vector offsets(num_experts_, 0); + { + int32_t acc = 0; + for (size_t e = 0; e < num_experts_; ++e) { + offsets[e] = acc; + acc += counts[e]; + } + } - if (k == 0) { - final_hidden_states_i = expert_out; - } else { - infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, expert_out); + auto perm_token_cpu = infinicore::Tensor::empty( + {total_assignments}, infinicore::DataType::I32, infinicore::Device(infinicore::Device::Type::CPU)); + auto perm_weight_cpu = infinicore::Tensor::empty( + {total_assignments}, dtype_, infinicore::Device(infinicore::Device::Type::CPU)); + int32_t *perm_token = reinterpret_cast(perm_token_cpu->data()); + void *perm_weight_raw = perm_weight_cpu->data(); + + { + std::vector cursor(offsets); + for (size_t a = 0; a < total_assignments; ++a) { + int32_t e = idx_host[a]; + int32_t row = cursor[e]++; + perm_token[row] = int32_t(a / k); + switch (dtype_) { + case infinicore::DataType::BF16: + reinterpret_cast(perm_weight_raw)[row] = f32_to_bf16(w_host[a]); + break; + case infinicore::DataType::F16: + reinterpret_cast(perm_weight_raw)[row] = f32_to_f16(w_host[a]); + break; + case infinicore::DataType::F32: + reinterpret_cast(perm_weight_raw)[row] = w_host[a]; + break; + default: + throw std::runtime_error("Qwen3MoeExperts: unsupported dtype for routing weight broadcast."); } } + } - final_hidden_states->narrow({{0, itok, 1}})->copy_from(final_hidden_states_i); + auto counts_cpu = infinicore::Tensor::empty( + {num_experts_}, infinicore::DataType::I32, infinicore::Device(infinicore::Device::Type::CPU)); + std::memcpy(counts_cpu->data(), counts.data(), counts.size() * sizeof(int32_t)); + + // ---- Step 3: upload bucketing results to the active device. + auto perm_token_dev = perm_token_cpu->to(device_); + auto perm_weight_dev = perm_weight_cpu->to(device_); + auto counts_dev = counts_cpu->to(device_); + + // ---- Step 4: gather hidden states by permuted token indices. + auto permuted_hidden = infinicore::op::embedding(perm_token_dev, hidden_states); + + // ---- Step 5: per-projection grouped GEMMs. + // `counts` already lives on the host (computed in Step 2), so we hand it to + // grouped_gemm directly: device backends then skip the per-call device->host + // copy + stream sync of the group sizes. On MetaX that copy is effectively a + // hard sync, so this removes 3 hard syncs per layer on the decode path. + // Under graph recording the op runs at replay time, after this local + // `counts` vector is gone; pass nullptr there so the device-sync path is used. + const int32_t *counts_host = infinicore::context::isGraphRecording() ? nullptr : counts.data(); + auto gate_out = infinicore::op::grouped_gemm(permuted_hidden, gate_proj_stacked_, counts_dev, 1.0f, 0.0f, counts_host); + auto up_out = infinicore::op::grouped_gemm(permuted_hidden, up_proj_stacked_, counts_dev, 1.0f, 0.0f, counts_host); + auto intermediate = infinicore::op::swiglu(up_out, gate_out); + auto down_out = infinicore::op::grouped_gemm(intermediate, down_proj_stacked_, counts_dev, 1.0f, 0.0f, counts_host); + + // ---- Step 6: scale each row by its routing weight (broadcast over hidden). + auto weight_broadcast = perm_weight_dev->as_strided( + {total_assignments, hidden_size_}, {1, 0}); + auto scaled = infinicore::op::mul(down_out, weight_broadcast); + + // ---- Step 7: scatter-add back into a zero-initialized per-token buffer. + auto output = infinicore::Tensor::zeros({ntoken, hidden_size_}, dtype_, device_); + infinicore::op::index_add_(output, output, /*dim=*/0, perm_token_dev, scaled, 1.0f); + + // ---- Step 8: TP all-reduce on the partial down_proj output. + if (tp_size_ > 1 && communicator_ != nullptr) { + infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); + infinicore::context::syncStream(); } - return final_hidden_states; + + return output; } } // namespace infinilm::models::qwen3_moe diff --git a/csrc/models/qwen3_moe/qwen3_moe_experts.hpp b/csrc/models/qwen3_moe/qwen3_moe_experts.hpp index 2f470ca88..560ad35f1 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_experts.hpp +++ b/csrc/models/qwen3_moe/qwen3_moe_experts.hpp @@ -1,16 +1,21 @@ #pragma once -#include "../../layers/moe/legacy/moe_mlp.hpp" -#include "infinicore/nn/module.hpp" -#include "infinicore/tensor.hpp" - -#include +#include #include namespace infinilm::models::qwen3_moe { -using Qwen3MoeMLP = infinilm::layers::moe::legacy::MoeMLP; - +// Stacked-weight experts that dispatch a single `grouped_gemm` per projection +// instead of one MLP forward per (token, expert) pair. +// +// Weight layout (after TP slicing): +// gate_proj : [num_experts, moe_intermediate_size / tp_size, hidden_size] +// up_proj : [num_experts, moe_intermediate_size / tp_size, hidden_size] +// down_proj : [num_experts, hidden_size, moe_intermediate_size / tp_size] +// +// HF state_dict keys `experts.{i}.{gate,up,down}_proj.weight` are aliased to +// views into the corresponding slab so the existing safetensors loader fills +// them directly -- no Python remapper required. class Qwen3MoeExperts : public infinicore::nn::Module { public: Qwen3MoeExperts(std::shared_ptr model_config, @@ -21,9 +26,23 @@ class Qwen3MoeExperts : public infinicore::nn::Module { const infinicore::Tensor &top_k_weights) const; protected: - INFINICORE_NN_MODULE_VEC(Qwen3MoeMLP, experts); - size_t num_experts_per_tok_{0}; + // Stacked weight parameters. We register them under non-HF names so the + // load_state_dict walk never picks them up directly; only the per-expert + // aliases below are populated by the loader. + INFINICORE_NN_PARAMETER(gate_proj_stacked); + INFINICORE_NN_PARAMETER(up_proj_stacked); + INFINICORE_NN_PARAMETER(down_proj_stacked); + size_t num_experts_{0}; + size_t num_experts_per_tok_{0}; + size_t hidden_size_{0}; + size_t moe_intermediate_size_per_rank_{0}; + infinicore::DataType dtype_; + infinicore::Device device_; + + int tp_size_{1}; + int tp_rank_{0}; + infinicclComm_t communicator_{nullptr}; }; } // namespace infinilm::models::qwen3_moe diff --git a/test/models/qwen3_moe/CORRECTNESS_README.md b/test/models/qwen3_moe/CORRECTNESS_README.md new file mode 100644 index 000000000..1353f6682 --- /dev/null +++ b/test/models/qwen3_moe/CORRECTNESS_README.md @@ -0,0 +1,59 @@ +# Qwen3-MoE 端到端正确性验证（prefill 首步对齐 HF） + +验证 InfiniLM 适配的 Qwen3-30B-A3B-Thinking-2507 与 HuggingFace 参照在 **prefill 首步** +的一致性。核心指标是**首个生成 token（= prefill logits 的 argmax）与 HF 一致**，并顺带校验 +**prompt 分词（chat template + `` 前缀）与 HF 一致**。 + +## 为什么是"首 token"而不是"整段 logits bit-exact" + +- 本模型 **run-to-run 非确定**（TP all-reduce 规约顺序 + BF16 `index_add_` 原子加）， + 同一二进制两次贪心生成会从中途分叉 → 整条 token 序列 / 逐元素 logit 对齐不可靠。 +- **首 token 的 argmax 对微小数值噪声鲁棒**，是贪心解码真正依赖的量，也正是赛题要求的 + "输出 token 一致性"最本质的一步。脚本额外报告"前缀一致长度"作为更强的参考信号。 + +## 为什么两阶段 + +30B 模型放不下两份（InfiniLM TP=2 + HF 参照）。所以： +1. 先单独用 HF 生成参照并存盘（`dump_hf_reference.py`）； +2. 再单独用 InfiniLM 加载并对比（`check_prefill_logits.py`）。 + +HF 仅作**测试参照**，不进入 InfiniLM 推理路径（符合赛题规则）。 + +## 用法 + +### 1) 生成 HF 参照（HF 单独占用显存/内存） + +```bash +cd /data/InfiniLM +pip install transformers torch # 若未装 + +# GPU（device_map=auto 可跨多卡）： +python3 test/models/qwen3_moe/dump_hf_reference.py \ + --model /data/huggingface_home/Qwen3-30B-A3B-Thinking-2507 \ + --device cuda --out /tmp/qwen3_ref.json +# 或 CPU（~60GB 内存，慢但不占 GPU）：--device cpu +``` + +### 2) InfiniLM 对比（确保 HF 进程已退出、显存已释放） + +```bash +python3 test/models/qwen3_moe/check_prefill_logits.py \ + --model /data/huggingface_home/Qwen3-30B-A3B-Thinking-2507 \ + --device metax --tp 2 --ref /tmp/qwen3_ref.json +``` + +退出码 0 = PASS，非 0 = FAIL（可用于 CI）。 + +## 判读 + +- **全部 PASS**：first_tok 与 prompt_tok 都对齐 → 权重加载 / chat template / 路由正确。 +- **first_tok 或 prompt_tok 不一致 → 真 bug**（权重/模板/路由适配错误）。 +- 仅个别 **prefix 早停但 first_tok 全 ==** → 通常是非确定性或近似平票，非适配错误。 + +## 参数 + +- `--max-new-tokens`（dump 侧，默认 16）：生成多少 token 供比对前缀。 +- `--prefix-threshold`（check 侧，默认 1）：PASS 要求每条前缀一致的最少 token 数； + 设 1 即"首 token 对齐"。 +- prompts 固定在 `dump_hf_reference.py` 的 `DEFAULT_PROMPTS`，并写入参照文件， + checker 复用同一批 messages，保证两侧 prompt 完全一致。 diff --git a/test/models/qwen3_moe/check_prefill_logits.py b/test/models/qwen3_moe/check_prefill_logits.py new file mode 100644 index 000000000..3f2421699 --- /dev/null +++ b/test/models/qwen3_moe/check_prefill_logits.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Stage 2/2 of the Qwen3-MoE prefill correctness check: compare InfiniLM vs HF. + +Loads the model through the InfiniLM engine (no transformers dependency in the +inference path), greedily generates for the same prompts saved by +`dump_hf_reference.py`, and checks alignment with the HF reference: + + 1. prompt tokenization matches HF (validates chat template + thinking prefix), + 2. **first generated token matches** — the prefill next-token = argmax of the + prefill logits; this is the meaningful "prefill 首步 logits 对齐" signal and + is robust to the model's run-to-run non-determinism (TP all-reduce order + + BF16 atomic scatter), which makes a full token-sequence / bit-exact logit + comparison unreliable, + 3. matching-prefix length is reported as a stronger (best-effort) signal. + +PASS when, for every prompt, the prompt tokens match and the first generated +token matches (and prefix >= --prefix-threshold). + +Usage: + python3 check_prefill_logits.py --model /path/Qwen3-30B-A3B-Thinking-2507 \ + --device metax --tp 2 --ref qwen3_moe_prefill_reference.json +""" +import argparse +import json +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../python")) + +from infinilm.llm.llm import LLM # noqa: E402 + +# Same mapping as BaseConfig.get_device_str: the LLM engine only accepts backend +# strings (cpu/cuda/mlu/musa/...), so e.g. "metax"/"iluvatar" must map to "cuda". +_DEVICE_STR_MAP = { + "cpu": "cpu", "nvidia": "cuda", "qy": "cuda", "cambricon": "mlu", + "ascend": "ascend", "metax": "cuda", "moore": "musa", "iluvatar": "cuda", + "kunlun": "kunlun", "hygon": "cuda", "ali": "cuda", + # already-backend strings pass through: + "cuda": "cuda", "mlu": "mlu", "musa": "musa", +} + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", required=True) + ap.add_argument("--ref", required=True, help="reference json from dump_hf_reference.py") + ap.add_argument("--device", default="metax") + ap.add_argument("--tp", type=int, default=2) + ap.add_argument("--prefix-threshold", type=int, default=1, + help="min matching leading tokens required per prompt to PASS (default 1)") + ap.add_argument("--cache-type", default="static", choices=["static", "paged"], + help="KV cache type; static matches examples/test_infer.py's proven path") + args = ap.parse_args() + + with open(args.ref, encoding="utf-8") as f: + ref = json.load(f) + entries = ref["entries"] + max_new = ref.get("max_new_tokens", 16) + + device_str = _DEVICE_STR_MAP.get(args.device.lower(), "cpu") + # Mirror examples/test_infer.py's proven invocation (static cache, default + # attn, no graph) so the engine initializes the same way it does there. + model = LLM( + model_path=os.path.expanduser(args.model), + device=device_str, + tensor_parallel_size=args.tp, + cache_type=args.cache_type, + max_batch_size=len(entries), + max_tokens=max_new, + temperature=1.0, + top_k=1, # greedy + top_p=1.0, + enable_graph=False, + attn_backend="default", + ) + + conversations = [e["messages"] for e in entries] + outputs = model.chat(messages=conversations) + + all_pass = True + print(f"\n{'result':6} | prompt_tok | first_tok (ref vs this) | prefix | prompt") + print("-" * 90) + for e, out in zip(entries, outputs): + ref_gen = e["gen_token_ids"] + this_gen = list(out.outputs[0].token_ids) + this_prompt = list(out.prompt_token_ids or []) + + prompt_match = (this_prompt == e["prompt_token_ids"]) if this_prompt else None + + prefix = 0 + for a, b in zip(ref_gen, this_gen): + if a == b: + prefix += 1 + else: + break + first_match = bool(ref_gen and this_gen and ref_gen[0] == this_gen[0]) + + ok = first_match and prefix >= args.prefix_threshold and (prompt_match is not False) + all_pass &= ok + + rt = ref_gen[0] if ref_gen else None + tt = this_gen[0] if this_gen else None + pm = {True: "ok", False: "MISMATCH", None: "n/a"}[prompt_match] + print(f"{'PASS' if ok else 'FAIL':6} | {pm:10} | {str(rt):>7} {'==' if first_match else '!='} " + f"{str(tt):<7} | {prefix:>2}/{min(len(ref_gen), len(this_gen))} | " + f"{e['messages'][0]['content'][:24]!r}") + + print("-" * 90) + print(f"=== OVERALL: {'PASS' if all_pass else 'FAIL'} " + f"({sum(1 for _ in entries)} prompts) ===") + if not all_pass: + print("提示：若仅个别 prefix 早停而 first_tok 全 ==，通常是非确定性/近似平票，非适配错误；" + "若 first_tok 或 prompt_tok 不一致，才是真 bug（权重加载 / chat template / 路由）。") + sys.exit(0 if all_pass else 1) + + +if __name__ == "__main__": + main() diff --git a/test/models/qwen3_moe/dump_hf_reference.py b/test/models/qwen3_moe/dump_hf_reference.py new file mode 100644 index 000000000..e779d63e6 --- /dev/null +++ b/test/models/qwen3_moe/dump_hf_reference.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Stage 1/2 of the Qwen3-MoE prefill correctness check: dump the HF reference. + +Loads Qwen3-30B-A3B-Thinking-2507 with HuggingFace transformers (**reference +only** — the InfiniLM adaptation itself must not depend on transformers), greedily +generates a few tokens for a fixed set of prompts, and saves the prompt token ids ++ generated token ids to JSON. Stage 2 (`check_prefill_logits.py`) then loads the +InfiniLM engine and compares against this file. + +Two-phase design because the 30B model does not fit in memory twice: run this +first (HF alone), then run the checker (InfiniLM alone). + +Usage: + # single GPU / multi-GPU (device_map=auto splits across visible GPUs): + python3 dump_hf_reference.py --model /path/Qwen3-30B-A3B-Thinking-2507 --device cuda + # CPU (needs ~60GB RAM for bf16; slow but avoids GPU contention): + python3 dump_hf_reference.py --model /path/... --device cpu +""" +import argparse +import json + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +# Fixed prompts covering zh/en + reasoning. Kept in the reference file so the +# checker replays the exact same messages (single source of truth). +DEFAULT_PROMPTS = [ + [{"role": "user", "content": "你好，请介绍一下自己"}], + [{"role": "user", "content": "用一句话解释什么是量子纠缠。"}], + [{"role": "user", "content": "What is the capital of France?"}], + [{"role": "user", "content": "计算 17 乘以 23 等于多少？请给出计算步骤。"}], +] + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", required=True) + ap.add_argument("--out", default="qwen3_moe_prefill_reference.json") + ap.add_argument("--device", default="cuda", choices=["cuda", "cpu"]) + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--max-new-tokens", type=int, default=16) + args = ap.parse_args() + + tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + dtype = getattr(torch, args.dtype) + model = AutoModelForCausalLM.from_pretrained( + args.model, + torch_dtype=dtype, + trust_remote_code=True, + device_map="auto" if args.device == "cuda" else None, + ) + if args.device == "cpu": + model = model.to("cpu") + model.eval() + + entries = [] + for messages in DEFAULT_PROMPTS: + # Newer transformers return a BatchEncoding (dict) here instead of a bare + # tensor; handle both. Also pass attention_mask when available. + tpl = tok.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt" + ) + if torch.is_tensor(tpl): + input_ids = tpl.to(model.device) + attn = None + else: + input_ids = tpl["input_ids"].to(model.device) + attn = tpl.get("attention_mask") + attn = attn.to(model.device) if attn is not None else None + with torch.no_grad(): + out = model.generate( + input_ids, attention_mask=attn, + max_new_tokens=args.max_new_tokens, do_sample=False, + ) + gen_ids = out[0, input_ids.shape[1]:].tolist() + entries.append({ + "messages": messages, + "prompt_token_ids": input_ids[0].tolist(), + "gen_token_ids": gen_ids, + "gen_text": tok.decode(gen_ids), + }) + first = gen_ids[0] if gen_ids else None + print(f"[hf] {messages[0]['content'][:20]!r} -> first_tok={first} " + f"{tok.decode(gen_ids[:1])!r}") + + payload = { + "model": args.model, + "max_new_tokens": args.max_new_tokens, + "note": "greedy (do_sample=False) HF reference; compare with check_prefill_logits.py", + "entries": entries, + } + with open(args.out, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + print(f"saved {len(entries)} entries -> {args.out}") + + +if __name__ == "__main__": + main()