diff --git a/tpu_sync/benchmarks/BUILD b/tpu_sync/benchmarks/BUILD index 03a283663..6148342a7 100644 --- a/tpu_sync/benchmarks/BUILD +++ b/tpu_sync/benchmarks/BUILD @@ -22,6 +22,14 @@ filegroup( visibility = ["//:internal"], ) +filegroup( + name = "gating_baselines", + srcs = [ + "h2d_d2h_gating_jax_baselines.json", + "h2d_d2h_gating_torch_baselines.json", + ], +) + py_library( name = "bap_metrics", testonly = True, @@ -59,7 +67,7 @@ py_binary( name = "h2d_d2h_benchmark_gating", testonly = True, srcs = ["h2d_d2h_benchmark_gating.py"], - data = ["h2d_d2h_gating_baselines.json"], + data = [":gating_baselines"], deps = [ ":bap_metrics", ":perf_core", diff --git a/tpu_sync/benchmarks/H2D_D2H_GATING_TEST.md b/tpu_sync/benchmarks/H2D_D2H_GATING_TEST.md index 433b0bcad..e2ade70be 100644 --- a/tpu_sync/benchmarks/H2D_D2H_GATING_TEST.md +++ b/tpu_sync/benchmarks/H2D_D2H_GATING_TEST.md @@ -17,6 +17,19 @@ bandwidth number**. This gate blocks a PR that regresses it. ## Setup +```mermaid +flowchart LR + J[("baselines.json")] + J -->|"read: configs, iters, warmup"| M["perf_core.measure()"] + M --> R["measured samples"] + R -->|"--record mode"| W["floor = median - 3.5*MADσ"] + W -->|"overwrite baseline_* and floor_*"| J + R -->|"gate mode (default)"| C{"median >= floor_*?"} + J -->|"read: floor_*"| C + C -->|yes| P["PASS"] + C -->|no| F["exit(1) — blocks the CL"] +``` + The gate binary (`h2d_d2h_benchmark_gating.py`) measures, via `perf_core`, the aggregate d2h and h2d bandwidth over the node's 8 chips. @@ -75,13 +88,19 @@ A perf-floor failure blocks the PR, unless the author adds `[skip-perf-gate]` ## Running it ```bash +# Full gate across both JAX and PyTorch (default): bazel run -c opt --config=oss --config=ci \ //tpu_sync/benchmarks:h2d_d2h_benchmark_gating + +# Select a single framework: +bazel run -c opt --config=oss --config=ci \ + //tpu_sync/benchmarks:h2d_d2h_benchmark_gating -- --framework=torch # or --framework=jax ``` -In CI it runs as the `h2d_d2h_gating` workload in `benchmark_registry.pbtxt` via -the `run_benchmarks` workflow (on pull requests). Re-record baselines with the -`--record` flag (or the record workflow). +In CI it runs as the unified `h2d_d2h_gating` workload in `benchmark_registry.pbtxt` +via the `run_benchmarks` workflow (on pull requests), evaluating both frameworks +on the assigned Cloud TPU runner. Re-record baselines for all frameworks with +the `--record` flag (or the `h2d_d2h_record` workflow). ## Scope diff --git a/tpu_sync/benchmarks/benchmark_registry.pbtxt b/tpu_sync/benchmarks/benchmark_registry.pbtxt index cfa72856a..7bb360db8 100644 --- a/tpu_sync/benchmarks/benchmark_registry.pbtxt +++ b/tpu_sync/benchmarks/benchmark_registry.pbtxt @@ -1,6 +1,6 @@ benchmarks { name: "h2d_d2h_gating" - description: "Perf gate: 5 stable configs, fail if median Gbps drops >3% vs gating_baselines.json" + description: "Perf gate: stable configs across JAX and PyTorch, fail if median Gbps drops below statistical floor vs gating_baselines.json" owner: "raiden-dev" workload { @@ -16,11 +16,23 @@ benchmarks { workload_action_inputs { key: "bazel_run_flags" value: "-c opt --config=oss --config=ci" } } + # Symmetrical framework metrics: JAX + metrics { name: "jax_int32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "jax_int32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "jax_float32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "jax_float32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } + + # Symmetrical framework metrics: PyTorch + metrics { name: "torch_int32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "torch_int32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "torch_float32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } + metrics { name: "torch_float32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } + + # Legacy un-prefixed aliases for MLCompass monitoring (go/tpu-sync-oss-mlcompass) metrics { name: "int32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } metrics { name: "int32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } metrics { name: "float32_L1_8x128x1024x128/d2h_gbps" unit: "Gbps" stats { stat: MEAN } } metrics { name: "float32_L1_8x128x1024x128/h2d_gbps" unit: "Gbps" stats { stat: MEAN } } - } benchmarks { diff --git a/tpu_sync/benchmarks/h2d_d2h_benchmark_gating.py b/tpu_sync/benchmarks/h2d_d2h_benchmark_gating.py index da286b27c..be71a9a73 100644 --- a/tpu_sync/benchmarks/h2d_d2h_benchmark_gating.py +++ b/tpu_sync/benchmarks/h2d_d2h_benchmark_gating.py @@ -13,11 +13,12 @@ # limitations under the License. import json +import multiprocessing import os import subprocess import sys -from absl import app +from absl import app # pylint: disable=unused-import # used by the OSS entry point below from absl import flags import numpy as np @@ -35,29 +36,54 @@ 'Override iters from the baselines file. Use a large value when recording ' '(the floor depends on a MAD/sigma estimate, which needs many samples to be ' 'stable); the gate itself runs the smaller value in the baselines file.') +_FRAMEWORK = flags.DEFINE_enum( + 'framework', + 'jax', + ['all', 'jax', 'torch'], + 'Framework to benchmark: "jax", "torch", or "all" (both JAX and PyTorch).', +) +_ENFORCE = flags.DEFINE_bool( + 'enforce', + None, + 'Override the "enforce" setting in the baselines file.', +) # The perf floor (NOT the correctness check) can be bypassed per-PR by putting # one of these tags in the CL description / commit message. _SKIP_TAGS = ('[skip-perf-gate]', '[skip-h2d-d2h-gating]') -def _baselines_path(): +def _baselines_path(framework): + """Baselines file for one framework. + + JAX and torch exercise different transfer code, so their floors are not + interchangeable and never share a file. + """ if _BASELINES.value: return _BASELINES.value - return os.path.join(os.path.dirname(os.path.abspath(__file__)), - 'h2d_d2h_gating_baselines.json') + return os.path.join( + os.path.dirname(os.path.abspath(__file__)), + f'h2d_d2h_gating_{framework}_baselines.json') def _opted_out(): - """True if the HEAD commit message asks to skip the perf floor for this PR.""" - try: - msg = subprocess.run(['git', 'log', '-1', '--format=%B'], - capture_output=True, text=True).stdout.lower() - except Exception: # pylint: disable=broad-exception-caught - return None - for tag in _SKIP_TAGS: - if tag in msg: - return tag + """True if the CL description, commit message, or environment requests bypass.""" + for var in ('SKIP_PERF_GATE', 'TPU_RAIDEN_SKIP_PERF_GATE'): + if os.environ.get(var, '').lower() in ('1', 'true', 'yes'): + return var + for cmd in ( + ['hg', 'log', '-r', '.', '-T', '{desc}'], + ['git', 'log', '-1', '--format=%B'], + ): + try: + res = subprocess.run(cmd, capture_output=True, text=True, timeout=5) + if res.returncode == 0: + msg = res.stdout.lower() + for tag in _SKIP_TAGS: + if tag in msg: + return tag + except Exception: # pylint: disable=broad-exception-caught + pass return None @@ -75,68 +101,96 @@ def _core_floor(samples, k): return float(med - k * sigma) -def _write_tb_metrics(results): +def _write_tb_metrics(results, framework): """Log per-config throughput to TENSORBOARD_OUTPUT_DIR so BAP ingests it. - Each tag MUST have a matching metrics{name:...} in benchmark_registry.pbtxt.""" + + Each tag MUST have a matching metrics{name:...} in benchmark_registry.pbtxt. + """ scalars = {} for c, r in results: - label = f"{c['dtype']}_L{c['num_layers']}_{'x'.join(map(str, c['shape']))}" - scalars[f'{label}/d2h_gbps'] = r['d2h_gbps'] - scalars[f'{label}/h2d_gbps'] = r['h2d_gbps'] + tag_suffix = ( + f"{c['dtype']}_L{c['num_layers']}_{'x'.join(map(str, c['shape']))}" + ) + # Symmetrical framework-qualified metric tags + scalars[f'{framework}_{tag_suffix}/d2h_gbps'] = r['d2h_gbps'] + scalars[f'{framework}_{tag_suffix}/h2d_gbps'] = r['h2d_gbps'] + + # Legacy un-prefixed alias for JAX to maintain continuity with existing + # MLCompass regression monitoring (go/tpu-sync-oss-mlcompass). + if framework == 'jax': + scalars[f'{tag_suffix}/d2h_gbps'] = r['d2h_gbps'] + scalars[f'{tag_suffix}/h2d_gbps'] = r['h2d_gbps'] bap_metrics.emit(scalars) -def main(_): - path = _baselines_path() - with open(path) as f: - cfg = json.load(f) - sigma_k = float(cfg.get('sigma_k', 3.5)) # #robust-sigmas below median - iters = int(cfg.get('iters', 20)) - if _ITERS.value is not None: - iters = _ITERS.value - warmup = int(cfg.get('warmup', 3)) - configs = cfg['configs'] - - # NOTE: the correctness check (a d2h->h2d round-trip byte-equality assertion) - # runs INSIDE perf_core.measure(), before any floor comparison. A corrupt or - # no-op transfer raises there and fails the run regardless of the skip tag or - # the "enforce" switch below -- correctness is never bypassable. +def _measure_configs(configs, fw, iters, warmup): + """Measure every config once under one framework. + + The round-trip byte-equality check runs inside perf_core.measure(), before + any floor comparison, so a corrupt or no-op transfer raises here regardless + of the skip tag or "enforce". + """ results = [] for c in configs: - r = perf_core.measure(shape=c['shape'], num_layers=c['num_layers'], - dtype=c['dtype'], shard_axis=c.get('shard_axis', 2), - iters=iters, warmup=warmup) + r = perf_core.measure( + shape=c['shape'], + num_layers=c['num_layers'], + dtype=c['dtype'], + shard_axis=c.get('shard_axis', 2), + iters=iters, + warmup=warmup, + framework=fw, + ) results.append((c, r)) - print(f"[measured] {c['dtype']} L{c['num_layers']} " - f"{'x'.join(map(str, c['shape']))} " - f"d2h {r['d2h_gbps']:.1f} h2d {r['h2d_gbps']:.1f} Gbps") + print( + f"[{fw} measured] {c['dtype']} L{c['num_layers']} " + f"{'x'.join(map(str, c['shape']))} " + f"d2h {r['d2h_gbps']:.1f} h2d {r['h2d_gbps']:.1f} Gbps" + ) + return results + + +def _record_output_path(path): + """Where recorded baselines can actually be retrieved from. + + Under a test runner the source tree is a sandbox copy that is deleted + afterwards, so the result has to go where the harness collects artifacts. + """ + for var in ('WORKLOAD_ARTIFACTS_DIR', 'TEST_UNDECLARED_OUTPUTS_DIR'): + directory = os.environ.get(var) + if directory: + return os.path.join(directory, os.path.basename(path)) + return path - # --- record mode: overwrite baselines + floors, no gating --- - if _RECORD.value: - for c, r in results: - c['baseline_d2h'] = round(r['d2h_gbps'], 1) - c['baseline_h2d'] = round(r['h2d_gbps'], 1) - c['floor_d2h'] = round(_core_floor(r['d2h_gbps_all'], sigma_k), 1) - c['floor_h2d'] = round(_core_floor(r['h2d_gbps_all'], sigma_k), 1) - out_path = path - adir = os.environ.get('WORKLOAD_ARTIFACTS_DIR') - if adir: - out_path = os.path.join(adir, 'h2d_d2h_gating_baselines.json') - with open(out_path, 'w') as f: - json.dump(cfg, f, indent=2) - print(f'Recorded {len(results)} baselines+floors ' - f'(iters={iters}, sigma_k={sigma_k}) -> {out_path}') - return - # emit per-config throughput to TB for the BAP dashboard (gate mode only) - _write_tb_metrics(results) +def _record_baselines(cfg, results, sigma_k, path, fw, iters): + """Overwrite this framework's baselines and floors from these measurements.""" + for c, r in results: + for d in ('d2h', 'h2d'): + c[f'baseline_{d}'] = round(r[f'{d}_gbps'], 1) + c[f'floor_{d}'] = round(_core_floor(r[f'{d}_gbps_all'], sigma_k), 1) + out_path = _record_output_path(path) + with open(out_path, 'w') as f: + json.dump(cfg, f, indent=2) + print( + f'Recorded {fw} baselines+floors ' + f'(iters={iters}, sigma_k={sigma_k}) -> {out_path}' + ) + # Also echo it: on a remote runner the log is the one artifact always kept. + print(json.dumps(cfg, indent=2)) - # --- gate mode: compare the median of `iters` runs against the recorded floor --- + +def _gate(results, fw, iters, sigma_k): + """Print the comparison table; return how many directions are below floor.""" + print( + f'\n{fw.upper()} perf gate: median of {iters} iters vs per-config floor' + f' (median - {sigma_k} robust-sigmas / MAD)\n' + ) + print( + f"{'config':30}{'dir':4}{'baseline':>9}{'floor':>9}{'median':>9}" + f"{'drop':>7} verdict" + ) fails = 0 - print(f'\nperf gate: median of {iters} iters vs per-config floor ' - f'(median - {sigma_k} robust-sigmas / MAD)\n') - print(f"{'config':30}{'dir':4}{'baseline':>9}{'floor':>9}{'median':>9}" - f"{'drop':>7} verdict") for c, r in results: label = f"{c['dtype']} L{c['num_layers']} {'x'.join(map(str, c['shape']))}" for d in ('d2h', 'h2d'): @@ -145,29 +199,96 @@ def main(_): med = r[f'{d}_gbps'] ok = med >= floor fails += not ok - print(f"{label:30}{d:4}{base:9.1f}{floor:9.1f}{med:9.1f}" - f"{(base-med)/base*100:6.1f}% {'PASS' if ok else 'FAIL <-- REGRESSION'}") + print( + f'{label:30}{d:4}{base:9.1f}{floor:9.1f}{med:9.1f}' + f'{(base-med)/base*100:6.1f}% ' + f" {'PASS' if ok else 'FAIL <-- REGRESSION'}" + ) + return fails + + +def main(_): + fw_choice = _FRAMEWORK.value.lower() + frameworks = ['jax', 'torch'] if fw_choice == 'all' else [fw_choice] + if _BASELINES.value and len(frameworks) > 1: + raise ValueError( + '--baselines names a single file but --framework=all needs one per ' + 'framework. Run each framework separately.' + ) + + total_configs = 0 + blocking_fails = 0 + reported_fails = 0 + + for fw in frameworks: + path = _baselines_path(fw) + with open(path) as f: + cfg = json.load(f) + sigma_k = float(cfg.get('sigma_k', 3.5)) # robust sigmas below median + warmup = int(cfg.get('warmup', 3)) + iters = int(cfg.get('iters', 20)) + if _ITERS.value is not None: + iters = _ITERS.value + + results = _measure_configs(cfg['configs'], fw, iters, warmup) + + if _RECORD.value: + _record_baselines(cfg, results, sigma_k, path, fw, iters) + continue + + # per-config throughput to TB for the BAP dashboard (gate mode only) + _write_tb_metrics(results, framework=fw) + total_configs += len(results) + fails = _gate(results, fw, iters, sigma_k) + enforce = ( + _ENFORCE.value + if _ENFORCE.value is not None + else cfg.get('enforce', True) + ) + if enforce: + blocking_fails += fails + else: + reported_fails += fails + + if _RECORD.value: + return print() - if not fails: - print(f'GATE PASS: all {len(results)} configs at/above their floor') + if reported_fails: + print( + f'{reported_fails} direction(s) below the floor in a framework with' + ' "enforce": false -- reported, NOT blocking' + ) + if not blocking_fails: + print( + f'GATE PASS: all {total_configs} configs across {frameworks}' + ' at/above their floor' + ) return - # A perf-floor regression. It is blocking unless the maintainer switched the - # gate to report-only ("enforce": false) or the author opted out via a tag. - msg = f'GATE FAIL: {fails} direction(s) below the floor' - enforce = bool(cfg.get('enforce', True)) + # A perf-floor regression, blocking unless the author opted out via a tag. + msg = f'GATE FAIL: {blocking_fails} direction(s) below the floor' opt = _opted_out() - if not enforce: - print(msg + ' [report-only: "enforce": false in baselines, NOT blocking]') - return if opt: - print(msg + f' [report-only: {opt} in commit message, NOT blocking]') + print(msg + f' [report-only: {opt} bypass detected, NOT blocking]') return - print(msg + f' (to bypass, add {_SKIP_TAGS[0]} to your CL description)', - file=sys.stderr) + print( + msg + + f' (to bypass: add {_SKIP_TAGS[0]} to your CL description, or pass' + ' --enforce=false / set SKIP_PERF_GATE=1)', + file=sys.stderr, + ) sys.exit(1) +def _parse_known_flags(args): + """Ignore flags owned by the test runner rather than this binary.""" + return flags.FLAGS(args, known_only=True) + + if __name__ == '__main__': - app.run(main, flags_parser=lambda args: flags.FLAGS(args, known_only=True)) + # The torch path measures from worker processes (perf_core._measure_torch), + # and hermetic Python has no interpreter for spawn to re-exec. handle_main + # supplies one and routes spawned children back into the worker. + multiprocessing.set_start_method('spawn', force=True) + app.run(main, flags_parser=_parse_known_flags) diff --git a/tpu_sync/benchmarks/h2d_d2h_gating_baselines.json b/tpu_sync/benchmarks/h2d_d2h_gating_jax_baselines.json similarity index 96% rename from tpu_sync/benchmarks/h2d_d2h_gating_baselines.json rename to tpu_sync/benchmarks/h2d_d2h_gating_jax_baselines.json index 117423784..ba43fbbf7 100644 --- a/tpu_sync/benchmarks/h2d_d2h_gating_baselines.json +++ b/tpu_sync/benchmarks/h2d_d2h_gating_jax_baselines.json @@ -1,7 +1,7 @@ { "_comment": [ "Regenerate with: bazel run //tpu_sync/benchmarks:h2d_d2h_benchmark_gating --", - " --record --iters=", + " --record --iters= [--framework=jax|torch]", "Use a large --iters when recording: the floor depends on a MAD estimate, which", "needs many samples to be stable. Record mode rewrites this file in place and", "preserves this comment.", diff --git a/tpu_sync/benchmarks/h2d_d2h_gating_torch_baselines.json b/tpu_sync/benchmarks/h2d_d2h_gating_torch_baselines.json new file mode 100644 index 000000000..02684e4ce --- /dev/null +++ b/tpu_sync/benchmarks/h2d_d2h_gating_torch_baselines.json @@ -0,0 +1,67 @@ +{ + "_comment": [ + "Recorded on 8 ghostlite chips, --iters=1000.", + "http://sponge2/2d116850-0a4d-421d-bcf5-25b74434b9c6", + "", + "Re-record with: bazel run //tpu_sync/benchmarks:h2d_d2h_benchmark_gating --", + " --record --framework=torch --iters=", + "Use a large --iters when recording: the floor depends on a MAD estimate, which", + "needs many samples to be stable. At --iters=20 these same floors came out up to", + "30% different run to run. Record mode rewrites this file in place and preserves", + "this comment.", + "", + "Top level", + " sigma_k : robust sigmas below the median used to derive each floor (see below).", + " iters : timed iterations the GATE runs per config; their median is what is", + " compared against the floor. Smaller than the value used to record.", + " warmup : untimed iterations run before the timed ones.", + " enforce : optional, defaults to true. Set false to report regressions without", + " blocking the build.", + "", + "Per config", + " dtype / num_layers / shape / shard_axis : the KV-cache arrays being transferred.", + " shape is per layer; shape[shard_axis] must be divisible by the chip", + " count (8 on v5e) or the sharded shape silently mismatches.", + " baseline_d2h / baseline_h2d : median throughput in Gbps observed when this file", + " was recorded. REFERENCE ONLY -- the gate never compares against it; it", + " is here so a drop can be read as a percentage.", + " floor_d2h / floor_h2d : THE GATE THRESHOLD, in Gbps. The build FAILS if the", + " measured median falls below it. Computed at record time as", + " floor = median - sigma_k * (1.4826 * MAD)", + " where MAD is the median absolute deviation of the recorded samples --", + " i.e. the lower edge of the normal run-to-run spread. MAD is used rather", + " than a plain stddev because it is outlier-resistant: one slow sample", + " cannot drag the floor down and quietly weaken the gate.", + "", + " These floors sit ~30% below their baselines because ghostlite is", + " shared, noisy hardware: one robust sigma is ~9% of the median. So this", + " gate catches gross breakage, not a 5% regression. Tightening it needs", + " quieter hardware, not a smaller sigma_k." + ], + "sigma_k": 3.5, + "iters": 20, + "warmup": 3, + "enforce": true, + "configs": [ + { + "dtype": "int32", + "num_layers": 1, + "shape": [8, 128, 1024, 128], + "shard_axis": 2, + "baseline_d2h": 1519.7, + "baseline_h2d": 1416.9, + "floor_d2h": 1028.3, + "floor_h2d": 1013.3 + }, + { + "dtype": "float32", + "num_layers": 1, + "shape": [8, 128, 1024, 128], + "shard_axis": 2, + "baseline_d2h": 1528.6, + "baseline_h2d": 1414.5, + "floor_d2h": 1079.4, + "floor_h2d": 994.7 + } + ] +} diff --git a/tpu_sync/benchmarks/perf_core.py b/tpu_sync/benchmarks/perf_core.py index b2314b095..6934ac434 100644 --- a/tpu_sync/benchmarks/perf_core.py +++ b/tpu_sync/benchmarks/perf_core.py @@ -23,18 +23,115 @@ import ctypes import gc import os +import pathlib +import socket import sys import time +from typing import Optional -import jax -import jax.numpy as jnp import numpy as np -from tpu_sync.frameworks.jax import _tpu_raiden_jax as kv_cache_manager -DTYPE_MAP = {'float32': jnp.float32, 'bfloat16': jnp.bfloat16, 'float16': jnp.float16, - 'int32': jnp.int32, 'float8_e4m3fn': jnp.float8_e4m3fn} -ITEMSIZE = {'float32': 4, 'bfloat16': 2, 'float16': 2, 'int32': 4, 'float8_e4m3fn': 1} +ITEMSIZE = { + "float32": 4, + "bfloat16": 2, + "float16": 2, + "int32": 4, + "float8_e4m3fn": 1, +} + +_GOOGLE_PCI_VENDOR_ID = "0x1ae0" + +# Mirrors _TOPOLOGY_BY_TPU_PCI_DEVICE_ID in +# third_party/py/torch_tpu/_internal/utils/hardware.py +_V4_TOPOLOGY = {1: "1,1,1", 2: "1,2,1", 4: "2,2,1", 8: "2,2,2"} +_V5P_TOPOLOGY = {1: "1,1,1", 2: "1,2,1", 4: "2,2,1", 8: "2,2,2"} +# 8-chip v5e is discovered as a 2x2x2 cube internally, but Cloud TPU VMs wire it +# as a 2x4 mesh. Set TORCH_TPU_TOPOLOGY=2,4,1 to override there. +_V5E_TOPOLOGY = {1: "1,1,1", 4: "2,2,1", 8: "2,2,2"} +_V6E_TOPOLOGY = {1: "1,1,1", 4: "2,2,1", 8: "2,4,1"} +_V7_TOPOLOGY = {2: "1,1,1,2", 4: "1,2,1,2", 8: "2,2,1,2"} +_V8I_TOPOLOGY = {2: "1,1,1,2", 4: "2,1,1,2", 8: "4,1,1,2"} + +_TOPOLOGY_BY_TPU_PCI_DEVICE_ID = { + "0x005e": _V4_TOPOLOGY, + "0x0062": _V5P_TOPOLOGY, + "0x0063": _V5E_TOPOLOGY, + "0x006f": _V6E_TOPOLOGY, + "0x0076": _V7_TOPOLOGY, + "0x0083": _V8I_TOPOLOGY, +} + + +# ---------------- PCI / TPU Hardware Helpers ---------------- +def _scan_pci_tpus() -> tuple[int, Optional[dict[int, str]]]: + """Scans PCI bus to identify local physical TPU device IDs and topology.""" + count = 0 + topology_map = None + pci_devices = pathlib.Path("/sys/bus/pci/devices") + if not pci_devices.exists(): + return 0, None + for device_path in pci_devices.iterdir(): + try: + vendor_id = (device_path / "vendor").read_text().strip() + if vendor_id != _GOOGLE_PCI_VENDOR_ID: + continue + device_id = (device_path / "device").read_text().strip() + if device_id in _TOPOLOGY_BY_TPU_PCI_DEVICE_ID: + try: + group_id = (device_path / "iommu_group").readlink().name + (pathlib.Path("/dev/vfio") / group_id).stat() + except OSError: + continue + count += 1 + if topology_map is None: + topology_map = _TOPOLOGY_BY_TPU_PCI_DEVICE_ID[device_id] + except OSError: + continue + return count, topology_map + + +def get_tpu_device_count() -> int: + count, _ = _scan_pci_tpus() + return count + + +def get_tpu_topology(world_size: int) -> str: + """Chip mesh for world_size devices on the attached TPU generation.""" + _, topology_map = _scan_pci_tpus() + if topology_map is None: + raise RuntimeError("No TPU devices found.") + if world_size not in topology_map: + raise RuntimeError( + f"No TPU topology for {world_size} chips on this host. Set" + " TORCH_TPU_TOPOLOGY to override." + ) + return topology_map[world_size] + + +def pick_unused_ports(count: int = 1) -> list[int]: + sockets = [] + ports = [] + for _ in range(count): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("localhost", 0)) + ports.append(s.getsockname()[1]) + sockets.append(s) + for s in sockets: + s.close() + return ports + + +def prepare_tpu_environment(world_size: int) -> None: + if "TORCH_TPU_XPROF_SESSION_ID" not in os.environ: + os.environ["TORCH_TPU_XPROF_SESSION_ID"] = str(time.time_ns()) + if "TORCH_TPU_SLICEBUILDER_ADDRESSES" not in os.environ: + ports = pick_unused_ports(world_size) + os.environ["TORCH_TPU_SLICEBUILDER_ADDRESSES"] = ",".join( + [f"localhost:{p}" for p in ports] + ) + if "TORCH_TPU_TOPOLOGY" not in os.environ: + os.environ["TORCH_TPU_TOPOLOGY"] = get_tpu_topology(world_size) # ---------------- NUMA (optional experiment knob) ---------------- @@ -81,35 +178,46 @@ def summarize(values): 'max': float(a.max()), 'stddev': float(a.std(ddof=1)) if len(a) > 1 else 0.0} -# ---------------- per-device placement (ported from the DMA benchmark) ---------------- -def create_sharded_array(shape, sharding, dtype, is_host=False, is_random=False): +# ---------------- JAX Backend Implementation ---------------- +def _create_sharded_array_jax( + shape, sharding, dtype, is_host=False, is_random=False +): """Places each shard on its own device (pinned_host if is_host) -> NUMA-local per chip.""" + import jax + import jax.numpy as jnp + mesh, spec = sharding.mesh, sharding.spec devices = list(mesh.devices.flat) shard_shape = list(shape) shard_axis = None for i, axis in enumerate(spec): if axis is not None: - shard_axis = i; break + shard_axis = i + break if shard_axis is not None: shard_shape[shard_axis] = shape[shard_axis] // len(devices) shards = [] for idx, device in enumerate(devices): - sd = (jax.sharding.SingleDeviceSharding(device, memory_kind='pinned_host') - if is_host else jax.sharding.SingleDeviceSharding(device)) + sd = ( + jax.sharding.SingleDeviceSharding(device, memory_kind="pinned_host") + if is_host + else jax.sharding.SingleDeviceSharding(device) + ) if is_random: shard_np = np.random.uniform(0, 1, shard_shape).astype(np.float32) elif dtype == jnp.int32: start = idx * int(np.prod(shard_shape)) - shard_np = (np.arange(np.prod(shard_shape), dtype=np.int32) + start).reshape(shard_shape) + shard_np = ( + np.arange(np.prod(shard_shape), dtype=np.int32) + start + ).reshape(shard_shape) else: shard_np = np.zeros(shard_shape, dtype=np.float32) shards.append(jax.device_put(shard_np, sd).astype(dtype)) return jax.make_array_from_single_device_arrays(shape, sharding, shards) -def verify_roundtrip(manager, dev_arrs, num_blocks): +def _verify_roundtrip_jax(manager, dev_arrs, num_blocks): """One-shot d2h->h2d data-integrity check (ported from V2 verify_device_cache). Uses the SAME manager + d2h/h2d as the benchmark, but with disjoint offsets @@ -138,54 +246,74 @@ def verify_roundtrip(manager, dev_arrs, num_blocks): for s in a.addressable_shards: d = np.asarray(s.data) np.testing.assert_array_equal( - d[half:2 * half], d[0:half], - err_msg=f'CORRUPTION: layer {li} d2h/h2d round-trip mismatch') + d[half : 2 * half], + d[0:half], + err_msg=f"CORRUPTION: layer {li} d2h/h2d round-trip mismatch", + ) -def measure(shape, num_layers, dtype, shard_axis=2, iters=20, warmup=3, - lock_buffers=True, verify=True): # verify: run one integrity check before timing - """Runs the d2h/h2d transfer benchmark for one config and returns a result dict. - `shape` may be a "a,b,c" string (flag) or a list/tuple (json). Pure compute + - timing, no flag reads and no file I/O -> reusable from any binary. - """ - if isinstance(shape, str): - shape = shape.split(',') - shape = tuple(int(x) for x in shape) +def _measure_jax( + shape, + num_layers, + dtype, + shard_axis=2, + iters=20, + warmup=3, + lock_buffers=True, + verify=True, +): + """Executes the JAX transfer benchmark.""" + import jax + import jax.numpy as jnp + from tpu_sync.frameworks.jax import _tpu_raiden_jax as kv_cache_manager - devices = jax.devices('tpu') + dtype_map = { + "float32": jnp.float32, + "bfloat16": jnp.bfloat16, + "float16": jnp.float16, + "int32": jnp.int32, + "float8_e4m3fn": jnp.float8_e4m3fn, + } + + devices = jax.devices("tpu") if not devices: - raise RuntimeError('No TPU devices found.') + raise RuntimeError("No TPU devices found.") num_devices = len(devices) - num_blocks = shape[0] # major dim copied - dt = DTYPE_MAP.get(dtype, jnp.float32) + num_blocks = shape[0] + dt = dtype_map.get(dtype, jnp.float32) itemsize = ITEMSIZE.get(dtype, 4) - # mesh (1, num_devices) ("data","model"); shard `shard_axis` across the device axis - mesh = jax.sharding.Mesh(np.array(devices).reshape(1, num_devices), ('data', 'model')) + mesh = jax.sharding.Mesh( + np.array(devices).reshape(1, num_devices), ("data", "model") + ) spec = jax.sharding.PartitionSpec( - *[('model' if i == shard_axis else None) for i in range(len(shape))]) + *[("model" if i == shard_axis else None) for i in range(len(shape))] + ) tpu_sharding = jax.sharding.NamedSharding(mesh, spec) - src_arrs = [create_sharded_array(shape, tpu_sharding, dt, is_host=False, - is_random=(dt != jnp.int32)) - for _ in range(num_layers)] + src_arrs = [ + _create_sharded_array_jax( + shape, tpu_sharding, dt, is_host=False, is_random=(dt != jnp.int32) + ) + for _ in range(num_layers) + ] jax.block_until_ready(src_arrs) - # invalidate device shadow copies before each measured d2h (like the DMA benchmark) - mutate = jax.jit(lambda x: x + jnp.array(1 if x.dtype == jnp.int32 else 0.01, dtype=x.dtype)) + mutate = jax.jit( + lambda x: x + + jnp.array(1 if x.dtype == jnp.int32 else 0.01, dtype=x.dtype) + ) manager = kv_cache_manager.KVCacheManager( device_arrays=src_arrs, host_blocks_to_allocate=num_blocks, - unsafe_skip_buffer_lock=not lock_buffers) + unsafe_skip_buffer_lock=not lock_buffers, + ) - # Verify the transfer is byte-correct BEFORE timing (and before any mutate), so - # src_arrs still holds the known init values and the timed loop below is - # untouched. A corrupt/no-op transfer raises here and fails the gate. if verify: - verify_roundtrip(manager, src_arrs, num_blocks) + _verify_roundtrip_jax(manager, src_arrs, num_blocks) - offsets, sizes = [0], [num_blocks] # full-major-dim copy + offsets, sizes = [0], [num_blocks] total_bytes = num_layers * int(np.prod(shape)) * itemsize def once(): @@ -194,16 +322,24 @@ def once(): jax.block_until_ready(src_arrs) gc.disable() t0 = time.perf_counter() - manager.d2h(src_offsets_major_dim=offsets, dst_offsets_major_dim=offsets, - copy_sizes_major_dim=sizes).Await() + manager.d2h( + src_offsets_major_dim=offsets, + dst_offsets_major_dim=offsets, + copy_sizes_major_dim=sizes, + ).Await() d2h = time.perf_counter() - t0 - gc.enable(); gc.collect() + gc.enable() + gc.collect() gc.disable() t0 = time.perf_counter() - manager.h2d(src_offsets_major_dim=offsets, dst_offsets_major_dim=offsets, - copy_sizes_major_dim=sizes).Await() + manager.h2d( + src_offsets_major_dim=offsets, + dst_offsets_major_dim=offsets, + copy_sizes_major_dim=sizes, + ).Await() h2d = time.perf_counter() - t0 - gc.enable(); gc.collect() + gc.enable() + gc.collect() return d2h, h2d for _ in range(warmup): @@ -211,19 +347,366 @@ def once(): d2h_times, h2d_times = [], [] for _ in range(iters): d, h = once() - d2h_times.append(d); h2d_times.append(h) + d2h_times.append(d) + h2d_times.append(h) d2h_gbps_all = [(total_bytes * 8) / (t * 1e9) for t in d2h_times] h2d_gbps_all = [(total_bytes * 8) / (t * 1e9) for t in h2d_times] - d2h_med_t, h2d_med_t = float(np.median(d2h_times)), float(np.median(h2d_times)) + d2h_med_t, h2d_med_t = float(np.median(d2h_times)), float( + np.median(h2d_times) + ) return { - 'shape': list(shape), 'num_layers': num_layers, 'dtype': dtype, - 'total_bytes': total_bytes, - 'd2h_times_sec': d2h_times, 'h2d_times_sec': h2d_times, - 'd2h_gbps_all': d2h_gbps_all, 'h2d_gbps_all': h2d_gbps_all, - 'd2h_med_t': d2h_med_t, 'h2d_med_t': h2d_med_t, - 'd2h_gbps': (total_bytes * 8) / (d2h_med_t * 1e9), - 'h2d_gbps': (total_bytes * 8) / (h2d_med_t * 1e9), - 'd2h_gbps_summary': summarize(d2h_gbps_all), - 'h2d_gbps_summary': summarize(h2d_gbps_all), + "shape": list(shape), + "num_layers": num_layers, + "dtype": dtype, + "total_bytes": total_bytes, + "d2h_times_sec": d2h_times, + "h2d_times_sec": h2d_times, + "d2h_gbps_all": d2h_gbps_all, + "h2d_gbps_all": h2d_gbps_all, + "d2h_med_t": d2h_med_t, + "h2d_med_t": h2d_med_t, + "d2h_gbps": (total_bytes * 8) / (d2h_med_t * 1e9), + "h2d_gbps": (total_bytes * 8) / (h2d_med_t * 1e9), + "d2h_gbps_summary": summarize(d2h_gbps_all), + "h2d_gbps_summary": summarize(h2d_gbps_all), } + + +# ---------------- PyTorch Backend Implementation ---------------- +def _torch_worker_fn( + rank: int, + world_size: int, + master_port: int, + shape: tuple[int, ...], + num_layers: int, + dtype: str, + shard_axis: int, + iters: int, + warmup: int, + lock_buffers: bool, + verify: bool, + result_queue, +): + """Single-process worker driving local TPU shard transfers in PyTorch.""" + # pylint: disable=g-import-not-at-top + import torch + import torch.distributed as dist + import torch_tpu + from tpu_sync.api.torch import torch_tpu_common_loader + from tpu_sync.frameworks.torch import _tpu_raiden_torch as kv_cache_manager_torch + # pylint: enable=g-import-not-at-top + + if world_size > 1: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(master_port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["PJRT_LOCAL_PROCESS_RANK"] = str(rank) + os.environ["GROUP_RANK"] = "0" + os.environ["LOCAL_WORLD_SIZE"] = str(world_size) + + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{master_port}", + rank=rank, + world_size=world_size, + ) + + try: + torch_tpu_common_loader.load_torch_tpu_common() + device = torch.device("tpu") + + shard_shape = list(shape) + if world_size > 1: + shard_shape[shard_axis] = shape[shard_axis] // world_size + shard_shape = tuple(shard_shape) + num_blocks = shape[0] + + torch_dt_map = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "int32": torch.int32, + "float8_e4m3fn": torch.float8_e4m3fn, + } + dt = torch_dt_map.get(dtype, torch.float32) + itemsize = ITEMSIZE.get(dtype, 4) + total_bytes = num_layers * int(np.prod(shape)) * itemsize + + torch.manual_seed(42 + rank) + src_tensors = [] + for _ in range(num_layers): + if dtype == "int32": + start = rank * int(np.prod(shard_shape)) + np_data = ( + np.arange(np.prod(shard_shape), dtype=np.int32) + start + ).reshape(shard_shape) + t = torch.from_numpy(np_data).to(device) + else: + t = torch.randn(shard_shape, dtype=dt, device=device) + src_tensors.append([t]) + torch.tpu.synchronize() + + manager = kv_cache_manager_torch.KVCacheManager( + device_tensors=src_tensors, + local_port=0, + host_blocks_to_allocate=num_blocks, + unsafe_skip_buffer_lock=not lock_buffers, + ) + + if verify: + half = num_blocks // 2 + if half > 0: + manager.D2h( + src_offsets_major_dim=[0], + dst_offsets_major_dim=[0], + copy_sizes_major_dim=[half], + ).Await() + manager.H2d( + src_offsets_major_dim=[0], + dst_offsets_major_dim=[half], + copy_sizes_major_dim=[half], + ).Await() + torch.tpu.synchronize() + for li, shards in enumerate(src_tensors): + d = shards[0].cpu().numpy() + np.testing.assert_array_equal( + d[half : 2 * half], + d[0:half], + err_msg=( + f"CORRUPTION: rank {rank} layer {li} d2h/h2d round-trip" + " mismatch" + ), + ) + if world_size > 1: + dist.barrier() + + offsets, sizes = [0], [num_blocks] + + def once(): + for shards in src_tensors: + shards[0].add_(1 if shards[0].dtype == torch.int32 else 0.01) + torch.tpu.synchronize() + if world_size > 1: + dist.barrier() + gc.disable() + t0 = time.perf_counter() + manager.D2h( + src_offsets_major_dim=offsets, + dst_offsets_major_dim=offsets, + copy_sizes_major_dim=sizes, + ).Await() + torch.tpu.synchronize() + t1 = time.perf_counter() + gc.enable() + gc.collect() + + if world_size > 1: + dist.barrier() + gc.disable() + t2 = time.perf_counter() + manager.H2d( + src_offsets_major_dim=offsets, + dst_offsets_major_dim=offsets, + copy_sizes_major_dim=sizes, + ).Await() + torch.tpu.synchronize() + t3 = time.perf_counter() + gc.enable() + gc.collect() + + if world_size > 1: + t0_t = torch.tensor([t0], dtype=torch.float64) + t1_t = torch.tensor([t1], dtype=torch.float64) + t2_t = torch.tensor([t2], dtype=torch.float64) + t3_t = torch.tensor([t3], dtype=torch.float64) + dist.all_reduce(t0_t, op=dist.ReduceOp.MIN) + dist.all_reduce(t1_t, op=dist.ReduceOp.MAX) + dist.all_reduce(t2_t, op=dist.ReduceOp.MIN) + dist.all_reduce(t3_t, op=dist.ReduceOp.MAX) + return (t1_t - t0_t).item(), (t3_t - t2_t).item() + return (t1 - t0), (t3 - t2) + + for _ in range(warmup): + once() + + d2h_times, h2d_times = [], [] + for _ in range(iters): + d, h = once() + d2h_times.append(d) + h2d_times.append(h) + + if rank == 0: + d2h_gbps_all = [(total_bytes * 8) / (t * 1e9) for t in d2h_times] + h2d_gbps_all = [(total_bytes * 8) / (t * 1e9) for t in h2d_times] + d2h_med_t, h2d_med_t = float(np.median(d2h_times)), float( + np.median(h2d_times) + ) + res = { + "shape": list(shape), + "num_layers": num_layers, + "dtype": dtype, + "total_bytes": total_bytes, + "d2h_times_sec": d2h_times, + "h2d_times_sec": h2d_times, + "d2h_gbps_all": d2h_gbps_all, + "h2d_gbps_all": h2d_gbps_all, + "d2h_med_t": d2h_med_t, + "h2d_med_t": h2d_med_t, + "d2h_gbps": (total_bytes * 8) / (d2h_med_t * 1e9), + "h2d_gbps": (total_bytes * 8) / (h2d_med_t * 1e9), + "d2h_gbps_summary": summarize(d2h_gbps_all), + "h2d_gbps_summary": summarize(h2d_gbps_all), + } + if result_queue is not None: + result_queue.put(res) + return res + finally: + if world_size > 1 and dist.is_initialized(): + dist.destroy_process_group() + + +def _measure_torch( + shape, + num_layers, + dtype, + shard_axis=2, + iters=20, + warmup=3, + lock_buffers=True, + verify=True, +): + """Executes the PyTorch transfer benchmark across local TPU chips.""" + import torch.distributed as dist + import torch.multiprocessing as mp + + # If already inside a distributed worker process, run worker directly: + if dist.is_available() and dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + return _torch_worker_fn( + rank=rank, + world_size=world_size, + master_port=0, + shape=shape, + num_layers=num_layers, + dtype=dtype, + shard_axis=shard_axis, + iters=iters, + warmup=warmup, + lock_buffers=lock_buffers, + verify=verify, + result_queue=None, + ) + + count = get_tpu_device_count() + world_size = count if count > 0 else 1 + + if world_size <= 1: + return _torch_worker_fn( + rank=0, + world_size=1, + master_port=0, + shape=shape, + num_layers=num_layers, + dtype=dtype, + shard_axis=shard_axis, + iters=iters, + warmup=warmup, + lock_buffers=lock_buffers, + verify=verify, + result_queue=None, + ) + + # Multi-device SPMD execution on host + prepare_tpu_environment(world_size) + master_port = pick_unused_ports(1)[0] + ctx = mp.get_context("spawn") + result_queue = ctx.SimpleQueue() + mp.spawn( + _torch_worker_fn, + args=( + world_size, + master_port, + shape, + num_layers, + dtype, + shard_axis, + iters, + warmup, + lock_buffers, + verify, + result_queue, + ), + nprocs=world_size, + join=True, + ) + return result_queue.get() + + +# ---------------- Unified Public Measurement Interface ---------------- +def measure( + shape, + num_layers, + dtype, + shard_axis=2, + iters=20, + warmup=3, + lock_buffers=True, + verify=True, + framework="jax", +): + """Runs the d2h/h2d transfer benchmark for one config and returns a result dict. + + Args: + shape: Per-layer shape string ('a,b,c') or tuple/list. + num_layers: Number of cache arrays (layers). + dtype: Element data type ('float32', 'bfloat16', 'int32', etc.). + shard_axis: Axis to partition across device chips. + iters: Number of timed iterations. + warmup: Number of untimed warmup iterations. + lock_buffers: Whether to lock host buffers. + verify: Run one round-trip byte integrity assertion before timing. + framework: 'jax' or 'torch'. + + Returns: + Result dictionary containing throughput and latency statistics. + """ + if isinstance(shape, str): + shape = shape.split(",") + shape = tuple(int(x) for x in shape) + + fw = framework.lower() + if fw == "torch": + return _measure_torch( + shape, + num_layers, + dtype, + shard_axis=shard_axis, + iters=iters, + warmup=warmup, + lock_buffers=lock_buffers, + verify=verify, + ) + elif fw == "jax": + return _measure_jax( + shape, + num_layers, + dtype, + shard_axis=shard_axis, + iters=iters, + warmup=warmup, + lock_buffers=lock_buffers, + verify=verify, + ) + else: + raise ValueError( + f'Unsupported framework: {framework}. Expected "jax" or "torch".' + ) + + +# Backward-compatibility aliases +create_sharded_array = _create_sharded_array_jax +verify_roundtrip = _verify_roundtrip_jax