From 1da4880262ab5febbc55aa690e72e446e6b1eb42 Mon Sep 17 00:00:00 2001 From: justheuristic Date: Sun, 18 Sep 2022 00:19:22 +0300 Subject: change typecast behavior --- bitsandbytes/autograd/_functions.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) (limited to 'bitsandbytes') diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index b54ac24..5499db9 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -356,7 +356,7 @@ class MatMul8bitLt(torch.autograd.Function): if req_gradBias: # compute grad_bias first before changing grad_output dtype - grad_bias = grad_output.sum(0).to(ctx.dtype_bias) + grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) # Cast grad_output to fp16 if len(grad_output.shape) == 3: @@ -385,9 +385,8 @@ class MatMul8bitLt(torch.autograd.Function): elif state.CB is not None: CB = state.CB.to(ctx.B_dtype) - SCB = (state.SCB.unsqueeze(1) / 127.0).half() - CB *= SCB - grad_A = torch.mm(grad_output, CB).view(ctx.grad_shape).to(ctx.A_dtype) + CB.mul_(state.SCB.unsqueeze(1).div_(127.0).to(ctx.B_dtype)) + grad_A = torch.matmul(grad_output, CB).view(ctx.grad_shape).to(ctx.A_dtype) else: raise Exception('State must contain either CBt or CB matrix for backward') -- cgit v1.2.3