Add device_id parameter to all torch distributed initialization calls

benchmark/examples/benchmark_all_gather_gemm_pull.py

@@ -75,7 +75,9 @@ def worker(rank: int, world_size: int, init_url: str, args: argparse.Namespace):
     This function will be executed by each spawned process.
     """
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
 
     shmem = iris.iris(args.heap_size)
     torch.cuda.set_device(rank)

benchmark/examples/benchmark_all_gather_gemm_push.py

@@ -74,7 +74,9 @@ def worker(rank: int, world_size: int, init_url: str, args: argparse.Namespace):
     This function will be executed by each spawned process.
     """
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
 
     shmem = iris.iris(args.heap_size)
     torch.cuda.set_device(rank)

benchmark/examples/benchmark_flash_decode.py

@@ -106,7 +106,9 @@ def prepare_perf_data(cfg, num_query_heads, num_kv_heads):
 
 def run_benchmark(rank, world_size, init_url, args):
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
     # Set the correct GPU for this specific process
     torch.cuda.set_device(rank)
 

examples/00_load/load_bench.py

@@ -235,7 +235,13 @@ def print_bandwidth_matrix(matrix, label="Unidirectional LOAD bandwidth GiB/s [R
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/01_store/store_bench.py

@@ -208,7 +208,13 @@ def print_bandwidth_matrix(matrix, label="Unidirectional STORE bandwidth GiB/s [
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/02_all_load/all_load_bench.py

@@ -316,7 +316,13 @@ def print_bandwidth_matrix(
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     heap_size = args["heap_size"]

examples/03_all_store/all_store_bench.py

@@ -245,7 +245,13 @@ def print_bandwidth_matrix(
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     heap_size = args["heap_size"]

examples/04_atomic_add/atomic_add_bench.py

@@ -206,7 +206,13 @@ def print_bandwidth_matrix(
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/05_atomic_xchg/atomic_xchg_bench.py

@@ -212,7 +212,13 @@ def print_bandwidth_matrix(
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/06_message_passing/message_passing_load_store.py

@@ -137,7 +137,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/06_message_passing/message_passing_put.py

@@ -125,7 +125,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/07_gemm_all_scatter/benchmark.py

@@ -61,7 +61,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/08_gemm_atomics_all_reduce/benchmark.py

@@ -80,7 +80,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     # Main benchmark logic
     shmem = iris.iris(args["heap_size"])

examples/09_gemm_one_shot_all_reduce/benchmark.py

@@ -77,7 +77,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     shmem = iris.iris(args["heap_size"])
     rank = shmem.get_rank()

examples/10_gemm_all_scatter_wg_specialization/benchmark.py

@@ -65,7 +65,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     shmem = iris.iris(args["heap_size"])
     rank = shmem.get_rank()

examples/11_gemm_all_scatter_producer_consumer/benchmark.py

@@ -66,7 +66,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     shmem = iris.iris(args["heap_size"])
     rank = shmem.get_rank()

examples/12_gemm_all_scatter_bulk_synchronous/benchmark.py

@@ -66,7 +66,13 @@ def parse_args():
 def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     """Worker function for PyTorch distributed execution."""
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=init_url,
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(f"cuda:{local_rank}"),
+    )
 
     shmem = iris.iris(args["heap_size"])
     rank = shmem.get_rank()

examples/13_flash_decode/example_run.py

@@ -82,7 +82,9 @@ def setup_example_data(rank, world_size, args, dtype):
 
 def example_run(rank: int, world_size: int, init_url: str, args: dict):
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
 
     # 1. Initialize Iris for distributed communication
     shmem = iris.iris()

examples/14_all_gather_gemm/example_run_pull.py

@@ -74,7 +74,9 @@ def setup_example_data(rank, world_size, args, dtype):
 
 def example_run(rank: int, world_size: int, init_url: str, args: argparse.Namespace):
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
 
     # Initialize Iris for distributed communication
     shmem = iris.iris()

examples/14_all_gather_gemm/example_run_push.py

@@ -72,7 +72,9 @@ def setup_example_data(rank, world_size, args, dtype):
 
 def example_run(rank: int, world_size: int, init_url: str, args: argparse.Namespace):
     backend = "nccl" if torch.cuda.is_available() else "gloo"
-    dist.init_process_group(backend=backend, init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend=backend, init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
 
     shmem = iris.iris()
     torch.manual_seed(42)

examples/benchmark/reference/all_gather.py

@@ -8,6 +8,7 @@
 import random
 import iris
 import argparse
+import os
 
 from examples.common.utils import JSONWriter
 
@@ -53,7 +54,8 @@ def main():
     validate = args["validate"]
     benchmark = args["benchmark"]
 
-    dist.init_process_group("nccl")
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    dist.init_process_group("nccl", device_id=torch.device(f"cuda:{local_rank}"))
 
     rank = dist.get_rank()
     world_size = dist.get_world_size()

examples/benchmark/reference/all_gather_gemm/benchmark_rccl_torch.py

@@ -46,7 +46,9 @@ def parse_args():
 
 
 def worker(rank: int, world_size: int, init_url: str, args: argparse.Namespace):
-    dist.init_process_group(backend="nccl", init_method=init_url, world_size=world_size, rank=rank)
+    dist.init_process_group(
+        backend="nccl", init_method=init_url, world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{rank}")
+    )
     torch.cuda.set_device(rank)
 
     output_dir = args.output_dir

examples/benchmark/reference/all_reduce.py

@@ -8,6 +8,7 @@
 import random
 import iris
 import argparse
+import os
 
 from examples.common.utils import JSONWriter
 
@@ -52,7 +53,8 @@ def main():
     validate = args["validate"]
     benchmark = args["benchmark"]
 
-    dist.init_process_group("nccl")
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    dist.init_process_group("nccl", device_id=torch.device(f"cuda:{local_rank}"))
 
     rank = dist.get_rank()
     world_size = dist.get_world_size()

examples/benchmark/reference/flash_decode_rccl/benchmark_flash_decode_rccl.py

@@ -91,7 +91,8 @@ def prepare_perf_data(config, num_query_heads, num_kv_heads, page_size, datatype
 
 
 def run_benchmark(args):
-    dist.init_process_group(backend="nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    dist.init_process_group(backend="nccl", device_id=torch.device(f"cuda:{local_rank}"))
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

examples/benchmark/reference/reduce_scatter.py

@@ -8,6 +8,7 @@
 import random
 import iris
 import argparse
+import os
 
 from examples.common.utils import JSONWriter
 
@@ -43,7 +44,8 @@ def main():
     m, n, k = args["m"], args["n"], args["k"]
     validate, benchmark = args["validate"], args["benchmark"]
 
-    dist.init_process_group("nccl")
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    dist.init_process_group("nccl", device_id=torch.device(f"cuda:{local_rank}"))
     rank = dist.get_rank()
     world_size = dist.get_world_size()
     torch.cuda.set_device(rank)