dask serialization for cuda pytorch tensors
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
**What happened**:
Serializer fails when data are tensors allocated in cuda memory. Curiously the deserializer is implemented to handle this situation as it recreates the tensor in cpu and then move it to gpu. https://github.com/dask/distributed/blob/bd3f47ecf2be443e1e5f19e5d3cadb3a30fd45ed/distributed/protocol/torch.py#L46
```
Traceback (most recent call last):
File "dask_bug.py", line 3, in
serialize_torch_Tensor(torch.tensor([3]).cuda())
File "/home/sanchezg/app/anaconda3/envs/pytorch/lib/python3.7/site-packages/distributed/protocol/torch.py", line 14, in serialize_torch_Tensor
header, frames = serialize(t.numpy())
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
```
The issue is similar to #2619
**What you expected to happen**:
The serializer code should move the tensor to the cpu before the serialization process. It should also store the device in which it was stored
**Minimal Complete Verifiable Example**:
```python
from distributed.protocol.torch import serialize_torch_Tensor
import torch
serialize_torch_Tensor(torch.tensor([3]).cuda())
```
**Anything else we need to know?**:
Proposed solution for https://github.com/dask/distributed/blob/bd3f47ecf2be443e1e5f19e5d3cadb3a30fd45ed/distributed/protocol/torch.py#L13:
```
def serialize_torch_Tensor(t):
device = t.device
requires_grad_ = t.requires_grad
t = t.cpu()
if requires_grad_:
sub_header, frames = serialize(t.detach().numpy())
else:
sub_header, frames = serialize(t.numpy())
header = {"sub-header": sub_header}
if t.grad is not None:
grad_header, grad_frames = serialize(t.grad.numpy())
header["grad"] = {"header": grad_header, "start": len(frames)}
frames += grad_frames
header["requires_grad"] = requires_grad_
header["device"] = device.type
return header, frames
```
Contributor guide
Assessment
This issue has not been assessed yet.