From 8d34d36f150b0fd4914cdb56d4e3bda34c029ccc Mon Sep 17 00:00:00 2001 From: dbaranchuk Date: Mon, 29 Aug 2022 00:56:08 +0300 Subject: req_gradA for casted & more efficient and accurate fp16 backward --- bitsandbytes/autograd/_functions.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) (limited to 'bitsandbytes/autograd/_functions.py') diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index e266d69..3bd39a9 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -213,10 +213,6 @@ class MatMul8bitLt(torch.autograd.Function): else: return torch.empty(A.shape[:-1]+B.shape[:1], dtype=torch.float16, device=A.device) - # Cast A to fp16 - A_dtype = A.dtype - A = A.to(torch.float16) - # 1. Quantize A # 2. Quantize B # 3. Matmul @@ -229,6 +225,11 @@ class MatMul8bitLt(torch.autograd.Function): input_shape = A.shape if state.outlier_pool is None: state.outlier_pool = GlobalOutlierPooler.get_instance() + + # Cast A to fp16 + A_dtype = A.dtype + A = A.to(torch.float16) + assert ( A.dtype == torch.float16 ), f"The input data type needs to be fp16 but {A.dtype} was found!" @@ -337,14 +338,14 @@ class MatMul8bitLt(torch.autograd.Function): 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 - # Cast grad_output to fp16 - grad_output_dtype = grad_output.dtype - grad_output.to(torch.float16) - req_gradA, req_gradB, req_gradBias = ctx.req_grads assert not req_gradB, "TODO: support weight updates as well" state = ctx.state + # Cast grad_output to fp16 + grad_output_dtype = grad_output.dtype + grad_output = grad_output.to(torch.float16) + if len(grad_output.shape) == 3: grad_output = grad_output.reshape( -1, grad_output.shape[-1] @@ -354,9 +355,9 @@ class MatMul8bitLt(torch.autograd.Function): if req_gradA: CB = state.CB.half() - SCB = state.SCB.unsqueeze(1).half() - B = (CB * SCB) / 127.0 - grad_A = torch.mm(grad_output, B).view(ctx.grad_shape) + SCB = (state.SCB.unsqueeze(1) / 127.0).half() + CB *= SCB + grad_A = torch.mm(grad_output, CB).view(ctx.grad_shape) if req_gradBias: grad_bias = grad_output.sum(0) -- cgit v1.2.3