Replace ColossalAIStrategy with external implementation #16757

awaelchli · 2023-02-15T02:39:41Z

What does this PR do?

The Colossal AI Strategy moves to an external repo where it will be maintained: Lightning-Universe/lightning-ColossalAI#1. A corresponding package will be published to PyPi from which users can install the integration.

Replaces the implementation with the one in the lightning-colossalai package.
Registeres the strategy="colossal" as a strategy if the package is installed
Produces an error message if you request the strategy but don't have the package installed

Before:

trainer = Trainer(strategy="colossalai")

# OR

from lightning.pytorch.strategies import ColossalAIStrategy
trainer = Trainer(strategy=ColossalAIStrategy(...))

Now:

pip install lightning-colossalai

trainer = Trainer(strategy="colossalai")  # only works if  `lightning-colossalai` is installed.

# OR

from lightning_colossalai import ColossalAIStrategy
trainer = Trainer(strategy=ColossalAIStrategy(...))

Follow up:

Update docs once the lightning-colossalai package is finalized and public

cc @Borda @justusschock

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

…olossal

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

…olossal

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

…olossal

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

requirements/pytorch/strategies.txt

github-actions · 2023-02-15T13:19:55Z

⚡ 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.8, 1.11, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 1.11)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.12)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 1.11)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.12)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.13)	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 requirements/pytorch/strategies.txt, src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py, tests/tests_pytorch/helpers/runif.py, tests/tests_pytorch/strategies/test_colossalai.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt, src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py, tests/tests_pytorch/helpers/runif.py, tests/tests_pytorch/strategies/test_colossalai.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
pytorch-lightning.Benchmark	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt.

🟢 pytorch_lightning: Azure HPU

Check ID	Status
pytorch-lightning (HPUs)	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt, src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py, tests/tests_pytorch/helpers/runif.py, tests/tests_pytorch/strategies/test_colossalai.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py.

🟢 pytorch_lightning: Azure IPU

Check ID	Status
pytorch-lightning (IPUs)	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt, src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py, tests/tests_pytorch/helpers/runif.py, tests/tests_pytorch/strategies/test_colossalai.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.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/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, docs/source-pytorch/api_references.rst, docs/source-pytorch/extensions/plugins.rst, requirements/pytorch/strategies.txt.

🟢 pytorch_lightning: Docker

Check ID	Status
build-cuda (3.9, 1.11, 11.3.1)	success	✅
build-cuda (3.9, 1.12, 11.6.1)	success	✅
build-cuda (3.9, 1.13, 11.7.1)	success	✅
build-hpu (1.5.0, 1.11.0)	success	✅
build-ipu (3.9, 1.13)	success	✅
build-NGC	success	✅
build-pl (3.9, 1.11, 11.3.1)	success	✅
build-pl (3.9, 1.12, 11.6.1)	success	✅
build-pl (3.9, 1.13, 11.7.1)	success	✅
build-xla (3.8, 1.12)	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to requirements/pytorch/strategies.txt, src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.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/callbacks/gradient_accumulation_scheduler.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/colossalai.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/colossalai.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, src/lightning/pytorch/utilities/imports.py, requirements/pytorch/strategies.txt.

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.

justusschock

lgtm. Will the external strategies be registered to pypi?

src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py

awaelchli · 2023-02-15T14:06:30Z

lgtm. Will the external strategies be registered to pypi?

That's the plan yes. Right now there is a chicken and egg problem but once the two PRs land we can complete the integration.

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

src/lightning/pytorch/utilities/imports.py

colossalai

7890cb2

awaelchli added breaking change Includes a breaking change strategy: colossalai (removed) Colossal-AI labels Feb 15, 2023

github-actions bot added the pl Generic label for PyTorch Lightning package label Feb 15, 2023

pre-commit-ci bot and others added 11 commits February 15, 2023 02:40

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

29cbf01

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

test

dd94853

update

88fe5ac

Merge remote-tracking branch 'origin/removal/colossal' into removal/c…

e7ffe79

…olossal

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

0ef9135

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

fix

b6be2a6

Merge remote-tracking branch 'origin/removal/colossal' into removal/c…

ff25058

…olossal

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

a2734b8

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

add git repo

270fecc

Merge remote-tracking branch 'origin/removal/colossal' into removal/c…

aed69b6

…olossal

token

c06cd0e

awaelchli added this to the 2.0 milestone Feb 15, 2023

awaelchli and others added 6 commits February 15, 2023 04:45

add test

c613678

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

9a433f5

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

update test

a743f7b

avoid circular import

03b60c2

debug

e09a920

register them early

5acf0bf

awaelchli commented Feb 15, 2023

View reviewed changes

requirements/pytorch/strategies.txt Outdated Show resolved Hide resolved

awaelchli marked this pull request as ready for review February 15, 2023 13:19

awaelchli requested review from carmocca, justusschock, williamFalcon, lantiga, Borda and tchaton as code owners February 15, 2023 13:19

awaelchli changed the title ~~Replace ColossalAIStrategy with external implementation~~ WIP: Replace ColossalAIStrategy with external implementation Feb 15, 2023

awaelchli added 2 commits February 15, 2023 14:21

docs

3d571ee

docs

0402bbd

awaelchli requested a review from edenlightning as a code owner February 15, 2023 13:31

awaelchli changed the title ~~WIP: Replace ColossalAIStrategy with external implementation~~ Replace ColossalAIStrategy with external implementation Feb 15, 2023

justusschock approved these changes Feb 15, 2023

View reviewed changes

src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py Show resolved Hide resolved

Borda approved these changes Feb 15, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Feb 15, 2023

remove debug statement

2bea54f

carmocca approved these changes Feb 15, 2023

View reviewed changes

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

src/lightning/pytorch/utilities/imports.py Outdated Show resolved Hide resolved

try requirement cache

cc9caf5

awaelchli enabled auto-merge (squash) February 15, 2023 14:39

awaelchli merged commit 83f4c83 into master Feb 15, 2023

awaelchli deleted the removal/colossal branch February 15, 2023 15:11

Borda mentioned this pull request Feb 16, 2023

Remove the BaguaStrategy #16746

Merged

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Replace ColossalAIStrategy with external implementation #16757

Replace ColossalAIStrategy with external implementation #16757

awaelchli commented Feb 15, 2023 •

edited

Loading

github-actions bot commented Feb 15, 2023 •

edited

Loading

justusschock left a comment

awaelchli commented Feb 15, 2023 •

edited

Loading

Replace ColossalAIStrategy with external implementation #16757

Replace ColossalAIStrategy with external implementation #16757

Conversation

awaelchli commented Feb 15, 2023 • edited Loading

What does this PR do?

github-actions bot commented Feb 15, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

justusschock left a comment

Choose a reason for hiding this comment

awaelchli commented Feb 15, 2023 • edited Loading

awaelchli commented Feb 15, 2023 •

edited

Loading

github-actions bot commented Feb 15, 2023 •

edited

Loading

awaelchli commented Feb 15, 2023 •

edited

Loading