diff --git a/README.md b/README.md index 5d9c49a..d71529b 100644 --- a/README.md +++ b/README.md @@ -179,7 +179,17 @@ PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash WARNING 1 config key(s) got different winners across shards — parallel interference reached the sweep ``` -**冲突数就是这次加速的可信度指标**：0 可放心用；偏多说明 config 选择已被污染，出正式结论前不设该变量重跑一遍。其余行为（`REPLAY_FROM` 已设或可见卡不足 2 张时跳过、分片失败回退串行、轮转分片而非按块切、用子进程而非 xdist 的原因）见插件 docstring。分片绑卡遵守当前后端的可见性变量（CUDA 为 `CUDA_VISIBLE_DEVICES`）——用它选过卡时，分片只会落在你选的那几张上。 +**冲突数就是这次加速的可信度指标**：0 可放心用；偏多说明 config 选择已被污染，出正式结论前不设该变量重跑一遍。 + +插件会在三种情况下**主动放弃分片、退回串行 sweep**，日志各有独立措辞，都不是静默的： + +- **`BENCH_LEVEL` 不是 `core`**。`comprehensive` 会让上游在 `set_shapes()` 里并入 `set_more_shapes()`，那部分是各子类自己的 Python（硬编码元组加 `models_shapes/*.yaml`），插件只看得到 yaml 里的 shape 列表。按子集分片会让未覆盖的 key 全部现场自调优，而这一轮却自称在 replay——`mm` 上实测是 29 个 shape 只扫到 5 个、90 次 `AUTOTUNE_REPLAY_FALLBACK reason=key_missing`。 +- **任一分片非零退出**。死掉的分片的 shape 从未被扫过，合并出来的 config 集是带洞的，publish 之后每个洞都会静默自调优。因此整轮 sweep 直接作废，worker 日志保留在临时目录里备查。 +- yaml 里查不到该算子的 shape（见下）、`REPLAY_FROM` 已设、或可见卡不足 2 张。 + +yaml key 的查找与上游对齐：先按 `op_name`，查不到再沿 Benchmark 类链逐级尝试类名（`mm` 因此命中 `BlasBenchmark`）。类链是静态解析测试模块里 `op_name=` 那次构造调用得到的，不 import FlagGems。日志会打出命中的 key；一个都没命中时会列出尝试过的候选。 + +其余行为（轮转分片而非按块切、用子进程而非 xdist 的原因）见插件 docstring。分片绑卡遵守当前后端的可见性变量（CUDA 为 `CUDA_VISIBLE_DEVICES`）——用它选过卡时，分片只会落在你选的那几张上。 ## 批量多算子测试（run_batch.py） diff --git a/_parallel_warmup_plugin.py b/_parallel_warmup_plugin.py index eed03ad..76640d0 100644 --- a/_parallel_warmup_plugin.py +++ b/_parallel_warmup_plugin.py @@ -33,7 +33,7 @@ import sys import tempfile import time from pathlib import Path -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple import pytest @@ -88,15 +88,114 @@ def _visible_devices() -> List[str]: return [] -def _read_shape_yaml(path: str, op: str) -> Tuple[List[Any], Dict[str, Any]]: - """Shapes for this op, plus the sibling keys to carry into each shard.""" +def _attr_name(node: Any) -> str: + """Last component of a Name/Attribute node: `base.BlasBenchmark` -> `BlasBenchmark`.""" + import ast + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return "" + + +def _class_bases(py_path: Path) -> Dict[str, List[str]]: + """{class name: [base class names]} parsed statically from one module.""" + import ast + try: + tree = ast.parse(py_path.read_text(errors="replace")) + except Exception: + return {} + return { + node.name: [n for n in (_attr_name(b) for b in node.bases) if n] + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) + } + + +def _class_chain(test_py: Path, op: str) -> List[str]: + """Benchmark class names to try for `op`, nearest ancestor first. + + FlagGems' `Benchmark.set_shapes` looks up `op_name` and, failing that, walks + `type(self).__mro__` trying each class name as a yaml key -- which is how + BLAS ops reach the shared `BlasBenchmark` entry in core_shapes.yaml. This + plugin runs at pytest_configure, before any Benchmark exists, so the real + MRO is unavailable; the class is recovered statically instead, from the + `op_name=` keyword of the constructor call inside the test module, then its + bases from that module plus benchmark/base.py. + + Breadth-first over the inheritance graph, not a true C3 linearisation: the + benchmark classes are single-inheritance in practice, and a miss only costs + the sharding (see the caller's skip path), never correctness. + """ + import ast + try: + tree = ast.parse(test_py.read_text(errors="replace")) + except Exception: + return [] + start = "" + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + for kw in node.keywords: + if (kw.arg == "op_name" and isinstance(kw.value, ast.Constant) + and kw.value.value == op): + start = _attr_name(node.func) + break + if start: + break + if not start: + return [] + + graph = _class_bases(test_py) + for name, bases in _class_bases(test_py.parent / "base.py").items(): + graph.setdefault(name, bases) + + chain: List[str] = [] + seen: set = set() + queue = [start] + while queue: + name = queue.pop(0) + if not name or name in seen: + continue + seen.add(name) + chain.append(name) + queue.extend(graph.get(name, [])) + return chain + + +def _test_module_path(argv: List[str]) -> Optional[Path]: + """The benchmark module under test, from pytest's own argv.""" + for a in argv: + if a.startswith("-"): + continue + p = Path(a.split("::", 1)[0]) + if p.suffix == ".py" and p.is_file(): + return p + return None + + +def _read_shape_yaml(path: str, op: str, test_py: Optional[Path] = None + ) -> Tuple[List[Any], Dict[str, Any], str]: + """Shapes for this op, the sibling keys to carry into each shard, and the + yaml key they came from ("" when nothing matched). + + Tries `op` first, then the benchmark class chain, mirroring FlagGems. Without + the class fallback every BLAS op (mm/addmm/bmm/baddbmm, whose shapes live + under `BlasBenchmark`) resolved to zero shapes and sharding silently turned + itself off. + """ import yaml with open(path) as f: doc = yaml.safe_load(f) or {} - entry = doc.get(op) or {} - shapes = entry.get("shapes") or [] - extra = {k: v for k, v in entry.items() if k != "shapes"} - return list(shapes), extra + keys = [op] + if not isinstance(doc.get(op), dict) and test_py is not None: + keys += _class_chain(test_py, op) + for key in keys: + entry = doc.get(key) + if isinstance(entry, dict) and entry.get("shapes"): + extra = {k: v for k, v in entry.items() if k != "shapes"} + return list(entry["shapes"]), extra, key + return [], {}, "" def _shard(shapes: List[Any], n: int) -> List[List[Any]]: @@ -226,6 +325,23 @@ def pytest_configure(config): return op = _op_name() + + # Only --level core is shardable. At comprehensive, FlagGems' set_shapes() + # merges set_more_shapes() on top of whatever the yaml holds, and that + # addition is arbitrary per-subclass Python (hardcoded tuples plus + # models_shapes/*.yaml) which this plugin cannot see. Sharding the yaml list + # alone would sweep a strict subset -- measured on mm: 5 of 29 shapes, 90 + # AUTOTUNE_REPLAY_FALLBACK reason=key_missing -- and every uncovered key + # then self-tunes while the run reports itself as replaying. Refusing costs + # wall time; sharding a subset costs the same-config premise. + level = str(getattr(config.option, "level", "") or "").strip().lower() + if level and level != "core": + print(f"{tag('[parallel-warmup-plugin]')} --level {level} merges shapes " + "this plugin cannot enumerate (set_more_shapes); sharding a subset " + "would break the same-config premise, so sweeping serially instead", + flush=True) + return + shape_file = getattr(config.option, "shape_file", "") or "" if not shape_file or not Path(shape_file).is_file(): print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; " @@ -236,15 +352,24 @@ def pytest_configure(config): "already pinned, nothing to sweep, skipping", flush=True) return + test_py = _test_module_path(list(config.invocation_params.args)) try: - shapes, extra = _read_shape_yaml(shape_file, op) + shapes, extra, key = _read_shape_yaml(shape_file, op, test_py) except Exception as exc: print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes " f"({type(exc).__name__}: {exc}); skipping", flush=True) return + if not key: + # Distinct from "found a short list": nothing in the yaml describes this + # op at all, so the shard files would be built from an empty set. + print(f"{tag('[parallel-warmup-plugin]')} no yaml entry for {op!r} " + f"(tried: {', '.join([op] + _class_chain(test_py, op)) if test_py else op}); " + "skipping", flush=True) + return if len(shapes) < 2: - print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s); " - "sharding would not pay off, skipping", flush=True) + print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s) " + f"under yaml key {key!r}; sharding would not pay off, skipping", + flush=True) return devices = _visible_devices() @@ -255,8 +380,9 @@ def pytest_configure(config): shards = _shard(shapes, min(gpus, len(shapes), len(devices))) scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_")) - print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes on " - f"{len(shards)} GPUs (shards: {[len(s) for s in shards]}); this phase's " + print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes " + f"(yaml key {key!r}) on {len(shards)} GPUs " + f"(shards: {[len(s) for s in shards]}); this phase's " "latency is discarded, only configs are kept", flush=True) keep_scratch = False @@ -264,9 +390,23 @@ def pytest_configure(config): try: recs, bad = _spawn(shards, op, extra, scratch, list(config.invocation_params.args), devices) - merged, total, conflicts = _merge(recs, op) elapsed = time.time() - t0 + if bad: + # A dead shard means its shapes were never swept, so publishing the + # merge would hand the measurement a config set with holes; every + # hole then silently self-tunes under a run that claims to replay. + # Discard the whole sweep rather than publish a partial one. + keep_scratch = True + print(f"{tag('[parallel-warmup-plugin]')} shard(s) {bad} exited " + f"non-zero after {elapsed:.0f}s; their shapes were never swept, " + "so the merged configs would have holes. Discarding this sweep " + f"and falling back to normal serial autotune (worker logs kept " + f"in {scratch})", flush=True) + return + + merged, total, conflicts = _merge(recs, op) + if not total: keep_scratch = True print(f"{tag('[parallel-warmup-plugin]')} no configs recovered in " @@ -291,9 +431,6 @@ def pytest_configure(config): msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: " f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; " "measurement continues serially on one GPU") - if bad: - msg += f" (WARNING shard(s) {bad} exited non-zero; missing keys will " - msg += "fall back to live autotune)" print(msg, flush=True) if conflicts: print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config "