Skip to content

[BugFix] Support bf16 in zero-copy tensor serialization #16860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[BugFix] Support bf16 in zero-copy tensor serialization #16860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions tests/v1/test_serial_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ def test_encode_decode():
list_of_tensors=[
torch.rand((1, 10), dtype=torch.float32),
torch.rand((3, 5, 4000), dtype=torch.float64),
# Make sure to test bf16 which numpy doesn't support.
torch.rand((3, 5, 1000), dtype=torch.bfloat16),
torch.tensor([float("-inf"), float("inf")] * 1024,
dtype=torch.bfloat16),
torch.tensor(1984), # test scalar too
],
numpy_array=np.arange(512),
Expand All @@ -64,7 +68,7 @@ def test_encode_decode():
# There should be the main buffer + 4 large tensor buffers
# + 1 large numpy array. "large" is <= 512 bytes.
# The two small tensors are encoded inline.
assert len(encoded) == 6
assert len(encoded) == 8

decoded: MyType = decoder.decode(encoded)

Expand All @@ -76,7 +80,7 @@ def test_encode_decode():

encoded2 = encoder.encode_into(obj, preallocated)

assert len(encoded2) == 6
assert len(encoded2) == 8
assert encoded2[0] is preallocated

decoded2: MyType = decoder.decode(encoded2)
Expand Down
31 changes: 27 additions & 4 deletions vllm/v1/serial_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,11 @@

bytestr = Union[bytes, bytearray, memoryview, zmq.Frame]

NP_FP16_STR = torch.tensor(0, dtype=torch.float16,
device="cpu").numpy().dtype.str
# Special dtype string for bf16 which numpy doesn't support.
ENC_BF16_STR = "!bf16"


class MsgpackEncoder:
"""Encoder with custom torch tensor and numpy array serialization.
Expand Down Expand Up @@ -80,7 +85,7 @@ def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:

def enc_hook(self, obj: Any) -> Any:
if isinstance(obj, torch.Tensor):
return self._encode_ndarray(obj.numpy())
return self._encode_tensor(obj)

# Fall back to pickle for object or void kind ndarrays.
if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
Expand Down Expand Up @@ -113,6 +118,15 @@ def enc_hook(self, obj: Any) -> Any:
return msgpack.Ext(CUSTOM_TYPE_PICKLE,
pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL))

def _encode_tensor(
self, obj: torch.Tensor
) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
if obj.dtype != torch.bfloat16:
return self._encode_ndarray(obj.numpy())
# Numpy doesn't support as bf16 so send as fp16 view.
_, shape, data = self._encode_ndarray(obj.view(torch.float16).numpy())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be dumb question, but is this conversion lossless? I thought we need to upcast it to fp32 to preserve precision since bf16 and fp16 have different exponents

Copy link
Contributor

@p88h p88h Apr 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

torch.view doesn't alter the data, it's more like a type cast in C.

@njhill -- What I thought of doing here is also this, but perhaps simplified / generalized more ->

  1. check if tensor is_contiguous
  2. store shape & dtype on the tensor
  3. create a view of the tensor as one-dimensional nbytes 8-bit int array
  4. encode that using numpy (since torch doesn't expose it's buffers directly)

And do that in reverse on the decode path. Torch has a bunch of weird 'bits' types that Numpy doesn't support, so probably doing this encoding for every single type doesn't make much sense. WDYT ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ywang96 yes like @p88h said this isn't actually a conversion, it's just temporarily pretending that the raw bf16 data is raw fp16 data.

@p88h I like that idea for a general case but I'm not sure there's any other unsupported types that we care about here tbh, and it means some additional conversions for the common/supported types.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@p88h shall we get this merged to address the immediate issue and then you could open a follow-on PR for the above? Then we could microbenchmark it for the supported type cases...

Copy link
Contributor

@p88h p88h Apr 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the more generic variant: #16866

There isn't any more conversions - view() is basically free, the only real difference is that this will now send a longer dtype string since it will use torch naming.

As for what types are there it is quite a bunch - https://pytorch.org/docs/stable/tensors.html
All of the quantized types, custom 8bit floats are not supported by numpy.

return ENC_BF16_STR, shape, data

def _encode_ndarray(
self, obj: np.ndarray
) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
Expand All @@ -135,7 +149,7 @@ def _encode_ndarray(

def _encode_nested_tensors(self, nt: NestedTensors) -> Any:
if isinstance(nt, torch.Tensor):
return self._encode_ndarray(nt.numpy())
return self._encode_tensor(nt)
if isinstance(nt, (int, float)):
# Although it violates NestedTensors type, MultiModalKwargs
# values are sometimes floats.
Expand Down Expand Up @@ -186,7 +200,7 @@ def dec_hook(self, t: type, obj: Any) -> Any:
if issubclass(t, np.ndarray):
return self._decode_ndarray(obj)
if issubclass(t, torch.Tensor):
return torch.from_numpy(self._decode_ndarray(obj))
return self._decode_tensor(obj)
if issubclass(t, MultiModalKwargs):
if isinstance(obj, list):
return MultiModalKwargs.from_items(
Expand All @@ -197,6 +211,15 @@ def dec_hook(self, t: type, obj: Any) -> Any:
})
return obj

def _decode_tensor(self, arr: Any) -> torch.Tensor:
dtype, shape, data = arr
if dtype != ENC_BF16_STR:
return torch.from_numpy(self._decode_ndarray(arr))
# Numpy doesn't support as bf16 so convert from fp16 view.
arr = NP_FP16_STR, shape, data
tensor = torch.from_numpy(self._decode_ndarray(arr))
return tensor.view(torch.bfloat16)

def _decode_ndarray(self, arr: Any) -> np.ndarray:
dtype, shape, data = arr
# Copy from inline representation, otherwise Torch is unhappy since
Expand Down Expand Up @@ -228,7 +251,7 @@ def _decode_nested_tensors(self, obj: Any) -> NestedTensors:
if not isinstance(obj, list):
raise TypeError(f"Unexpected NestedTensors contents: {type(obj)}")
if obj and isinstance(obj[0], str):
return torch.from_numpy(self._decode_ndarray(obj))
return self._decode_tensor(obj)
return [self._decode_nested_tensors(x) for x in obj]

def ext_hook(self, code: int, data: memoryview) -> Any:
Expand Down