257 lines
8.1 KiB
Python
257 lines
8.1 KiB
Python
#!/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()
|