feat: add intranode examples (#31)

Signed-off-by: chang-ning <spiderpower02@gmail.com>

Author: crazyguitar

README.md

@@ -66,6 +66,7 @@ to when you pressed `r`, so the trace spans exactly your record window.
 Use `rdmatop` to monitor RDMA traffic while running GPU
 communication benchmarks:
 
+- [PyTorch](examples/pytorch/) — intranode NVLink/XGMI traffic
 - [IB Perftest](examples/ib/) — two-node `ib_write_bw` benchmark
 - [UCX Perftest](examples/ucx/) — two-node `ucx_perftest` bandwidth / latency
 - [NCCL](examples/nccl/) — collective communication

examples/pytorch/README.md

@@ -0,0 +1,62 @@
+# PyTorch Collectives
+
+Small `torch.distributed` loops that drive sustained traffic over the
+intranode GPU interconnect — **XGMI** on AMD (RCCL) or **NVLink** on
+NVIDIA (NCCL). Run one in a shell and watch the `amdgpu<N>` /
+`nvidia<N>` rows in `rdmatop` from another.
+
+Collectives are symmetric: every GPU talks to every peer, so all
+interconnect rows light up with roughly equal TX/RX, and the per-peer
+detail pane shows traffic on every link.
+
+| Script             | Collective              | TX per GPU per iteration    |
+|--------------------|-------------------------|-----------------------------|
+| `allgather.py`     | `all_gather`            | `(world-1) * buffer`        |
+| `allreduce.py`     | `all_reduce`            | `~2*(world-1)/world * buffer` |
+| `alltoall.py`      | `all_to_all_single`     | `(world-1)/world * buffer`  |
+| `reducescatter.py` | `reduce_scatter_tensor` | `(world-1)/world * buffer`  |
+
+## Prerequisites
+
+- 2+ GPUs on one node with PyTorch matching your GPU stack:
+  ```bash
+  # AMD (ROCm)
+  pip install torch --index-url https://download.pytorch.org/whl/rocm6.4
+  # NVIDIA (CUDA)
+  pip install torch
+  ```
+- `rdmatop` built with the matching feature to see the GPU rows:
+  ```bash
+  cargo install --path . --features xgmi     # AMD
+  cargo install --path . --features nvlink   # NVIDIA
+  ```
+
+## Usage
+
+```bash
+cd examples/pytorch
+torchrun --nproc_per_node=8 allgather.py                # 60s, 256 MiB/rank
+torchrun --nproc_per_node=8 allreduce.py --mb 512 --secs 300
+torchrun --nproc_per_node=4 alltoall.py                 # only 4 GPUs
+torchrun --nproc_per_node=8 reducescatter.py
+```
+
+Rank 0 prints an estimated per-GPU transmit rate each batch
+(`~N Gbps TX/GPU`, from the ring-algorithm model in the table above).
+Compare it against the per-row TX in `rdmatop`.
+
+## Notes
+
+- The printed rate is a ring-algorithm estimate; RCCL/NCCL may pick
+  other algorithms, so treat it as a sanity reference. Ground truth on
+  AMD is `amd-smi xgmi` accumulator deltas, which read the same
+  counters `rdmatop` samples.
+- The `nccl` backend name is historical — on ROCm builds of PyTorch it
+  is backed by RCCL, no code changes needed.
+
+## Related Links
+
+- [torch.distributed](https://pytorch.org/docs/stable/distributed.html)
+  — collective API used by the scripts
+- [rccl](https://github.com/ROCm/rccl) — AMD collective library
+- [nccl](https://github.com/NVIDIA/nccl) — NVIDIA collective library

examples/pytorch/allgather.py

@@ -0,0 +1,16 @@
+import torch
+import torch.distributed as dist
+
+import common
+
+
+def make(world, n):
+    x = torch.randn(n, dtype=torch.float16, device="cuda")
+    out = [torch.empty_like(x) for _ in range(world)]
+    # ring all-gather: each GPU transmits (world-1) shards per iteration
+    tx = x.numel() * x.element_size() * (world - 1)
+    return lambda: dist.all_gather(out, x), tx
+
+
+if __name__ == "__main__":
+    common.run("all_gather", make)

examples/pytorch/allreduce.py

@@ -0,0 +1,15 @@
+import torch
+import torch.distributed as dist
+
+import common
+
+
+def make(world, n):
+    x = torch.randn(n, dtype=torch.float16, device="cuda")
+    # ring all-reduce: each GPU transmits ~2*(world-1)/world of its buffer
+    tx = int(x.numel() * x.element_size() * 2 * (world - 1) / world)
+    return lambda: dist.all_reduce(x), tx
+
+
+if __name__ == "__main__":
+    common.run("all_reduce", make)

examples/pytorch/alltoall.py

@@ -0,0 +1,19 @@
+import torch
+import torch.distributed as dist
+
+import common
+
+
+def make(world, n):
+    # all_to_all_single sends one equal slice to each rank, so the buffer
+    # is built as `world` slices of n//world elements.
+    slice_len = n // world
+    x = torch.randn(slice_len * world, dtype=torch.float16, device="cuda")
+    out = torch.empty_like(x)
+    # each GPU transmits (world-1) slices per iteration (own slice stays)
+    tx = slice_len * x.element_size() * (world - 1)
+    return lambda: dist.all_to_all_single(out, x), tx
+
+
+if __name__ == "__main__":
+    common.run("all_to_all", make)

examples/pytorch/common.py

@@ -0,0 +1,47 @@
+"""Shared harness for the intranode collective benchmarks.
+
+Each benchmark defines its collective and a per-iteration TX-bytes model;
+run() handles setup, warmup, the timed loop, and rank-0 rate reporting.
+"""
+
+import argparse
+import time
+
+import torch
+import torch.distributed as dist
+
+
+def run(name, make):
+    p = argparse.ArgumentParser(description=f"{name} loop for rdmatop monitoring")
+    p.add_argument("--mb", type=int, default=256, help="per-rank buffer size in MiB")
+    p.add_argument("--secs", type=float, default=60, help="run duration in seconds")
+    args = p.parse_args()
+
+    dist.init_process_group("nccl")  # NCCL on CUDA, RCCL on ROCm
+    rank = dist.get_rank()
+    world = dist.get_world_size()
+    torch.cuda.set_device(rank % torch.cuda.device_count())
+
+    n = args.mb * (1 << 20) // 2  # fp16 elements per rank
+    step, tx_bytes_per_iter = make(world, n)
+
+    for _ in range(3):  # warmup
+        step()
+    torch.cuda.synchronize()
+
+    t0 = time.time()
+    iters = 0
+    while time.time() - t0 < args.secs:
+        for _ in range(10):
+            step()
+            iters += 1
+        torch.cuda.synchronize()
+        if rank == 0:
+            elapsed = time.time() - t0
+            gbps = iters * tx_bytes_per_iter * 8 / 1e9 / elapsed
+            print(
+                f"{elapsed:6.1f}s  {iters:5d} iters  ~{gbps:7.1f} Gbps TX/GPU",
+                flush=True,
+            )
+
+    dist.destroy_process_group()

examples/pytorch/reducescatter.py

@@ -0,0 +1,19 @@
+import torch
+import torch.distributed as dist
+
+import common
+
+
+def make(world, n):
+    # reduce_scatter_tensor reduces the input into one shard per rank, so
+    # the buffer is built as `world` shards of n//world elements.
+    shard_len = n // world
+    x = torch.randn(shard_len * world, dtype=torch.float16, device="cuda")
+    out = torch.empty(shard_len, dtype=torch.float16, device="cuda")
+    # ring reduce-scatter: each GPU transmits (world-1) shards per iteration
+    tx = shard_len * x.element_size() * (world - 1)
+    return lambda: dist.reduce_scatter_tensor(out, x), tx
+
+
+if __name__ == "__main__":
+    common.run("reduce_scatter", make)