Add Qwen235B NCCL profile runner
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user