Skip to content

fix: wrong absmax indexing in AVX512 dequantizeBlockwise4bitCpu - #2096

Open
NanoCode012 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
NanoCode012:patch-1
Open

NanoCode012 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
NanoCode012:patch-1

Conversation

@NanoCode012

Copy link
Copy Markdown

Was adding a mini test in Axolotl for merge qlora and found an edge case on GitHub CI that didn't appear on local Mac testing. Claude helped solve this, but repro is ran and verified on a L40s with AVX512.

Bug:

On an x86 CPU with AVX512 (like Github CPU CI), there's a guard missing in dequantizeBlockwise4bitCpu unlike in the ARM path (n % blocksize == 0).

This PR adds the same condition to the AVX512 branch, so those shapes fall through to the scalar path, which indexes absmax by flat offset.

Repro:

On an AVX512 CPU with no visible GPU:

# CUDA_VISIBLE_DEVICES="" python repro.py
import torch
import bitsandbytes.functional as F

for shape in [(64, 16), (64, 80), (64, 64)]:
    w = torch.randn(*shape, dtype=torch.bfloat16)
    q2d, s2d = F.quantize_4bit(w, quant_type="nf4")
    q1d, s1d = F.quantize_4bit(w.flatten(), quant_type="nf4")
    d2d = F.dequantize_4bit(q2d, s2d).flatten()
    d1d = F.dequantize_4bit(q1d, s1d).flatten()
    print(shape, torch.equal(d2d, d1d), (d2d.float() - d1d.float()).abs().max().item())

The quantized bytes are identical for the 2D and flattened inputs, so both dequantizations should match.

AMD EPYC 9554 with bnb 0.50.2 and torch 2.13.0:

Shape Before After
(64, 16) mismatch, max abs error about 1.5 match
(64, 80) mismatch, max abs error about 1.7 match
(64, 64) match match

The same results hold with compress_statistics=True.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant