Category: tpu
Question: How to use TPUs in PyTorch?
Answer: Once you have flipped the "Accelerator" switch in your notebook to "TPU v3-8", this is how to enable TPU training in PyTorch:


# Step 1: Install Torch-XLA (PyTorch with Accelerated Linear Algebra (XLA) support)
```shell
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev
```

# Step 2: Run your PyTorch code
TPUs (TPU v3-8) have 8 cores, and each core is itself an XLA device.You can run code on a single XLA device, but to take full advantage of the TPU you will want to run your code on all 8 cores simultaneously. For examples that demonstrate how to do this, you can refer to 
- [The Ultimate PyTorch TPU Tutorial](https://www.kaggle.com/tanlikesmath/the-ultimate-pytorch-tpu-tutorial-jigsaw-xlm-r),
- [I Like Clean TPU Training Kernels and I Can Not Lie](https://www.kaggle.com/abhishek/i-like-clean-tpu-training-kernels-i-can-not-lie),
- [Super Duper Fast PyTorch TPU Kernel](https://www.kaggle.com/abhishek/super-duper-fast-pytorch-tpu-kernel),
- [XLM Roberta Large Pytorch TPU](https://www.kaggle.com/philippsinger/xlm-roberta-large-pytorch-pytorch-tpu?scriptVersionId=38462589).

You should also note the following when using TPUs with PyTorch:

#1: Startup Script 
[https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py](https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py)

#2: Distributed training function mp_fn
```python
xmp.spawn(_mp_fn, nprocs=8, start_method='fork')
```

#3: Instantiate model outside of mp_fn and use MpModelWrapper
```python
MX = JigsawModel()   =>    MX = xmp.MpModelWrapper(JigsawModel())
```

#4: Send model to TPU device
```python
device = xm.xla_device()
model = MX.to(device)
```

#5: Changes to training loop: send data to device
```python
ids = ids.to(device, dtype=torch.long)
token_type_ids = token_type_ids.to(device, dtype=torch.long)
mask = mask.to(device, dtype=torch.long)
targets = targets.to(device, dtype=torch.float)
```

#6: Printing messages
```python
xm.master_print
```

#7: Loading data
```python
train_dataset = … # user-defined, can be outside of mp_fn
# in mp_fn:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
            num_replicas=xm.xrt_world_size(),rank=xm.get_ordinal(), …)
train_data_loader = torch.utils.data.DataLoader(train_dataset,
                            sampler=train_sampler, …)
```

#8: Training on data
```python
for epoch in range(EPOCHS):
    para_loader = pl.ParallelLoader(train_data_loader, [device])
    train_fn(para_loader.per_device_loader(device), …)
```

#9: Results from TPU
```python
xm.mesh_reduce
```

#10: Model save / restore (memory-optimized)
```python
import torch_xla.utils.serialization as xser
xser.save(model.state_dict(), f"model.bin", master_only=True)
model.load_state_dict(xser.load(f"model.bin"))
```

#11: Model save / restore (PyTorch standard)
```python
torch_xla.core.xla_model.save
torch.load(...)
```

#12: Out of memory datasets:
```python
Can be loaded from localhost
Of loaded from GCS in TFRecord format, a TFRecords PyTorch loader exists
```