Recommended PyTorch version for JetPack 5.1

virtual.ramblings · November 21, 2023, 5:46am

Hi, I am trying to compare performance across bare-metal and docker on my AGX Orin devkit (JetPack version: 5.1).

For bare-metal, the official PyTorch 2.0 wheel seems to be torch-2.0.0+nv23.05-cp38 from here

For L4T-ML container, the PyTorch 2.0 version seems to be torch-2.0.0a0+ec3941ad.nv23.2-cp38 from here

How different are these 2 PyTorch versions? Is there any reason to choose different versions on bare-metal v/s container?

AastaLLL · November 21, 2023, 8:02am

Hi,

These two are all based on the PyTorch 2.0.0 but in different snapshots.
Since we have ~ monthly PyTorch releases, they just pick up different packages.

Thanks.

virtual.ramblings · November 21, 2023, 8:06am

Thanks for your response. A couple of follow-up questions:

I’m seeing major performance differences across these two versions (2-2.5x slowdown). Is there a way I could compare the code changes that went into both of these?

torch-2.0.0+nv23.05-cp38 is newer than torch-2.0.0a0+ec3941ad.nv23.2-cp38 right?

dusty_nv · November 21, 2023, 2:15pm

Hi @virtual.ramblings, yes I believe that torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl is newer. Which is the version that is faster for you?

Also, the containers have been updated to use these updated wheels (and PyTorch 2.1) - you can rebuild them or pull the updated images from DockerHub:

github.com

dusty-nv/jetson-containers/blob/086fae178ea6c900fdf0ca56e5373c8c15ac8747/packages/pytorch/config.py#L64


      
                  'TORCH_CUDA_ARCH_LIST': cuda_arch,
              }
          
              pkg['dockerfile'] = dockerfile
          
              pkg['depends'] = requires
              
              return pkg
              
          package = [
              # JetPack 5
              pytorch('2.1', 'torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl', 'https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl', '>=34.1.0'),
              pytorch('2.0', 'torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl', 'https://nvidia.box.com/shared/static/i8pukc49h3lhak4kkn67tg9j4goqm0m7.whl', '>=34.1.0', default=True),
              pytorch('1.13', 'torch-1.13.0a0+340c4120.nv22.06-cp38-cp38-linux_aarch64.whl', 'https://developer.download.nvidia.com/compute/redist/jp/v50/pytorch/torch-1.13.0a0+340c4120.nv22.06-cp38-cp38-linux_aarch64.whl', '>=34.1.0'),
              pytorch('1.12', 'torch-1.12.0a0+8a1a93a9.nv22.5-cp38-cp38-linux_aarch64.whl', 'https://developer.download.nvidia.com/compute/redist/jp/v50/pytorch/torch-1.12.0a0+8a1a93a9.nv22.5-cp38-cp38-linux_aarch64.whl', '>=34.1.0'),
              pytorch('1.11', 'torch-1.11.0-cp38-cp38-linux_aarch64.whl', 'https://nvidia.box.com/shared/static/ssf2v7pf5i245fk4i0q926hy4imzs2ph.whl', '>=34.1.0'),
              
              # JetPack 4
              #pytorch('1.11', 'torch-1.11.0a0+17540c5-cp36-cp36m-linux_aarch64.whl', 'https://developer.download.nvidia.com/compute/redist/jp/v461/pytorch/torch-1.11.0a0+17540c5+nv22.01-cp36-cp36m-linux_aarch64.whl', '==32.*'),  # (built without LAPACK support)
              pytorch('1.10', 'torch-1.10.0-cp36-cp36m-linux_aarch64.whl', 'https://nvidia.box.com/shared/static/fjtbno0vpo676a25cgvuqc1wty0fkkg6.whl', '==32.*', default=True),
              pytorch('1.9', 'torch-1.9.0-cp36-cp36m-linux_aarch64.whl', 'https://nvidia.box.com/shared/static/h1z9sw4bb1ybi0rm3tu8qdj8hs05ljbm.whl', '==32.*'),

virtual.ramblings · November 22, 2023, 4:35am

Thanks for your reply!

23.02 is much faster than 23.05 on the workload we are looking at. (BERT for question answering)

Do you recommend moving to 23.05? Is this a more stable version even if slower?

AastaLLL · November 22, 2023, 6:13am

Hi,

As above mentioned, we have a container with the newer PyTorch version (PyTorch 2.1 nv23.06).
It’s recommended to try our latest release.

Thanks.