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

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

@ -355,17 +355,22 @@ class LowBitLinear(nn.Linear):
if x_2d.is_contiguous() is False: if x_2d.is_contiguous() is False:
x_2d = x_2d.contiguous() x_2d = x_2d.contiguous()
# 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] input_seq_size = x_shape[1]
if self.training and x_2d.requires_grad: if self.training and x_2d.requires_grad:
result = MatMulLowBit.apply(x_2d, self.weight, input_seq_size) result = MatMulLowBit.apply(x_2d, self.weight, input_seq_size)
else: else:
result = linear_q4_0.forward_new(x_2d, self.weight.data, self.weight.qtype, # current workaround to reduce first token latency of fp32 input
input_seq_size) # 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()
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)
new_shape = x_shape[:-1] + (self.out_len,) new_shape = x_shape[:-1] + (self.out_len,)
result = result.view(new_shape) result = result.view(new_shape)
if self.bias is not None: if self.bias is not None: