32 bit optimizer update error despite gradients being the same about bitsandbytes

[Edenzzzz]

@matthewdouglas @TimDettmers any insights? Thanks!

[matthewdouglas]

Hi @Edenzzzz,

Make sure that this chunk is contiguous as F.optimizer_update_32bit ultimately treats it as 1D.

dist_low_rank_grad = torch.load("dist_low_rank_grad.pt").contiguous()

I was able to reproduce your results, and after this change I believe I'm seeing the desired result.

RTX 3060, CUDA 12.4, torch==2.2.2+cu121, bitsandbytes==0.43.1

[Edenzzzz]

Thanks a lot! This worked. The non-contiguous tensor came from torch.chunk and torch.distributed.all_gather.
I wonder if that's due to the c++ kernel not considering the reshaped strides and assuming row-major format? I can file in a PR to make it contiguous if you feel that's helpful.

[matthewdouglas]

I wonder if that's due to the c++ kernel not considering the reshaped strides and assuming row-major format?

Yes, exactly. The C++ kernel assumes it's row-major and only knows the total number of elements.

I can file in a PR to make it contiguous if you feel that's helpful.

That seems like a reasonable check to me, so a PR to add that sounds good!