Unexpected behaviour in inference with merged QLoRA weights

Hi, a few weeks ago @morettif and I finetuned the `Llama70B` with QLoRA on a H100:

- `r=32`
- `alpha=64`
- `quantize=bnb.nf4-dq`
- `precision=bf16-true`
- `weight_decay=0`
- `batch_size=32`
- `micro_batch_size=2`
- `lora_dropout=0.05`
- All LoRA layers

The validation loss after the finetuning was about `0.5`, but during inference on the validation set, we obtained very strange results. We checked the loss on the validation set and it was ~2.5 (as before the finetuning).

We used the `generate/lora.py` script and we noticed that after commenting this line: https://github.com/Lightning-AI/lit-gpt/blob/8a101b633dfeafd378f8fbaba6a80a4417c33576/generate/lora.py#L120 the model generated meaningful output with a loss coherent with the one observed during training (~0.5).

We analyzed the dtype of the pre-trained model weights at the moment of the call to the `merge_lora_weights()` function, some of them are `torch.uint8`, this causes that during merging the following if (in the `merge()` function) is evaluated as `True`: https://github.com/Lightning-AI/lit-gpt/blob/8a101b633dfeafd378f8fbaba6a80a4417c33576/lit_gpt/lora.py#L151

Could there be an error in this piece of code related to the management of quantization?

Start with generate/lora.py at the merge_lora_weights() call and inspect lit_gpt/lora.py in merge(), especially the torch.uint8 quantization branch. Reproduce the Llama70B QLoRA validation inference case with the reported settings and compare the loss before and after merging. Done means merged weights produce meaningful output and validation loss consistent with the approximately 0.5 training loss.

Written by the indexing model from the issue text.