3 Commits
Author SHA1 Message Date
Aharrypotter 2786eb9886 bench: orchestrate a multi-card campaign without touching the measurement path
campaign.sh fans a screened operator list across cards, one ab_repeat.sh per
card, serial within a card. It does no measuring: ab_repeat.sh is the path that
produced the four operator verdicts already on record, and it stays byte for
byte as it was.

Each card gets its own AB_ROOT. Eight ab_repeat.sh started in the same second
would otherwise compute the same timestamped directory name and two cards would
write into one.

ZL_ISOLATE_TUNER_CACHE=1 is forced, since eight processes on one sqlite is the
thing to avoid. The cost is that pinned configs differ from a shared-cache
batch, so absolute latencies cannot be quoted across the two; the caliber guard
refuses such comparisons on its own, which is why this is a note and not a
warning someone has to remember.

partition.py's shell output gained the per-card weight, because EXTRA_OPS has
to land on the lightest card and counting operators gets that wrong — one heavy
operator outweighs three light ones. It writes through stdout.buffer: Python
turns "\n" into CRLF on Windows, and the trailing \r then rides into the last
field and breaks the numeric comparison that picks the card.
2026-09-10 19:43:37 +08:00
Aharrypotter 03fdcbab28 bench: partition operators across cards by measured cost, and rescore in place
Two pieces the parallel campaign needs.

partition.py assigns screened operators to cards longest-first. run_batch.py's
existing hash partition ignores weight: on this batch it puts the critical path
at 120 minutes and leaves a card idle, where greedy gets 70. The floor is set
by the single heaviest operator — topk alone is 70.6 minutes against an ideal
share of 45.4 — so partitioning cannot beat it, only keep the other cards from
waiting. Weights come from screen.tsv's elapsed_before + elapsed_after, which
is one pair, i.e. 2 cells; N rounds is 2N cells, so the estimate is N times the
weight. Getting that wrong reported topk as 11.8 minutes.

RESCORE=<root> re-judges an existing screening root against the current
fingerprint and rewrites its screen.tsv, keeping the original alongside. The
criterion has already changed twice — index.tsv used to be hashed, and the
switch's own ttg.rlc-* annotations had to be stripped — and the second change
moved 24 operators from 23 IR_CHANGED to 16. Rescoring the artifacts already on
disk takes seconds; re-running that screen takes 79 minutes.
2026-09-10 19:25:34 +08:00
Aharrypotter c6a7f834cb bench: ignore the switch's own annotations when fingerprinting the IR
FlagTree's _apply_musa_rlc_policy(mod) stamps the whole RLC policy onto the
module as ttg.rlc-* attributes whenever the switch is on, regardless of whether
the pass then changes anything. Hashing the raw text therefore carried a
constant offset: every operator that compiles a kernel differed on that one
line and screened as IR_CHANGED. A filter that excludes nothing is not
filtering.

The null control could not catch it. Both of its sides run with the switch off,
so both carry the same attributes and the fingerprints match — it proved the
dump is deterministic, which was true, and said nothing about this.

Stripping ttg.rlc-* before hashing turns the screen back into a filter. Over
the same 24 operators, recomputed from the artifacts already on disk: 23
IR_CHANGED becomes 16, with 7 falling back to NO_IR_DIFF — tanh, gelu, mul,
softmax, mm, bmm, replication_pad3d, each byte-identical apart from that line.
replication_pad3d is the PR's largest NVIDIA win at +124.6%; on MUSA this pass
does not touch it.

The counts also show why the convert_layout count fails as a criterion in both
directions: addmm keeps cvt at 15 -> 15 while all 15 of its ttgir files differ
in substance, and tanh has no conversions and no change.
2026-09-10 19:13:51 +08:00
7 changed files with 442 additions and 8 deletions

No files matched your search

+31 -1
View File
@@ -137,6 +137,29 @@ ttgir 关键设计:
- **命名 = 组内有区分度的参数**:同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名;多词参数缩写为首字母(`BLOCK_SIZE_M` → `BSM`,图例见 naming.md);同名冲突依次用 dtype(`__fp16`)、参数对齐特化(`__EMdiv16`)、hash 前缀消歧。 - **命名 = 组内有区分度的参数**:同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名;多词参数缩写为首字母(`BLOCK_SIZE_M` → `BSM`,图例见 naming.md);同名冲突依次用 dtype(`__fp16`)、参数对齐特化(`__EMdiv16`)、hash 前缀消歧。
- **崩溃也能拿到 IR**:dump 挂在进程 atexit 上,CUDA crash 后仍会落盘已编译部分。 - **崩溃也能拿到 IR**:dump 挂在进程 atexit 上,CUDA crash 后仍会落盘已编译部分。
## 多卡跑一批:campaign.sh
`ab_repeat.sh` 是单卡串行的。一批算子要并行时用 `campaign.sh` 编排——它只分活,
测量仍然全部由 `ab_repeat.sh` 做,那个脚本一行不动。
```bash
SWITCH=FLAGTREE_MUSA_RLC_ENHANCE DEVICES=0,1,2,3,4,5,6,7 SCREEN=runs/ab_rlc_XXX/screen.tsv ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems DRY_RUN=1 bash campaign.sh
```
`DRY_RUN=1` 先看方案,确认了再去掉它起跑。
- **只取筛选表里 `IR_CHANGED` 的算子**——第一步的判定在这里生效
- **按实测耗时贪心分片**(`partition.py`),不是 hash 取模。S5000 上 16 个算子实测:
hash 分片关键路径 120 分钟、有一张卡空转;贪心 70 分钟。**下界由最重的单个算子
决定**(`topk` 一个就 70.6 分钟),加卡不解决,分片只保证别的卡不空等
- **每卡独立 `AB_ROOT`**——八个 `ab_repeat.sh` 同一秒启动,默认的时间戳目录名会撞
- **强制 `ZL_ISOLATE_TUNER_CACHE=1`**。8 个进程写同一个 sqlite 是要避开的事。代价是
pin 的 config 与共享缓存批次不同,**两批的绝对延迟不可互引**——口径守卫会自己拒绝
- `EXTRA_OPS` 追加筛选表里没有的算子(如 benchmark 刚修好的),放到**负载最小**的卡上
- 跑完合并各卡 `status.tsv`(首列是卡号),并统计非零退出的格与三个标记的总数
设备可见性变量名从 backends 层取,不写死 `MUSA_VISIBLE_DEVICES`。
## 两步:先筛选,再实测 ## 两步:先筛选,再实测
跑一个算子的六轮 A/B 要 13~44 分钟(S5000 实测:`rand` 15 shape 13.0 分钟, 跑一个算子的六轮 A/B 要 13~44 分钟(S5000 实测:`rand` 15 shape 13.0 分钟,
@@ -164,7 +187,14 @@ SWITCH=FLAGTREE_MUSA_RLC_ENHANCE ROUNDS=6 OPS="topk:topk" bash ab_repeat.sh
两件容易做错的: 两件容易做错的:
- **判据是 ttgir 指纹,不是 `convert_layout` 计数。** cvt 数相同不等于生成代码 - **判据是 ttgir 指纹,不是 `convert_layout` 计数。** cvt 数相同不等于生成代码
相同,而要测的是"生成代码变了没有"。cvt 数照样记,它解释的是变化的性质。 相同,而要测的是"生成代码变了没有"。两个方向都会骗人:`addmm` 的 cvt 是 15→15
没变,15 个 ttgir 文件却实质全不同;`tanh`/`softmax` 的 cvt 是 0→0,生成代码也
真的没动。cvt 数照样记,它解释的是变化的性质。
- **指纹先抹掉开关自己盖的注解。** FlagTree 的 `_apply_musa_rlc_policy(mod)` 在开关
打开时把整套 RLC 策略旋钮盖成 module 属性(`ttg.rlc-*`),跟这个 pass 后面有没有
真的改动什么无关。不抹的话判据有个常量偏移——**凡是能编出 kernel 的算子都必然判成
"变了",筛选器谁都不排除**。24 个算子实测:抹之前 23 个 `IR_CHANGED`,抹之后 7 个
落回 `NO_IR_DIFF`。零对照抓不到这个,因为两侧同值时属性也相同。
- **不要全局写死一个 `DTYPES`。** `OPS` 元素支持 `OP:OP_FILE:DTYPES`,留空表示用 - **不要全局写死一个 `DTYPES`。** `OPS` 元素支持 `OP:OP_FILE:DTYPES`,留空表示用
该算子自己的全套。`fft` 只吃 complex64,被一刀切的 `DTYPES=float16` 判死过一次 该算子自己的全套。`fft` 只吃 complex64,被一刀切的 `DTYPES=float16` 判死过一次
——那是 harness 的问题,被记成了算子的问题。 ——那是 harness 的问题,被记成了算子的问题。
+15 -6
View File
@@ -2,11 +2,18 @@
# 由 ab_repeat.sh 与 screen.sh 共用;两处必须用同一个判据,否则筛选放行的算子 # 由 ab_repeat.sh 与 screen.sh 共用;两处必须用同一个判据,否则筛选放行的算子
# 到了 A/B 那步又被拦下,或者反过来。 # 到了 A/B 那步又被拦下,或者反过来。
# 一个 run 生成代码的指纹:只取 .ttgir 的相对路径 + 内容。 # 一个 run 生成代码的指纹:只取 .ttgir 的相对路径 + 内容,并先抹掉开关自己盖上去的
# 路径进指纹是因为"变体集合变了"与"某个变体内容变了"同样算改变,而只 cat 内容 # 注解。
# 会把前者洗掉。只认 .ttgir 则是因为 dump 目录里还有 index.tsv 和 naming.md: #
# 一个 kernel 都没编出来的 run 只剩这两个文件、两侧完全相同,会把"什么都没生成" # 只认 .ttgir:dump 目录里还有 index.tsv 和 naming.md,一个 kernel 都没编出来的 run
# 报成"开关没改变生成代码"——router_gemm 在 MUSA 上就是这样。 # 只剩这两个文件、两侧完全相同,会把"什么都没生成"报成"开关没改变生成代码"。
#
# 抹掉 ttg.rlc-*:FlagTree 的 _apply_musa_rlc_policy(mod) 在开关打开时把整套 RLC 策略
# 旋钮盖成 module 属性,跟这个 pass 后面有没有真的改动什么无关。不抹的话判据有个常量
# 偏移——凡是能编出 kernel 的算子都必然判成"变了",筛选器谁都不排除。实测 24 个算子
# 里有 7 个(tanh gelu mul softmax mm bmm replication_pad3d)除这行属性外逐字节相同。
# 零对照抓不到它:两侧同值时属性也相同。
#
# 零 kernel 返回非零,交给调用方当作"没有可比的 IR",而不是"一致"。 # 零 kernel 返回非零,交给调用方当作"没有可比的 IR",而不是"一致"。
ir_fingerprint() { ir_fingerprint() {
local dir="$1/ttgir" local dir="$1/ttgir"
@@ -16,7 +23,9 @@ ir_fingerprint() {
(( n > 0 )) || return 1 (( n > 0 )) || return 1
local out local out
out=$( { find "$dir" -name '*.ttgir' -type f | LC_ALL=C sort | while IFS= read -r f; do out=$( { find "$dir" -name '*.ttgir' -type f | LC_ALL=C sort | while IFS= read -r f; do
printf '%s\n' "${f#"$dir"/}"; cat "$f" printf '%s
' "${f#"$dir"/}"
sed -E 's/"ttg[.]rlc-[^"]*" = [0-9-]+ : i32, ?//g' "$f"
done; } | sha256sum | cut -d' ' -f1 ) done; } | sha256sum | cut -d' ' -f1 )
[[ -n "$out" ]] && printf '%s' "$out" [[ -n "$out" ]] && printf '%s' "$out"
} }
+3 -1
View File
@@ -42,7 +42,9 @@ if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
export FLAGGEMS_PERF_COLOR=always export FLAGGEMS_PERF_COLOR=always
fi fi
ROOT="runs/ab_repeat_${SWITCH}_$(date +%Y%m%d-%H%M%S)" # AB_ROOT 供 campaign.sh 指定每卡独立的产物根:八个 ab_repeat.sh 同一秒启动时,
# 默认的时间戳目录名会撞在一起,两张卡往同一个目录里写。
ROOT="${AB_ROOT:-runs/ab_repeat_${SWITCH}_$(date +%Y%m%d-%H%M%S)}"
mkdir -p "$ROOT" mkdir -p "$ROOT"
LOG="$ROOT/summary.log" LOG="$ROOT/summary.log"
STATUS="$ROOT/status.tsv" STATUS="$ROOT/status.tsv"
+45
View File
@@ -39,6 +39,51 @@ set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "./_ir_fingerprint.sh" source "./_ir_fingerprint.sh"
# RESCORE=<已有的 root>:判据改了之后,用当前的 ir_fingerprint 对盘上已有的 ttgir
# 重新打分并改写该 root 的 screen.tsv,不重跑任何测量。
#
# 需要它是因为判据确实改过:_ir_fingerprint.sh 起初把 index.tsv 也算进哈希,后来又
# 发现要抹掉开关自己盖的 ttg.rlc-* 属性。同一批 24 个算子,结论从 23 个 IR_CHANGED
# 变成 16 个。重打分是秒级的,重跑那批筛选要 79 分钟。
#
# 原表存为 screen.tsv.before-rescore,判定有变化的算子逐个打出来。
RESCORE="${RESCORE:-}"
if [[ -n "$RESCORE" ]]; then
[[ -d "$RESCORE" ]] || { echo "!!! 目录不存在: $RESCORE" >&2; exit 2; }
TSV="$RESCORE/screen.tsv"
[[ -f "$TSV" ]] || { echo "!!! 没有 screen.tsv: $TSV" >&2; exit 2; }
cp "$TSV" "$TSV.before-rescore"
hdr=$(head -1 "$TSV.before-rescore")
# 列名 -> 序号,不写死列序
idx() { echo "$hdr" | tr '\t' '\n' | grep -nxF "$1" | cut -d: -f1; }
c_op=$(idx op); c_rb=$(idx rows_before); c_ra=$(idx rows_after); c_v=$(idx verdict)
[[ -n "$c_op" && -n "$c_v" ]] || { echo "!!! screen.tsv 缺 op/verdict 列" >&2; exit 2; }
printf '%s\n' "$hdr" > "$TSV"
changed=0
while IFS= read -r line; do
op=$(printf '%s' "$line" | cut -f"$c_op")
old=$(printf '%s' "$line" | cut -f"$c_v")
rb=$(printf '%s' "$line" | cut -f"${c_rb:-0}"); ra=$(printf '%s' "$line" | cut -f"${c_ra:-0}")
fa=$(ir_fingerprint "$RESCORE/${op}_before") || fa=""
fb=$(ir_fingerprint "$RESCORE/${op}_after") || fb=""
if [[ "${rb:-0}" == 0 || "${ra:-0}" == 0 ]]; then v=NO_ROWS
elif [[ -z "$fa" || -z "$fb" ]]; then v=NO_TRITON
elif [[ "$fa" == "$fb" ]]; then v=NO_IR_DIFF
else v=IR_CHANGED; fi
printf '%s\n' "$line" | awk -v FS='\t' -v OFS='\t' -v i="$c_v" -v v="$v" '{$i=v; print}' >> "$TSV"
if [[ "$old" != "$v" ]]; then
printf ' %-22s %-12s -> %s\n' "$op" "$old" "$v" >&2
changed=$((changed+1))
fi
done < <(tail -n +2 "$TSV.before-rescore")
echo ">>> 重打分完成,$changed 个算子的判定变了;原表存为 $TSV.before-rescore"
awk -v i="$c_v" -F'\t' 'NR>1{c[$i]++} END{for (k in c) printf(" %-12s %d\n", k, c[k])}' "$TSV"
exit 0
fi
VENDOR="${VENDOR:-MUSA}" VENDOR="${VENDOR:-MUSA}"
ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE" ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"
+128
View File
@@ -0,0 +1,128 @@
#!/bin/bash
# 多卡跑一批算子的六轮 A/B:每卡一个 ab_repeat.sh,卡内串行。
#
# 它只做编排,不做测量。ab_repeat.sh 一行不动——那是已经跑出结论的、验证过的串行
# 路径,不该为了并行去改它。分片交给 partition.py(按实测耗时贪心,比 hash 分片在
# 这批活上快 50 分钟)。
#
# 并行必须隔离 tuner cache:8 个进程写同一个 sqlite 是要避开的事。代价是 pin 下来
# 的 config 与共享缓存下跑的批次不同,两批的绝对延迟不可互引——batch_summary.py
# 的口径守卫会自己拒绝这种对比,不用记着。
#
# 用法:
# SWITCH=FLAGTREE_MUSA_RLC_ENHANCE DEVICES=0,1,2,3,4,5,6,7 \
# SCREEN=runs/ab_rlc_XXX/screen.tsv \
# ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash campaign.sh
#
# DRY_RUN=1 只打印分片方案,不起跑
# EXTRA_OPS=".." 追加筛选表里没有的算子(如 benchmark 刚修好的),塞给最空的卡
# ROUNDS=6 CELL_TIMEOUT=2400 VERDICT=IR_CHANGED
set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
: "${SCREEN:?SCREEN is required, 指向 ab_rlc_test.sh 的 screen.tsv}"
DEVICES="${DEVICES:-0}"
ROUNDS="${ROUNDS:-6}"
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"
VERDICT="${VERDICT:-IR_CHANGED}"
EXTRA_OPS="${EXTRA_OPS:-}"
DRY_RUN="${DRY_RUN:-0}"
TAB=$'\t'
[[ -f "$SCREEN" ]] || { echo "!!! 筛选表不存在: $SCREEN" >&2; exit 2; }
# 设备可见性变量名从 backends 层拿,不写死 MUSA_VISIBLE_DEVICES
VIS_KEY=$(python -c "import sys; sys.path.insert(0, '.')
from backends import get_backend
print(get_backend().visibility_env_keys[0])" 2>/dev/null) || VIS_KEY=""
[[ -n "$VIS_KEY" ]] || { echo "!!! 取不到后端的可见性变量名,检查 ZL_BACKEND" >&2; exit 2; }
echo ">>> switch=$SWITCH rounds=$ROUNDS devices=$DEVICES 可见性变量=$VIS_KEY"
python partition.py --screen "$SCREEN" --devices "$DEVICES" \
--verdict "$VERDICT" --rounds "$ROUNDS" || exit 2
# 每卡一行 "<device><TAB><OPS><TAB><权重秒>"
mapfile -t PLAN < <(python partition.py --screen "$SCREEN" --devices "$DEVICES" \
--verdict "$VERDICT" --format shell) || exit 2
if [[ -n "$EXTRA_OPS" ]]; then
# 追加到当前负载最小的卡。按算子个数选是错的——一个重算子能顶三个轻的,
# partition.py 的第三列给的就是权重秒数
lightest=0; least=""
for i in "${!PLAN[@]}"; do
IFS="$TAB" read -r _d _o _w <<< "${PLAN[$i]}"
if [[ -z "$least" ]] || (( ${_w:-0} < least )); then least=${_w:-0}; lightest=$i; fi
done
IFS="$TAB" read -r d o w <<< "${PLAN[$lightest]}"
PLAN[$lightest]="${d}${TAB}${o:+$o }${EXTRA_OPS}${TAB}${w}"
echo ">>> EXTRA_OPS 追加到卡${d}(当前最空,$(( ${w:-0} * ROUNDS / 60 )) 分钟): $EXTRA_OPS"
fi
if [[ "$DRY_RUN" == "1" ]]; then
echo ">>> DRY_RUN=1,不起跑。逐卡将执行:"
for row in "${PLAN[@]}"; do
IFS="$TAB" read -r d o w <<< "$row"
printf ' 卡%s %4s 分钟 OPS="%s"\n' "$d" "$(( ${w:-0} * ROUNDS / 60 ))" "$o"
done
exit 0
fi
ROOT="runs/campaign_${SWITCH}_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
printf '%s\n' "${PLAN[@]}" > "$ROOT/plan.tsv"
echo ">>> 产物: $(pwd)/$ROOT"
declare -a PIDS=() CARDS=()
for row in "${PLAN[@]}"; do
IFS="$TAB" read -r dev ops _w <<< "$row"
[[ -n "$ops" ]] || { echo ">>> 卡$dev 没有分到活,跳过"; continue; }
out="$ROOT/card${dev}"
mkdir -p "$out"
# AB_ROOT 让每卡有独立产物根:八个 ab_repeat.sh 同一秒启动,默认的时间戳目录名会撞
env "$VIS_KEY=$dev" ZL_DEVICE="$dev" ZL_ISOLATE_TUNER_CACHE=1 \
AB_ROOT="$out" SWITCH="$SWITCH" ROUNDS="$ROUNDS" \
CELL_TIMEOUT="$CELL_TIMEOUT" OPS="$ops" NO_COLOR=1 \
bash ab_repeat.sh > "$ROOT/card${dev}.log" 2>&1 &
PIDS+=($!); CARDS+=("$dev")
echo ">>> 卡$dev pid=$! ops=$ops"
done
(( ${#PIDS[@]} )) || { echo "!!! 没有一张卡分到活" >&2; exit 2; }
echo ">>> ${#PIDS[@]} 张卡已启动,等待全部结束……"
fail=0
for i in "${!PIDS[@]}"; do
if wait "${PIDS[$i]}"; then
echo ">>> 卡${CARDS[$i]} 结束 rc=0"
else
rc=$?; fail=$((fail+1))
echo "!!! 卡${CARDS[$i]} 结束 rc=$rc —— 看 $ROOT/card${CARDS[$i]}.log"
fi
done
# 合并逐格记录:卡号加在最前面,便于事后按卡追
{
first=1
for d in "${CARDS[@]}"; do
f="$ROOT/card${d}/status.tsv"
[[ -f "$f" ]] || continue
if (( first )); then printf 'device%s%s\n' "$TAB" "$(head -1 "$f")"; first=0; fi
tail -n +2 "$f" | sed "s/^/${d}${TAB}/"
done
} > "$ROOT/status.tsv"
echo
echo "########## 汇总 ##########"
echo ">>> 失败的卡: $fail / ${#PIDS[@]}"
if [[ -s "$ROOT/status.tsv" ]]; then
echo ">>> 逐格记录: $(pwd)/$ROOT/status.tsv ($(( $(wc -l < "$ROOT/status.tsv") - 1 )) 格)"
awk -F"$TAB" 'NR>1 && ($6+0)!=0 {n++} END{printf ">>> 非零退出的格: %d\n", n+0}' "$ROOT/status.tsv"
awk -F"$TAB" 'NR>1 {fb+=$10; ng+=$11; po+=$12}
END{printf ">>> 三个标记合计: replay_fallback=%d no_cudagraph=%d capture_poisoned=%d\n",
fb+0, ng+0, po+0}' "$ROOT/status.tsv"
fi
echo ">>> 逐算子报告在各卡的 summary.log 里;单独重出:"
echo " python batch_summary.py --ab-rounds $ROOT/card<N>/<op> --label <op>"
echo ">>> 注意:本批 ZL_ISOLATE_TUNER_CACHE=1,绝对延迟不可与共享缓存的批次互引"
+145
View File
@@ -0,0 +1,145 @@
#!/usr/bin/env python3
"""把筛选通过的算子按实测耗时分到各张卡上。
为什么不用 `run_batch.py` 的 hash 分片:它按算子名的 hash 取模,不看重量。S5000 上
17 个算子的六轮 campaign,hash 分片关键路径 120 分钟、有一张卡空转;按耗时贪心是
70 分钟。同一批活,差 50 分钟。
贪心(最长优先塞进当前最空的卡)在这里够用,也不可能更好:`topk` 一个算子就要 70.6
分钟,而理想均分是 45.4——**关键路径的下界由最重的那个算子决定,加卡不解决**。分片
只能保证除它之外的卡不空等。
权重来自 `screen.tsv` 的 `elapsed_before + elapsed_after`(**一对 = 2 格**的秒数)。
N 轮是 2N 格,所以预估耗时是权重的 N 倍——不是 N/6 倍,这里错过一次,把 topk 的 70.6
分钟报成了 11.8。缺列时退化成等权重,并在 stderr 说明——等权重的分片仍然
比不分好,但别让人以为它按耗时排过。
用法:
python partition.py --screen runs/ab_rlc_XXX/screen.tsv --devices 0,1,2,3
python partition.py --screen ... --devices 0,1 --verdict IR_CHANGED --format shell
"""
from __future__ import annotations
import argparse
import csv
import sys
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple
Item = Tuple[str, float]
def plan(items: Sequence[Item], slots: Sequence[str]) -> Dict[str, List[Item]]:
"""最长优先贪心:每次把最重的活放进当前总量最小的卡。
卡少于 1 张是调用方的错,直接抛;活比卡少时后面的卡就是空的,这不是错误。
"""
if not slots:
raise ValueError("需要至少一张卡")
out: Dict[str, List[Item]] = {s: [] for s in slots}
load: Dict[str, float] = {s: 0.0 for s in slots}
# 同重量时按名字排,保证同一份输入每次得到同一个方案
for name, weight in sorted(items, key=lambda kv: (-kv[1], kv[0])):
target = min(slots, key=lambda s: (load[s], s))
out[target].append((name, weight))
load[target] += weight
return out
def imbalance(assignment: Dict[str, List[Item]]) -> float:
"""最慢的卡 / 理想均分。1.0 是完美;下界由最重的单个算子决定。"""
loads = [sum(w for _, w in v) for v in assignment.values()]
total = sum(loads)
if total <= 0:
return 1.0
return max(loads) / (total / len(loads))
def read_screen(path: Path, verdict: str = "IR_CHANGED"
) -> Tuple[List[Item], Dict[str, str], List[str]]:
"""(算子, 一对耗时秒), op -> op_file, 被判定挡掉的算子。"""
with path.open(newline="", encoding="utf-8") as f:
rows = list(csv.DictReader(f, delimiter="\t"))
if not rows:
raise SystemExit(f"筛选表没有数据行: {path}")
items: List[Item] = []
files: Dict[str, str] = {}
skipped: List[str] = []
warned = False
for r in rows:
op = r.get("op", "")
if not op:
continue
if verdict and r.get("verdict") != verdict:
skipped.append(f"{op}({r.get('verdict', '?')})")
continue
try:
w = float(r["elapsed_before"]) + float(r["elapsed_after"])
except (KeyError, TypeError, ValueError):
w = 1.0
if not warned:
print("!!! 筛选表没有 elapsed_* 列,退化成等权重分片——"
"分片仍然比不分好,但它没有按耗时排过", file=sys.stderr)
warned = True
items.append((op, w))
files[op] = r.get("op_file") or op
return items, files, skipped
def main(argv: Optional[List[str]] = None) -> int:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0],
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--screen", type=Path, required=True, help="ab_rlc_test.sh 的 screen.tsv")
p.add_argument("--devices", required=True, help="逗号分隔的设备列表,如 0,1,2,3")
p.add_argument("--verdict", default="IR_CHANGED",
help="只取这个判定的算子;空字符串表示全要")
# 权重是一对(2 格)的秒数;N 轮是 2N 格,所以预估耗时是权重的 N 倍
p.add_argument("--rounds", type=int, default=6, help="仅用于打印预估耗时")
p.add_argument("--format", choices=("human", "shell"), default="human",
help="shell 输出每卡一行 `<device>\\t<OPS>`,供驱动脚本读")
args = p.parse_args(argv)
if not args.screen.is_file():
raise SystemExit(f"筛选表不存在: {args.screen}")
devices = [d.strip() for d in args.devices.split(",") if d.strip()]
if not devices:
raise SystemExit("--devices 为空")
items, files, skipped = read_screen(args.screen, args.verdict)
if not items:
raise SystemExit(f"筛选表里没有 verdict={args.verdict} 的算子——"
"先确认筛选跑过,且判据没有把所有算子都挡掉")
assignment = plan(items, devices)
if args.format == "shell":
# 第三列是该卡的权重秒数,供驱动脚本决定往哪张卡追加算子——按负载选,不是
# 按算子个数:一个重算子能顶三个轻的。
# 显式写 "\n" 而不用 print:Windows 上 print 会发 CRLF,读的一侧会把 \r
# 带进最后一个字段,数值比较就废了。
for d in devices:
ops = " ".join(f"{op}:{files[op]}" for op, _ in assignment[d])
load = sum(w for _, w in assignment[d])
sys.stdout.buffer.write(f"{d}\t{ops}\t{load:.0f}\n".encode())
sys.stdout.buffer.flush()
return 0
total = sum(w for _, w in items) * args.rounds
print(f">>> {len(items)} 个算子 / {len(devices)} 张卡 / {args.rounds} 轮")
if skipped:
print(f">>> 判定挡掉 {len(skipped)} 个: {', '.join(skipped)}")
for d in devices:
load = sum(w for _, w in assignment[d]) * args.rounds
names = " ".join(op for op, _ in assignment[d]) or "(空)"
print(f" 卡{d} {load/60:6.1f} 分钟 {names}")
crit = max(sum(w for _, w in v) for v in assignment.values()) * args.rounds
heaviest = max(items, key=lambda kv: kv[1])
print(f">>> 单卡串行 {total/3600:.2f} 小时;{len(devices)} 卡关键路径 "
f"{crit/60:.1f} 分钟,失衡 {imbalance(assignment):.2f}x")
print(f">>> 下界由最重的算子决定:{heaviest[0]} 单独就要 "
f"{heaviest[1]*args.rounds/60:.1f} 分钟,加卡不解决")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+75
View File
@@ -659,5 +659,80 @@ class ForeignVendorModuleTests(unittest.TestCase):
self.assertEqual(foreign_vendor_modules("mthreads", mods), []) self.assertEqual(foreign_vendor_modules("mthreads", mods), [])
class PartitionTests(unittest.TestCase):
"""按耗时把算子分到各卡。hash 分片在这批活上比它慢 50 分钟。"""
HEADER = "op\top_file\telapsed_before\telapsed_after\tverdict"
def _load(self, assignment):
return sorted(sum(w for _, w in v) for v in assignment.values())
def _screen(self, tmp, rows, header=None):
p = Path(tmp, "screen.tsv")
p.write_text((header or self.HEADER) + "\n" + "\n".join(rows) + "\n",
encoding="utf-8")
return p
def test_heavy_first_keeps_cards_even(self):
from partition import plan
items = [("a", 10), ("b", 9), ("c", 8), ("d", 7), ("e", 6), ("f", 5)]
self.assertEqual(self._load(plan(items, ["0", "1", "2"])), [15, 15, 15])
def test_one_item_heavier_than_the_ideal_share_sets_the_floor(self):
# topk 单个 70 分钟、理想均分 45——加卡不解决,分片只保证别人不空等
from partition import imbalance, plan
items = [("topk", 70)] + [(f"op{i}", 5) for i in range(10)]
a = plan(items, [str(i) for i in range(8)])
self.assertEqual(max(sum(w for _, w in v) for v in a.values()), 70)
self.assertGreater(imbalance(a), 1.0)
def test_more_cards_than_operators_leaves_some_empty(self):
from partition import plan
a = plan([("a", 1), ("b", 1)], ["0", "1", "2", "3"])
self.assertEqual(sum(1 for v in a.values() if not v), 2)
def test_same_input_gives_the_same_plan(self):
from partition import plan
items = [("a", 5), ("b", 5), ("c", 5), ("d", 5)]
first = plan(items, ["0", "1"])
for _ in range(5):
self.assertEqual(plan(list(reversed(items)), ["0", "1"]), first)
def test_no_cards_is_the_callers_error(self):
from partition import plan
with self.assertRaises(ValueError):
plan([("a", 1)], [])
def test_perfect_balance_is_one(self):
from partition import imbalance, plan
self.assertEqual(imbalance(plan([("a", 5), ("b", 5)], ["0", "1"])), 1.0)
def test_verdict_filters_and_reports_what_it_dropped(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\ta\t10\t10\tIR_CHANGED",
"b\tbb\t5\t5\tNO_IR_DIFF",
"c\tc\t1\t1\tNO_TRITON"])
items, files, skipped = read_screen(f)
self.assertEqual(items, [("a", 20.0)])
self.assertEqual(files["a"], "a")
self.assertEqual(len(skipped), 2)
def test_op_file_defaults_to_op_name(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\t\t1\t1\tIR_CHANGED"])
_, files, _ = read_screen(f)
self.assertEqual(files["a"], "a")
def test_missing_elapsed_falls_back_to_equal_weights(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\ta\tIR_CHANGED", "b\tb\tIR_CHANGED"],
header="op\top_file\tverdict")
items, _, _ = read_screen(f)
self.assertEqual([w for _, w in items], [1.0, 1.0])
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()