Torch.distributed is not available on Jetson PyTorch

[zergzzlun]

Hi there, I’m working on Jetson Orin NX and PyTorch, I followed instructions here (Installing PyTorch for Jetson Platform - NVIDIA Docs) to install PyTorch and found that “torch.distributed.is_available()” is “False”. If I directly install PyTorch using conda (conda install pytorch -c pytorch -c nvidia) I will have “torch.distributed.is_available()” to be “True” but “torch.cuda.is_available()” to be “False”.

Unfortunately I have to use both CUDA and distributed module. I’m trying to deploy a mmpose (GitHub - open-mmlab/mmpose: OpenMMLab Pose Estimation Toolbox and Benchmark.) module and its dependencies (mmengine, mmdeploy, mmcv) are largely depend on both CUDA and torch.distributed.

How can I fix this problem? Many thanks.

[AastaLLL]

Hi,

Could you try the container mentioned below:

Which nvcr.io/nvidia/l4t-pytorch images have torch.distributed enabled?

I’m running Jetpack 5.0.2 (l4t R35.1.0), and need a docker image with torch distributed enabled. I’ve found the repo, GitHub - dusty-nv/jetson-containers: Machine Learning Containers for NVIDIA Jetson and JetPack-L4T, but is there a way to tell if a given image is compiled with torch.distributed enabled?

Thanks.

[dusty_nv]

@zergzzlun if you prefer installing the wheel as opposed to using the container, I believe this was the last PyTorch wheel built with USE_DISTRIBUTED enabled:

PyTorch v1.11.0

• JetPack 5.0 (L4T R34.1) / JetPack 5.0.2 (L4T R35.1) / JetPack 5.1 (L4T R35.2.1) / JetPack 5.1.1 (L4T R35.3.1)
  • Python 3.8 - torch-1.11.0-cp38-cp38-linux_aarch64.whl

For newer versions, you would need to build PyTorch from source to enable distributed mode (see this post for build instructions: PyTorch for Jetson)

Also, the reason that you get torch.cuda.is_available = False when you install it with conda, is because those upstream wheels from PyPi/conda/pip were not built for JetPack or with CUDA enabled.