Skip to content

Commit 6463723

Browse files
author
Verdi March
committed
Script to probe the nccl stack that PyTorch actually uses
1 parent 091d536 commit 6463723

3 files changed

Lines changed: 188 additions & 0 deletions

File tree

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
# Runtime sanity checks: which NCCL is loaded by PyTorch
2+
3+
The scripts in this folder disambiguate the exact NCCL library that a PyTorch application actually
4+
uses, in the presence of potentially multiple NCCL libraries installed.
5+
6+
## 1. Motivation
7+
8+
Knowing the precise NCCL version is important for performance tuning. The exact NCCL library used by
9+
PyTorch depends on how the PyTorch was installed. On DLAMI which pre-installs several different
10+
versions of NCCL, it is misleading to assume that PyTorch will use one of these libraries.
11+
12+
Suppose your EC2 instance runs DLAMI `Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04)
13+
20240314`. This AMI provides `libnccl.so` version 2.18.5 under `/usr/lib/cuda` ([Section
14+
3.1](#31-system-wide-nccl-provided-by-dlami)). You then install PyTorch to a conda or a regular
15+
Python virtual environment, using the [prebuilt PyTorch
16+
installer](https://pytorch.org/get-started/locally/) from the pytorch
17+
[channel](https://anaconda.org/pytorch/repo) or [PyPI](https://pypi.org/project/torch/),
18+
respectively. We can easily observe that our PyTorch installs its own NCCL and CUDA runtime.
19+
20+
```console
21+
$ conda list | egrep 'torch|nvidia'
22+
nvidia-cublas-cu12 12.1.3.1 pypi_0 pypi
23+
nvidia-cuda-cupti-cu12 12.1.105 pypi_0 pypi
24+
nvidia-cuda-nvrtc-cu12 12.1.105 pypi_0 pypi
25+
nvidia-cuda-runtime-cu12 12.1.105 pypi_0 pypi
26+
nvidia-cudnn-cu12 8.9.2.26 pypi_0 pypi
27+
nvidia-cufft-cu12 11.0.2.54 pypi_0 pypi
28+
nvidia-curand-cu12 10.3.2.106 pypi_0 pypi
29+
nvidia-cusolver-cu12 11.4.5.107 pypi_0 pypi
30+
nvidia-cusparse-cu12 12.1.0.106 pypi_0 pypi
31+
nvidia-nccl-cu12 2.19.3 pypi_0 pypi
32+
nvidia-nvjitlink-cu12 12.4.99 pypi_0 pypi
33+
nvidia-nvtx-cu12 12.1.105 pypi_0 pypi
34+
torch 2.2.1 pypi_0 pypi
35+
torchaudio 2.2.1 pypi_0 pypi
36+
torchvision 0.17.1 pypi_0 pypi
37+
```
38+
39+
We can verify that our PyTorch indeed loads its NCCL version `(2, 19, 3)` from the conda
40+
environment, instead of the system-wide version `2.18.5`.
41+
42+
```console
43+
# Library search paths, without conda environment
44+
$ env | grep ^LD
45+
LD_LIBRARY_PATH=/opt/amazon/efa/lib:/opt/amazon/openmpi/lib:/opt/aws-ofi-nccl/lib:/usr/local/cuda-12.1/lib:/usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1:/usr/local/cuda-12.1/targets/x86_64-linux/lib/:/usr/local/cuda-12.1/extras/CUPTI/lib64:/usr/local/lib:/usr/lib
46+
47+
# Print version of the system-wide NCCL, pre-installed with DLAMI
48+
$ strings /usr/local/cuda/lib/libnccl.so | grep -m1 -i '^NCCL version .*\+cuda.*$'
49+
NCCL version 2.18.5+cuda12.2
50+
51+
# Activate the conda environment where PyTorch is located.
52+
$ source miniconda3/bin/activate ./conda_env_pytorch/
53+
54+
# Conda environment does not change the library search paths.
55+
(conda_env_name) $ env | grep ^LD
56+
LD_LIBRARY_PATH=/opt/amazon/efa/lib:/opt/amazon/openmpi/lib:/opt/aws-ofi-nccl/lib:/usr/local/cuda-12.1/lib:/usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1:/usr/local/cuda-12.1/targets/x86_64-linux/lib/:/usr/local/cuda-12.1/extras/CUPTI/lib64:/usr/local/lib:/usr/lib
57+
58+
# Query the nccl version to PyTorch
59+
(conda_env_pytorch) $ python -c 'import torch; print(f"{torch.cuda.nccl.version()=}")'
60+
torch.cuda.nccl.version()=(2, 19, 3)
61+
62+
# Similar to above, but also pinpoint the exact `libnccl.so` file.
63+
(conda_env_name) $ LD_DEBUG=libs python -c 'import torch; print(f"{torch.cuda.nccl.version()=}")' 2>&1 | egrep 'torch.cuda.nccl.version|libnccl.so'
64+
47352: find library=libnccl.so.2 [0]; searching
65+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cublas/lib/libnccl.so.2
66+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cuda_cupti/lib/libnccl.so.2
67+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cuda_nvrtc/lib/libnccl.so.2
68+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cuda_runtime/lib/libnccl.so.2
69+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cudnn/lib/libnccl.so.2
70+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cufft/lib/libnccl.so.2
71+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/curand/lib/libnccl.so.2
72+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cusolver/lib/libnccl.so.2
73+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/cusparse/lib/libnccl.so.2
74+
47352: trying file=/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2
75+
47352: calling init: /fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2
76+
torch.cuda.nccl.version()=(2, 19, 3)
77+
47352: calling fini: /fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2 [0]
78+
```
79+
80+
By now, hopefully you're convinced on the need for a runtime probe to pinpoint the version of NCCL
81+
loaded by PyTorch.
82+
83+
## 2. Howto
84+
85+
Pre-requisites:
86+
87+
- one GPU.
88+
- the command `python` will invoke the binary from your PyTorch environment. Typically, you must
89+
activate your environment (conda, virtual env, etc.).
90+
91+
Direct invocation: `./probe-pt-nccl-aws-libs.sh`
92+
93+
Via Slurm: `srun -l -N1 ./probe-pt-nccl-aws-libs.sh`
94+
95+
## 3. Appendix
96+
97+
### 3.1. System-wide NCCL provided by DLAMI
98+
99+
See the [DLAMI release
100+
notes](https://aws.amazon.com/releasenotes/aws-deep-learning-base-gpu-ami-ubuntu-20-04/) for the
101+
versions of pre-installed [CUDA runtime](https://docs.nvidia.com/cuda/),
102+
[NCCL](https://github.com/NVIDIA/nccl), and [aws-ofi-nccl](https://github.com/aws/aws-ofi-nccl).
103+
104+
Optionally, follow below steps to validate the versions on a live EC2 instance.
105+
106+
```console
107+
# DLAMI provides multiple versions of CUDA runtime
108+
$ ls -ald /usr/local/cuda-*/ /usr/local/cuda
109+
lrwxrwxrwx 1 root root 21 Mar 14 22:22 /usr/local/cuda -> /usr/local/cuda-12.1/
110+
drwxr-xr-x 17 root root 4096 Mar 14 22:01 /usr/local/cuda-11.7/
111+
drwxr-xr-x 19 root root 4096 Mar 14 22:19 /usr/local/cuda-11.8/
112+
drwxr-xr-x 19 root root 4096 Mar 14 22:13 /usr/local/cuda-12.0/
113+
drwxr-xr-x 19 root root 4096 Mar 14 22:25 /usr/local/cuda-12.1/
114+
drwxr-xr-x 19 root root 4096 Mar 14 22:07 /usr/local/cuda-12.2/
115+
116+
# Each CUDA version brings its own NCCL library
117+
$ find /usr/local/cuda-* -name 'libnccl.so'
118+
/usr/local/cuda-11.7/lib/libnccl.so
119+
/usr/local/cuda-11.8/lib/libnccl.so
120+
/usr/local/cuda-12.0/lib/libnccl.so
121+
/usr/local/cuda-12.1/lib/libnccl.so
122+
/usr/local/cuda-12.2/lib/libnccl.so
123+
124+
# Print NCCL versions
125+
$ find /usr/local/cuda-* -name 'libnccl.so' | xargs -n1 -I{} bash -c "echo -n {} '=> ' ; strings {} | grep '^NCCL version' -m1"
126+
/usr/local/cuda-11.7/lib/libnccl.so => NCCL version 2.16.2+cuda11.8
127+
/usr/local/cuda-11.8/lib/libnccl.so => NCCL version 2.16.2+cuda11.8
128+
/usr/local/cuda-12.0/lib/libnccl.so => NCCL version 2.18.5+cuda12.2
129+
/usr/local/cuda-12.1/lib/libnccl.so => NCCL version 2.18.5+cuda12.2
130+
/usr/local/cuda-12.2/lib/libnccl.so => NCCL version 2.18.5+cuda12.2
131+
132+
# Print aws-ofi-nccl version
133+
$ strings /opt/aws-ofi-nccl/lib/libnccl-net.so | grep '^NET/OFI Initializing aws-ofi-nccl .*-aws'
134+
NET/OFI Initializing aws-ofi-nccl 1.7.4-aws
135+
```
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#!/usr/bin/env python
2+
3+
"""
4+
All-reduce on one GPU, to probe libraries opened. This script can be run without torchrun.
5+
6+
Usage:
7+
8+
./conda_env_pytorch/bin/python all_reduce_bench.py
9+
"""
10+
import os
11+
import torch
12+
13+
def main():
14+
os.environ.setdefault('MASTER_ADDR', 'localhost')
15+
os.environ.setdefault('MASTER_PORT', '0')
16+
os.environ.setdefault('RANK', '0')
17+
torch.cuda.set_device(0)
18+
torch.distributed.init_process_group('nccl', world_size=1)
19+
X = torch.rand(500, 2, dtype=torch.float32).cuda(0)
20+
torch.distributed.all_reduce(X)
21+
torch.cuda.synchronize()
22+
if torch.distributed.get_rank() == 0:
23+
print(f"Test completed.")
24+
25+
if __name__ == "__main__":
26+
main()
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
#!/bin/bash
2+
3+
set -uo pipefail
4+
## NOTE: you must activate the python environment, or make sure the right python is in PATH
5+
6+
# Call PyTorch collective on a single GPU, to force load libnccl.
7+
strace_output=$( strace -e trace=open,openat -e status=successful python ./all_reduce_single_gpu.py 2>&1 )
8+
retval="$?"
9+
[[ "retval" -eq 0 ]] || { echo "$strace_output" ; exit "$retval" ; }
10+
11+
set -e
12+
13+
## strace output expected to contain these lines:
14+
#openat(AT_FDCWD, "/fsx/marcverd/awsome-distributed-training/3.test_cases/10.FSDP/conda_env_pytorch/lib/python3.10/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2", O_RDONLY|O_CLOEXEC) = 3
15+
#openat(AT_FDCWD, "/opt/aws-ofi-nccl/lib/libnccl-net.so", O_RDONLY|O_CLOEXEC) = 85
16+
declare -a OPENED_LIB=( $(echo "$strace_output" | egrep 'libnccl.so|libnccl-net.so' | cut -d',' -f 2 | tr -d '"' ) )
17+
18+
set +o pipefail
19+
NCCL_VERSION=$(strings "${OPENED_LIB[0]}" | grep -m1 '^NCCL version' | cut -d' ' -f3 | cut -d'+' -f1)
20+
[[ ! "$NCCL_VERSION" == "" ]] || NCCL_VERSION="not found"
21+
echo "NCCL version :" "$NCCL_VERSION"
22+
echo "NCCL path :" "$(realpath ${OPENED_LIB[0]})"
23+
24+
AWS_OFI_NCCL_VERSION=$(strings "${OPENED_LIB[1]}" | grep -m1 '^NET/OFI Initializing aws-ofi-nccl ' | cut -d' ' -f4 | cut -d'-' -f1)
25+
[[ ! "$AWS_OFI_NCCL_VERSION" == "" ]] || AWS_OFI_NCCL_VERSION="not found"
26+
echo "aws-ofi-nccl version :" "$AWS_OFI_NCCL_VERSION"
27+
echo "aws-ofi-nccl path :" "$(realpath ${OPENED_LIB[1]})"

0 commit comments

Comments
 (0)