ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

[LLM]Solve the problem of calling bmm operator in BF16Linear (#9924)

Browse source 
* Solve the problem of calling bmm operator in BF16Linear
This commit is contained in:
Ziteng Zhang 2024-01-17 18:08:35 +08:00 committed by GitHub
parent e403e4a8b7
commit 18cd1f1432
1 changed files with 11 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
python/llm/src/bigdl/llm/transformers/low_bit_linear.py
View file

@ -665,5 +665,16 @@ class BF16Linear(nn.Linear):
if self.bias is not None and self.bias.dtype != x.dtype:
self.bias.data = self.bias.data.to(x.dtype)
# If x.shape>3, F.linear will use bmm, accounting for performance degradation.
original_shape = x.shape
# Convert to 2D shape
if len(original_shape) > 2:
x = x.reshape(-1, original_shape[-1])
result = F.linear(x, self.weight, self.bias)
# Convert to original shape
if len(original_shape) > 2:
result = result.reshape(*original_shape[:-1], result.shape[-1])
return result.to(x.dtype)