# Welcome to ⚡ Lightning Thunder
Lightning Thunder is a source-to-source compiler for PyTorch.
It makes PyTorch programs faster both on single accelerators or in distributed settings.
Thunder aims to be usable, understandable, and extensible.
## Performance
Thunder can achieve significant speedups over standard PyTorch eager code, through the compounding effects of optimizations and the use of best in class executors. Here is an example of the pretraining throughput for Llama 2 7B as implemented in [LitGPT](https://github.com/Lightning-AI/litgpt).
We achieve a 40% speedup in training throughput compared to eager code on H100 using a combination of executors including nvFuser, torch.compile, cuDNN, and TransformerEngine FP8.
Thunder supports distributed strategies like DDP and FSDP (ZeRO2 and ZeRO3). Here is the normalized throughput measured for Llama 2 7B (this time without FP8 mixed precision, support for FSDP is underway).
**NOTE: Lightning Thunder is alpha.** Feel free to get involved, expect a few bumps along the way.
## Start with Thunder
Try Thunder without installing by using our [Zero to Thunder Tutorial Studio](https://lightning.ai/lightning-ai/studios/zero-to-thunder-tutorial).
## Install Thunder
Install [nvFuser](https://github.com/NVIDIA/Fuser) nightly, which will also install the matching PyTorch nightly:
```bash
pip install --pre 'nvfuser-cu121[torch]' --extra-index-url https://pypi.nvidia.com
```
Install Thunder:
```bash
pip install git+https://github.com/Lightning-AI/lightning-thunder.git
```
or install from the local repo:
```bash
pip install .
```
## Hello World
Here is a simple example of how Thunder lets you compile and run PyTorch code:
```python
import torch
import thunder
def foo(a, b):
return a + b
jfoo = thunder.jit(foo)
a = torch.full((2, 2), 1)
b = torch.full((2, 2), 3)
result = jfoo(a, b)
print(result)
# prints
# tensor(
# [[4, 4]
# [4, 4]])
```
The compiled function `jfoo` takes and returns PyTorch tensors, just like the original function, so modules and functions compiled by Thunder can be used as part of larger PyTorch programs.
## Running training
Thunder is in its early stages, it should not be used for production runs yet.
However, it can already deliver outstanding performance on models supported by [LitGPT](https://github.com/Lightning-AI/lit-gpt), such as Mistral, Llama2, Gemma, Falcon, and derivatives.
Run training loop for Llama, single-GPU:
```bash
python examples/lit-gpt/train.py
```
Run training loop for Llama, multi-GPU, using FSDP:
```bash
python examples/lit-gpt/train_fsdp.py
```
See [README.md](examples/lit-gpt/README.md) for details on running LitGPT with Thunder.
## What's in the box
Given a python callable or PyTorch module, Thunder can generate an optimized program that:
- Computes its forward and backward passes
- Coalesces operations into efficient fusion regions
- Dispatches computations to optimized kernels
- Distributes computations optimally across machines
To do so, Thunder ships with:
- A JIT for acquiring Python programs targeting PyTorch and custom operations
- A multi-level IR to represent operations as a trace of a reduced op-set
- An extensible set of transformations on the trace, such as `grad`, fusions, distributed (like `ddp`, `fsdp`), functional (like `vmap`, `vjp`, `jvp`)
- A way to dispatch operations to an extensible collection of executors
Thunder is written entirely in Python. Even its trace is represented as valid Python at all stages of transformation. This allows unprecedented levels of introspection and extensibility.
Thunder doesn't generate code for accelerators directly. It acquires and transforms user programs so that it's possible to optimally select or generate device code using fast executors like:
- [torch.compile](https://pytorch.org/get-started/pytorch-2.0/)
- [nvFuser](https://github.com/NVIDIA/Fuser)
- [cuDNN](https://developer.nvidia.com/cudnn)
- [Apex](https://github.com/NVIDIA/apex)
- [TransformerEngine](https://github.com/NVIDIA/TransformerEngine)
- [PyTorch eager](https://github.com/pytorch/pytorch)
- custom kernels, including those written with [OpenAI Triton](https://github.com/openai/triton)
Modules and functions compiled with Thunder fully interoperate with vanilla PyTorch and support PyTorch's autograd. Also, Thunder works alongside torch.compile to leverage its state-of-the-art optimizations.
## Build the documentation
Docs are currently not hosted publicly. However you can build them locally really quickly:
```bash
make docs
```
and point your browser to the generated docs at `docs/build/index.html`.
## Develop and run tests
You can set up your environment for developing Thunder by installing the development requirements:
```bash
pip install -r requirements/devel.txt
```
Install Thunder as an editable package (optional):
```bash
pip install -e .
```
Now you run tests:
```bash
pytest thunder/tests
```
Thunder is very thoroughly tested, so expect this to take a while.
## License
Lightning Thunder is released under the [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) license.
See LICENSE file for details.
[](https://github.com/Lightning-AI/lightning-thunder/actions/workflows/ci-testing.yml)
[](https://github.com/Lightning-AI/lightning-thunder/actions/workflows/ci-checks.yml)
[](https://lightning-thunder.readthedocs.io/en/latest/?badge=latest)
[](https://results.pre-commit.ci/latest/github/Lightning-AI/lightning-thunder/main?badge_token=mqheL1-cTn-280Vx4cJUdg)
Raw data
{
"_id": null,
"home_page": "https://github.com/Lightning-AI/lightning-thunder",
"name": "lightning-thunder",
"maintainer": null,
"docs_url": null,
"requires_python": "<3.12,>=3.10",
"maintainer_email": null,
"keywords": "deep learning, AI",
"author": "Lightning-AI et al",
"author_email": "community@lightning.ai",
"download_url": "https://files.pythonhosted.org/packages/90/21/26d3cdf52c8322c9fb6b2324d9060b8349b8369881106c8e09cafd7b32a8/lightning-thunder-0.1.0.tar.gz",
"platform": null,
"description": "\n\n# Welcome to \u26a1 Lightning Thunder\n\nLightning Thunder is a source-to-source compiler for PyTorch.\n\nIt makes PyTorch programs faster both on single accelerators or in distributed settings.\n\nThunder aims to be usable, understandable, and extensible.\n\n## Performance\n\nThunder can achieve significant speedups over standard PyTorch eager code, through the compounding effects of optimizations and the use of best in class executors. Here is an example of the pretraining throughput for Llama 2 7B as implemented in [LitGPT](https://github.com/Lightning-AI/litgpt).\n\n\n\nWe achieve a 40% speedup in training throughput compared to eager code on H100 using a combination of executors including nvFuser, torch.compile, cuDNN, and TransformerEngine FP8.\n\nThunder supports distributed strategies like DDP and FSDP (ZeRO2 and ZeRO3). Here is the normalized throughput measured for Llama 2 7B (this time without FP8 mixed precision, support for FSDP is underway).\n\n\n\n**NOTE: Lightning Thunder is alpha.** Feel free to get involved, expect a few bumps along the way.\n\n## Start with Thunder\n\nTry Thunder without installing by using our [Zero to Thunder Tutorial Studio](https://lightning.ai/lightning-ai/studios/zero-to-thunder-tutorial).\n\n## Install Thunder\n\nInstall [nvFuser](https://github.com/NVIDIA/Fuser) nightly, which will also install the matching PyTorch nightly:\n\n```bash\npip install --pre 'nvfuser-cu121[torch]' --extra-index-url https://pypi.nvidia.com\n```\n\nInstall Thunder:\n\n```bash\npip install git+https://github.com/Lightning-AI/lightning-thunder.git\n```\n\nor install from the local repo:\n\n```bash\npip install .\n```\n\n## Hello World\n\nHere is a simple example of how Thunder lets you compile and run PyTorch code:\n\n```python\nimport torch\nimport thunder\n\n\ndef foo(a, b):\n return a + b\n\n\njfoo = thunder.jit(foo)\n\na = torch.full((2, 2), 1)\nb = torch.full((2, 2), 3)\n\nresult = jfoo(a, b)\n\nprint(result)\n\n# prints\n# tensor(\n# [[4, 4]\n# [4, 4]])\n```\n\nThe compiled function `jfoo` takes and returns PyTorch tensors, just like the original function, so modules and functions compiled by Thunder can be used as part of larger PyTorch programs.\n\n## Running training\n\nThunder is in its early stages, it should not be used for production runs yet.\n\nHowever, it can already deliver outstanding performance on models supported by [LitGPT](https://github.com/Lightning-AI/lit-gpt), such as Mistral, Llama2, Gemma, Falcon, and derivatives.\n\nRun training loop for Llama, single-GPU:\n\n```bash\npython examples/lit-gpt/train.py\n```\n\nRun training loop for Llama, multi-GPU, using FSDP:\n\n```bash\npython examples/lit-gpt/train_fsdp.py\n```\n\nSee [README.md](examples/lit-gpt/README.md) for details on running LitGPT with Thunder.\n\n## What's in the box\n\nGiven a python callable or PyTorch module, Thunder can generate an optimized program that:\n\n- Computes its forward and backward passes\n- Coalesces operations into efficient fusion regions\n- Dispatches computations to optimized kernels\n- Distributes computations optimally across machines\n\nTo do so, Thunder ships with:\n\n- A JIT for acquiring Python programs targeting PyTorch and custom operations\n- A multi-level IR to represent operations as a trace of a reduced op-set\n- An extensible set of transformations on the trace, such as `grad`, fusions, distributed (like `ddp`, `fsdp`), functional (like `vmap`, `vjp`, `jvp`)\n- A way to dispatch operations to an extensible collection of executors\n\nThunder is written entirely in Python. Even its trace is represented as valid Python at all stages of transformation. This allows unprecedented levels of introspection and extensibility.\n\nThunder doesn't generate code for accelerators directly. It acquires and transforms user programs so that it's possible to optimally select or generate device code using fast executors like:\n\n- [torch.compile](https://pytorch.org/get-started/pytorch-2.0/)\n- [nvFuser](https://github.com/NVIDIA/Fuser)\n- [cuDNN](https://developer.nvidia.com/cudnn)\n- [Apex](https://github.com/NVIDIA/apex)\n- [TransformerEngine](https://github.com/NVIDIA/TransformerEngine)\n- [PyTorch eager](https://github.com/pytorch/pytorch)\n- custom kernels, including those written with [OpenAI Triton](https://github.com/openai/triton)\n\nModules and functions compiled with Thunder fully interoperate with vanilla PyTorch and support PyTorch's autograd. Also, Thunder works alongside torch.compile to leverage its state-of-the-art optimizations.\n\n## Build the documentation\n\nDocs are currently not hosted publicly. However you can build them locally really quickly:\n\n```bash\nmake docs\n```\n\nand point your browser to the generated docs at `docs/build/index.html`.\n\n## Develop and run tests\n\nYou can set up your environment for developing Thunder by installing the development requirements:\n\n```bash\npip install -r requirements/devel.txt\n```\n\nInstall Thunder as an editable package (optional):\n\n```bash\npip install -e .\n```\n\nNow you run tests:\n\n```bash\npytest thunder/tests\n```\n\nThunder is very thoroughly tested, so expect this to take a while.\n\n## License\n\nLightning Thunder is released under the [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) license.\nSee LICENSE file for details.\n\n[](https://github.com/Lightning-AI/lightning-thunder/actions/workflows/ci-testing.yml)\n[](https://github.com/Lightning-AI/lightning-thunder/actions/workflows/ci-checks.yml)\n[](https://lightning-thunder.readthedocs.io/en/latest/?badge=latest)\n[](https://results.pre-commit.ci/latest/github/Lightning-AI/lightning-thunder/main?badge_token=mqheL1-cTn-280Vx4cJUdg)\n",
"bugtrack_url": null,
"license": "Apache 2.0",
"summary": "Lightning Thunder project.",
"version": "0.1.0",
"project_urls": {
"Bug Tracker": "https://github.com/Lightning-AI/lightning-thunder/issues",
"Documentation": "https://lightning-thunder.rtfd.io/en/latest/",
"Download": "https://github.com/Lightning-AI/lightning-thunder",
"Homepage": "https://github.com/Lightning-AI/lightning-thunder",
"Source Code": "https://github.com/Lightning-AI/lightning-thunder"
},
"split_keywords": [
"deep learning",
" ai"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "336ecfaffd25f946b0db76fb813d91125871b84b1498c1e00e72169012cf0cc1",
"md5": "cab06adf92a0c61cd36701c5fda0e2ae",
"sha256": "dd95bbfc2fbd81a1e8cafc34d28ed94b33ae77d6656c617887af3dc1091cd550"
},
"downloads": -1,
"filename": "lightning_thunder-0.1.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "cab06adf92a0c61cd36701c5fda0e2ae",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": "<3.12,>=3.10",
"size": 483953,
"upload_time": "2024-03-20T19:19:05",
"upload_time_iso_8601": "2024-03-20T19:19:05.922720Z",
"url": "https://files.pythonhosted.org/packages/33/6e/cfaffd25f946b0db76fb813d91125871b84b1498c1e00e72169012cf0cc1/lightning_thunder-0.1.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "902126d3cdf52c8322c9fb6b2324d9060b8349b8369881106c8e09cafd7b32a8",
"md5": "bbf08857b2d912d10a9dc26dd209fc46",
"sha256": "80ccab2d9c9b1ad73f3fc3c05dd87bd996e25b07f215c82603106ffabe5b8216"
},
"downloads": -1,
"filename": "lightning-thunder-0.1.0.tar.gz",
"has_sig": false,
"md5_digest": "bbf08857b2d912d10a9dc26dd209fc46",
"packagetype": "sdist",
"python_version": "source",
"requires_python": "<3.12,>=3.10",
"size": 447089,
"upload_time": "2024-03-20T19:19:08",
"upload_time_iso_8601": "2024-03-20T19:19:08.389585Z",
"url": "https://files.pythonhosted.org/packages/90/21/26d3cdf52c8322c9fb6b2324d9060b8349b8369881106c8e09cafd7b32a8/lightning-thunder-0.1.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-03-20 19:19:08",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "Lightning-AI",
"github_project": "lightning-thunder",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [],
"lcname": "lightning-thunder"
}
JSON
