Repository: ray-project/ray
Issue category: Feature Request

Issue:
[core][compiled graphs] Support all torch.dtypes for tensors sent through shared memory channels
### What happened + What you expected to happen

Got the following error when experimenting with OpenRLHF:

**TypeError: Got unsupported ScalarType BFloat16** :

```
(CriticModelRayActor pid=33730)   cpu_tensor = torch.from_numpy(np_array) [repeated 2x across cluster]
(RewardModelRayActor pid=33731) Traceback (most recent call last): [repeated 4x across cluster]
(RewardModelRayActor pid=33731) 2024-10-21 09:10:49 ERROR    Compiled DAG task exited with exception
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/shared_memory_channel.py", line 469, in write [repeated 5x across cluster]
(RewardModelRayActor pid=33731)     serialized_value = self._worker.get_serialization_context().serialize(
(RewardModelRayActor pid=33731)                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/torch_tensor_type.py", line 108, in serialize [repeated 2x across cluster]
(RewardModelRayActor pid=33731)     return self._serialize_to_msgpack(value)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 497, in _serialize_to_msgpack
(RewardModelRayActor pid=33731)     pickle5_serialized_object = self._serialize_to_pickle5(
(RewardModelRayActor pid=33731)                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 439, in _serialize_to_pickle5 [repeated 2x across cluster]
(RewardModelRayActor pid=33731)     raise e
(RewardModelRayActor pid=33731)     inband = pickle.dumps(
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^ [repeated 2x across cluster]
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/cloudpickle/cloudpickle.py", line 1479, in dumps
(RewardModelRayActor pid=33731)     cp.dump(obj)
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/cloudpickle/cloudpickle.py", line 1245, in dump
(RewardModelRayActor pid=33731)     return super().dump(obj)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 175, in _CloudPicklerReducer
(RewardModelRayActor pid=33731)     return custom_deserializer, (custom_serializer(obj),)
(RewardModelRayActor pid=33731)                                  ^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)     return ctx.serialization_context.serialize_tensor(t)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/serialization_context.py", line 77, in serialize_tensor
(RewardModelRayActor pid=33731)     return self.serialize_to_numpy(tensor)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/serialization_context.py", line 88, in serialize_to_numpy
(RewardModelRayActor pid=33731)     return tensor.numpy()
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731) TypeError: Got unsupported ScalarType BFloat16
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 118, in do_exec_tasks
(RewardModelRayActor pid=33731)     done = tasks[operation.exec_task_idx].exec_operation(
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 547, in exec_operation
(RewardModelRayActor pid=33731)     return self._write()
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 517, in _write
(RewardModelRayActor pid=33731)     self.output_writer.write(output_val)
(RewardModelRayActor pid=33731)     channel.write(val_i, timeout)
(RewardModelRayActor pid=33731)     channel.write(value, timeout)
(RewardModelRayActor pid=33731)     self._buffers[self._next_write_index].write(value, timeout)
```

### Versions / Dependencies

head

### Reproduction script

will add later

### Issue Severity

None


Hints:
Another data point is that Pytorch has a number of other data types that aren't natively supported in numpy (like float8, etc). I'm not particularly sure of the details for which `tensor` was converted to a numpy array here but it would be best to not assume any interoperability with numpy and keep tensors as tensors during serialization. 
Hmm really what we need is to be able to zero-copy deserialize torch.Tensors. Using numpy was easiest at the time but this is indeed an issue. Maybe [arrow](https://arrow.apache.org/docs/cpp/api/tensor.html#tensors) is an option?
Hmm actually seems like this could be supported by using torch.Tensor.view with a uint8 dtype:

```python
In [99]: t
Out[99]:
tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
         0.8750, -1.7500], dtype=torch.float8_e4m3fn)

In [104]: torch.as_tensor(t.view(torch.uint8).numpy()).view(torch.float8_e4m3fn)
Out[104]:
tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
         0.8750, -1.7500], dtype=torch.float8_e4m3fn)

```
Hi @stephanie-wang, 

I am interested in taking up this issue.
> Hi @stephanie-wang,
> 
> I am interested in taking up this issue.

Thanks, feel free to open a PR. Will assign it to you for now and revisit in a week or so.
Thanks, do you have suggestions on reproducing the issue on my local?
You want a DAG that looks something like:
```python
t = A.return_tensor(...).with_type_hint(TorchTensorType())
dag = B.read_tensor(t)
```
And A should return a tensor with a dtype unsupported by numpy.
I am new to the eco-system. I'm sorry if this is an obvious question. What do `A` and `B` represent in your previous comment? I tried creating an example using the "Getting Started" docs. However, I am struggling to reproduce the issue.

```python
import ray
from ray import workflow
from ray.experimental.channel.torch_tensor_type import TorchTensorType
import torch

@ray.remote
def return_tensor():
    return torch.Tensor([1.0, 2.0])

@ray.remote
def read_tensor(a):
    return a

dag = read_tensor.bind(return_tensor.bind().with_type_hint(TorchTensorType()))

print(workflow.run(dag))
```

Am I on the right track? Also, if I try to specify the `dtype` param in the `torch.Tensor` init, the program fails with a weird torch data type error. I am not sure if that's related.
Hi @sahilgupta2105, yes we don't have great docs for compiled graphs yet. Please work through [the developer guide](https://docs.google.com/document/d/1xDs-yIBFLXuqUDLou-_KELERkzWrmGWri6p_xG18SH4/edit?tab=t.0#heading=h.kzrrbhon06yz) first and comment here if you still have questions.
Thanks, the developer guide is helpful. I was able to reproduce the issue.

```python
import ray
import ray.dag
from ray.experimental.channel.torch_tensor_type import TorchTensorType
import torch

@ray.remote
class Actor:
  def process(self, tensor: torch.Tensor):
    return tensor.shape

actor = Actor.remote()

with ray.dag.InputNode() as inp:
  inp = inp.with_type_hint(TorchTensorType())
  dag = actor.process.bind(inp)

  dag = dag.experimental_compile()
  print(ray.get(dag.execute(torch.zeros(10, dtype=torch.float8_e4m3fn))))
```
> Hmm actually seems like this could be supported by using torch.Tensor.view with a uint8 dtype:
> 
> ```python
> In [99]: t
> Out[99]:
> tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
>          0.8750, -1.7500], dtype=torch.float8_e4m3fn)
> 
> In [104]: torch.as_tensor(t.view(torch.uint8).numpy()).view(torch.float8_e4m3fn)
> Out[104]:
> tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
>          0.8750, -1.7500], dtype=torch.float8_e4m3fn)
> ```

@stephanie-wang to be sure, you meant type-casting the torch tensor to the unsigned int type with the appropriate precision, right? E.g. float8 -> uint8, float16 -> uint16, ...
No, cast everything to uint8 (because it's a supported np dtype and the lowest possible precision) and then cast back on the receiving side.
Gotcha. For my understanding, wouldn't casting everything to uint8 cause a loss of information if the input tensor is of higher precision? 
@stephanie-wang It seems like the casting operations only update the metadata, which ensures that information is not lost when `floatXX` is converted to `uint8`. However, to ensure data integrity, we need information about the `dtype` of the original array. Numpy arrays support [custom metadata](https://numpy.org/doc/stable/reference/generated/numpy.dtype.metadata.html). Shall we use that to store the `dtype` on serialization for deserialization to ensure the correct format?
It would be better to not rely on a third-party API for passing the dtype. You can pass it through Ray instead. Check out how [TorchTensorType](https://github.com/ray-project/ray/blob/master/python/ray/experimental/channel/torch_tensor_type.py) is used.
I dug more into the code. Can you confirm my understanding before I send out a PR?

- `TorchTensorType` uses a serializer from the serialization context that internally calls the `.numpy()`.
- Each `TorchTensorType` hint is unique per DAG node and contains the type in `_dtype` attribute.
- We can declare a custom serializer (skip calling the serializer from the serialization context) and use the `_dtype` attribute to maintain data integrity.

If you agree so far, do I need special handling for the "AUTO" dtype?
I don't think you need to touch TorchTensorType at all. Should be enough to just modify the existing custom serialization function.

Patch:
diff --git a/python/ray/experimental/channel/serialization_context.py b/python/ray/experimental/channel/serialization_context.py
index 613f923b7a6c..4aa5f280b51d 100644
--- a/python/ray/experimental/channel/serialization_context.py
+++ b/python/ray/experimental/channel/serialization_context.py
@@ -97,7 +97,11 @@ def serialize_tensor(self, tensor: "torch.Tensor") -> Union[int, "np.ndarray"]:
 
         return self.serialize_to_numpy(tensor)
 
-    def serialize_to_numpy(self, tensor: "torch.Tensor") -> "np.ndarray":
+    def serialize_to_numpy(
+        self, tensor: "torch.Tensor"
+    ) -> Tuple["np.ndarray", "torch.dtype"]:
+        import torch
+
         # Transfer through Ray's shared memory store for now.
         # TODO(swang): This requires two copies, one to transfer from GPU to
         # CPU and another from CPU to shared memory. Ideally we should elide
@@ -106,7 +110,10 @@ def serialize_to_numpy(self, tensor: "torch.Tensor") -> "np.ndarray":
         if tensor.device.type == "cuda":
             tensor = tensor.to("cpu")
 
-        return tensor.numpy()
+        # Numpy does not have an equivalent dtype for all torch dtypes, so
+        # instead of casting directly to numpy, we first use a view with a
+        # common dtype and then view as numpy array.
+        return (tensor.view(torch.uint8).numpy(), tensor.dtype)
 
     def deserialize_tensor(self, val: Union["np.ndarray", int]):
         # Found a placeholder for a tensor that was serialized via NCCL.
@@ -119,13 +126,17 @@ def deserialize_tensor(self, val: Union["np.ndarray", int]):
 
         return self.deserialize_from_numpy(val)
 
-    def deserialize_from_numpy(self, np_array: "np.ndarray"):
+    def deserialize_from_numpy(
+        self, np_array_dtype: Tuple["np.ndarray", "torch.dtype"]
+    ):
         import torch
 
         from ray.experimental.channel import ChannelContext
 
         ctx = ChannelContext.get_current()
 
+        np_array, dtype = np_array_dtype
+
         # TODO(swang): Support local P2P transfers if available.
         # If there is a GPU assigned to this worker, move it there.
         if ctx.torch_device is not None and ctx.torch_device.type == "cuda":
@@ -133,7 +144,7 @@ def deserialize_from_numpy(self, np_array: "np.ndarray"):
             def convert_numpy_to_tensor(np_array, ctx):
                 # It does zero-copy convert np_array inside shared memroy to
                 # a tensor. Since we move data to GPU immediately, it is safe.
-                cpu_tensor = torch.from_numpy(np_array)
+                cpu_tensor = torch.from_numpy(np_array).view(dtype)
                 return cpu_tensor.to(device=ctx.torch_device)
 
             global _TORCH_WARNING_FILTER_ACTIVATE
@@ -160,4 +171,4 @@ def convert_numpy_to_tensor(np_array, ctx):
         # TODO(swang): Use zero-copy from_numpy() if np_array.flags.writeable
         # is True. This is safe to set when deserializing np_array if the
         # upstream task has num_readers=1.
-        return torch.tensor(np_array, device=ctx.torch_device)
+        return torch.tensor(np_array, device=ctx.torch_device).view(dtype)


Test patch:
diff --git a/python/ray/tests/test_channel.py b/python/ray/tests/test_channel.py
index 69127ee157a8..95fc659c5a32 100644
--- a/python/ray/tests/test_channel.py
+++ b/python/ray/tests/test_channel.py
@@ -8,11 +8,13 @@
 
 import numpy as np
 import pytest
+import torch
 
 import ray
 import ray.cluster_utils
 import ray.exceptions
 import ray.experimental.channel as ray_channel
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
 from ray.exceptions import RayChannelError, RayChannelTimeoutError
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from ray.dag.compiled_dag_node import CompiledDAG
@@ -1407,6 +1409,20 @@ def write(self, i, timeout=None) -> bool:
         )
 
 
+def test_torch_dtype():
+    typ = TorchTensorType()
+    typ.register_custom_serializer()
+
+    t = torch.randn(5, 5, dtype=torch.bfloat16)
+    with pytest.raises(TypeError):
+        t.numpy()
+
+    ref = ray.put(t)
+    t_out = ray.get(ref)
+    assert (t == t_out).all()
+    assert t_out.dtype == t.dtype
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))