[LLM]Fix Arc falcon abnormal output issue (#9096)
* update * update * fix error & style * fix style * update train * to input_seq_size
This commit is contained in:
		
							parent
							
								
									548e4dd5fe
								
							
						
					
					
						commit
						4c4f8d1663
					
				
					 1 changed files with 8 additions and 6 deletions
				
			
		| 
						 | 
					@ -288,10 +288,10 @@ def ggml_matmul_src1_x_src0_t(src0: torch.Tensor,
 | 
				
			||||||
class MatMulLowBit(torch.autograd.Function):
 | 
					class MatMulLowBit(torch.autograd.Function):
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    @staticmethod
 | 
					    @staticmethod
 | 
				
			||||||
    def forward(ctx, A, weight):
 | 
					    def forward(ctx, A, weight, input_seq_size):
 | 
				
			||||||
        ctx.is_empty = False
 | 
					        ctx.is_empty = False
 | 
				
			||||||
        import linear_q4_0
 | 
					        import linear_q4_0
 | 
				
			||||||
        result = linear_q4_0.forward_new(A, weight.data, weight.qtype)
 | 
					        result = linear_q4_0.forward_new(A, weight.data, weight.qtype, input_seq_size)
 | 
				
			||||||
        if any(ctx.needs_input_grad[:2]):
 | 
					        if any(ctx.needs_input_grad[:2]):
 | 
				
			||||||
            ctx.tensors = (A, weight)
 | 
					            ctx.tensors = (A, weight)
 | 
				
			||||||
        else:
 | 
					        else:
 | 
				
			||||||
| 
						 | 
					@ -304,14 +304,14 @@ class MatMulLowBit(torch.autograd.Function):
 | 
				
			||||||
        if ctx.is_empty:
 | 
					        if ctx.is_empty:
 | 
				
			||||||
            bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias)
 | 
					            bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias)
 | 
				
			||||||
            return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None
 | 
					            return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None
 | 
				
			||||||
        req_gradA, _ = ctx.needs_input_grad
 | 
					        req_gradA, _, _ = ctx.needs_input_grad
 | 
				
			||||||
        A, weight = ctx.tensors
 | 
					        A, weight = ctx.tensors
 | 
				
			||||||
        grad_A, grad_weight = None, None
 | 
					        grad_A, grad_weight = None, None
 | 
				
			||||||
        if req_gradA:
 | 
					        if req_gradA:
 | 
				
			||||||
            dequant_weight = linear_q4_0.dequant(A, weight.data, weight.qtype)
 | 
					            dequant_weight = linear_q4_0.dequant(A, weight.data, weight.qtype)
 | 
				
			||||||
            grad_A = torch.matmul(grad_output, dequant_weight.reshape(weight._shape))
 | 
					            grad_A = torch.matmul(grad_output, dequant_weight.reshape(weight._shape))
 | 
				
			||||||
 | 
					
 | 
				
			||||||
        return grad_A, grad_weight
 | 
					        return grad_A, grad_weight, None
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
class LowBitLinear(nn.Linear):
 | 
					class LowBitLinear(nn.Linear):
 | 
				
			||||||
| 
						 | 
					@ -353,10 +353,12 @@ class LowBitLinear(nn.Linear):
 | 
				
			||||||
            # disable the conversion when training
 | 
					            # disable the conversion when training
 | 
				
			||||||
            if self.conver_to_half and x_2d.shape[0] > 1 and x_2d.dtype == torch.float32:
 | 
					            if self.conver_to_half and x_2d.shape[0] > 1 and x_2d.dtype == torch.float32:
 | 
				
			||||||
                x_2d = x_2d.half()
 | 
					                x_2d = x_2d.half()
 | 
				
			||||||
 | 
					            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)
 | 
					                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)
 | 
					                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:
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
		Reference in a new issue