FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric

Original PR #20440 by lantiga
Original: Lightning-AI/pytorch-lightning#20440

Author: Shallow Copy Bot

.actions/assistant.py

@@ -18,11 +18,10 @@
 import shutil
 import tempfile
 import urllib.request
-from collections.abc import Iterable, Iterator, Sequence
 from itertools import chain
 from os.path import dirname, isfile
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple
 
 from packaging.requirements import Requirement
 from packaging.version import Version
@@ -128,7 +127,7 @@ def _parse_requirements(lines: Iterable[str]) -> Iterator[_RequirementWithCommen
         pip_argument = None
 
 
-def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str = "all") -> list[str]:
+def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str = "all") -> List[str]:
     """Loading requirements from a file.
 
     >>> path_req = os.path.join(_PROJECT_ROOT, "requirements")
@@ -223,7 +222,7 @@ def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requireme
         fp.writelines([ln + os.linesep for ln in requires] + [os.linesep])
 
 
-def _retrieve_files(directory: str, *ext: str) -> list[str]:
+def _retrieve_files(directory: str, *ext: str) -> List[str]:
     all_files = []
     for root, _, files in os.walk(directory):
         for fname in files:
@@ -233,7 +232,7 @@ def _retrieve_files(directory: str, *ext: str) -> list[str]:
     return all_files
 
 
-def _replace_imports(lines: list[str], mapping: list[tuple[str, str]], lightning_by: str = "") -> list[str]:
+def _replace_imports(lines: List[str], mapping: List[Tuple[str, str]], lightning_by: str = "") -> List[str]:
     """Replace imports of standalone package to lightning.
 
     >>> lns = [
@@ -321,7 +320,7 @@ def copy_replace_imports(
             fo.writelines(lines)
 
 
-def create_mirror_package(source_dir: str, package_mapping: dict[str, str]) -> None:
+def create_mirror_package(source_dir: str, package_mapping: Dict[str, str]) -> None:
     """Create a mirror package with adjusted imports."""
     # replace imports and copy the code
     mapping = package_mapping.copy()

.github/workflows/_legacy-checkpoints.yml

@@ -60,7 +60,7 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           # Python version here needs to be supported by all PL versions listed in back-compatible-versions.txt.
-          python-version: "3.9"
+          python-version: 3.8
 
       - name: Install PL from source
         env:

.github/workflows/call-clear-cache.yml

@@ -23,7 +23,7 @@ on:
 jobs:
   cron-clear:
     if: github.event_name == 'schedule' || github.event_name == 'pull_request'
-    uses: Lightning-AI/utilities/.github/workflows/[email protected].9
+    uses: Lightning-AI/utilities/.github/workflows/[email protected].8
     with:
       scripts-ref: v0.11.8
       dry-run: ${{ github.event_name == 'pull_request' }}
@@ -32,7 +32,7 @@ jobs:
 
   direct-clear:
     if: github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request'
-    uses: Lightning-AI/utilities/.github/workflows/[email protected].9
+    uses: Lightning-AI/utilities/.github/workflows/[email protected].8
     with:
       scripts-ref: v0.11.8
       dry-run: ${{ github.event_name == 'pull_request' }}

.github/workflows/ci-check-md-links.yml

@@ -14,7 +14,7 @@ on:
 
 jobs:
   check-md-links:
-    uses: Lightning-AI/utilities/.github/workflows/[email protected].9
+    uses: Lightning-AI/utilities/.github/workflows/[email protected].8
     with:
       config-file: ".github/markdown-links-config.json"
       base-branch: "master"

.github/workflows/ci-schema.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   check:
-    uses: Lightning-AI/utilities/.github/workflows/[email protected].9
+    uses: Lightning-AI/utilities/.github/workflows/[email protected].8
     with:
       # skip azure due to the wrong schema file by MSFT
       # https://github.com/Lightning-AI/lightning-flash/pull/1455#issuecomment-1244793607

.pre-commit-config.yaml

@@ -74,7 +74,7 @@ repos:
     hooks:
       # try to fix what is possible
       - id: ruff
-        args: ["--fix", "--unsafe-fixes"]
+        args: ["--fix"]
       # perform formatting updates
       - id: ruff-format
       # validate if all is fine with preview mode

pyproject.toml

@@ -44,7 +44,7 @@ ignore-words-list = "te, compiletime"
 
 [tool.ruff]
 line-length = 120
-target-version = "py39"
+target-version = "py38"
 # Exclude a variety of commonly ignored directories.
 exclude = [
     ".git",

setup.py

@@ -45,10 +45,9 @@
 import logging
 import os
 import tempfile
-from collections.abc import Generator, Mapping
 from importlib.util import module_from_spec, spec_from_file_location
 from types import ModuleType
-from typing import Optional
+from typing import Generator, Mapping, Optional
 
 import setuptools
 import setuptools.command.egg_info

src/lightning/__setup__.py

@@ -3,7 +3,7 @@
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
 from types import ModuleType
-from typing import Any
+from typing import Any, Dict
 
 from setuptools import find_namespace_packages
 
@@ -26,7 +26,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:
 _ASSISTANT = _load_py_module(name="assistant", location=os.path.join(_PROJECT_ROOT, ".actions", "assistant.py"))
 
 
-def _prepare_extras() -> dict[str, Any]:
+def _prepare_extras() -> Dict[str, Any]:
     # https://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-extras
     # Define package extras. These are only installed if you specify them.
     # From remote, use like `pip install "lightning[dev, docs]"`
@@ -63,7 +63,7 @@ def _prepare_extras() -> dict[str, Any]:
     return extras
 
 
-def _setup_args() -> dict[str, Any]:
+def _setup_args() -> Dict[str, Any]:
     about = _load_py_module("about", os.path.join(_PACKAGE_ROOT, "__about__.py"))
     version = _load_py_module("version", os.path.join(_PACKAGE_ROOT, "__version__.py"))
     long_description = _ASSISTANT.load_readme_description(

src/lightning/fabric/accelerators/cpu.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Union
+from typing import List, Union
 
 import torch
 from typing_extensions import override
@@ -45,7 +45,7 @@ def parse_devices(devices: Union[int, str]) -> int:
 
     @staticmethod
     @override
-    def get_parallel_devices(devices: Union[int, str]) -> list[torch.device]:
+    def get_parallel_devices(devices: Union[int, str]) -> List[torch.device]:
         """Gets parallel devices for the Accelerator."""
         devices = _parse_cpu_cores(devices)
         return [torch.device("cpu")] * devices