GSPMD + PyTorch Compile + TPU crash

[agemagician]

Hi,

I am trying to combine both GSPMD + PyTorch Compile, but it doesn't work.
I took a copy of the test script "test_train_spmd_imagenet.py" and test it in colab, and it started normally. However, after I added the compile line :

device = xm.xla_device()
  model = get_model_property('model_fn')().to(device)

  model = torch.compile(
        model, backend='aot_torchxla_trace_once')

It crashed.

Here is a Colab example to reproduce the results:
https://colab.research.google.com/drive/1KNcBydAfZXLATpSo-CXILxtHJkK8JD-2?usp=sharing

[yeounoh]

Hi @agemagician torch.Compile doesn't support GSPMD out of the box. Looking at the title, does it run on non-TPU device? Also could you share some more context, like your use case?

[agemagician]

Hi @agemagician torch.Compile doesn't support GSPMD out of the box. Looking at the title, does it run on non-TPU device? Also could you share some more context, like your use case?

1. I didn't test it with multi-GPU. I only have access to a TPU pod.
2. We train a large language model based on transformer architecture. We currently use Jax but want to switch to PyTorch for easier debugging and testing. Our models don't fit on a single TPU/GPU core, so we have to use either fsdp or GSPMD, but neither of them works with torch.compile. We must use torch compile to get close to Jax training speed on the TPU.
3. We have opened another issue for fsdp here : Pytorch 2 compile + fsdp + transformers crash #4823

[yeounoh]

Thank you @agemagician, we plan to support GSPMD with Dynamo, and first for the inference workloads and then for training. I am looking at sometime in H1. Do you have a contact from Google whom you've been working with, asking since you already have access to a TPU pod. Maybe we can chat offline together.

cc @miladm @shauheen @alanwaketan @JackCaoG

[agemagician]

Yes, I am working with several Googlers + I am in a good contact with HuggingFace :)

Sure, let me know what is the best way to chat offline.