lightning 2.3.0.dev20240318__tar.gz → 2.3.0.dev20240505__tar.gz

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/.actions/assistant.py

@@ -359,17 +359,19 @@ def copy_replace_imports(
 
 
 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()
     mapping.pop("lightning", None)  # pop this key to avoid replacing `lightning` to `lightning.lightning`
 
     mapping = {f"lightning.{sp}": sl for sp, sl in mapping.items()}
     for pkg_from, pkg_to in mapping.items():
+        source_imports, target_imports = zip(*mapping.items())
         copy_replace_imports(
             source_dir=os.path.join(source_dir, pkg_from.replace(".", os.sep)),
             # pytorch_lightning uses lightning_fabric, so we need to replace all imports for all directories
-            source_imports=mapping.keys(),
-            target_imports=mapping.values(),
+            source_imports=source_imports,
+            target_imports=target_imports,
             target_dir=os.path.join(source_dir, pkg_to.replace(".", os.sep)),
             lightning_by=pkg_from,
         )

{lightning-2.3.0.dev20240318/src/lightning.egg-info → lightning-2.3.0.dev20240505}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.3.0.dev20240318
+Version: 2.3.0.dev20240505
 Summary: The Deep Learning framework to train, deploy, and ship AI products Lightning fast.
 Home-page: https://github.com/Lightning-AI/lightning
 Author: Lightning AI et al.
@@ -27,18 +27,18 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Provides-Extra: fabric-strategies
-Provides-Extra: fabric-test
-Provides-Extra: fabric-examples
-Provides-Extra: app-ui
-Provides-Extra: app-test
 Provides-Extra: app-components
 Provides-Extra: app-cloud
-Provides-Extra: pytorch-extra
+Provides-Extra: app-ui
+Provides-Extra: app-test
+Provides-Extra: store-test
 Provides-Extra: pytorch-strategies
 Provides-Extra: pytorch-test
+Provides-Extra: pytorch-extra
 Provides-Extra: pytorch-examples
-Provides-Extra: store-test
+Provides-Extra: fabric-strategies
+Provides-Extra: fabric-test
+Provides-Extra: fabric-examples
 Provides-Extra: fabric-all
 Provides-Extra: fabric-dev
 Provides-Extra: pytorch-all
@@ -46,15 +46,15 @@ Provides-Extra: pytorch-dev
 Provides-Extra: app-extra
 Provides-Extra: app-all
 Provides-Extra: app-dev
-Provides-Extra: data
-Provides-Extra: strategies
-Provides-Extra: test
-Provides-Extra: examples
-Provides-Extra: ui
 Provides-Extra: components
 Provides-Extra: cloud
+Provides-Extra: ui
 Provides-Extra: app
+Provides-Extra: test
+Provides-Extra: strategies
 Provides-Extra: extra
+Provides-Extra: examples
+Provides-Extra: data
 Provides-Extra: all
 Provides-Extra: dev
 Provides-Extra: store
@@ -78,7 +78,7 @@ ______________________________________________________________________
   <a href="https://lightning.ai/docs/pytorch/stable/">PyTorch Lightning</a> •
   <a href="https://lightning.ai/docs/fabric/stable/">Fabric</a> •
   <a href="https://lightning.ai/docs/app/stable/">Lightning Apps</a> •
-  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240318">Docs</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240505">Docs</a> •
   <a href="#community">Community</a> •
   <a href="https://lightning.ai/docs/pytorch/stable/generated/CONTRIBUTING.html">Contribute</a> •
 </p>
@@ -623,7 +623,7 @@ Lightning is rigorously tested across multiple CPUs, GPUs and TPUs and against m
 
 |       System / PyTorch ver.        | 1.13                                                                                                                                                                                                                            | 2.0                                                                                                                                                                                                                             |                                                                                                               2.1                                                                                                               |
 | :--------------------------------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: |:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
-|        Linux py3.9 \[GPUs\]        |  |  | [![Build Status](https://dev.azure.com/Lightning-AI/lightning/_apis/build/status%2Fpytorch-lightning%20%28GPUs%29?branchName=refs%2Ftags%2F2.3.0.dev20240318)](https://dev.azure.com/Lightning-AI/lightning/_build/latest?definitionId=24&branchName=master) |
+|        Linux py3.9 \[GPUs\]        |  |  | [![Build Status](https://dev.azure.com/Lightning-AI/lightning/_apis/build/status%2Fpytorch-lightning%20%28GPUs%29?branchName=refs%2Ftags%2F2.3.0.dev20240505)](https://dev.azure.com/Lightning-AI/lightning/_build/latest?definitionId=24&branchName=master) |
 |        Linux py3.9 \[TPUs\]        |                                                                                                                                                                                                                                 |  [![Test PyTorch - TPU](https://github.com/Lightning-AI/lightning/actions/workflows/tpu-tests.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/tpu-tests.yml)     |      |
 |  Linux (multiple Python versions)  | [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                                 | [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                                 |                 [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                 |
 |   OSX (multiple Python versions)   | [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                                 | [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                                 |                 [![Test PyTorch](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml/badge.svg)](https://github.com/Lightning-AI/lightning/actions/workflows/ci-tests-pytorch.yml)                 |

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/app/app.txt

@@ -1,4 +1,4 @@
-lightning-cloud == 0.5.65  # Must be pinned to ensure compatibility
+lightning-cloud == 0.5.68  # Must be pinned to ensure compatibility
 packaging
 typing-extensions >=4.4.0, <4.10.0
 deepdiff >=5.7.0, <6.6.0
@@ -6,7 +6,7 @@ fsspec[http] >=2022.5.0, <2023.11.0
 croniter >=1.3.0, <1.5.0  # strict; TODO: for now until we find something more robust.
 traitlets >=5.3.0, <5.12.0
 arrow >=1.2.0, <1.3.0
-lightning-utilities >=0.8.0, <0.11.0
+lightning-utilities >=0.8.0, <0.12.0
 beautifulsoup4 >=4.8.0, <4.13.0
 inquirer >=2.10.0, <3.2.0
 psutil <5.9.6

lightning-2.3.0.dev20240505/requirements/base.txt

@@ -0,0 +1,10 @@
+PyYAML<7.0,>=5.4
+fsspec[http]<2025.0,>=2022.5.0
+lightning-utilities<1.0,>=0.8.0
+numpy<2.0,>=1.17.2
+packaging<24.0,>=20.0
+torch<3.0,>=2.0.0
+torchmetrics<2.0,>=0.7.0
+tqdm<5.0,>=4.57.0
+typing-extensions<5.0,>=4.4.0
+

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/fabric/base.txt

@@ -2,8 +2,8 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <=2.2.0
-fsspec[http] >=2022.5.0, <2023.11.0
+torch >=2.0.0, <2.4.0
+fsspec[http] >=2022.5.0, <2024.4.0
 packaging >=20.0, <=23.1
 typing-extensions >=4.4.0, <4.10.0
-lightning-utilities >=0.8.0, <0.11.0
+lightning-utilities >=0.8.0, <0.12.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/fabric/examples.txt

@@ -1,6 +1,6 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torchvision >=0.14.0, <0.18.0
+torchvision >=0.15.0, <0.19.0
 torchmetrics >=0.10.0, <1.3.0
-lightning-utilities >=0.8.0, <0.11.0
+lightning-utilities >=0.8.0, <0.12.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/fabric/strategies.txt

@@ -5,5 +5,5 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.8.2, <=0.9.3; platform_system != "Windows"  # strict
+deepspeed >=0.8.2, <=0.9.3; platform_system != "Windows" and platform_system != "Darwin"  # strict
 bitsandbytes >=0.42.0,<0.43.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/pytorch/base.txt

@@ -2,11 +2,11 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <=2.2.0
+torch >=2.0.0, <2.4.0
 tqdm >=4.57.0, <4.67.0
 PyYAML >=5.4, <6.1.0
-fsspec[http] >=2022.5.0, <2023.11.0
+fsspec[http] >=2022.5.0, <2024.4.0
 torchmetrics >=0.7.0, <1.3.0  # needed for using fixed compare_version
 packaging >=20.0, <=23.1
 typing-extensions >=4.4.0, <4.10.0
-lightning-utilities >=0.8.0, <0.11.0
+lightning-utilities >=0.8.0, <0.12.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/pytorch/examples.txt

@@ -2,8 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 requests <2.32.0
-torchvision >=0.14.0, <0.18.0
-gym[classic_control] >=0.17.0, <0.27.0
+torchvision >=0.15.0, <0.19.0
 ipython[all] <8.15.0
 torchmetrics >=0.10.0, <1.3.0
-lightning-utilities >=0.8.0, <0.11.0
+lightning-utilities >=0.8.0, <0.12.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/pytorch/extra.txt

@@ -5,7 +5,7 @@
 matplotlib>3.1, <3.9.0
 omegaconf >=2.0.5, <2.4.0
 hydra-core >=1.0.5, <1.4.0
-jsonargparse[signatures] >=4.27.5, <4.28.0
+jsonargparse[signatures] >=4.27.7, <4.28.0
 rich >=12.3.0, <13.6.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.42.0,<0.43.0

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/requirements/pytorch/strategies.txt

@@ -3,4 +3,4 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.8.2, <=0.9.3; platform_system != "Windows"  # strict
+deepspeed >=0.8.2, <=0.9.3; platform_system != "Windows" and platform_system != "Darwin"  # strict

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/app/utilities/network.py

@@ -89,18 +89,32 @@ _CONNECTION_RETRY_BACKOFF_FACTOR = 0.5
 _DEFAULT_REQUEST_TIMEOUT = 30  # seconds
 
 
+def create_retry_strategy():
+    return Retry(
+        # wait time between retries increases exponentially according to: backoff_factor * (2 ** (retry - 1))
+        # but the the maximum wait time is 120 secs. By setting a large value (2880), we'll make sure clients
+        # are going to be alive for a very long time (~ 4 days) but retries every 120 seconds
+        total=_CONNECTION_RETRY_TOTAL,
+        backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
+        status_forcelist={
+            408,  # Request Timeout
+            429,  # Too Many Requests
+            *range(500, 600),  # Any 5xx Server Error status
+        },
+        allowed_methods={
+            "POST",  # Default methods are idempotent, add POST here
+            *Retry.DEFAULT_ALLOWED_METHODS,
+        },
+    )
+
+
 def _configure_session() -> Session:
     """Configures the session for GET and POST requests.
 
     It enables a generous retrial strategy that waits for the application server to connect.
 
     """
-    retry_strategy = Retry(
-        # wait time between retries increases exponentially according to: backoff_factor * (2 ** (retry - 1))
-        total=_CONNECTION_RETRY_TOTAL,
-        backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
-        status_forcelist=[429, 500, 502, 503, 504],
-    )
+    retry_strategy = create_retry_strategy()
     adapter = HTTPAdapter(max_retries=retry_strategy)
     http = requests.Session()
     http.mount("https://", adapter)
@@ -157,21 +171,7 @@ class HTTPClient:
         self, base_url: str, auth_token: Optional[str] = None, log_callback: Optional[Callable] = None
     ) -> None:
         self.base_url = base_url
-        retry_strategy = Retry(
-            # wait time between retries increases exponentially according to: backoff_factor * (2 ** (retry - 1))
-            # but the the maximum wait time is 120 secs. By setting a large value (2880), we'll make sure clients
-            # are going to be alive for a very long time (~ 4 days) but retries every 120 seconds
-            total=_CONNECTION_RETRY_TOTAL,
-            backoff_factor=_CONNECTION_RETRY_BACKOFF_FACTOR,
-            status_forcelist=[
-                408,  # Request Timeout
-                429,  # Too Many Requests
-                500,  # Internal Server Error
-                502,  # Bad Gateway
-                503,  # Service Unavailable
-                504,  # Gateway Timeout
-            ],
-        )
+        retry_strategy = create_retry_strategy()
         adapter = CustomRetryAdapter(max_retries=retry_strategy, timeout=_DEFAULT_REQUEST_TIMEOUT)
         self.session = requests.Session()
 

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/fabric/CHANGELOG.md

@@ -11,7 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Enabled consolidating distributed checkpoints through `fabric consolidate` in the new CLI [#19560](https://github.com/Lightning-AI/pytorch-lightning/pull/19560))
 
--
+- Added support for PyTorch 2.3 ([#19708](https://github.com/Lightning-AI/pytorch-lightning/pull/19708))
 
 -
 
@@ -39,7 +39,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Removed
 
--
+- Removed support for PyTorch 1.13 ([#19706](https://github.com/Lightning-AI/lightning/pull/19706))
 
 -
 
@@ -47,13 +47,22 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- Fixed an issue causing a TypeError when using `torch.compile` as a decorator ([#19627](https://github.com/Lightning-AI/pytorch-lightning/pull/19627))
+-
 
 -
 
 -
 
 
+## [2.2.2] - 2024-04-11
+
+### Fixed
+
+- Fixed an issue causing a TypeError when using `torch.compile` as a decorator ([#19627](https://github.com/Lightning-AI/pytorch-lightning/pull/19627))
+- Fixed issue where some model methods couldn't be monkeypatched after being Fabric wrapped ([#19705](https://github.com/Lightning-AI/pytorch-lightning/pull/19705))
+- Fixed an issue causing weights to be reset in `Fabric.setup()` when using FSDP ([#19755](https://github.com/Lightning-AI/pytorch-lightning/pull/19755))
+
+
 ## [2.2.1] - 2024-03-04
 
 ### Fixed

lightning-2.3.0.dev20240505/src/lightning/fabric/accelerators/cuda.py

@@ -0,0 +1,181 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 functools import lru_cache
+from typing import List, Optional, Union
+
+import torch
+from typing_extensions import override
+
+from lightning.fabric.accelerators.accelerator import Accelerator
+from lightning.fabric.accelerators.registry import _AcceleratorRegistry
+from lightning.fabric.utilities.rank_zero import rank_zero_info
+
+
+class CUDAAccelerator(Accelerator):
+    """Accelerator for NVIDIA CUDA devices."""
+
+    @override
+    def setup_device(self, device: torch.device) -> None:
+        """
+        Raises:
+            ValueError:
+                If the selected device is not of type CUDA.
+        """
+        if device.type != "cuda":
+            raise ValueError(f"Device should be CUDA, got {device} instead.")
+        _check_cuda_matmul_precision(device)
+        torch.cuda.set_device(device)
+
+    @override
+    def teardown(self) -> None:
+        _clear_cuda_memory()
+
+    @staticmethod
+    @override
+    def parse_devices(devices: Union[int, str, List[int]]) -> Optional[List[int]]:
+        """Accelerator device parsing logic."""
+        from lightning.fabric.utilities.device_parser import _parse_gpu_ids
+
+        return _parse_gpu_ids(devices, include_cuda=True)
+
+    @staticmethod
+    @override
+    def get_parallel_devices(devices: List[int]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        return [torch.device("cuda", i) for i in devices]
+
+    @staticmethod
+    @override
+    def auto_device_count() -> int:
+        """Get the devices when set to auto."""
+        return num_cuda_devices()
+
+    @staticmethod
+    @override
+    def is_available() -> bool:
+        return num_cuda_devices() > 0
+
+    @classmethod
+    @override
+    def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> None:
+        accelerator_registry.register(
+            "cuda",
+            cls,
+            description=cls.__name__,
+        )
+
+
+def find_usable_cuda_devices(num_devices: int = -1) -> List[int]:
+    """Returns a list of all available and usable CUDA GPU devices.
+
+    A GPU is considered usable if we can successfully move a tensor to the device, and this is what this function
+    tests for each GPU on the system until the target number of usable devices is found.
+
+    A subset of GPUs on the system might be used by other processes, and if the GPU is configured to operate in
+    'exclusive' mode (configurable by the admin), then only one process is allowed to occupy it.
+
+    Args:
+        num_devices: The number of devices you want to request. By default, this function will return as many as there
+            are usable CUDA GPU devices available.
+
+    Warning:
+        If multiple processes call this function at the same time, there can be race conditions in the case where
+        both processes determine that the device is unoccupied, leading into one of them crashing later on.
+
+    """
+    if num_devices == 0:
+        return []
+    visible_devices = _get_all_visible_cuda_devices()
+    if not visible_devices:
+        raise ValueError(
+            f"You requested to find {num_devices} devices but there are no visible CUDA devices on this machine."
+        )
+    if num_devices > len(visible_devices):
+        raise ValueError(
+            f"You requested to find {num_devices} devices but this machine only has {len(visible_devices)} GPUs."
+        )
+
+    available_devices = []
+    unavailable_devices = []
+
+    for gpu_idx in visible_devices:
+        try:
+            torch.tensor(0, device=torch.device("cuda", gpu_idx))
+        except RuntimeError:
+            unavailable_devices.append(gpu_idx)
+            continue
+
+        available_devices.append(gpu_idx)
+        if len(available_devices) == num_devices:
+            # exit early if we found the right number of GPUs
+            break
+
+    if num_devices != -1 and len(available_devices) != num_devices:
+        raise RuntimeError(
+            f"You requested to find {num_devices} devices but only {len(available_devices)} are currently available."
+            f" The devices {unavailable_devices} are occupied by other processes and can't be used at the moment."
+        )
+    return available_devices
+
+
+def _get_all_visible_cuda_devices() -> List[int]:
+    """Returns a list of all visible CUDA GPU devices.
+
+    Devices masked by the environment variabale ``CUDA_VISIBLE_DEVICES`` won't be returned here. For example, assume you
+    have 8 physical GPUs. If ``CUDA_VISIBLE_DEVICES="1,3,6"``, then this function will return the list ``[0, 1, 2]``
+    because these are the three visible GPUs after applying the mask ``CUDA_VISIBLE_DEVICES``.
+
+    """
+    return list(range(num_cuda_devices()))
+
+
+def num_cuda_devices() -> int:
+    """Returns the number of available CUDA devices."""
+    return torch.cuda.device_count()
+
+
+def is_cuda_available() -> bool:
+    """Returns a bool indicating if CUDA is currently available."""
+    # We set `PYTORCH_NVML_BASED_CUDA_CHECK=1` in lightning.fabric.__init__.py
+    return torch.cuda.is_available()
+
+
+def _is_ampere_or_later(device: Optional[torch.device] = None) -> bool:
+    major, _ = torch.cuda.get_device_capability(device)
+    return major >= 8  # Ampere and later leverage tensor cores, where this setting becomes useful
+
+
+@lru_cache(1)  # show the warning only ever once
+def _check_cuda_matmul_precision(device: torch.device) -> None:
+    if not torch.cuda.is_available() or not _is_ampere_or_later(device):
+        return
+    # check that the user hasn't changed the precision already, this works for both `allow_tf32 = True` and
+    # `set_float32_matmul_precision`
+    if torch.get_float32_matmul_precision() == "highest":  # default
+        rank_zero_info(
+            f"You are using a CUDA device ({torch.cuda.get_device_name(device)!r}) that has Tensor Cores. To properly"
+            " utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off"
+            " precision for performance. For more details, read https://pytorch.org/docs/stable/generated/"
+            "torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision"
+        )
+    # note: no need change `torch.backends.cudnn.allow_tf32` as it's enabled by default:
+    # https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+
+
+def _clear_cuda_memory() -> None:
+    # strangely, the attribute function be undefined when torch.compile is used
+    if hasattr(torch._C, "_cuda_clearCublasWorkspaces"):
+        # https://github.com/pytorch/pytorch/issues/95668
+        torch._C._cuda_clearCublasWorkspaces()
+    torch.cuda.empty_cache()

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/fabric/accelerators/mps.py

@@ -11,6 +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.
+import os
 import platform
 from functools import lru_cache
 from typing import List, Optional, Union
@@ -70,7 +71,8 @@ class MPSAccelerator(Accelerator):
     @lru_cache(1)
     def is_available() -> bool:
         """MPS is only available on a machine with the ARM-based Apple Silicon processors."""
-        return torch.backends.mps.is_available() and platform.processor() in ("arm", "arm64")
+        mps_disabled = os.getenv("DISABLE_MPS", "0") == "1"
+        return not mps_disabled and torch.backends.mps.is_available() and platform.processor() in ("arm", "arm64")
 
     @classmethod
     @override

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/fabric/fabric.py

@@ -51,7 +51,6 @@ from lightning.fabric.strategies import (
     FSDPStrategy,
     SingleDeviceStrategy,
     Strategy,
-    XLAFSDPStrategy,
     XLAStrategy,
 )
 from lightning.fabric.strategies.fsdp import _has_meta_device_parameters
@@ -67,7 +66,6 @@ from lightning.fabric.utilities.data import (
 )
 from lightning.fabric.utilities.device_dtype_mixin import _update_properties
 from lightning.fabric.utilities.distributed import DistributedSamplerWrapper, _InfiniteBarrier
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
 from lightning.fabric.utilities.rank_zero import rank_zero_deprecation, rank_zero_warn
 from lightning.fabric.utilities.registry import _load_external_callbacks
 from lightning.fabric.utilities.seed import seed_everything
@@ -699,26 +697,14 @@ class Fabric:
 
     def init_tensor(self) -> ContextManager:
         """Tensors that you instantiate under this context manager will be created on the device right away and have
-        the right data type depending on the precision setting in Fabric.
-
-        The automatic device placement under this context manager is only supported with PyTorch 2.0 and newer.
-
-        """
-        if not _TORCH_GREATER_EQUAL_2_0 and self.device.type != "cpu":
-            rank_zero_warn(
-                "`Fabric.init_tensor()` can't place tensors on the device directly"
-                " with PyTorch < 2.0. Parameters will remain on CPU until `Fabric.setup()` is called."
-                " Upgrade to PyTorch >= 2.0 to fully utilize this feature.",
-                category=PossibleUserWarning,
-            )
+        the right data type depending on the precision setting in Fabric."""
         return self._strategy.tensor_init_context()
 
     def init_module(self, empty_init: Optional[bool] = None) -> ContextManager:
         """Instantiate the model and its parameters under this context manager to reduce peak memory usage.
 
         The parameters get created on the device and with the right data type right away without wasting memory being
-        allocated unnecessarily. The automatic device placement under this context manager is only supported with
-        PyTorch 2.0 and newer.
+        allocated unnecessarily.
 
         Args:
             empty_init: Whether to initialize the model with empty weights (uninitialized memory).
@@ -727,13 +713,6 @@ class Fabric:
 
         """
         self._validate_launched()
-        if not _TORCH_GREATER_EQUAL_2_0 and self.device.type != "cpu":
-            rank_zero_warn(
-                "`Fabric.init_module()` can't place the model parameters on the device directly"
-                " with PyTorch < 2.0. Parameters will remain on CPU until `Fabric.setup()` is called."
-                " Upgrade to PyTorch >= 2.0 to fully utilize this feature.",
-                category=PossibleUserWarning,
-            )
         return self._strategy.module_init_context(empty_init=empty_init)
 
     def save(
@@ -1036,12 +1015,6 @@ class Fabric:
         if any(isinstance(opt, _FabricOptimizer) for opt in optimizers):
             raise ValueError("An optimizer should be passed only once to the `setup` method.")
 
-        if isinstance(self._strategy, (FSDPStrategy, XLAFSDPStrategy)) and not _TORCH_GREATER_EQUAL_2_0:
-            raise RuntimeError(
-                f"The `{type(self).__name__}` requires the model and optimizer(s) to be set up separately."
-                " Create and set up the model first through `model = self.setup_module(model)`. Then create the"
-                " optimizer and set it up: `optimizer = self.setup_optimizer(optimizer)`."
-            )
         if isinstance(self._strategy, FSDPStrategy) and any(
             _has_meta_device_parameters(optimizer) for optimizer in optimizers
         ):

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/fabric/plugins/precision/amp.py

@@ -20,7 +20,6 @@ from torch.nn import Module
 from torch.optim import LBFGS, Optimizer
 from typing_extensions import override
 
-from lightning.fabric.accelerators.cuda import _patch_cuda_is_available
 from lightning.fabric.plugins.precision.precision import Precision
 from lightning.fabric.plugins.precision.utils import _convert_fp_tensor
 from lightning.fabric.utilities.types import Optimizable
@@ -50,9 +49,7 @@ class MixedPrecision(Precision):
 
         self.precision = precision
         if scaler is None and self.precision == "16-mixed":
-            with _patch_cuda_is_available():
-                # if possible, we defer CUDA initialization to support strategies that will attempt forks
-                scaler = torch.cuda.amp.GradScaler()
+            scaler = torch.cuda.amp.GradScaler()
         if scaler is not None and self.precision == "bf16-mixed":
             raise ValueError(f"`precision='bf16-mixed'` does not use a scaler, found {scaler}.")
         self.device = device

{lightning-2.3.0.dev20240318 → lightning-2.3.0.dev20240505}/src/lightning/fabric/plugins/precision/fsdp.py

@@ -23,7 +23,6 @@ from typing_extensions import get_args, override
 from lightning.fabric.plugins.precision.amp import _optimizer_handles_unscaling
 from lightning.fabric.plugins.precision.precision import Precision
 from lightning.fabric.plugins.precision.utils import _convert_fp_tensor, _DtypeContextManager
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
 from lightning.fabric.utilities.types import Optimizable
 
 if TYPE_CHECKING:
@@ -78,21 +77,18 @@ class FSDPPrecision(Precision):
     def mixed_precision_config(self) -> "TorchMixedPrecision":
         from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision as TorchMixedPrecision
 
-        # With PyTorch < 2.0, FSDP uses the noneness of `param_dtype` as a proxy for the `_uses_param_mixed_precision`
-        # property. In order to avoid FSDP assertion failures, we therefore avoid setting `param_dtype` to
-        # `torch.float32` here with PyTorch < 2.0.
         if self.precision == "16-mixed":
-            param_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
+            param_dtype = torch.float32
             reduce_dtype = buffer_dtype = torch.float16
         elif self.precision == "bf16-mixed":
-            param_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
+            param_dtype = torch.float32
             reduce_dtype = buffer_dtype = torch.bfloat16
         elif self.precision == "16-true":
             param_dtype = reduce_dtype = buffer_dtype = torch.float16
         elif self.precision == "bf16-true":
             param_dtype = reduce_dtype = buffer_dtype = torch.bfloat16
         elif self.precision == "32-true":
-            param_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
+            param_dtype = torch.float32
             reduce_dtype = buffer_dtype = torch.float32
         else:
             raise ValueError(f"Was unable to infer precision type, received {self.precision!r}.")