Merge branch 'main' into crpa/lookaside_autograd-function-apply

Lightning-AI · Nov 5, 2024 · 627845d · 627845d
2 parents 71db6cd + 3d42c10
commit 627845d
Show file tree

Hide file tree

Showing 78 changed files with 4,121 additions and 595 deletions.
diff --git a/.azure/docker-build.yml b/.azure/docker-build.yml
@@ -40,8 +40,8 @@ jobs:
       #maxParallel: "3"
       matrix:
         # CUDA 12.1
-        "cuda 12.1 | torch 2.4.0 | cudnn FE v1.5.2":
-          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.4.0", TRITON_VERSION: "3.0.0", CUDNN_FRONTEND_VERSION: "1.5.2" }
+        "cuda 12.1 | torch 2.5.1 | cudnn FE v1.5.2":
+          { CUDA_VERSION: "12.1.1", TORCH_VERSION: "2.5.1", TRITON_VERSION: "3.1.0", CUDNN_FRONTEND_VERSION: "1.5.2" }
         "cuda 12.1 | torch 2.5 /nightly | cudnn FE v1.5.2":
           { CUDA_VERSION: "12.1.1", TORCH_VERSION: "main", TORCH_INSTALL: "source", CUDNN_FRONTEND_VERSION: "1.5.2" }
         #'cuda 12.1': # this version - '8.9.5.29-1+cuda12.1' for 'libcudnn8' was not found

diff --git a/.azure/gpu-tests.yml b/.azure/gpu-tests.yml
@@ -16,11 +16,11 @@ jobs:
     strategy:
       matrix:
         # CUDA 12.1
-        "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.4.0 | regular":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.4.0-dev"
+        "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.5.1 | regular":
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.5.1-dev"
           CUDA_VERSION_MM: "121"
-        "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.4.0 | distributed":
-          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.4.0-dev"
+        "ubuntu22.04 | cuda 12.1 | python 3.10 | torch 2.5.1 | distributed":
+          docker-image: "ubuntu22.04-cuda12.1.1-cudnn-fe1.5.2-py3.10-pt_2.5.1-dev"
           CUDA_VERSION_MM: "121"
           testing: "distributed"
         "ubuntu22.04 | cuda 12.1 | python 3.10 | torch-nightly | regular":
@@ -84,7 +84,6 @@ jobs:
             pytest thunder/tests/ \
               -m "not standalone" \
               -v --datefmt="%Y%m%d-%H:%M:%S.%f" \
-              --timeout=240 \
               --random-order-seed=42 \
               --durations=250 \
               --timeout=240 \
@@ -97,7 +96,7 @@ jobs:
           ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
             --flags=gpu,pytest,regular --name="GPU-coverage" --env=linux,azure
         condition: ne(variables['testing'], 'distributed')
-        timeoutInMinutes: "30"
+        timeoutInMinutes: "40"
         displayName: "Testing: regular"
 
       - bash: |

diff --git a/.github/workflows/ci-checks.yml b/.github/workflows/ci-checks.yml
@@ -11,19 +11,19 @@ concurrency:
 
 jobs:
   precommit-run:
-    uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.11.7
+    uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.11.8
     with:
       python-version: "3.10"
 
   check-schema:
-    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.11.7
+    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.11.8
     with:
       azure-dir: ".azure"
 
   check-package:
-    uses: Lightning-AI/utilities/.github/workflows/check-package.yml@v0.11.7
+    uses: Lightning-AI/utilities/.github/workflows/check-package.yml@v0.11.8
     with:
-      actions-ref: v0.11.3.post0
+      actions-ref: v0.11.8
       import-name: "thunder"
       artifact-name: dist-packages-${{ github.sha }}
       testing-matrix: |

diff --git a/.github/workflows/release-nightly.yml b/.github/workflows/release-nightly.yml
@@ -33,15 +33,15 @@ jobs:
       # We do this, since failures on test.pypi aren't that bad
       - name: Publish to Test PyPI
         if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'
-        uses: pypa/gh-action-pypi-publish@v1.10.2
+        uses: pypa/gh-action-pypi-publish@v1.11.0
         with:
           user: __token__
           password: ${{ secrets.test_pypi_password }}
           repository_url: https://test.pypi.org/legacy/
 
       - name: Publish distribution 📦 to PyPI
         if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'
-        uses: pypa/gh-action-pypi-publish@v1.10.2
+        uses: pypa/gh-action-pypi-publish@v1.11.0
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}
diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml
@@ -26,15 +26,15 @@ jobs:
       # We do this, since failures on test.pypi aren't that bad
       - name: Publish to Test PyPI
         if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
-        uses: pypa/gh-action-pypi-publish@v1.10.2
+        uses: pypa/gh-action-pypi-publish@v1.11.0
         with:
           user: __token__
           password: ${{ secrets.test_pypi_password }}
           repository_url: https://test.pypi.org/legacy/
 
       - name: Publish distribution 📦 to PyPI
         if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
-        uses: pypa/gh-action-pypi-publish@v1.10.2
+        uses: pypa/gh-action-pypi-publish@v1.11.0
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -8,7 +8,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: end-of-file-fixer
       - id: trailing-whitespace
@@ -23,7 +23,7 @@ repos:
       - id: detect-private-key
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.17.0
+    rev: v3.19.0
     hooks:
       - id: pyupgrade
         args: ["--py310-plus"]
@@ -38,14 +38,14 @@ repos:
         #args: ["--write-changes"] # uncomment if you want to get automatic fixing
 
   - repo: https://github.com/psf/black
-    rev: 24.8.0
+    rev: 24.10.0
     hooks:
       - id: black
         name: Black code
         exclude: "examples"
 
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.17
+    rev: 0.7.18
     hooks:
       - id: mdformat
         additional_dependencies:
@@ -55,7 +55,7 @@ repos:
         exclude: "examples"
 
   - repo: https://github.com/sphinx-contrib/sphinx-lint
-    rev: v0.9.1
+    rev: v1.0.0
     hooks:
       - id: sphinx-lint
 

diff --git a/README.md b/README.md
@@ -73,13 +73,19 @@ The easiest way to get started with Thunder, requiring no extra installations or
 
 ## Install Thunder
 
-To use Thunder on your local machine:
+Thunder is in alpha and the latest development is happening on the `main` branch. You can install the latest version of Thunder from the `main` branch as follows:
 
-- install [nvFuser](https://github.com/NVIDIA/Fuser) and PyTorch stable together as follows:
+```bash
+pip install git+https://github.com/Lightning-AI/lightning-thunder.git@main
+```
+
+To achieve the best performance, you can install Thunder with the following additional dependencies:
+
+- install prerelease [nvFuser](https://github.com/NVIDIA/Fuser) built for PyTorch 2.5.1 as follows:
 
 ```bash
-# install nvFuser which installs the matching stable PyTorch
-pip install --pre nvfuser-cu121-torch24
+# install nvFuser built for the matching stable PyTorch
+pip install --pre nvfuser-cu121-torch25
 ```
 
 - install [cudnn](https://gitlab-master.nvidia.com/cudnn/cudnn_frontend) as follows:
@@ -89,35 +95,12 @@ pip install --pre nvfuser-cu121-torch24
 pip install nvidia-cudnn-frontend
 ```
 
-- Finally, install Thunder as follows:
-
-```
-# install thunder
-pip install lightning-thunder
-```
-
 <details>
   <summary>Advanced install options</summary>
     <!-- following section will be skipped from PyPI description -->
 
 &#160;
 
-### Install from main
-
-Alternatively, you can install the latest version of Thunder directly from this GitHub repository as follows:
-
-```
-# 1) Install nvFuser and PyTorch dependencies:
-pip install --pre nvfuser-cu121-torch24
-```
-
-```bash
-# 2) Install Thunder itself
-pip install git+https://github.com/Lightning-AI/lightning-thunder.git
-```
-
-&#160;
-
 ### Install to tinker and contribute
 
 If you are interested in tinkering with and contributing to Thunder, we recommend cloning the Thunder repository and installing it in pip's editable mode:

diff --git a/docs/source/fundamentals/installation.rst b/docs/source/fundamentals/installation.rst
@@ -6,12 +6,12 @@ Minimal dependencies
 
 Follow these instructions to install PyTorch, nvFuser, and finally Thunder.
 
-Install PyTorch and nvFuser with pip (command shown is for CUDA 12.1 and PyTorch 2.4.x)::
+Install PyTorch and nvFuser with pip (command shown is for CUDA 12.1 and PyTorch 2.5.x)::
 
-  pip install --pre nvfuser-cu121-torch24
+  pip install --pre nvfuser-cu121-torch25
 
 cu121 can be replaced with cu118 depending on your CUDA version. NVFuser builds typically support the latest point release of PyTorch stable versions.
-For torch 2.4, cu124 is also supported. For nightly versions and more detailed instructions, please see https://github.com/NVIDIA/Fuser/#installation
+For torch 2.5, cu124 is also supported. For nightly versions and more detailed instructions, please see https://github.com/NVIDIA/Fuser/#installation
 
 You're all set with minimal dependencies, so you can follow `Install Thunder`_.
 

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -102,6 +102,7 @@ The compiled function ``jitted_foo`` takes and returns PyTorch tensors, just lik
    FSDP Under the Hood Tutorial <notebooks/dev_tutorials/fsdp_tutorial>
    Benchmarking Thunder <intermediate/benchmarking>
    Writing a Transform <notebooks/writing_a_trace_transform_cpu_offloading>
+   Using Liger-Kernel with Thunder <notebooks/liger_kernel>
 
 .. toctree::
    :maxdepth: 1
@@ -143,6 +144,7 @@ API reference
    reference/torch/index
    reference/extend/index
    reference/transforms/index
+   reference/dynamo/index
 
 
 Indices and tables

diff --git a/docs/source/reference/dynamo/index.rst b/docs/source/reference/dynamo/index.rst
@@ -0,0 +1,9 @@
+.. module:: thunder.dynamo
+
+thunder.dynamo
+==============
+
+.. autosummary::
+    :toctree:
+
+    ThunderCompiler