56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import json
|
|
import os
|
|
import statistics
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def main() -> None:
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
dist.init_process_group(backend="nccl")
|
|
|
|
# 16 tokens x 4096 hidden values in BF16: 128 KiB per rank.
|
|
tensor = torch.ones((16, 4096), dtype=torch.bfloat16, device="cuda")
|
|
for _ in range(10):
|
|
dist.all_reduce(tensor)
|
|
torch.cuda.synchronize()
|
|
|
|
samples_ms = []
|
|
for _ in range(50):
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
dist.all_reduce(tensor)
|
|
end.record()
|
|
end.synchronize()
|
|
samples_ms.append(float(start.elapsed_time(end)))
|
|
|
|
if dist.get_rank() == 0:
|
|
ordered = sorted(samples_ms)
|
|
result = {
|
|
"backend": "nccl",
|
|
"collective": "all_reduce",
|
|
"dtype": "bfloat16",
|
|
"elements_per_rank": tensor.numel(),
|
|
"bytes_per_rank": tensor.numel() * tensor.element_size(),
|
|
"world_size": dist.get_world_size(),
|
|
"warmup_iterations": 10,
|
|
"measured_iterations": len(samples_ms),
|
|
"mean_ms": statistics.fmean(samples_ms),
|
|
"p50_ms": statistics.median(samples_ms),
|
|
"p95_ms": ordered[int(0.95 * (len(ordered) - 1))],
|
|
"min_ms": min(samples_ms),
|
|
"max_ms": max(samples_ms),
|
|
}
|
|
print(json.dumps(result, sort_keys=True))
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|