deftdawg/bitsandbytes-rocm
1
0
Fork 0
forked from mrq/bitsandbytes-rocm
Code Pull Requests Activity

Matmullt with direct outlier extraction for 8-bit inference.

Browse Source
This commit is contained in:
Tim Dettmers 2022-07-26 19:15:35 -07:00
parent 32fa459ed7
commit 47a73d94c3
1 changed files with 33 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
bitsandbytes/autograd/_functions.py
View File

@ -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 # 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 # we also need to convert it to the turing/ampere format
state.CxB, state.SB = F.transform(state.CB, to_order=formatB) 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: #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 # # generate outlier index and subB
outlier_idx = torch.unique(coo_tensorA.colidx).long() # outlier_idx = torch.unique(coo_tensorA.colidx).long()
state.outlier_pool.add_outliers(outlier_idx, A.shape[-1]) # state.outlier_pool.add_outliers(outlier_idx, A.shape[-1])
if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]: # if state.use_pool and state.outlier_pool.model_dim == A.shape[-1]:
# do not use pool for 2nd FFN layer # # do not use pool for 2nd FFN layer
state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device) # state.idx = state.outlier_pool.get_current_outlier_idx().to(A.device)
else: # else:
state.idx = outlier_idx # state.idx = outlier_idx
state.subB = (state.CB[:, state.idx].float().t().contiguous()*(state.SCB/127)).half() # state.subB = (state.CB[:, state.idx].float().t().contiguous()*(state.SCB/127)).half()
if state.idx is not None: #if state.idx is not None:
# extract outliers # # extract outliers
CA[:, state.idx] = 0 # CA[:, state.idx] = 0
CAt[:, state.idx] = 0 # CAt[:, state.idx] = 0
subA = A[:, state.idx] # subA = A[:, state.idx]
else: #else:
subA = None # subA = None
else: else:
if not state.has_fp16_weights and state.CxB is None: if not state.has_fp16_weights and state.CxB is None:
state.CxB, state.SB = F.transform(state.CB, to_order=formatB) state.CxB, state.SB = F.transform(state.CB, to_order=formatB)
@ -229,6 +229,22 @@ class MatMul8bitLt(torch.autograd.Function):
else: else:
has_grad = False 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] shapeB = state.SB[0]
if len(input_shape) == 3: if len(input_shape) == 3: