Xla: How to save/load a model trained with `torch_xla_py.data_parallel`

[ibeltagy]

A torch_xla_py.data_parallel model doesn't have an implementation for the function state_dict() which is required to save/load the model. Is there a way around this?

Thanks

[dlibenzi]

The object returned by DataParallel() is not a model, if that's what you mean.
The models (the model argument) that are passed to the loop functions are fully flagged torch.nn.Module instances with their own tensors having XLA devices.

[tweetyone]

Do you mean the model saving code of XLA is same as the original pytorch code? still use state_dict()?

[ibeltagy]

IIIC, the model object inside the train_loop_fn (https://github.com/pytorch/xla/blob/master/test/test_train_mnist.py#L96) is a regular pytorch model that can be saved with state_dict or torch.save. What I didn't quite get is, which model should be saved, knowing that we have multiple copies of the model, one on each TPU core?

[jysohn23]

Yes, I believe something like torch.save(model_parallel._models[0].state_dict(), PATH) should work just fine. If you check the data_parallel.py implementation here, it's just a list of standard PyTorch torch.nn.Modules.

[jysohn23]

Looks like our comments crossed :)

But per your question, each of the cores should have the exact same model as at the end of each step we sync all the gradients and take a global step on each of the cores with the exact same gradients. Thus, you should be able to chose any of them. We're doing synchronous gradient descent here so all models should be pretty much equal.

[dlibenzi]

We had some issues with saving models which reside on XLA devices.
I am not sure whether that was fixed.
@ailzhang Do you remember that?

[asuhan]

We carry a patch for torch.save and checkpoints worked in the past (for BERT at least), not sure there's anything to do.

[ibeltagy]

torch.save(model_parallel._models[0].state_dict(), PATH) works. Thanks @jysohn23.