diff --git a/examples/microbenchmarks/torch_dma_kv_cache_benchmark.py b/examples/microbenchmarks/torch_dma_kv_cache_benchmark.py index fb7f1b574..f154f09eb 100644 --- a/examples/microbenchmarks/torch_dma_kv_cache_benchmark.py +++ b/examples/microbenchmarks/torch_dma_kv_cache_benchmark.py @@ -12,14 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Single-host KV-cache DMA benchmark: KVCacheManager vs PyTorch host copies. + +Each case times D2H and H2D of a sharded KV cache on one TPU host, one process +per local TPU device, three ways: KVCacheManager zero-copy DMA, PyTorch pinned +host memory, and PyTorch pageable host memory. +""" + +import dataclasses import gc import json import os import pathlib +import queue import socket import sys import time -from typing import Sequence +import traceback +from typing import Any, Sequence from absl import flags from absl.testing import absltest @@ -62,7 +72,8 @@ def log_telemetry( """Appends benchmark telemetry stats to the JSONL file at telemetry_log_path. Takes its settings as arguments rather than reading FLAGS: it runs inside - mp.spawn workers, where absl flags are not parsed in open-source builds. + spawned worker processes, where absl flags are not parsed in open-source + builds. Args: telemetry_log_path: JSONL file to append the record to. @@ -365,19 +376,124 @@ def _report_results( sys.stderr.flush() -def _mpmd_benchmark_worker( +_CASE_TIMEOUT_S = 1800.0 +_RESULT_POLL_S = 5.0 +_WORKER_EXIT_TIMEOUT_S = 120.0 + + +def _join_all(processes: Sequence[Any], timeout_s: float) -> None: + """Waits for every process to exit, for at most `timeout_s` in total.""" + deadline = time.monotonic() + timeout_s + for process in processes: + process.join(timeout=max(0.0, deadline - time.monotonic())) + + +@dataclasses.dataclass(frozen=True) +class _BenchmarkCase: + """One benchmark configuration, run by every rank of the worker group.""" + + test_name: str + dtype: torch.dtype + num_layers: int + shape: tuple[int, ...] + benchmark_runs: int + locality: str + telemetry_log_path: str + + +def _run_case(rank: int, world_size: int, case: _BenchmarkCase) -> None: + """Runs one benchmark case on this rank; rank 0 reports the results.""" + device = torch.device("tpu") + shard_shape = list(case.shape) + if len(case.shape) > 2: + # shard axis 2 (num_kv_heads) across TPU ranks. + shard_shape[2] //= world_size + + src_tensors = create_tpu_tensors( + case.num_layers, shard_shape, case.dtype, device + ) + num_blocks = shard_shape[0] + + # Branch 1: Raiden Zero-Copy DMA + mgr_d2h, mgr_h2d = _benchmark_raiden( + src_tensors, num_blocks, case.benchmark_runs + ) + if rank == 0: + _report_results( + "KVCacheManager", + case.dtype, + case.num_layers, + case.shape, + mgr_d2h, + mgr_h2d, + ) + log_telemetry( + case.telemetry_log_path, + case.locality, + case.test_name, + case.dtype, + case.num_layers, + case.shape, + mgr_d2h, + mgr_h2d, + ) + + # Branch 2: PyTorch Pinned Host Baseline + pin_d2h, pin_h2d = _benchmark_torch_copy( + src_tensors, + shard_shape, + case.dtype, + device, + case.benchmark_runs, + pinned=True, + ) + if rank == 0: + _report_results( + "PyTorch Pinned", + case.dtype, + case.num_layers, + case.shape, + pin_d2h, + pin_h2d, + ) + + # Branch 3: PyTorch Standard Host Baseline + std_d2h, std_h2d = _benchmark_torch_copy( + src_tensors, + shard_shape, + case.dtype, + device, + case.benchmark_runs, + pinned=False, + ) + if rank == 0: + _report_results( + "PyTorch Standard", + case.dtype, + case.num_layers, + case.shape, + std_d2h, + std_h2d, + ) + + +def _benchmark_worker_main( rank: int, world_size: int, master_port: int, - test_name: str, - dtype: torch.dtype, - num_layers: int, - shape: Sequence[int], - benchmark_runs: int, - locality: str, - telemetry_log_path: str, + cases: Any, + results: Any, ) -> None: - """Runs the benchmark worker for one distributed rank.""" + """Worker process: joins the TPU slice once, then runs cases until None. + + Args: + rank: This worker's rank; also the local TPU device it owns. + world_size: Number of workers, one per local TPU device. + master_port: Port of the gloo rendezvous shared by all workers. + cases: Queue this worker reads `_BenchmarkCase`s from; None means exit. + results: Queue shared by all workers; each case puts (rank, error), where + error is None on success or the formatted traceback on failure. + """ os.environ.pop("TPU_PREMAPPED_BUFFER_SIZE", None) os.environ.pop("stairways_ifrt_tpu_premapped_buffer_size", None) os.environ.pop("pathways_tpu_premapped_buffer_size", None) @@ -397,60 +513,132 @@ def _mpmd_benchmark_worker( rank=rank, world_size=world_size, ) - try: - device = torch.device("tpu") - shard_shape = list(shape) - if len(shape) > 2: - # shard axis 2 (num_kv_heads) across TPU ranks. - shard_shape[2] //= world_size - - src_tensors = create_tpu_tensors(num_layers, shard_shape, dtype, device) - num_blocks = shard_shape[0] - - # Branch 1: Raiden Zero-Copy DMA - mgr_d2h, mgr_h2d = _benchmark_raiden( - src_tensors, num_blocks, benchmark_runs - ) - if rank == 0: - _report_results( - "KVCacheManager", dtype, num_layers, shape, mgr_d2h, mgr_h2d - ) - log_telemetry( - telemetry_log_path, - locality, - test_name, - dtype, - num_layers, - shape, - mgr_d2h, - mgr_h2d, - ) + while (case := cases.get()) is not None: + try: + _run_case(rank, world_size, case) + except Exception: + # Report, then exit: the other ranks may be stuck in a collective, so + # this group cannot run more cases. + results.put((rank, traceback.format_exc())) + raise + gc.collect() + results.put((rank, None)) + finally: + dist.destroy_process_group() - # Branch 2: PyTorch Pinned Host Baseline - pin_d2h, pin_h2d = _benchmark_torch_copy( - src_tensors, shard_shape, dtype, device, benchmark_runs, pinned=True - ) - if rank == 0: - _report_results( - "PyTorch Pinned", dtype, num_layers, shape, pin_d2h, pin_h2d - ) - # Branch 3: PyTorch Standard Host Baseline - std_d2h, std_h2d = _benchmark_torch_copy( - src_tensors, shard_shape, dtype, device, benchmark_runs, pinned=False - ) - if rank == 0: - _report_results( - "PyTorch Standard", dtype, num_layers, shape, std_d2h, std_h2d - ) +class _BenchmarkWorkers: + """Runs benchmark cases on one long-lived process per local TPU device. - finally: - dist.destroy_process_group() + The processes form one multi-process TPU slice: it is created when the + first case starts and torn down once by `close()`. A fresh process group + per case would re-create the slice every time, and on TPU v7 the new + slice's START_SESSION fails for minutes whenever the previous slice's + teardown timed out (b/568095515). + + When a case fails, its workers are stopped (the surviving ranks may be stuck + in a collective) and the next `run()` starts a fresh group. + """ + + def __init__(self, world_size: int): + self._world_size = world_size + self._processes = [] + self._cases = [] + self._results = None + + def run(self, case: _BenchmarkCase) -> None: + """Runs `case` on every rank. + + Args: + case: The benchmark configuration to run. + + Raises: + RuntimeError: if a rank fails, exits, or the case exceeds + _CASE_TIMEOUT_S. + """ + if not self._processes: + self._start() + for cases in self._cases: + cases.put(case) + pending = set(range(self._world_size)) + deadline = time.monotonic() + _CASE_TIMEOUT_S + while pending: + try: + rank, error = self._results.get(timeout=_RESULT_POLL_S) + except queue.Empty: + exited = { + r: self._processes[r].exitcode + for r in sorted(pending) + if not self._processes[r].is_alive() + } + timed_out = time.monotonic() > deadline + if exited or timed_out: + self._stop() + if exited: + reason = ( + f"workers exited without a result (rank: exitcode {exited})" + ) + else: + reason = ( + f"ranks {sorted(pending)} timed out after" + f" {_CASE_TIMEOUT_S:.0f}s" + ) + raise RuntimeError(f"{case.test_name}: {reason}") from None + continue + if error is not None: + self._stop() + raise RuntimeError(f"{case.test_name}: rank {rank} failed:\n{error}") + pending.discard(rank) + + def close(self) -> None: + """Asks every worker to exit, which tears the TPU slice down once.""" + for cases in self._cases: + cases.put(None) + _join_all(self._processes, _WORKER_EXIT_TIMEOUT_S) + self._stop() + + def _start(self) -> None: + """Spawns one worker per device; they join the TPU slice on first use.""" + prepare_tpu_environment(self._world_size) + master_port = pick_unused_ports(1)[0] + ctx = mp.get_context("spawn") + self._results = ctx.Queue() + self._cases = [ctx.Queue() for _ in range(self._world_size)] + self._processes = [ + ctx.Process( + target=_benchmark_worker_main, + args=( + rank, + self._world_size, + master_port, + self._cases[rank], + self._results, + ), + ) + for rank in range(self._world_size) + ] + for process in self._processes: + process.start() + + def _stop(self) -> None: + """Terminates any worker still running and forgets the group.""" + for process in self._processes: + if process.is_alive(): + process.terminate() + _join_all(self._processes, _WORKER_EXIT_TIMEOUT_S) + for process in self._processes: + if process.is_alive(): + process.kill() + process.join() + self._processes = [] + self._cases = [] + self._results = None class KVCacheManagerPerfTest(parameterized.TestCase): device_count: int = 0 + workers: _BenchmarkWorkers | None = None @classmethod def setUpClass(cls): @@ -459,6 +647,14 @@ def setUpClass(cls): os.environ.pop("stairways_ifrt_tpu_premapped_buffer_size", None) os.environ.pop("pathways_tpu_premapped_buffer_size", None) cls.device_count, _ = _scan_pci_tpus() + if cls.device_count: + cls.workers = _BenchmarkWorkers(cls.device_count) + + @classmethod + def tearDownClass(cls): + if cls.workers is not None: + cls.workers.close() + super().tearDownClass() def setUp(self): super().setUp() @@ -489,24 +685,16 @@ def _run_benchmark( f" ({world_size})" ) - prepare_tpu_environment(world_size) - master_port = pick_unused_ports(1)[0] - - mp.spawn( - _mpmd_benchmark_worker, - args=( - world_size, - master_port, - test_name, - dtype, - num_layers, - shape, - FLAGS.benchmark_runs, - FLAGS.locality, - FLAGS.telemetry_log_path, - ), - nprocs=world_size, - join=True, + self.workers.run( + _BenchmarkCase( + test_name=test_name, + dtype=dtype, + num_layers=num_layers, + shape=tuple(shape), + benchmark_runs=FLAGS.benchmark_runs, + locality=FLAGS.locality, + telemetry_log_path=FLAGS.telemetry_log_path, + ) ) @parameterized.named_parameters( diff --git a/examples/microbenchmarks/torch_weight_sync_benchmark_runner.py b/examples/microbenchmarks/torch_weight_sync_benchmark_runner.py index 65e4ec471..1cb8205fa 100644 --- a/examples/microbenchmarks/torch_weight_sync_benchmark_runner.py +++ b/examples/microbenchmarks/torch_weight_sync_benchmark_runner.py @@ -26,27 +26,35 @@ are directly comparable. Per-stage metrics mirror the JAX runner too: D2H on the source, Net transfer and H2D on the destination. -Start the SOURCE first; with --parallelism_list it hosts an in-process -`RaidenControllerServer` on port (29550 + p) for each rung and waits for the -destination to drive and then shut each rung down: +Start the SOURCE first. For each rung p in --parallelism_list it hosts an +in-process `RaidenControllerServer` on port (29550 + p) and waits for the +destination to drive and then shut that rung down: cd PYTHONUNBUFFERED=1 PYTHONPATH=$PWD python3 \ examples/microbenchmarks/torch_weight_sync_benchmark_runner.py \ --role=source \ --parallelism_list=1,4,8 \ - --src_log_dir=/tmp/raiden_ws + --results_dir=/tmp/raiden_ws -Then, for each rung P, the DESTINATION pointed at the source host: +Then start the DESTINATION, pointed at the source host, with the same list. It +drives every rung from one set of processes, connecting to each rung's +controller in turn: cd PYTHONUNBUFFERED=1 PYTHONPATH=$PWD python3 \ examples/microbenchmarks/torch_weight_sync_benchmark_runner.py \ --role=destination \ - --controller_address=10.128.0.241:$((29550 + P)) \ - --parallelism=$P + --controller_host=10.128.0.241 \ + --parallelism_list=1,4,8 \ + --results_dir=/tmp/raiden_ws + +Both roles keep one set of processes for the whole ladder: starting fresh +processes per rung re-creates the multi-process TPU slice, and on TPU v7 the +new slice's START_SESSION can fail for minutes after the previous slice's +teardown timed out (b/568095515). -The payload flags must match on both sides. +The payload flags and --parallelism_list must match on both sides. """ import asyncio @@ -82,10 +90,10 @@ required=True, ) flags.DEFINE_string( - "controller_address", - "localhost:29551", - "Host:port of the RaidenController (destination only). The source hosts" - " the controller itself on port (29550 + parallelism).", + "controller_host", + "localhost", + "Destination only: host of the source, which serves each rung's" + " RaidenController on port (29550 + parallelism).", ) flags.DEFINE_integer( "num_processes", @@ -116,13 +124,15 @@ flags.DEFINE_integer( "parallelism", 8, - "Number of parallel transport streams for weight synchronization.", + "Number of parallel transport streams for a single-rung run.", ) flags.DEFINE_list( "parallelism_list", [], - "Source only: comma-separated parallelism values to sweep. For each" - " value p the source hosts a RaidenControllerServer on port (29550 + p).", + "Comma-separated parallelism values to sweep, in order; overrides" + " --parallelism. Must match on both roles: for each value p the source" + " hosts a RaidenControllerServer on port (29550 + p) and the destination" + " drives that rung.", ) flags.DEFINE_integer( "group_size", @@ -130,9 +140,10 @@ "Number of variables grouped per transfer request. 0 = all variables.", ) flags.DEFINE_string( - "src_log_dir", + "results_dir", "", - "Optional directory to write torch_ws_src_p{p}.log files on the source.", + "Optional directory for per-rung result files: torch_ws_src_p{p}.log on" + " the source, torch_ws_dst_p{p}.log on the destination.", ) flags.DEFINE_string( "bind_ip", @@ -395,6 +406,36 @@ def _ping_port(addr: str) -> bool: return False +def _controller_port(parallelism: int) -> int: + """Port the source serves rung `parallelism`'s RaidenController on.""" + return _CONTROLLER_BASE_PORT + parallelism + + +def _wait_for_port(addr: str, timeout_s: float) -> None: + """Blocks until addr accepts TCP connections; raises after timeout_s.""" + deadline = time.time() + timeout_s + while not _ping_port(addr): + if time.time() > deadline: + raise RuntimeError(f"Timed out after {timeout_s:.0f}s waiting for {addr}") + time.sleep(1.0) + + +def _emit_rung_report( + lines: abc.Sequence[str], + file_prefix: str, + parallelism: int, + results_dir: str, +) -> None: + """Prints a rung's result block; also saves it under results_dir if set.""" + report = "\n".join(lines) + "\n" + print(report, end="", flush=True) + if results_dir: + os.makedirs(results_dir, exist_ok=True) + path = os.path.join(results_dir, f"{file_prefix}_p{parallelism}.log") + with open(path, "w") as f: + f.write(report) + + # ---------------------------------------------------------------------------- # Source # ---------------------------------------------------------------------------- @@ -404,13 +445,13 @@ def _report_source( d2h_times: abc.Sequence[float], total_bytes: int, parallelism: int, - src_log_dir: str, + results_dir: str, ) -> None: total_mb = total_bytes / (1024.0 * 1024.0) avg_d2h_time = float(np.mean(d2h_times)) avg_d2h_bw_gbs = (total_bytes / 1e9) / max(avg_d2h_time, 1e-9) - report = ( - "\n".join([ + _emit_rung_report( + [ "==================================================", "SOURCE BENCHMARK RESULTS", "==================================================", @@ -422,16 +463,11 @@ def _report_source( f" ({avg_d2h_bw_gbs * 8.0:.2f} Gbps)" ), "==================================================", - ]) - + "\n" + ], + "torch_ws_src", + parallelism, + results_dir, ) - print(report, end="", flush=True) - if src_log_dir: - os.makedirs(src_log_dir, exist_ok=True) - with open( - os.path.join(src_log_dir, f"torch_ws_src_p{parallelism}.log"), "w" - ) as f: - f.write(report) def _run_source_worker(rank: int, world_size: int, cfg: dict[str, Any]) -> None: @@ -476,7 +512,7 @@ def _run_source_worker(rank: int, world_size: int, cfg: dict[str, Any]) -> None: worker_rpc_client = None try: for p in p_values: - ctrl_port = _CONTROLLER_BASE_PORT + p + ctrl_port = _controller_port(p) server = None if rank == 0: worker_rpc_client = raiden_controller.WeightSyncWorkerRpcClient( @@ -513,7 +549,7 @@ async def _noop_shutdown(*unused_args: Any, **unused_kw: Any): _barrier() d2h_times.append(time.perf_counter() - t0) if rank == 0: - _report_source(d2h_times, total_bytes, p, cfg["src_log_dir"]) + _report_source(d2h_times, total_bytes, p, cfg["results_dir"]) # Register this rank's shard with the rung's controller. ctrl_client = raiden_controller.RaidenControllerClientFacade( @@ -560,16 +596,42 @@ async def _noop_shutdown(*unused_args: Any, **unused_kw: Any): def _run_destination_worker( rank: int, world_size: int, cfg: dict[str, Any] ) -> None: - """Destination process: holds column shards, drives one rung, verifies.""" + """Destination process: holds column shards and drives every rung in turn. + + Like the source, one process (and one TPU client) serves the whole ladder; + each rung gets fresh tensors and a fresh WeightSynchronizer. + + Args: + rank: This process's rank among the destination's local processes. + world_size: Number of destination processes, one per local TPU device. + cfg: Benchmark configuration built by main(); cfg["p_values"] lists the + rungs to drive, in order. + """ if world_size > 1: dist.init_process_group(backend="tpu_dist") + try: + for parallelism in cfg["p_values"]: + _run_destination_rung(rank, world_size, cfg, parallelism) + finally: + if world_size > 1: + dist.destroy_process_group() + + +def _run_destination_rung( + rank: int, world_size: int, cfg: dict[str, Any], parallelism: int +) -> None: + """Drives one rung: registers, transfers, verifies, then stops its controller.""" device = torch.device("tpu") torch_dtype, item_size = _resolve_torch_dtype(cfg["dtype"]) var_shape: tuple[int, int] = cfg["var_shape"] num_variables: int = cfg["num_variables"] - parallelism: int = cfg["parallelism"] self_ip: str = cfg["self_ip"] + controller_address = ( + f"{cfg['controller_host']}:{_controller_port(parallelism)}" + ) + # Freshly zeroed per rung, so the warmup verification proves this rung's + # transfer rather than data left over from the previous rung. tensors = _allocate_tensors( "destination", rank, @@ -598,8 +660,10 @@ def _run_destination_worker( while not _ping_port(ctrl_addr): time.sleep(0.2) + # The source opens a rung's controller only after its own D2H stage. + _wait_for_port(controller_address, _REGISTRATION_TIMEOUT_S) ctrl_client = raiden_controller.RaidenControllerClientFacade( - cfg["controller_address"], name_resolver=None + controller_address, name_resolver=None ) ctrl_client.register_work_unit( _unit(_DST_JOB, rank), @@ -680,26 +744,33 @@ def _to_id(m: Any) -> raiden_controller.RaidenId: avg_h2d = float(np.mean(h2d_times)) net_gbs = (total_bytes / 1e9) / max(avg_net, 1e-9) h2d_gbs = (total_bytes / 1e9) / max(avg_h2d, 1e-9) - print("==================================================") - print("DESTINATION BENCHMARK RESULTS") - print("==================================================") - print(f"Total Size: {total_mb:.2f} MB") - print(f"Parallelism: {parallelism}") - print(f"Avg Net Transfer Time:{avg_net:.6f} s") - print( - f"Avg Net Transfer BW: {net_gbs:.2f} GB/s ({net_gbs * 8.0:.2f} Gbps)" - ) - print(f"Avg H2D Time: {avg_h2d:.6f} s") - print( - f"Avg H2D BW: {h2d_gbs:.2f} GB/s ({h2d_gbs * 8.0:.2f} Gbps)" + _emit_rung_report( + [ + "==================================================", + "DESTINATION BENCHMARK RESULTS", + "==================================================", + f"Total Size: {total_mb:.2f} MB", + f"Parallelism: {parallelism}", + f"Avg Net Transfer Time:{avg_net:.6f} s", + ( + f"Avg Net Transfer BW: {net_gbs:.2f} GB/s" + f" ({net_gbs * 8.0:.2f} Gbps)" + ), + f"Avg H2D Time: {avg_h2d:.6f} s", + ( + f"Avg H2D BW: {h2d_gbs:.2f} GB/s" + f" ({h2d_gbs * 8.0:.2f} Gbps)" + ), + "==================================================", + ], + "torch_ws_dst", + parallelism, + cfg["results_dir"], ) - print("==================================================", flush=True) ctrl_client.shutdown() _barrier() finally: - del ws - if world_size > 1: - dist.destroy_process_group() + ws.shutdown() # ---------------------------------------------------------------------------- @@ -764,6 +835,17 @@ def main(argv: abc.Sequence[str]) -> None: var_shape = tuple(int(x) for x in FLAGS.variable_shape) if len(var_shape) != 2: raise app.UsageError("--variable_shape must be 2D (rows,cols).") + p_values = ( + [int(x) for x in FLAGS.parallelism_list] + if FLAGS.parallelism_list + else [FLAGS.parallelism] + ) + # Rungs are keyed by parallelism (controller port, the source's last-rung + # check), so a repeated value would collide with an earlier rung. + if len(set(p_values)) != len(p_values): + raise app.UsageError( + f"--parallelism_list must not repeat a value: {p_values}" + ) world_size = _local_world_size() cfg: dict[str, Any] = { "role": FLAGS.role, @@ -773,14 +855,9 @@ def main(argv: abc.Sequence[str]) -> None: "num_iterations": FLAGS.num_iterations, "group_size": FLAGS.group_size, "self_ip": _resolve_local_ip(), - "src_log_dir": FLAGS.src_log_dir, - "controller_address": FLAGS.controller_address, - "parallelism": FLAGS.parallelism, - "p_values": ( - [int(x) for x in FLAGS.parallelism_list] - if FLAGS.parallelism_list - else [FLAGS.parallelism] - ), + "results_dir": FLAGS.results_dir, + "controller_host": FLAGS.controller_host, + "p_values": p_values, # Source and destination hosts are assumed symmetric (same chip count); # the destination only needs a lower bound on registered source shards. "num_src_shards": world_size,