Fix parameter count in ModelSummary when parameters are DTensors #20163

awaelchli · 2024-08-04T23:23:22Z

What does this PR do?

Fixes #20151

Running

python examples/pytorch/tensor_parallel/train.py

Before:

  | Name  | Type            | Params | Mode 
--------------------------------------------------
0 | model | FSDPTransformer | 0      | train
--------------------------------------------------
0         Trainable params
0         Non-trainable params
0         Total params
0.000     Total estimated model params size (MB)
421       Modules in train mode
0         Modules in eval mode

(shows 0 parameters)

After:

  | Name  | Type            | Params | Mode 
--------------------------------------------------
0 | model | FSDPTransformer | 6.7 B  | train
--------------------------------------------------
6.7 B     Trainable params
0         Non-trainable params
6.7 B     Total params
26,953.662Total estimated model params size (MB)
421       Modules in train mode
0         Modules in eval mode

(shows correct counts)

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

cc @carmocca @justusschock @awaelchli

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

github-actions · 2024-08-04T23:33:18Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-13, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (macOS-14, lightning, 3.10, 2.1)	success	✅
pl-cpu (macOS-14, lightning, 3.11, 2.2)	success	✅
pl-cpu (macOS-14, lightning, 3.11, 2.3)	success	✅
pl-cpu (macOS-14, lightning, 3.12, 2.4)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.11, 2.2)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.11, 2.3)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.12, 2.4)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.1)	success	✅
pl-cpu (windows-2022, lightning, 3.11, 2.2)	success	✅
pl-cpu (windows-2022, lightning, 3.11, 2.3)	success	✅
pl-cpu (windows-2022, lightning, 3.12, 2.4)	success	✅
pl-cpu (macOS-14, pytorch, 3.9, 2.1)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.9, 2.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.9, 2.1)	success	✅
pl-cpu (macOS-12, pytorch, 3.10, 2.1)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.10, 2.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.10, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/utilities/distributed.py, src/lightning/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py, tests/tests_pytorch/utilities/test_model_summary.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/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py, tests/tests_pytorch/utilities/test_model_summary.py, src/lightning/fabric/utilities/distributed.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/utilities/distributed.py, src/lightning/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.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/utilities/distributed.py.

🟢 pytorch_lightning: Docs

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

These checks are required after the changes to src/lightning/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-13, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (macOS-14, lightning, 3.10, 2.1)	success	✅
fabric-cpu (macOS-14, lightning, 3.11, 2.2)	success	✅
fabric-cpu (macOS-14, lightning, 3.11, 2.3)	success	✅
fabric-cpu (macOS-14, lightning, 3.12, 2.4)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 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 (ubuntu-20.04, lightning, 3.12, 2.4)	success	✅
fabric-cpu (windows-2022, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 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 (windows-2022, lightning, 3.12, 2.4)	success	✅
fabric-cpu (macOS-14, fabric, 3.9, 2.1)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.9, 2.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.9, 2.1)	success	✅
fabric-cpu (macOS-12, fabric, 3.10, 2.1)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.10, 2.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.10, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/utilities/distributed.py, tests/tests_fabric/utilities/test_distributed.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/utilities/distributed.py, tests/tests_fabric/utilities/test_distributed.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/utilities/distributed.py, src/lightning/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py.

🟢 install

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

These checks are required after the changes to src/lightning/fabric/utilities/distributed.py, src/lightning/pytorch/utilities/model_summary/model_summary.py, src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.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-08-04T23:50:07Z

Codecov Report

Attention: Patch coverage is 94.11765% with 1 line in your changes missing coverage. Please review.

Project coverage is 89%. Comparing base (d4de8e2) to head (d4cdde4).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #20163   +/-   ##
=======================================
- Coverage      89%      89%   -0%     
=======================================
  Files         267      267           
  Lines       23063    23071    +8     
=======================================
+ Hits        20565    20572    +7     
- Misses       2498     2499    +1

…htning-AI#20163)

awaelchli added 4 commits August 4, 2024 18:24

update

0fc4dde

update

0a69fcf

add simple test

cf1ad7d

add test

a3a408b

github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Aug 4, 2024

pre-commit-ci bot and others added 3 commits August 4, 2024 23:23

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

9d6dfc8

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

update

9e70ae2

update

2aed8db

awaelchli added this to the 2.4 milestone Aug 4, 2024

awaelchli added strategy: fsdp Fully Sharded Data Parallel fun Staff contributions outside working hours - to differentiate from the "community" label callback: model summary bug Something isn't working labels Aug 4, 2024

awaelchli marked this pull request as ready for review August 4, 2024 23:32

awaelchli requested review from lantiga, Borda, tchaton and justusschock as code owners August 4, 2024 23:32

add comment

d4cdde4

Borda approved these changes Aug 5, 2024

View reviewed changes

awaelchli merged commit 345450b into master Aug 5, 2024
97 of 98 checks passed

awaelchli deleted the bugfix/model-summary-dtensor-shapes branch August 5, 2024 14:57

ammyk9 pushed a commit to ammyk9/pytorch-lightning that referenced this pull request Aug 6, 2024

Fix parameter count in ModelSummary when parameters are DTensors (Lig…

25954a1

…htning-AI#20163)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix parameter count in ModelSummary when parameters are DTensors #20163

Fix parameter count in ModelSummary when parameters are DTensors #20163

awaelchli commented Aug 4, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Aug 4, 2024 •

edited

Loading

codecov bot commented Aug 4, 2024 •

edited

Loading

Fix parameter count in ModelSummary when parameters are DTensors #20163

Fix parameter count in ModelSummary when parameters are DTensors #20163

Conversation

awaelchli commented Aug 4, 2024 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented Aug 4, 2024 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Aug 4, 2024 • edited Loading

Codecov Report

awaelchli commented Aug 4, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Aug 4, 2024 •

edited

Loading

codecov bot commented Aug 4, 2024 •

edited

Loading