Repository: pytorch
Issue: Saving & Loading Model Across [In-tree and out-of-tree] Devices
Labels: module: serialization,triaged,module: backend,module: PrivateUse1

### 🚀 The feature, motivation and pitch

# Problem

1. User trains a model with a out-of-tree backend. The backend is registered through [privateuseone interface](https://pytorch.org/tutorials/advanced/privateuseone.html).
2. The training code involves saving a torch.Device object.
    ```
        device = next(model.parameters()).device
        self.device = device
    ```
3. Later user tries to load the model on CPU or GPU platform which doesn't have the device plugin installed. A error is raised:
    ```
    Traceback (most recent call last):
      File "/torch/code/ultralytics/test.py", line 10, in <module>
        model = YOLO("./best-torch2.1.pt").to(device)
      File "/torch/code/ultralytics/ultralytics/models/yolo/model.py", line 23, in __init__
        super().__init__(model=model, task=task, verbose=verbose)
      File "/torch/code/ultralytics/ultralytics/engine/model.py", line 142, in __init__
        self._load(model, task=task)
      File "/torch/code/ultralytics/ultralytics/engine/model.py", line 229, in _load
        self.model, self.ckpt = attempt_load_one_weight(weights)
      File "/torch/code/ultralytics/ultralytics/nn/tasks.py", line 788, in attempt_load_one_weight
        ckpt, weight = torch_safe_load(weight)  # load ckpt
      File "/torch/code/ultralytics/ultralytics/nn/tasks.py", line 714, in torch_safe_load
        ckpt = torch.load(file, map_location="cpu")
      File "/torch/code/camb-pytorch/catch/venv/pytorch/lib/python3.10/site-packages/torch/serialization.py", line 1025, in load
        return _load(opened_zipfile,
      File "/torch/code/camb-pytorch/catch/venv/pytorch/lib/python3.10/site-packages/torch/serialization.py", line 1446, in _load
        result = unpickler.load()
    RuntimeError: Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: mlu
    ```


Below is a minimum case to illustrate the problem. It's derived from the [YoloV8 project.](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/loss.py#L152)

The training process runs smoothly but a error is raised when `torch.load` is called.

The problem is straightforward.  However you'll need a device plugin environment to reproduce it.

```python
import torch
import torch.nn as nn
import torch.optim as optim
import os
 
class SimpleLoss(nn.MSELoss):
    def __init__(self, model):
        super(SimpleLoss, self).__init__()
        # Problematic code
        device = next(model.parameters()).device
        self.device = device
 
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(1, 1)
        self.criterion = None
 
    def forward(self, x):
        x = self.fc(x)
        return x
 
    def init_criterion(self, criterion):
        self.criterion = criterion
 
 
_SAVE_MODEL = os.environ.get('SAVE_MODEL', 0)
if int(_SAVE_MODEL) == 1:
    import torch_mlu
    model = SimpleModel().to('mlu:0')
    criterion = SimpleLoss(model)
    model.init_criterion(criterion)
    optimizer = optim.SGD(model.parameters(), lr=0.01)
     
    x = torch.tensor([[1.0], [2.0], [3.0], [4.0]], requires_grad=True).to('mlu:0')
    y = torch.tensor([[2.0], [4.0], [6.0], [8.0]]).to('mlu:0')
     
    for epoch in range(5):
        outputs = model(x)
        loss = criterion(outputs, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
     
    model = model.cpu()
 
    torch.save(model, './model.pt')
    print('save model is success')
else:
    print('try load model')
    model = torch.load('./model.pt').to("cpu")
```

# Analysis

`torch.save` serializes the model objects which contain `torch.Device` which type is the plugined device type.

`torch.load` deserializes the model objects and tries to instantiate `torch.Device` with the plugined device type.

If the device plugin is installed,  the parse_type check runs happily and returns `DeviceType::PrivateUse1;` on [line 49](https://github.com/pytorch/pytorch/blob/main/c10/core/Device.cpp#L48). Otherwise a error is reported on [line 57](https://github.com/pytorch/pytorch/blob/main/c10/core/Device.cpp#L57).

It's not practical to ask the user to install the training device plugin when doing inference on a total irrevelant hardware platform. 



### Alternatives

# Workarounds

## Workaround 1

In the training code, tell users "hey if you want your model to work across in tree and out-of-tree backends then please do not save a device object in your model. Change `self.device = device` to `self.device = device.type` instead".

This is kinda weird.

## Workaround 2

In the inference code, before `torch.load` call `torch.utils.rename_privateuse1_backend("mlu") ` first.
This can fool around the privateuse1 check.

This is also a little bit weird from user's perspective since he/she has no idea about the model training context.

### Additional context

Not sure how to fix this on technical side. Or maybe a new chapter about working across in tree and out tree devices?
[saving-loading-model-across-devices](https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-across-devices)

cc @mruberry @mikaylagawarecki @bdhirsh