Lightning-AI / Lightning-AI/pytorch-lightning
Fabric FSDP with bitsandbytes plugin is not supported
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
Description & Motivation
While the transformers library supports FSDP with bitsandbytes 4bit quantization (https://huggingface.co/docs/bitsandbytes/main/en/fsdp_qlora), Fabric appears unable to combine the two:
mode = "nf4"
plugin = BitsandbytesPrecision(mode=mode, dtype=torch.bfloat16)
policy = partial(size_based_auto_wrap_policy, min_num_params=1E8)
# use data parallel with quantization rather than FSDP:
fabric = Fabric(accelerator="auto", devices="auto",
strategy=FSDPStrategy(
auto_wrap_policy=policy, cpu_offload=False,
state_dict_type="sharded"),
plugins=plugin)
gemma-fine-tuning, pid=3426) File "/home/ubuntu/miniconda3/envs/fabric/lib/python3.10/site-packages/lightning/fabric/strategies/fsdp.py", line 242, in precision
raise TypeError(f"The FSDP strategy can only work with the `FSDPPrecision` plugin, found {precision}")
TypeError: The FSDP strategy can only work with the `FSDPPrecision` plugin, found <lightning.fabric.plugins.precision.bitsandbytes.BitsandbytesPrecision object at 0x78b7aa999ba0>
Pitch
Being able to fine-tuning large models with FSDP and quantization is becoming more important as large multi-modal LLMs require a lot of GPU memory for batches (e.g., multi-images per example and batch size >1).
Alternatives
No response
Additional context
No response
cc @lantiga @borda @justusschock
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the supplied Fabric configuration and inspect lightning/fabric/strategies/fsdp.py, especially the FSDPStrategy precision check and the BitsandbytesPrecision entry point. Done means the shown FSDP and bitsandbytes configuration is accepted without the reported TypeError and its behavior is covered by the project's relevant tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100