xnn.common.distributed

Process-group helpers shared by training and benchmarking.

xnn’s multi-GPU / multi-node support is native torch.distributed driven purely by the environment: a launcher such as torchrun (or Slurm plus torchrun) exports RANK / LOCAL_RANK / WORLD_SIZE into every process it spawns, and the code paths here switch on those variables. A plain python / xnn invocation sees WORLD_SIZE unset and runs the unchanged single-process pipeline. No framework wrapper is involved; these are the handful of primitives the Trainer and the Benchmark share so they cannot drift apart.

Functions

all_gather_cat(t, device)	Concatenate a flat tensor from every rank, in rank order.
all_reduce_sum(values, device)	Sum a list of scalars over all ranks (identity outside a distributed run).
barrier()	Synchronize all ranks (a no-op outside a distributed run).
destroy_process_group(owns)	Tear down the process group if owns says this caller created it.
init_process_group(device)	Join the launcher's process group, if there is one.
is_distributed()	Whether this process was spawned by a distributed launcher.
is_main()	Whether this process is rank 0 (the one that prints and writes).
merge_dicts(d)	Merge a small picklable dict over all ranks (identity outside a distributed run).
rank()	This process's global rank (0 when not launched distributed).
shard_indices(n)	The dataset indices this rank scores: every WORLD_SIZE-th one.
world_size()	Number of processes in the distributed launch (1 when not launched).