Merge branch 'main' into wjy/permute

.azure/gpu-tests.yml

@@ -65,8 +65,14 @@ jobs:
         # drop pt from requirements so not to interfere with the existing one
         bash .azure/remove-torch-lines.sh requirements/base.txt
         cat requirements/base.txt
+
         # double check on test requirements
         pip install -r requirements/test.txt
+
+        # https://docs.codecov.com/docs/codecov-uploader
+        curl -Os https://uploader.codecov.io/latest/linux/codecov
+        chmod +x codecov
+
         # install this package
         python setup.py develop
       displayName: 'Install package & ...'
@@ -85,6 +91,12 @@ jobs:
             --durations=250 \
             --numprocesses=9 \
             --ignore=thunder/tests/distributed --ignore=thunder/tests/test_networks.py
+        # compile coverage results
+        python -m coverage report
+        python -m coverage xml
+        # upload to codecov
+        ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
+          --flags=gpu,pytest,regular --name="GPU-coverage" --env=linux,azure
       condition: ne(variables['testing'], 'distributed')
       displayName: 'Testing: regular'
 
@@ -95,6 +107,12 @@ jobs:
              thunder/tests/test_networks.py \
              -m "not standalone" \
              -v --random-order-seed=42 --durations=0 --numprocesses=3
+        # compile coverage results
+        python -m coverage report
+        python -m coverage xml
+        # upload to codecov
+        ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
+          --flags=gpu,pytest,networks --name="GPU-coverage" --env=linux,azure
       condition: ne(variables['testing'], 'distributed')
       displayName: 'Testing: networks'
 
@@ -108,6 +126,12 @@ jobs:
     - bash: |
         # run all found tests in given past as standalone
         bash scripts/run_standalone_tests.sh "thunder/tests/distributed"
+        # compile coverage results
+        python -m coverage report
+        python -m coverage xml
+        # upload to codecov
+        ./codecov --token=$(CODECOV_TOKEN) --commit=$(Build.SourceVersion) \
+          --flags=gpu,pytest,distributed --name="GPU-coverage" --env=linux,azure
       condition: eq(variables['testing'], 'distributed')
       displayName: 'Testing: distributed'
 

.github/workflows/ci-checks.yml

@@ -10,10 +10,10 @@ concurrency:
   cancel-in-progress: ${{ ! (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/release/')) }}
 
 jobs:
-  #check-code:
-  #  uses: Lightning-AI/utilities/.github/workflows/check-code.yml@main
-  #  with:
-  #    actions-ref: main
+  precommit-run:
+    uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.10.1
+    with:
+      python-version: "3.10"
 
   check-schema:
     uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.11.0

.github/workflows/ci-testing.yml

@@ -114,15 +114,15 @@ jobs:
         coverage report
         coverage xml
 
-    #- name: Upload coverage to Codecov
-    #  uses: codecov/codecov-action@v3
-    #  with:
-    #    token: ${{ secrets.CODECOV_TOKEN }}
-    #    file: ./coverage.xml
-    #    flags: unittests
-    #    env_vars: OS,PYTHON
-    #    name: codecov-umbrella
-    #    fail_ci_if_error: false
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        file: ./coverage.xml
+        flags: unittests
+        env_vars: OS,PYTHON
+        name: codecov-umbrella
+        fail_ci_if_error: false
 
 
   testing-guardian:

.github/workflows/docs-build.yml

@@ -14,35 +14,54 @@ defaults:
     shell: bash
 
 jobs:
-  build-docs:
+  docs-make:
     uses: Lightning-AI/utilities/.github/workflows/check-docs.yml@v0.11.0
     with:
       python-version: "3.10"
       requirements-file: "requirements/docs.txt"
       install-tex: true
 
   deploy-docs:
-  # https://github.com/marketplace/actions/deploy-to-github-pages
-    needs: build-docs
-    if: github.event_name != 'pull_request'
+    needs: docs-make
+    if: github.repository_owner == 'Lightning-AI' && github.event_name == 'push'
     runs-on: ubuntu-latest
+    env:
+      GCP_TARGET: "gs://lightning-docs-thunder"
     steps:
-      - name: Checkout 🛎️
-        uses: actions/checkout@v4
-        with:
-          # If you're using actions/checkout@v4 you must set persist-credentials to false in most cases for the deployment to work correctly.
-          persist-credentials: false
       - uses: actions/download-artifact@v3
         with:
           name: docs-html-${{ github.sha }}
           path: docs/build/
 
-      - name: Deploy 🚀
-        uses: JamesIves/github-pages-deploy-action@v4.5.0
+      - name: Authenticate to Google Cloud
+        uses: google-github-actions/auth@v2
+        with:
+          credentials_json: ${{ secrets.GCS_SA_KEY }}
+
+      - name: Setup gcloud
+        uses: google-github-actions/setup-gcloud@v2
         with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          branch: gh-pages # The branch the action should deploy to.
-          folder: docs/build/html # The folder the action should deploy.
-          clean: true # Automatically remove deleted files from the deploy branch
-          target-folder: docs # If you'd like to push the contents of the deployment folder into a specific directory
-          single-commit: true # you'd prefer to have a single commit on the deployment branch instead of full history
+          project_id: ${{ secrets.GCS_PROJECT }}
+
+      # Uploading docs to GCS, so they can be served on lightning.ai
+      #- name: Upload docs/thunder/stable to GCS 🪣
+      #  if: startsWith(github.ref, 'refs/heads/release/')
+      #  run: gsutil -m rsync -d -R docs/build/html/ ${GCP_TARGET}/stable
+
+      # Uploading docs to GCS, so they can be served on lightning.ai
+      - name: Upload docs/thunder/latest to GCS 🪣
+        if: github.ref == 'refs/heads/main'
+        run: gsutil -m rsync -d -R docs/build/html/ ${GCP_TARGET}/latest
+
+      # Uploading docs to GCS, so they can be served on lightning.ai
+      #- name: Upload docs/thunder/release to GCS 🪣
+      #  if: startsWith(github.ref, 'refs/tags/')
+      #  run: gsutil -m rsync -d -R docs/build/html/ ${GCP_TARGET}/${{ github.ref_name }}
+
+      # Uploading docs as archive to GCS, so they can be as backup
+      #- name: Upload docs as archive to GCS 🪣
+      #  if: startsWith(github.ref, 'refs/tags/')
+      #  working-directory: docs/build
+      #  run: |
+      #    zip ${{ github.ref_name }}.zip -r html/
+      #    gsutil cp ${{ github.ref_name }}.zip ${GCP_TARGET}

.pre-commit-config.yaml

@@ -17,6 +17,9 @@ repos:
       - id: check-toml
       - id: check-json
       - id: check-added-large-files
+        with:
+          maxkb: 250
+          enforce-all: true
       - id: check-docstring-first
       - id: detect-private-key
 

README.md

@@ -125,19 +125,7 @@ Thunder is in its early stages and should not be used for production runs yet.
 
 However, it can already deliver outstanding performance on LLM model supported by [LitGPT](https://github.com/Lightning-AI/lit-gpt), such as Mistral, Llama 2, Gemma, Falcon, and others.
 
-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.
+Check out [the LitGPT integration](https://github.com/Lightning-AI/litgpt/tree/main/extensions/thunder) to learn about running LitGPT and Thunder together.
 
 ## Features
 

docs/source/conf.py

@@ -92,7 +92,7 @@ def _transform_changelog(path_in: str, path_out: str) -> None:
     "sphinx.ext.linkcode",
     "sphinx.ext.autosummary",
     "sphinx.ext.napoleon",
-    "sphinx.ext.imgmath",
+    "sphinx.ext.mathjax",
     "myst_parser",
     "nbsphinx",
     "sphinx_autodoc_typehints",
@@ -209,6 +209,11 @@ def _transform_changelog(path_in: str, path_out: str) -> None:
     (master_doc, project + ".tex", project + " Documentation", author, "manual"),
 ]
 
+# MathJax configuration
+mathjax3_config = {
+    "tex": {"packages": {"[+]": ["ams", "newcommand", "configMacros"]}},
+}
+
 # -- Options for manual page output ------------------------------------------
 
 # One entry per manual page. List of tuples