fix for itemsize => element_size() for torch backwards compat

[winglian]

This is a similar fix to huggingface/peft#1630.

.itemsize on a tensor is only supported on torch>=2.1

Fixes: #30304

Edit by @younesbelkada:

This PR makes sure the checks are compatible with earlier versions of pytorch that have been overlooked by myself in #30162

[winglian]

@pacman100 @BenjaminBossan

[muellerzr]

I meant to respond to the other one earlier, but may be good to have an if/else to know which version to call? E.g. param.numel() * 2 * elem_size where elemsize is defined earlier based on the torch version

[muellerzr]

Correction, as @BenjaminBossan pointed out this is a must for us to have

[peterdonnelly1]

Besides itemsize not being available in torch 2.0.1, I couldn't get element_size() to work either (not a dtype attribute error).

I put this 'helper' function at the top of the file to do it manually:

def get_dtype_size(dtype):
if dtype == torch.float32:
return 4
elif dtype == torch.float64:
return 8
elif dtype == torch.float16:
return 2
elif dtype == torch.uint8:
return 1
elif dtype == torch.int8:
return 1
elif dtype == torch.int16:
return 2
elif dtype == torch.int32:
return 4
elif dtype == torch.int64:
return 8
else:
raise ValueError("Unsupported dtype")

and then used this in place of the itemsize code:

if is_loaded_in_4bit and isinstance(param, bnb.nn.Params4bit):
quant_storage = self.hf_quantizer.quantization_config.bnb_4bit_quant_storage
nb_params = get_dtype_size(quant_storage)
total_numel.append(param.numel() * 2 * nb_params)
else:
total_numel.append(param.numel())

I've only tested it by using the qlora config on tiny-llama; and at least for this case it works:

accelerate launch -m axolotl.cli.train examples/tiny-llama/qlora.yml

I am using image 'winglian/axolotl-base:main-base-py3.11-cu121-2.2.1' (which actually uses torch 2.0.1 not 2.2.1 as the name implies.). Ubuntu 22.04 / laptop RTX4090 (16G)

[yand1ngs]

Thank you soooo much!!!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

Note that this PR fails on PEFT CI.

[winglian]

my last change isn't correct either. I'll take a look at this tomorrow.

[younesbelkada]

Hi @winglian , as pointed out by @hiyouga my fix did not cover all the cases, would you be happy to rebase your PR with main and we merge it ?

[younesbelkada]

Why not simply use:

quant_storage = self.hf_quantizer.quantization_config.bnb_4bit_quant_storage
num_bytes = quant_storage.element_size()

element_size seems present from 1.9.1: https://pytorch.org/docs/1.9.1/search.html?q=element_size&check_keywords=yes&area=default to latest: https://pytorch.org/docs/2.2/search.html?q=element_size&check_keywords=yes&area=default

[younesbelkada]

cc @winglian @BenjaminBossan @hiyouga

[hiyouga]

self.hf_quantizer.quantization_config.bnb_4bit_quant_storage is a torch.dtype instance while only the torch.tensor has element_size()

huggingface/peft#1635

[younesbelkada]

Thanks for taking care of the backward compatibility for previous torch versions !

[amyeroberts]

Thanks for handling this fix!

Just a small nit