Torch.distributed.pipelining can't be using

[s1wnd0702]

I’m a student, and I would like to refer to the link “PyTorch Distributed Overview — PyTorch Tutorials 2.5.0+cu124 documentation” to partition the model and then do the inference task. But I found that the code needs a command like “from torch.distributed.pipelining import pipeline, SplitPoint, PipelineStage, ScheduleGPipe”. But after I entered this command, it showed that there was no such package.
My jetson device is jetson AGX Orin, Jetpack version is 5.1.1, torch version is 1.11.0, and I installed it using the wheel file provided by nvidia.
Is torch.distributed.pipelining not working on jetson devices? Is there any alternative?
Any helpful suggestions would be greatly appreciated!

[AastaLLL]

Hi,

Please upgrade your device to JetPack 6.1 and install our latest PyTorch package below:
http://jetson.webredirect.org/jp6/cu126

$ python3
Python 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from torch.distributed.pipelining import pipeline, SplitPoint, PipelineStage, ScheduleGPipe
>>> 

Thanks.

[s1wnd0702]

Thank you for your suggestion! By the way, are there any software instructions that can be upgraded JetPack directly instead of upgrading it by hardware?

[AastaLLL]

Hi,

Do you want to upgrade JetPack without using a newer BSP?
If yes, please find the document below:

https://docs.nvidia.com/jetson/archives/jetpack-archived/jetpack-61/install-setup/index.html#upgradable-compute-stack

Thanks.