bitsandbytes: NaN error when using a GPU with no support for igemmlt

@0cc4m @Opdoop

I edit /site-packages/bitsandbytes/autograd/_functions.py

first at #406

        else:
            A_wo_outliers = A.clone()
            if state.idx is not None:
                A_wo_outliers[:, state.idx.long()] = 0
            #output = torch.nn.functional.linear(A_wo_outliers, state.CB.to(A.dtype))
            #output = output.mul_(state.SCB.unsqueeze(0).mul(1.0 / 127.0))
            CB = state.CB.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
            output = torch.nn.functional.linear(A_wo_outliers, CB)
            if bias is not None:
                output = output.add_(bias)

then at 468:

        if req_gradA:
            if state.CBt is not None:
                C32grad, Sgrad = F.transform(Cgrad, "col32")
                if state.CxBt is None:
                    state.CxBt, state.SBt = F.transform(state.CBt, to_order=formatB, transpose=True)
                gradA32, SgradA32 = F.igemmlt(C32grad, state.CxBt, Sgrad, state.SBt)
                CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
                grad_A = torch.matmul(grad_output, CB).view(ctx.grad_shape).to(ctx.dtype_A)
                #grad_A = F.mm_dequant(gradA32, SgradA32, SCgrad, state.SCBt).view(ctx.grad_shape).to(ctx.dtype_A)

and now pythia-12b in 8bits at 1.5 threshold no longer NaN on me.

I then switch to full 6.0 threshold and run inference again! It works!

@richardwth you are a hero, you fixed this bug and nobody noticed!

wahoo! https://github.com/TimDettmers/bitsandbytes/pull/335

No, I think it corresponds to the threshold mentioned in the bitsandbytes README, which defaults to 6.0. That explains why it works on older cards with 0.8, it doesn’t convert much to 8-bit anymore.

@zhaoqf123 Hi buddy, sorry for the late reply. I did not use any advanced methods as you used here. I manually inserted break points and used torch.isnan, torch.isinf or torch.isfinite to check which transformer layer and later on which line exactly gave the infinite results.

I came across a similar problem when finetuning Llama 7B: the hidden states became inf at LlamaMLP (specifically, down_proj). I used V100 with device_capability 7.0 so igemmlt is not supported naturally. Then I found the inf happens at this line of autograd._functions.MatMul8bitLt

# (line 390) 3. Matmul, else branch
output = torch.nn.functional.linear(A_wo_outliers, state.CB.to(A.dtype))
output = output.mul_(state.SCB.unsqueeze(0).mul(1.0 / 127.0))

The inf happens because output has some values larger than 65536 at F.linear.

As I understand, state.CB ranges between -127 and 127 and is relatively larger than A_wo_outliers (which is confined by threshold 6.0). Wouldn’t it be safer to calculate CB first then do F.linear? That is,

CB = state.CB.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
output = torch.nn.functional.linear(A_wo_outliers, CB)

Is it designed to prevent underflow? I also notice that CB is calculated first in the backward pass (line 455).

CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
grad_A = torch.matmul(grad_output, CB).view(ctx.grad_shape).to(ctx.dtype_A)

Only fix thus far is to lower the threshold for int8 like they did here: https://gist.github.com/whjms/2505ef082a656e7a80a3f663c16f4277

its still buggy and a bit slow.