Performance of Torchax

[yuanfz98]

❓ Questions and Help

Hello Community,

Will using torchAx be slower than native PyTorch? Is there any transformation layer of tensors which makes it slower ?

[yuanfz98]

@qihqi Hello, Does TorchAx support jax.distributed.initialize() (Distributed context and multi-host multi-devices) ?

[qihqi]

Hi @yuanfz98

If you run torchax eagerly, you will get eager Jax performance (which in general is not so great: jax-ml/jax#13960). If you pair you computation with jax.jit (via torchax.interop.jax_jit) you should get jax.jitt'd performance. That should be on par with Pytorch on GPU.

[qihqi]

@qihqi Hello, Does TorchAx support jax.distributed.initialize() (Distributed context and multi-host multi-devices) ?

Yes theoretically, I haven't tried it personally.

We tried to use it on Google TPUs on distributed setting (however, TPU doesn't need jax.distributed.initialize()). This blog post describes how to do tensor parallel using HuggingFace's model: https://huggingface.co/blog/qihqi/huggingface-jax-02