[docs] Clarify FSDP-QLoRA

[stevhliu]

Fix and clarify the docs around:

1. FSDP only supports floating data type in quant_storage
2. show a more common way to access and configure this option with BitsAndBytesConfig in Transformers
3. explain compute_dtype and the importance of matching quant_storage data type with the data types used everywhere else
4. link to code example from PEFT docs

cc @Titus-von-Koeller

[github-actions]

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.

[Titus-von-Koeller]

In the first !TIP box in the Training section, I would add a sentence (maybe in bold?) like "Note that FSDP is a distributed framework and therefore needs to be kicked off as a distributed training job, which is described in the following resources."

Otherwise it doesn't become clear for novices that these code examples alone (on this BNB doc page) don't suffice and FSDP is not run as standard single threaded Python that they can run in an interpreter or notebook, but needs to be kicked as a distributed way through Accelerate (or torchrun or alike if we weren't using HF libs).

[Titus-von-Koeller]

Hey @stevhliu, thanks a lot for taking the time to further polish this page, really appreciated! 🤗

I left a few suggestions. Please feel free to integrate my feedback at your discretion and then we'll merge.

[Titus-von-Koeller]

Suggested change

FSDP only supports sharding float data types which can be problematic because quantized weights are typically stored as integer data types (uint8). bitsandbytes doesn't have this problem because it uses `StoreChar` to read and write quantized weights regardless of the data type storage. This makes it simple to add a `quant_storage` parameter to the [`~nn.Linear4bit`] and [`~nn.Params4bit`] classes and set it to `torch.uint8` to maintain backward compatibility with the codebase. With the `quant_storage` parameter, you can select any of the FSDP supported data types to shard [`~nn.Linear4bit`] with such as bfloat16, float16 or float32.

FSDP only supports sharding float data types which can be problematic because quantized weights are typically stored as integer data types (uint8). bitsandbytes isn't limited by this convention because it uses `StoreChar` to read and write quantized weights regardless of the underlying data type used for storage of the quantized bytes. This made it possible to add a `quant_storage` parameter to the [`~nn.Linear4bit`] and [`~nn.Params4bit`] classes and set it to the default `torch.uint8` to maintain backward compatibility with previous iterations of the codebase. With the `quant_storage` parameter, you can select any of the FSDP supported data types to shard [`~nn.Linear4bit`] with such as bfloat16, float16 or float32.

[Titus-von-Koeller]

Suggested change

You'll typically access and configure this option from [`transformers.BitsAndBytesConfig`] by setting the `bnb_4bit_quant_storage` parameter. It is very **important** the `quant_storage` data type matches the data types used throughout the model because FSDP can only wrap layers and modules that have the *same floating data type*. Making sure the data types are aligned will ensure the model is correctly sharded.

You'll typically access and configure this option from [`transformers.BitsAndBytesConfig`] by setting the `bnb_4bit_quant_storage` parameter. It is very **important** that the `quant_storage` data type matches the data types used throughout the model because FSDP can only wrap layers and modules that have the *same floating point data type*. Making sure the data types are aligned will ensure the model is correctly sharded.

[Titus-von-Koeller]

Suggested change

> The `compute_dtype` is the data type used for computation inside the CUDA kernel, where the 4-bit quantized weights are unpacked from the data type in `quant_storage` and dequantized to `compute_dtype`. We recommend using torch.bfloat16 (if available on your hardware) for better numerical stability.

> Another dtype that you'll have to be sure to set correctly when using bitsandbytes, but that might sound confusing in this context without further explanation is the `compute_dtype`: It's unrelated to the storage format used for the quantized bytes, but instead is the data type used for computation inside the CUDA kernel. In the kernel, the 4-bit quantized weights are unpacked from the data type in `quant_storage` and dequantized to `compute_dtype` for the matrix multiplication operation. We recommend using torch.bfloat16 (if available on your hardware) for better numerical stability. What you set `compute_dtype` has no effect on FSDP and sharding: So numerical stability, hardware support and computational efficiency are the only relevant considerations for setting it.