Repository: ray-project/ray
Issue category: Bug Report

[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


Functions to edit:
python/ray/experimental/channel/serialization_context.py:_SerializationContext.serialize_to_numpy
python/ray/experimental/channel/serialization_context.py:_SerializationContext.deserialize_from_numpy

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__]))