Add opt-in multi-GPU sweep parallelism; fix GPU pinning and run dir races
For shape-heavy ops the autotune sweep dominates wall time, not the measurement. _parallel_warmup_plugin shards the sweep across N GPUs and hands the merged configs to the normal single-GPU serial measurement, so timing stays comparable. Activated by PARALLEL_WARMUP_GPUS=N (silent no-op when unset, so it can stay in the plugin list); output layout, REPLAY_FROM usage and the latency table are unchanged. Measurement is deliberately not parallelized: N processes saturating one box couple through the power/thermal budget, so per-card latency gets dragged by its neighbours by an amount that does not reproduce. Config keys whose winner differs across shards are counted and reported as a WARNING -- that count is how much to trust the run. Fixes found while auditing: - Shards no longer overwrite CUDA_VISIBLE_DEVICES with a bare shard index, which a caller who had selected idle cards (e.g. 6,7) would see re-interpreted as absolute ids 0,1 -- silently benchmarking on the busy cards they were avoiding. - Interrupting the sweep no longer leaves N subprocesses holding GPUs; children are terminated and reaped before the exception propagates (BaseException, since KeyboardInterrupt is not an Exception). - run_pytest.sh claims its output dir with a bare mkdir and retreats to a -2/-3 suffix on collision. The second-resolution timestamp meant two concurrent runs shared one directory and overwrote each other. - Replay eviction failures now emit a distinct compile_*_no_evict marker. LibTuner.cache is a sqlite-backed ConfigCache with no __delitem__, so the bad config could not be dropped and the retry re-read it, while the log still claimed a clean fallback to live autotune. Also trims comment density in both shell scripts and reworks the README: promotes the parallel and cudagraph sections out from under the A/B flow, groups the env table by purpose, documents that two record-mode runs are not comparable, and marks ab_fold_test.sh as a caliber example whose switch upstream has already removed.
This commit is contained in:
+30
-24
@@ -1,23 +1,19 @@
|
||||
#!/bin/bash
|
||||
# 单算子性能测试入口（手动调试用）。
|
||||
# 用法：改下方【配置区】，或用环境变量覆盖，例如：
|
||||
# 单算子性能测试入口。改配置区或用同名环境变量覆盖：
|
||||
# OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
|
||||
# 产物统一落在 runs/<op>_<时间戳>/：run.log、shapes.yaml、autotune_records/、ttgir/。
|
||||
# 产物落在 runs/<op>_<时间戳>/：run.log、shapes.yaml、autotune_records/、ttgir/
|
||||
set -euo pipefail
|
||||
|
||||
# ============================== 配置区 ==============================
|
||||
# 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。
|
||||
|
||||
# 被测算子：OP=测试函数名（去掉 test_ 前缀）；OP_FILE=benchmark 文件名
|
||||
# （对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>）
|
||||
# 被测算子，对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
|
||||
OP="${OP:-fused_marlin_moe_mxfp4}"
|
||||
OP_FILE="${OP_FILE:-fused_marlin_moe}"
|
||||
|
||||
# FlagGems 仓库路径
|
||||
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
|
||||
|
||||
# shape 来源：SHAPE_FILE 非空则用该 yaml；为空则用下方内置 yaml。
|
||||
# yaml 顶层 key 必须是 op 名。
|
||||
# shape 来源：非空则用该 yaml，否则用 INLINE_YAML。顶层 key 必须是 op 名。
|
||||
SHAPE_FILE="${SHAPE_FILE:-}"
|
||||
read -r -d '' INLINE_YAML <<'YAML' || true
|
||||
fused_marlin_moe_mxfp4:
|
||||
@@ -44,17 +40,17 @@ YAML
|
||||
# 调优空间：0=普通 autotune（默认，快速验证）；1=FlagTune 扩展空间（首跑全量搜索、慢）
|
||||
USE_FLAGTUNE="${USE_FLAGTUNE:-0}"
|
||||
|
||||
# autotune record/replay（见 _autotune_record_plugin.py）：
|
||||
# - 空：record 模式，本次选中的 config 记录到 $OUT_DIR/autotune_records/<op>.json；
|
||||
# - 指向某次历史 runs/<op>_<时间戳> 目录：replay 该次记录（A/B 两侧同 config）。
|
||||
# 空=record 模式，把本次选中的 config 记入 autotune_records/<op>.json；
|
||||
# 指向某次历史 run 目录则 replay 其记录，用于 A/B 两侧锁同一套 config。
|
||||
REPLAY_FROM="${REPLAY_FROM:-}"
|
||||
|
||||
# 终端颜色：always/never 强制开/关；为空则按 tty 自动判断（run.log 始终去色）
|
||||
# always/never 强制开关终端颜色，空则按 tty 判断（run.log 始终去色）
|
||||
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
|
||||
|
||||
# pytest 插件列表（可按需注释停用；各插件作用见 README「插件说明」）
|
||||
# 各插件作用见 README「插件说明」；注释掉某行即停用该插件
|
||||
PLUGINS=(
|
||||
-p _device_guard_plugin
|
||||
-p _parallel_warmup_plugin
|
||||
-p _seed_plugin
|
||||
-p _shape_inject_plugin
|
||||
-p _shape_iter_inject_plugin
|
||||
@@ -71,12 +67,24 @@ PLUGINS=(
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
|
||||
|
||||
OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
|
||||
# 产物目录，可用 OUT_DIR 指定。时间戳只到秒，并发启动会撞名，故用不带 -p 的
|
||||
# mkdir 抢占（已存在即失败），撞上就退让到 -2、-3……
|
||||
if [[ -n "${OUT_DIR:-}" ]]; then
|
||||
mkdir -p "$OUT_DIR"
|
||||
else
|
||||
mkdir -p "$SCRIPT_DIR/runs"
|
||||
BASE="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
|
||||
OUT_DIR="$BASE"
|
||||
n=1
|
||||
until mkdir "$OUT_DIR" 2>/dev/null; do
|
||||
n=$((n + 1))
|
||||
OUT_DIR="${BASE}-${n}"
|
||||
(( n > 99 )) && { echo "!!! 无法创建产物目录（$BASE 及后缀均被占用）" >&2; exit 1; }
|
||||
done
|
||||
fi
|
||||
LOG_FILE="$OUT_DIR/run.log"
|
||||
mkdir -p "$OUT_DIR"
|
||||
|
||||
# 颜色决策：下方 tee 管道会让 python 侧看不到 tty，所以在这里判断，
|
||||
# 并通过 FLAGGEMS_PERF_COLOR 下传给各插件（_term_style.py 消费）。
|
||||
# tee 管道会让 python 侧检测不到 tty，故在 shell 层判断后经环境变量下传
|
||||
if [[ -z "$FLAGGEMS_PERF_COLOR" && -t 1 && -z "${NO_COLOR:-}" ]]; then
|
||||
FLAGGEMS_PERF_COLOR=always
|
||||
fi
|
||||
@@ -94,7 +102,7 @@ export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
||||
export FLAGGEMS_PERF_CURRENT_OP="$OP"
|
||||
export PYTHONUNBUFFERED=1 # 实时输出不缓冲
|
||||
|
||||
# record/replay 二选一，通过环境变量激活对应模式
|
||||
# record/replay 互斥，各由自己的环境变量激活
|
||||
if [[ -n "$REPLAY_FROM" ]]; then
|
||||
AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records"
|
||||
[[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \
|
||||
@@ -106,7 +114,7 @@ else
|
||||
MODE_DESC=record
|
||||
fi
|
||||
|
||||
# 解析 shape 文件：SHAPE_FILE 为空时把内置 yaml 写到临时文件
|
||||
# 未指定 shape 文件时，把 INLINE_YAML 落到临时文件供 pytest 读取
|
||||
if [[ -z "$SHAPE_FILE" ]]; then
|
||||
SHAPE_FILE="$(mktemp --suffix=.yaml)"
|
||||
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
|
||||
@@ -119,9 +127,8 @@ status=0
|
||||
echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
|
||||
echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
|
||||
|
||||
# _ir_meta_plugin 在进程退出时把"实际被使用"的变体的 ttgir 按 shape 整理落盘
|
||||
# （挂在 atexit 上，CUDA crash 后已编译部分仍可拿到）。
|
||||
# 目录结构与命名图例见 <dump>/naming.md 和 index.tsv（后者也记录落选的 sweep 变体）。
|
||||
# 每次用独立的 Triton 缓存目录，跑完即删：保证编译过程可复现，且 ttgir 落盘
|
||||
# 只包含本次的变体（_ir_meta_plugin 在 atexit 里按 shape 整理）。
|
||||
CACHE_DIR="$OUT_DIR/.triton_cache"
|
||||
rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR"
|
||||
status=0
|
||||
@@ -143,7 +150,6 @@ status=0
|
||||
exit "$status"
|
||||
} 2>&1 | tee "$LOG_FILE" || status=$?
|
||||
|
||||
# 终端保留颜色；落盘的 run.log 去掉 ANSI 转义，保证 grep/diff 面对纯文本
|
||||
# （Ctrl-C 中断时会跳过去色，仅影响观感）。
|
||||
# run.log 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色；Ctrl-C 时会跳过这步）
|
||||
sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG_FILE"
|
||||
exit "$status"
|
||||
|
||||
Reference in New Issue
Block a user