Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized #19886

awaelchli · 2024-05-21T15:28:57Z

What does this PR do?

Fixes #19271
Since the Trainer directly imports the plugin from Fabric, this fix automatically applies there too.

Checklist:

Add test (fails on master)
Run against LitGPT

Internal discussion: https://pytorch-lightning.slack.com/archives/C05NRK0DH34/p1716296573578039

📚 Documentation preview 📚: https://pytorch-lightning--19886.org.readthedocs.build/en/19886/

cc @Borda @carmocca @justusschock @awaelchli

github-actions · 2024-05-21T15:29:23Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 2.0, oldest)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.1)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.2)	success	✅
pl-cpu (macOS-14, lightning, 3.10, 2.3)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 2.0, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.2)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.3)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 2.0, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.1)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.2)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.3)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 2.0)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 2.0)	success	✅
pl-cpu (windows-2022, pytorch, 3.8, 2.0)	success	✅
pl-cpu (macOS-12, pytorch, 3.11, 2.0)	success	✅
pl-cpu (macOS-12, pytorch, 3.11, 2.1)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.0)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.11, 2.0)	success	✅
pl-cpu (windows-2022, pytorch, 3.11, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs) (testing Lightning \| latest)	success	✅
pytorch-lightning (GPUs) (testing PyTorch \| latest)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

🟢 fabric: Docs

Check ID	Status
docs-make (fabric, doctest)	success	✅
docs-make (fabric, html)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, lightning, 3.8, 2.0, oldest)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
fabric-cpu (macOS-11, lightning, 3.11, 2.1)	success	✅
fabric-cpu (macOS-11, lightning, 3.11, 2.2)	success	✅
fabric-cpu (macOS-14, lightning, 3.10, 2.3)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 2.0, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.1)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.2)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.3)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 2.0, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
fabric-cpu (windows-2022, lightning, 3.11, 2.1)	success	✅
fabric-cpu (windows-2022, lightning, 3.11, 2.2)	success	✅
fabric-cpu (windows-2022, lightning, 3.11, 2.3)	success	✅
fabric-cpu (macOS-11, fabric, 3.8, 2.0)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 2.0)	success	✅
fabric-cpu (windows-2022, fabric, 3.8, 2.0)	success	✅
fabric-cpu (macOS-12, fabric, 3.11, 2.0)	success	✅
fabric-cpu (macOS-12, fabric, 3.11, 2.1)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.0)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.11, 2.0)	success	✅
fabric-cpu (windows-2022, fabric, 3.11, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py, tests/tests_fabric/plugins/precision/test_bitsandbytes.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs) (testing Fabric \| latest)	success	✅
lightning-fabric (GPUs) (testing Lightning \| latest)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py, tests/tests_fabric/plugins/precision/test_bitsandbytes.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

🟢 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.8)	success	✅
install-pkg (ubuntu-22.04, app, 3.11)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.11)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.11)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.11)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.11)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.11)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.11)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.11)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.11)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.11)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.11)	success	✅
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.11)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.11)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.11)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.11)	success	✅

These checks are required after the changes to src/lightning/fabric/plugins/precision/bitsandbytes.py.

Thank you for your contribution! 💜

Note
This comment is automatically generated and updates for 60 minutes every 180 seconds. If you have any other questions, contact carmocca for help.

codecov · 2024-05-21T15:44:26Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 59%. Comparing base (d76feef) to head (340d721).
Report is 122 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (d76feef) and HEAD (340d721). Click for more details.

HEAD has 294 uploads less than BASE

Flag BASE (d76feef) HEAD (340d721)

python3.10 35 16

cpu 118 48

lightning 73 32

pytest 74 28

python3.8 24 12

lightning_fabric 19 10

python3.11 40 20

examples 9 0

app 9 0

pytorch2.0 24 12

pytest-full 48 24

pytorch_lightning 19 10

pytorch2.3 6 3

pytorch2.1 12 6

pytorch2.2 6 3

lightning_app 2 0

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #19886     +/-   ##
=========================================
- Coverage      84%      59%    -25%     
=========================================
  Files         426      421      -5     
  Lines       35233    35135     -98     
=========================================
- Hits        29501    20716   -8785     
- Misses       5732    14419   +8687

for more information, see https://pre-commit.ci

src/lightning/fabric/plugins/precision/bitsandbytes.py

lantiga

Great find, thank you @awaelchli

…ady quantized (#19886)

bugfix

beb7358

awaelchli requested review from carmocca and justusschock as code owners May 21, 2024 15:28

github-actions bot added the fabric lightning.fabric.Fabric label May 21, 2024

awaelchli added bug Something isn't working precision: bnb Bitsandbytes quantization labels May 21, 2024

awaelchli added this to the 2.2.x milestone May 21, 2024

awaelchli and others added 3 commits May 21, 2024 12:19

add test

1865820

[pre-commit.ci] auto fixes from pre-commit.com hooks

3093a3d

for more information, see https://pre-commit.ci

update

ff28f0f

carmocca approved these changes May 21, 2024

View reviewed changes

src/lightning/fabric/plugins/precision/bitsandbytes.py Show resolved Hide resolved

awaelchli changed the title ~~WIP: Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized~~ Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized May 21, 2024

add chlog

340d721

awaelchli requested a review from williamFalcon as a code owner May 21, 2024 16:39

github-actions bot added the pl Generic label for PyTorch Lightning package label May 21, 2024

lantiga approved these changes May 21, 2024

View reviewed changes

mergify bot added the ready PRs ready to be merged label May 21, 2024

lantiga merged commit 7e87ce0 into master May 21, 2024
121 of 122 checks passed

lantiga deleted the bugfix/load-quantized-ckpt branch May 21, 2024 17:46

awaelchli added a commit that referenced this pull request May 22, 2024

Fix state dict loading in bitsandbytes plugin when checkpoint is alre…

e1ee230

…ady quantized (#19886)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized #19886

Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized #19886

awaelchli commented May 21, 2024 •

edited by github-actions bot

Loading

github-actions bot commented May 21, 2024 •

edited

Loading

codecov bot commented May 21, 2024 •

edited

Loading

lantiga left a comment

Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized #19886

Fix state dict loading in bitsandbytes plugin when checkpoint is already quantized #19886

Conversation

awaelchli commented May 21, 2024 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented May 21, 2024 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented May 21, 2024 • edited Loading

Codecov Report

lantiga left a comment

Choose a reason for hiding this comment

awaelchli commented May 21, 2024 •

edited by github-actions bot

Loading

github-actions bot commented May 21, 2024 •

edited

Loading

codecov bot commented May 21, 2024 •

edited

Loading