Lightning-AI / Lightning-AI/pytorch-lightning

TorchScript export incompatible with DeviceDtypeModuleMixin properties

Open
#13,887 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug lightningmodule priority: 2
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🐛 Bug

My model code makes extensive use of tensors created during the forward pass. As part of this, I have to specify the device associated with each tensor to avoid conflicts. I want to export my model in TorchScript for deployment elsewhere, but the DeviceDtypeModuleMixin properties do not appear to be compatible with TorchScript.

While this appears to affect the `dtype` property too, for my use case, using `self.device` in a LightningModule prevents successful TorchScript export.

### To Reproduce

```python
import pytorch_lightning as pl
import torch

class TestModel(pl.LightningModule):
def forward(self):
return torch.zeros((2, 2), device=self.device)

model = TestModel()
model.to_torchscript()
```

results in the following error:

```
RuntimeError:
Module 'TestModel' has no attribute 'device' :
File "/test.py", line 7
def forward(self):
return torch.zeros((2, 2), device=self.device)
~~~~~~~~~~~ <--- HERE
```

I am not sure how this can be fixed without changing the way DeviceDtypeModuleMixin works. It relies on the `@property` decorator, which uses descriptors. [TorchScript is not compatible with descriptors (3.3.2.2, 3.3.2.3)](https://pytorch.org/docs/stable/jit_python_reference.html#python-language-reference).

Additionally, the `device` property returns a `Union[str, torch.device]`, so even if it worked as a property in TorchScript, the return value of the function is incompatible with the type expected by the `device` keyword option (`Optional[Device]`). There appears to have been some discussion of this in [this issue](https://github.com/Lightning-AI/lightning/issues/1951), but I'm not sure how it was resolved.

I am happy to try contributing a solution, but I wanted to discuss here first. A possible solution seems to involve eliminating the `@property` decorator and save the current state of the device as a plain `device` attribute on the model object. The actions currently in the `device()` function body should be performed as the `device` attribute is being set rather than on the way out.

### Expected behavior

The code shown above is successfully saved as TorchScript.

### Environment

- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
- PyTorch Lightning Version (e.g., 1.5.0): 1.6.5
- PyTorch Version (e.g., 1.10): 1.12.0
- Python version (e.g., 3.9): 3.9.13
- OS (e.g., Linux): Linux
- CUDA/cuDNN version: 11.6
- How you installed PyTorch (`conda`, `pip`, source): pip

cc @carmocca @justusschock @awaelchli @borda @ananthsub @ninginthecloud @jjenniferdai @rohitgr7

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with LightningModule.to_torchscript and the DeviceDtypeModuleMixin device and dtype properties, then run the provided TestModel reproduction with the stated versions. The work is done when a model using self.device can be successfully exported as TorchScript, with coverage for the reported behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.