[BugFix] Support bf16 in zero-copy tensor serialization #16860
Conversation
Numpy doesn't support bfloat16 so we convert to/from a fp16 view. Signed-off-by: Nick Hill <[email protected]>
njhill
requested review from
WoosukKwon,
robertgshaw2-redhat,
ywang96,
comaniac and
alexm-redhat
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
| 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()) |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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 ->
- check if tensor is_contiguous
- store shape & dtype on the tensor
- create a view of the tensor as one-dimensional nbytes 8-bit int array
- 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 ?
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
@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...
There was a problem hiding this comment.
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.
Signed-off-by: Staszek Pasko <[email protected]>
|
Replaced by #16866 |
Follow-on from #13790 and #16432.
Numpy doesn't support bfloat16 so we convert to/from a fp16 view.