summaryrefslogtreecommitdiff
path: root/bitsandbytes/autograd
diff options
context:
space:
mode:
authorjustheuristic <justheuristic@gmail.com>2022-09-18 00:19:22 +0300
committerjustheuristic <justheuristic@gmail.com>2022-09-18 00:19:22 +0300
commit1da4880262ab5febbc55aa690e72e446e6b1eb42 (patch)
tree2173f553acaa064930c3468198a6bc86762b5794 /bitsandbytes/autograd
parent1145589f84d2ba4eb3b4a18fa33423298f5747c0 (diff)
change typecast behavior
Diffstat (limited to 'bitsandbytes/autograd')
-rw-r--r--bitsandbytes/autograd/_functions.py7
1 files changed, 3 insertions, 4 deletions
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')