From 47a73d94c3d3284f6073b0ff189ed5bc9e3a8762 Mon Sep 17 00:00:00 2001 From: Tim Dettmers Date: Tue, 26 Jul 2022 19:15:35 -0700 Subject: Matmullt with direct outlier extraction for 8-bit inference. --- bitsandbytes/autograd/_functions.py | 52 ++++++++++++++++++++++++------------- 1 file changed, 34 insertions(+), 18 deletions(-) (limited to 'bitsandbytes/autograd/_functions.py') diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 815a4f1..5503749 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -191,24 +191,24 @@ class MatMul8bitLt(torch.autograd.Function): # B in in 8-bit row-major, we can transform it back to 16-bit to extract outlier dimensions # we also need to convert it to the turing/ampere format state.CxB, state.SB = F.transform(state.CB, to_order=formatB) - if state.threshold > 0.0 and coo_tensorA is not None and state.idx is None and state.CB is not None: - # generate outlier index and subB - outlier_idx = torch.unique(coo_tensorA.colidx).long() - state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) - if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: - # do not use pool for 2nd FFN layer - state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) - else: - state.idx = outlier_idx - state.subB = (state.CB[:, state.idx].float().t().contiguous()*(state.SCB/127)).half() - - if state.idx is not None: - # extract outliers - CA[:, state.idx] = 0 - CAt[:, state.idx] = 0 - subA = A[:, state.idx] - else: - subA = None + #if state.threshold > 0.0 and coo_tensorA is not None and state.idx is None and state.CB is not None: + # # generate outlier index and subB + # outlier_idx = torch.unique(coo_tensorA.colidx).long() + # state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) + # if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: + # # do not use pool for 2nd FFN layer + # state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) + # else: + # state.idx = outlier_idx + # state.subB = (state.CB[:, state.idx].float().t().contiguous()*(state.SCB/127)).half() + + #if state.idx is not None: + # # extract outliers + # CA[:, state.idx] = 0 + # CAt[:, state.idx] = 0 + # subA = A[:, state.idx] + #else: + # subA = None else: if not state.has_fp16_weights and state.CxB is None: state.CxB, state.SB = F.transform(state.CB, to_order=formatB) @@ -229,6 +229,22 @@ class MatMul8bitLt(torch.autograd.Function): else: has_grad = False + if coo_tensorA is not None and not state.has_fp16_weights: + # extract outliers + + outlier_idx = torch.unique(coo_tensorA.colidx) + state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) + if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: + # do not use pool for 2nd FFN layer + state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) + else: + state.idx = outlier_idx + outliers = F.extract_outliers(state.CxB, state.SB, outlier_idx).half() + state.subB = (outliers*state.SCB.view(-1, 1).half()/127.0).t().contiguous() + CA[:, state.idx.long()] = 0 + CAt[:, state.idx.long()] = 0 + subA = A[:, state.idx.long()] + shapeB = state.SB[0] if len(input_shape) == 3: -- cgit v1.2.3 From 389f66ca5a737eb7f22f22fed420274ff622623e Mon Sep 17 00:00:00 2001 From: Tim Dettmers Date: Wed, 27 Jul 2022 01:46:35 -0700 Subject: Fixed direct extraction masking. --- bitsandbytes/autograd/_functions.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) (limited to 'bitsandbytes/autograd/_functions.py') diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 5503749..e641583 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -191,6 +191,7 @@ class MatMul8bitLt(torch.autograd.Function): # B in in 8-bit row-major, we can transform it back to 16-bit to extract outlier dimensions # we also need to convert it to the turing/ampere format state.CxB, state.SB = F.transform(state.CB, to_order=formatB) + #state.B = (state.CB.float()*(state.SCB.view(-1, 1)/127)).half() #if state.threshold > 0.0 and coo_tensorA is not None and state.idx is None and state.CB is not None: # # generate outlier index and subB # outlier_idx = torch.unique(coo_tensorA.colidx).long() @@ -214,7 +215,6 @@ class MatMul8bitLt(torch.autograd.Function): state.CxB, state.SB = F.transform(state.CB, to_order=formatB) subA = None - C32A, SA = F.transform(CA, 'col32') # 2. Quantize B if state.has_fp16_weights: @@ -233,14 +233,15 @@ class MatMul8bitLt(torch.autograd.Function): # extract outliers outlier_idx = torch.unique(coo_tensorA.colidx) - state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) - if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: - # do not use pool for 2nd FFN layer - state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) - else: - state.idx = outlier_idx - outliers = F.extract_outliers(state.CxB, state.SB, outlier_idx).half() - state.subB = (outliers*state.SCB.view(-1, 1).half()/127.0).t().contiguous() + state.idx = outlier_idx + #state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) + #if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: + # # do not use pool for 2nd FFN layer + # state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) + #else: + # state.idx = outlier_idx + outliers = F.extract_outliers(state.CxB, state.SB, state.idx.int()) + state.subB = (outliers*state.SCB.view(-1, 1)/127.0).t().contiguous().half() CA[:, state.idx.long()] = 0 CAt[:, state.idx.long()] = 0 subA = A[:, state.idx.long()] @@ -253,11 +254,12 @@ class MatMul8bitLt(torch.autograd.Function): output_shape = (input_shape[0], shapeB[0]) # 3. Matmul + C32A, SA = F.transform(CA, 'col32') out32, Sout32 = F.igemmlt(C32A, state.CxB, SA, state.SB) output = F.mm_dequant(out32, Sout32, SCA, state.SCB) # 4. Mixed-precision decomposition matmul - if state.threshold > 0.0 and coo_tensorA is not None and subA is not None: + if coo_tensorA is not None and subA is not None: output += torch.matmul(subA, state.subB) # 5. Save state -- cgit v1.2.3