From c71f37911079ea71b59f2fac445a794d10c07e0f Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 15 Jul 2026 19:10:52 +0800 Subject: [PATCH] Add Qwen235B NCCL profile runner --- .../best_effort/profile_allreduce.py | 256 ++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/profile_allreduce.py diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/profile_allreduce.py b/runs/frontier-multicase-sufficiency-v0/best_effort/profile_allreduce.py new file mode 100644 index 0000000..484a2e7 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/profile_allreduce.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Profile the NCCL all-reduce paths consumed by the Qwen235B experiment.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import os +import statistics +import sys +from datetime import datetime, timezone +from pathlib import Path + +import torch +import torch.distributed as dist + + +DEFAULT_TOKENS = [ + 1, + 2, + 3, + 4, + 6, + 8, + 12, + 16, + 24, + 32, + 48, + 63, + 64, + 65, + 96, + 127, + 128, + 129, + 192, + 255, + 256, + 257, + 384, + 511, + 512, + 513, + 768, + 1023, + 1024, + 1025, + 1536, + 2047, + 2048, + 2049, + 3072, + 4095, + 4096, + 4097, + 6144, + 8191, + 8192, + 8193, + 12288, + 16383, + 16384, +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--hidden-size", type=int, default=4096) + parser.add_argument("--world-sizes", type=int, nargs="+", default=[4, 8]) + parser.add_argument("--tokens", type=int, nargs="+", default=DEFAULT_TOKENS) + parser.add_argument("--warmup-iterations", type=int, default=5) + parser.add_argument("--measured-iterations", type=int, default=30) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def profile_one_size( + *, + group: dist.ProcessGroup, + group_size: int, + num_tokens: int, + hidden_size: int, + warmup_iterations: int, + measured_iterations: int, +) -> list[float]: + tensor = torch.ones( + (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda" + ) + for _ in range(warmup_iterations): + dist.all_reduce(tensor, group=group) + dist.barrier(group=group) + torch.cuda.synchronize() + + local_samples: list[float] = [] + for _ in range(measured_iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + dist.all_reduce(tensor, group=group) + end.record() + end.synchronize() + local_samples.append(float(start.elapsed_time(end))) + + samples_by_rank: list[list[float] | None] = [None] * group_size + dist.all_gather_object(samples_by_rank, local_samples, group=group) + del tensor + + # Collective completion is limited by the slowest rank. Preserve that + # critical-path statistic instead of reporting only rank 0. + return [ + max(float(rank_samples[index]) for rank_samples in samples_by_rank if rank_samples) + for index in range(measured_iterations) + ] + + +def profile_group( + *, + group: dist.ProcessGroup, + group_ranks: list[int], + args: argparse.Namespace, +) -> list[dict[str, object]]: + global_rank = dist.get_rank() + if global_rank not in group_ranks: + return [] + + group_size = len(group_ranks) + group_rank = dist.get_rank(group=group) + rows: list[dict[str, object]] = [] + trial_orders = (("forward", args.tokens), ("reverse", list(reversed(args.tokens)))) + for trial_id, (order_name, tokens) in enumerate(trial_orders): + for num_tokens in tokens: + samples = profile_one_size( + group=group, + group_size=group_size, + num_tokens=num_tokens, + hidden_size=args.hidden_size, + warmup_iterations=args.warmup_iterations, + measured_iterations=args.measured_iterations, + ) + if group_rank != 0: + continue + ordered = sorted(samples) + p95_index = int(0.95 * (len(ordered) - 1)) + rows.append( + { + "time_stats.all_reduce.min": min(samples), + "time_stats.all_reduce.max": max(samples), + "time_stats.all_reduce.mean": statistics.fmean(samples), + "time_stats.all_reduce.median": statistics.median(samples), + "time_stats.all_reduce.std": statistics.stdev(samples), + "time_stats.all_reduce.p95": ordered[p95_index], + "rank": 0, + "num_workers": group_size, + "size": num_tokens * args.hidden_size * 2, + "collective": "all_reduce", + "devices_per_node": group_size, + "max_devices_per_node": dist.get_world_size(), + "profiling_precision": "BF16", + "measurement_type": "CUDA_EVENT", + "backend": "nccl", + "num_tokens": num_tokens, + "hidden_size": args.hidden_size, + "trial_id": trial_id, + "order": order_name, + "warmup_iterations": args.warmup_iterations, + "measured_iterations": args.measured_iterations, + } + ) + return rows + + +def main() -> None: + args = parse_args() + if args.hidden_size <= 0: + raise ValueError("hidden-size must be positive") + if args.warmup_iterations <= 0 or args.measured_iterations < 2: + raise ValueError("warmup must be positive and measured iterations must be >= 2") + if sorted(set(args.tokens)) != sorted(args.tokens) or min(args.tokens) <= 0: + raise ValueError("tokens must be unique positive integers") + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + global_rank = dist.get_rank() + global_world_size = dist.get_world_size() + if max(args.world_sizes) > global_world_size: + raise ValueError( + f"requested world size {max(args.world_sizes)} exceeds torchrun world size " + f"{global_world_size}" + ) + + groups: dict[int, tuple[dist.ProcessGroup, list[int]]] = {} + for world_size in sorted(set(args.world_sizes)): + ranks = list(range(world_size)) + group = dist.group.WORLD if world_size == global_world_size else dist.new_group(ranks) + groups[world_size] = (group, ranks) + + all_rows: list[dict[str, object]] = [] + for world_size in args.world_sizes: + group, ranks = groups[world_size] + all_rows.extend(profile_group(group=group, group_ranks=ranks, args=args)) + dist.barrier() + + if global_rank == 0: + expected_rows = len(set(args.world_sizes)) * len(args.tokens) * 2 + if len(all_rows) != expected_rows: + raise RuntimeError( + f"profile row count mismatch: expected={expected_rows}, actual={len(all_rows)}" + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", newline="", encoding="utf-8") as output_file: + writer = csv.DictWriter(output_file, fieldnames=list(all_rows[0])) + writer.writeheader() + writer.writerows(all_rows) + + manifest = { + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "command": [sys.executable, *sys.argv], + "output": str(args.output.resolve()), + "output_sha256": sha256(args.output), + "rows": len(all_rows), + "world_sizes": args.world_sizes, + "tokens": args.tokens, + "hidden_size": args.hidden_size, + "dtype": "torch.bfloat16", + "backend": dist.get_backend(), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "nccl_version": list(torch.cuda.nccl.version()), + "gpu": torch.cuda.get_device_name(0), + "warmup_iterations": args.warmup_iterations, + "measured_iterations": args.measured_iterations, + "sample_semantics": "per-iteration maximum CUDA-event time across participating ranks", + } + manifest_path = args.output.with_suffix(args.output.suffix + ".manifest.json") + manifest_path.write_text( + json.dumps(manifest, indent=2) + "\n", encoding="utf-8" + ) + print(json.dumps(manifest, indent=2)) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main()