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

Fix removing convert dtype bug (#9216)

Browse source 
* Fix removing convert dtype bug

* fix style
This commit is contained in:
Yang Wang 2023-10-19 02:24:22 +08:00 committed by GitHub
parent 942d6418e7
commit b0ddde0410
1 changed files with 12 additions and 7 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

@ -355,14 +355,19 @@ class LowBitLinear(nn.Linear):
if x_2d.is_contiguous() is False:
x_2d = x_2d.contiguous()
input_seq_size = x_shape[1]
if self.training and x_2d.requires_grad:
result = MatMulLowBit.apply(x_2d, self.weight, input_seq_size)
else:
# current workaround to reduce first token latency of fp32 input
# sometimes fp16 cause nan and training instability
# disable the conversion when training
if self.conver_to_half and x_2d.shape[0] > 1 and x_2d.dtype == torch.float32:
x_2d = x_2d.half()
input_seq_size = x_shape[1]
if self.training and x_2d.requires_grad:
result = MatMulLowBit.apply(x_2d, self.weight, input_seq_size)
result = linear_q4_0.forward_new(x_2d, self.weight.data, self.weight.qtype,
input_seq_size)
result = result.to(x.dtype)
else:
result = linear_q4_0.forward_new(x_2d, self.weight.data, self.weight.qtype,
input_seq_size)