Add `Trainer.init_module` and `LightningModule.configure_model` #18004

carmocca · 2023-07-06T18:27:01Z

What does this PR do?

Port of #17462 for the Trainer

Key implementation differences:

No init_context support for DeepSpeedPrecisionPlugin because it doesn't support true precision (unlike Fabric)
Trainer only exposes init_module which functionally matches fabric.init_tensor because otherwise it would require that processes have been launched when it is called, but that won't happen until trainer.fit().
[NEW HOOK] Instead of Trainer.init_module, users are suggested LightningModule.configure_model with sharded strategies
LightningModule.configure_sharded_model is deprecated in favor of LightningModule.configure_model

Generated docs: https://pytorch-lightning--18004.org.readthedocs.build/

Enables Lightning-AI/litgpt#228

cc @Borda @awaelchli @carmocca @justusschock

docs/source-pytorch/advanced/model_parallel.rst

tests/tests_pytorch/strategies/test_common.py

github-actions · 2023-07-13T02:24:33Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 1.11)	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.8, 1.11, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11)	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.8, 1.11, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.11)	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.8, 1.11, oldest)	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	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py, tests/tests_pytorch/deprecated_api/test_no_removal_version.py, tests/tests_pytorch/models/test_hooks.py, tests/tests_pytorch/plugins/test_double_plugin.py, tests/tests_pytorch/strategies/test_common.py, tests/tests_pytorch/strategies/test_ddp.py, tests/tests_pytorch/strategies/test_deepspeed_strategy.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/trainer/logging_/test_logger_connector.py, tests/tests_pytorch/trainer/test_trainer.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py, tests/tests_pytorch/deprecated_api/test_no_removal_version.py, tests/tests_pytorch/models/test_hooks.py, tests/tests_pytorch/plugins/test_double_plugin.py, tests/tests_pytorch/strategies/test_common.py, tests/tests_pytorch/strategies/test_ddp.py, tests/tests_pytorch/strategies/test_deepspeed_strategy.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/trainer/logging_/test_logger_connector.py, tests/tests_pytorch/trainer/test_trainer.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/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py.

🟢 pytorch_lightning: Docs

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

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py, docs/source-pytorch/advanced/model_parallel.rst, docs/source-pytorch/common/lightning_module.rst, docs/source-pytorch/integrations/strategies/colossalai.rst.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, lightning, 3.8, 1.11)	success	✅
fabric-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
fabric-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.11)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.9, 1.12)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.11)	success	✅
fabric-cpu (windows-2022, lightning, 3.9, 1.12)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (macOS-11, fabric, 3.8, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 1.13)	success	✅
fabric-cpu (windows-2022, fabric, 3.8, 1.13)	success	✅

These checks are required after the changes to tests/tests_fabric/strategies/test_fsdp_integration.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to tests/tests_fabric/strategies/test_fsdp_integration.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.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.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.10)	success	✅
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.10)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.10)	success	✅

These checks are required after the changes to src/lightning/pytorch/core/hooks.py, src/lightning/pytorch/plugins/precision/double.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/strategy.py, src/lightning/pytorch/trainer/call.py, src/lightning/pytorch/trainer/configuration_validator.py, src/lightning/pytorch/trainer/connectors/logger_connector/fx_validator.py, src/lightning/pytorch/trainer/trainer.py.

🟢 link-check

Check ID	Status
check-md-links / markdown-link-check	success	✅

These checks are required after the changes to src/lightning/fabric/CHANGELOG.md, src/lightning/pytorch/CHANGELOG.md.

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.

awaelchli

Looks great. Minor questions and comments only

docs/source-pytorch/advanced/model_parallel.rst

src/lightning/fabric/CHANGELOG.md

src/lightning/pytorch/trainer/call.py

src/lightning/pytorch/trainer/trainer.py

tests/tests_pytorch/strategies/test_fsdp.py

carmocca added feature Is an improvement or enhancement strategy: deepspeed trainer strategy strategy: fsdp Fully Sharded Data Parallel pl Generic label for PyTorch Lightning package labels Jul 6, 2023

carmocca added this to the 2.1 milestone Jul 6, 2023

carmocca self-assigned this Jul 6, 2023

github-actions bot added fabric lightning.fabric.Fabric app Generic label for Lightning App package labels Jul 6, 2023

carmocca force-pushed the carmocca/trainer-init-module branch from ceb470f to 7ffddf2 Compare July 6, 2023 18:27

github-actions bot removed the app Generic label for Lightning App package label Jul 6, 2023

carmocca force-pushed the carmocca/trainer-init-module branch from 2535351 to f0939ea Compare July 6, 2023 18:57

carmocca commented Jul 6, 2023

View reviewed changes

docs/source-pytorch/advanced/model_parallel.rst Show resolved Hide resolved

tests/tests_pytorch/strategies/test_common.py Show resolved Hide resolved

carmocca force-pushed the carmocca/trainer-init-module branch from dbf2f54 to eb942ea Compare July 6, 2023 19:39

This was referenced Jul 6, 2023

Fabric always wraps the root module with FSDP #18005

Closed

Don't reapply activation checkpointing #18006

Merged

carmocca changed the title ~~Add trainer.init_module and trainer.init_tensor~~ Add trainer.init_module Jul 10, 2023

carmocca force-pushed the carmocca/trainer-init-module branch 7 times, most recently from 6fed2f8 to 4613d7c Compare July 11, 2023 00:53

carmocca mentioned this pull request Jul 11, 2023

Pretrain script with the Trainer Lightning-AI/litgpt#228

Merged

carmocca force-pushed the carmocca/trainer-init-module branch from b7562a3 to 49629af Compare July 11, 2023 01:18

carmocca added 3 commits July 13, 2023 02:27

init_module

b2f2382

CHANGELOG

8574d6a

mypy

3b78097

carmocca requested review from awaelchli, justusschock, williamFalcon, tchaton, edenlightning, lantiga and Borda as code owners July 13, 2023 02:24

Fix CPU

6f93a55

carmocca mentioned this pull request Jul 13, 2023

Minor miscellaneous fixes #18068

Merged

carmocca added 6 commits July 13, 2023 04:57

Update

c000843

Revert

c4f78d4

FSDP fix

b543c59

Fix

7d28150

Merge branch 'master' into carmocca/trainer-init-module

1a68d4c

Fix import

fd92683

carmocca changed the title ~~Add trainer.init_module~~ Add Trainer.init_module and LightningModule.configure_model Jul 13, 2023

carmocca added hooks Related to the hooks API and removed fabric lightning.fabric.Fabric strategy: deepspeed strategy: fsdp Fully Sharded Data Parallel labels Jul 13, 2023

awaelchli reviewed Jul 13, 2023

View reviewed changes

nit

73a2421

github-actions bot added the fabric lightning.fabric.Fabric label Jul 14, 2023

awaelchli approved these changes Jul 14, 2023

View reviewed changes

justusschock approved these changes Jul 14, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Jul 14, 2023

carmocca mentioned this pull request Jul 14, 2023

Revamp model parallel docs in PL #18083

Closed

carmocca merged commit 340eecd into master Jul 14, 2023
101 checks passed

carmocca deleted the carmocca/trainer-init-module branch July 14, 2023 17:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `Trainer.init_module` and `LightningModule.configure_model` #18004

Add `Trainer.init_module` and `LightningModule.configure_model` #18004

carmocca commented Jul 6, 2023 •

edited

github-actions bot commented Jul 13, 2023 •

edited

awaelchli left a comment

Add Trainer.init_module and LightningModule.configure_model #18004

Add Trainer.init_module and LightningModule.configure_model #18004

Conversation

carmocca commented Jul 6, 2023 • edited

What does this PR do?

github-actions bot commented Jul 13, 2023 • edited

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli left a comment

Choose a reason for hiding this comment

Add `Trainer.init_module` and `LightningModule.configure_model` #18004

Add `Trainer.init_module` and `LightningModule.configure_model` #18004

carmocca commented Jul 6, 2023 •

edited

github-actions bot commented Jul 13, 2023 •

edited