feat: all reduce bench slurm pyxis

[OrenLeung]

SLURM Pyxis Container Plugin allows for easy reproducible scripts where the environment is containerized as such the benchmark harness script does not rely on any host machine dependencies besides the standard SLURM, pyxis SLURM plugin and nvidia-drivers.

SLURM Pyxis Container Plugin is widely used across many companies and increasingly gaining adoption.

On CSPs that have enabled SLURM Pyxis Container Plugin, such as CoreWeave, Crusoe, Oracle, Azure, etc, all_reduce_bench.py can be easily ran & reproduced via the following command:

sbatch -n <num_of_nodes> ./all_reduce_bench_pyxis.sbatch

Note that on AWS & GCP, this launcher will work too once you swapped nvcr.io#nvidia/pytorch:25.02-py3 and AWS/GCP specific container image that has all of the env vars and AWS EFA or GCP GPUDirect-TCP or GPUDirect-TCPXO aka FastTrak or gIB nccl net plugin included

Testing

I did some quick smoke tests to make sure this script works properly on two HGX H100 700W SXM nodes (16 total GPUs) with 400G InfiniBand NDR connected between them at 1 hop distance.

[stas00]

Thanks a lot, Oren!

[OrenLeung]

Thanks for the PR review @stas00 !