dask / dask/distributed

dask serialization for cuda pytorch tensors

Open
#6,111 3 comments 2 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.