Max/4bit

[maxjeblick]

This PR adds 4bit training/inference.
I have tested with EleutherAI/gpt-neox-20b where 8 bit yields OOM. I have also compared EleutherAI/gpt-j-6B 4bit/8bit training and its corresponding chat responses.

There is an issue with merging LORA back which can yield to OOM issues: #130 I'll have a look and see if it can be fixed and added to this PR.

I'm also not sure about this line: https://github.com/h2oai/h2o-llmstudio/blob/main/llm_studio/src/utils/modeling_utils.py#L72
It's currently left as is, I haven't found any issues with it so far.

[psinger]

I'm also not sure about this line: https://github.com/h2oai/h2o-llmstudio/blob/main/llm_studio/src/utils/modeling_utils.py#L72
It's currently left as is, I haven't found any issues with it so far.

IIRC this was needed for training and loading weights in different precisions

[pascal-pfeiffer]

Thank you Max! Will need to test a bit more, but looking great so far.

let's also add some words in the README.md about 4bit training?

[psinger]

Getting errors when training in 4bit, and then doing inference on float16.
While loading the weights.
Score is correct though.

Might be related to my comment above.

[psinger]

One more note:
Inference in 4bit is extremely slow for me.

Would it maybe make sense to force 8bit in Chat window?
Or even better, make it somehow an option?
Might be a separate issue.

[psinger]

Found this:
https://github.com/huggingface/peft/blob/main/examples/fp4_finetuning/finetune_fp4_opt_bnb_peft.py

Here they are also setting llm_int8_threshold - does it have an impact in 4bit?

I think we might also be fine setting it to float16 instead of bfloat16.

[maxjeblick]

Here they are also setting llm_int8_threshold - does it have an impact in 4bit?

Should not affect 4bit training:
https://github.com/huggingface/transformers/blob/main/src/transformers/utils/bitsandbytes.py#L133

[psinger]

Not sure this was lucky or has impact, but I got better inference speed with these settings:

elif cfg.architecture.backbone_dtype == "int4":
        kwargs["device_map"] = {"": cfg.environment._device}
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            llm_int8_has_fp16_weight=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            llm_int8_threshold=0.0
        )
        # need to force pretrained
        cfg.architecture.pretrained = True
        kwargs["torch_dtype"] = torch.float16

[psinger]

Let's add this also to int8:

kwargs["torch_dtype"] = torch.float16

And do you think bnb_4bit_use_double_quant=True is useful?

[maxjeblick]

And do you think bnb_4bit_use_double_quant=True is useful?

From here:

Other options include bnb_4bit_use_double_quant which uses a second quantization after the first one to save an additional 0.4 bits per parameter

I haven't tested this option yet; in theory, it should allow you to finetune even larger models. I can compare infernece speed and enable it if it turns out to be similar.

[maxjeblick]

For h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v2 (evaluate only, that is train with 0 epochs) bnb_4bit_use_double_quant does not seem to have any noticeable impact, both time-wise, as well as max-GPU memory-wise.
Not sure how, if at all, it affects training performance. Let's maybe keep it disabled for now?

[maxjeblick]

Getting errors when training in 4bit, and then doing inference on float16.

Is this issue still present? I started a new experiment in fp16, using previous 4bit weights; that works.

[psinger]

Getting errors when training in 4bit, and then doing inference on float16.

Is this issue still present? I started a new experiment in fp16, using previous 4bit weights; that works.

I think after you made the changes above it works now.

[psinger]

@maxjeblick did you try pushing to HF and reloading?

[maxjeblick]

@maxjeblick did you try pushing to HF and reloading?

I tried downloading the model from UI and subsequently loading from the unzipped folder, following the model card example. That worked.

[psinger]

Cool, if possible lets also try the HF push.

[psinger]

Thanks!

[maxjeblick]

Checked HF hub, it also works.