Fix zero-grad behavior when entering the validation loop #18710

awaelchli · 2023-10-03T22:52:54Z

What does this PR do?

Fixes #18713

This PR does the following:

Fixes an issue in which the validation loop would erase gradients when validation happened in the middle of a gradient accumulation phase (see the added test).
The zero_grad() call is now wrapped by a new hook LightningModule.on_validation_model_zero_grad() that the user can override if they don't want the default behavior (Should the model's grads be cleared before entering the validation loop? #18713).
We still ensure that we clear the gradients before entering any of the trainer stages trainer.validate/test/predict as a precaution (no change there).

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

cc @carmocca @awaelchli @Borda @justusschock

src/lightning/pytorch/trainer/trainer.py

tests/tests_pytorch/loops/test_loops.py

github-actions · 2023-10-05T02:03:36Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 1.12, oldest)	success	✅
pl-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.12, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 1.12)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	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 (windows-2022, lightning, 3.8, 1.12, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 1.12)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.1)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 1.13)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.13)	success	✅
pl-cpu (windows-2022, pytorch, 3.8, 1.13)	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/pytorch/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py, tests/tests_pytorch/loops/test_loops.py, tests/tests_pytorch/models/test_hooks.py, tests/tests_pytorch/trainer/logging_/test_logger_connector.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
[pytorch-lightning (GPUs) (testing Lightning	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=178365&view=logs&jobId=b2def368-7fa8-5edf-f15e-38e7ac88d76c)	success
[pytorch-lightning (GPUs) (testing PyTorch	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=178365&view=logs&jobId=3f274fac-2e11-54ca-487e-194c91f3ae9f)	success

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py, tests/tests_pytorch/loops/test_loops.py, tests/tests_pytorch/models/test_hooks.py, tests/tests_pytorch/trainer/logging_/test_logger_connector.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.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/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.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/pytorch/core/hooks.py, src/lightning/pytorch/loops/evaluation_loop.py, src/lightning/pytorch/loops/prediction_loop.py, src/lightning/pytorch/loops/training_epoch_loop.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.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.

src/lightning/pytorch/core/hooks.py

codecov · 2023-10-06T21:40:48Z

Codecov Report

Merging #18710 (20b7220) into master (7fd5c02) will decrease coverage by 35%.
The diff coverage is 100%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #18710      +/-   ##
==========================================
- Coverage      84%      49%     -35%     
==========================================
  Files         439      431       -8     
  Lines       34363    34222     -141     
==========================================
- Hits        28795    16819   -11976     
- Misses       5568    17403   +11835

Borda

just minor formating, feel free to ignore

src/lightning/pytorch/core/hooks.py

src/lightning/pytorch/trainer/trainer.py

Co-authored-by: Jirka Borovec <[email protected]>

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

set to none

297ba71

github-actions bot added the pl Generic label for PyTorch Lightning package label Oct 3, 2023

awaelchli added the performance label Oct 3, 2023

awaelchli added this to the v1.9.x milestone Oct 3, 2023

awaelchli added 2 commits October 4, 2023 01:18

repro

83f4153

move the zero grad call

6463194

awaelchli changed the title ~~Set grads to None in PyTorch < 2.0 at validation start~~ Set grads to None only once at the beginning of Trainer.validate/test/predict Oct 3, 2023

awaelchli added 3 commits October 4, 2023 01:47

update tests

ea7ee7e

x

55cecc1

test

8141143

awaelchli mentioned this pull request Oct 4, 2023

Should the model's grads be cleared before entering the validation loop? #18713

Closed

carmocca reviewed Oct 4, 2023

View reviewed changes

src/lightning/pytorch/trainer/trainer.py Outdated Show resolved Hide resolved

awaelchli added 4 commits October 5, 2023 03:21

add test

4ac0ad3

add hook

660af78

update tests

65ddfb3

changelog

ee6dd16

awaelchli changed the title ~~Set grads to None only once at the beginning of Trainer.validate/test/predict~~ Fix zero-grad behavior when enterint the validation loop Oct 5, 2023

awaelchli changed the title ~~Fix zero-grad behavior when enterint the validation loop~~ Fix zero-grad behavior when entering the validation loop Oct 5, 2023

awaelchli added bug Something isn't working hooks Related to the hooks API trainer: validate loops Related to the Loop API labels Oct 5, 2023

awaelchli added 4 commits October 5, 2023 03:57

condition on pytorch 2.0

a15fcda

reset

f602b38

condition on torch 2.0

ab43576

make tests a bit ugly

333e793

awaelchli commented Oct 5, 2023

View reviewed changes

tests/tests_pytorch/loops/test_loops.py Show resolved Hide resolved

awaelchli marked this pull request as ready for review October 5, 2023 02:03

awaelchli requested review from williamFalcon and tchaton as code owners October 5, 2023 02:03

awaelchli requested review from justusschock and Borda as code owners October 5, 2023 02:03

carmocca reviewed Oct 5, 2023

View reviewed changes

src/lightning/pytorch/core/hooks.py Outdated Show resolved Hide resolved

awaelchli added 2 commits October 6, 2023 23:02

update test

e84a119

fix

28890bf

carmocca approved these changes Oct 7, 2023

View reviewed changes

Borda approved these changes Oct 9, 2023

View reviewed changes

src/lightning/pytorch/core/hooks.py Outdated Show resolved Hide resolved

src/lightning/pytorch/trainer/trainer.py Outdated Show resolved Hide resolved

mergify bot added the ready PRs ready to be merged label Oct 9, 2023

awaelchli and others added 4 commits October 9, 2023 12:10

Update src/lightning/pytorch/core/hooks.py

375aa1c

Co-authored-by: Jirka Borovec <[email protected]>

Update src/lightning/pytorch/trainer/trainer.py

b329019

Co-authored-by: Jirka Borovec <[email protected]>

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

dc06b15

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

Merge branch 'master' into bugfix/val-loop-zero-grad

20b7220

awaelchli merged commit a26424e into master Oct 9, 2023
91 checks passed

awaelchli deleted the bugfix/val-loop-zero-grad branch October 9, 2023 22:07

awaelchli modified the milestones: v1.9.x, 2.1 Oct 9, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix zero-grad behavior when entering the validation loop #18710

Fix zero-grad behavior when entering the validation loop #18710

awaelchli commented Oct 3, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Oct 5, 2023 •

edited

Loading

codecov bot commented Oct 6, 2023 •

edited

Loading

Borda left a comment

Fix zero-grad behavior when entering the validation loop #18710

Fix zero-grad behavior when entering the validation loop #18710

Conversation

awaelchli commented Oct 3, 2023 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented Oct 5, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Oct 6, 2023 • edited Loading

Codecov Report

Borda left a comment

Choose a reason for hiding this comment

awaelchli commented Oct 3, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Oct 5, 2023 •

edited

Loading

codecov bot commented Oct 6, 2023 •

edited

Loading