Support NVIDIA's Transformer Engine #17597

carmocca · 2023-05-09T15:34:30Z

What does this PR do?

Fixes #17172

Tested in Lightning-AI/litgpt#123

cc @Borda @carmocca @justusschock @awaelchli

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

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

carmocca

Thanks for the reviews!

I'm not sure if this will land, or at least with the conversion logic as it seems very flawed.
However, it's as flawed as for any technique that needs to do it.

We could still consider the plugin even if the conversion is not done automatically. Or we could just move on with this.

This is an important unsolved problem in pytorch. Layers are not composable at all.

What do you guys think?

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

github-actions · 2023-06-26T00:40:07Z

⚡ 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/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py.

🟢 fabric: Docs

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

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py, docs/source-fabric/api/fabric_args.rst, docs/source-fabric/fundamentals/precision.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 src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py, tests/tests_fabric/test_connector.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.py, tests/tests_fabric/test_connector.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.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/fabric/connector.py, src/lightning/fabric/plugins/precision/__init__.py, src/lightning/fabric/plugins/precision/precision.py, src/lightning/fabric/plugins/precision/transformer_engine.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.

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

Minor nits and questions. Great work!

docs/source-fabric/fundamentals/precision.rst

src/lightning/fabric/CHANGELOG.md

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

Support NVIDIA's Transformer Engine

3b325f0

carmocca added feature Is an improvement or enhancement fabric lightning.fabric.Fabric plugin pl Generic label for PyTorch Lightning package labels May 9, 2023

carmocca added this to the 2.1 milestone May 9, 2023

carmocca self-assigned this May 9, 2023

github-actions bot removed the pl Generic label for PyTorch Lightning package label May 9, 2023

carmocca commented May 9, 2023

View reviewed changes

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

Fix argument, changelog

c5956c7

carmocca mentioned this pull request May 9, 2023

H100 Transformer Engine implementation Lightning-AI/lit-llama#249

Open

Borda approved these changes May 9, 2023

View reviewed changes

carmocca added 2 commits May 10, 2023 06:33

convert module

2cdd233

Replacement

6290c7a

carmocca commented May 10, 2023

View reviewed changes

justusschock reviewed May 11, 2023

View reviewed changes

awaelchli reviewed May 11, 2023

View reviewed changes

src/lightning/fabric/plugins/precision/fp8_transformer_engine.py Outdated Show resolved Hide resolved

carmocca commented May 11, 2023

View reviewed changes

carmocca commented May 22, 2023

View reviewed changes

src/lightning/fabric/plugins/precision/fp8_transformer_engine.py Outdated Show resolved Hide resolved

Silly bug

5e56c2a

carmocca commented May 23, 2023

View reviewed changes

src/lightning/fabric/plugins/precision/fp8_transformer_engine.py Outdated Show resolved Hide resolved

carmocca added 9 commits June 9, 2023 07:01

Remove device validation. Done already by TE

6e6ac42

Merge branch 'master' into carmocca/transformer-engine

9ae4287

Merge branch 'master' into carmocca/transformer-engine

152c054

Mapping

c18cf10

condition

a9426c1

init_module

86d7136

Update test

9905a06

Cover recursion case

1ea5517

Update fp8_transformer_engine.py

8944c5e

carmocca added 6 commits June 9, 2023 11:52

Fix 270

c75dbfa

Fix layer inheritance

3e0b7c4

Docs

a288f21

mypy

4b57d20

Docs

f33d729

Merge branch 'master' into carmocca/transformer-engine

d541de7

carmocca marked this pull request as ready for review June 24, 2023 23:33

carmocca requested review from edenlightning and lantiga as code owners June 24, 2023 23:33

dtype arg

8d154dd

Remove

36e0370

Borda requested review from justusschock and awaelchli June 26, 2023 07:27

awaelchli approved these changes Jun 27, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Jun 27, 2023

Borda reviewed Jun 28, 2023

View reviewed changes

src/lightning/fabric/plugins/precision/fp8_transformer_engine.py Outdated Show resolved Hide resolved

carmocca requested a review from williamFalcon as a code owner July 4, 2023 02:03

github-actions bot added the pl Generic label for PyTorch Lightning package label Jul 4, 2023

Review comments

b81ffe6

carmocca force-pushed the carmocca/transformer-engine branch from 176d37a to b81ffe6 Compare July 4, 2023 02:04

Merge branch 'master' into carmocca/transformer-engine

91d42de

github-actions bot removed the pl Generic label for PyTorch Lightning package label Jul 4, 2023

Borda and others added 4 commits July 12, 2023 11:55

Merge branch 'master' into carmocca/transformer-engine

743ba10

Merge branch 'master' into carmocca/transformer-engine

ad1c67c

Fix import

7d14c2f

Resolve issue

4a131d7

carmocca merged commit 071f858 into master Jul 19, 2023

carmocca deleted the carmocca/transformer-engine branch July 19, 2023 16:22

MF-FOOM mentioned this pull request Aug 25, 2023

fix: support transformer-engine (h100 fp8) in speed_monitor Lightning-AI/litgpt#468

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support NVIDIA's Transformer Engine #17597

Support NVIDIA's Transformer Engine #17597

carmocca commented May 9, 2023 •

edited by github-actions bot

Loading

carmocca left a comment

github-actions bot commented Jun 26, 2023 •

edited

Loading

awaelchli left a comment

Support NVIDIA's Transformer Engine #17597

Support NVIDIA's Transformer Engine #17597

Conversation

carmocca commented May 9, 2023 • edited by github-actions bot Loading

What does this PR do?

carmocca left a comment

Choose a reason for hiding this comment

github-actions bot commented Jun 26, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli left a comment

Choose a reason for hiding this comment

carmocca commented May 9, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jun 26, 2023 •

edited

Loading