huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-06-06 17:32:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

add fix

Browse Source
This commit is contained in:
Mohit Sharma 2024-12-13 16:10:00 +00:00
parent ea7f4082c4
commit 1fa9ca2f16
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
server/text_generation_server/layers/fp8.py
View File

@ -395,6 +395,13 @@ class Fp8Linear(torch.nn.Module):
qinput, self.qweight.t(), scale, self.scale, input.dtype, self.bias qinput, self.qweight.t(), scale, self.scale, input.dtype, self.bias
) )
batched = False
input_shape = input.shape
if input.dim() == 3:
input = input.view(-1, input_shape[-1])
batched = True
qinput, scale = fp8_quantize( qinput, scale = fp8_quantize(
input, input,
self.input_scale, self.input_scale,
@ -438,6 +445,9 @@ class Fp8Linear(torch.nn.Module):
output = output.to(dtype=self.dtype) output = output.to(dtype=self.dtype)
if batched:
output = output.view(*input_shape[:-1], output.shape[-1])
return output return output