lightning 2.3.0.dev20240407__tar.gz → 2.3.0.dev20240526__tar.gz

{lightning-2.3.0.dev20240407/src/lightning.egg-info → lightning-2.3.0.dev20240526}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.3.0.dev20240407
+Version: 2.3.0.dev20240526
 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: store-test
-Provides-Extra: fabric-strategies
 Provides-Extra: fabric-examples
+Provides-Extra: fabric-strategies
 Provides-Extra: fabric-test
-Provides-Extra: pytorch-strategies
-Provides-Extra: pytorch-examples
-Provides-Extra: pytorch-test
-Provides-Extra: pytorch-extra
-Provides-Extra: app-cloud
+Provides-Extra: store-test
 Provides-Extra: app-ui
-Provides-Extra: app-components
+Provides-Extra: app-cloud
 Provides-Extra: app-test
+Provides-Extra: app-components
+Provides-Extra: pytorch-examples
+Provides-Extra: pytorch-strategies
+Provides-Extra: pytorch-extra
+Provides-Extra: pytorch-test
 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: test
 Provides-Extra: data
-Provides-Extra: strategies
 Provides-Extra: examples
-Provides-Extra: extra
+Provides-Extra: strategies
+Provides-Extra: test
+Provides-Extra: ui
 Provides-Extra: cloud
 Provides-Extra: app
-Provides-Extra: ui
 Provides-Extra: components
+Provides-Extra: extra
 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.dev20240407">Docs</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240526">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.dev20240407)](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.dev20240526)](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.dev20240407 → lightning-2.3.0.dev20240526}/requirements/app/app.txt

@@ -1,4 +1,4 @@
-lightning-cloud == 0.5.65  # Must be pinned to ensure compatibility
+lightning-cloud == 0.5.69  # Must be pinned to ensure compatibility
 packaging
 typing-extensions >=4.4.0, <4.10.0
 deepdiff >=5.7.0, <6.6.0

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/requirements/base.txt

@@ -1,9 +1,9 @@
 PyYAML<7.0,>=5.4
-fsspec[http]<2024.0,>=2022.5.0
+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,>=1.13.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.dev20240407 → lightning-2.3.0.dev20240526}/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.3.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.12.0

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/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.12.0

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/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.dev20240407 → lightning-2.3.0.dev20240526}/requirements/pytorch/base.txt

@@ -2,10 +2,10 @@
 #  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.3.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-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/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.12.0

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/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.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/app/core/app.py

@@ -30,6 +30,7 @@ from lightning.app import _console
 from lightning.app.api.request_types import _APIRequest, _CommandRequest, _DeltaRequest
 from lightning.app.core.constants import (
     BATCH_DELTA_COUNT,
+    CHECK_ERROR_QUEUE_INTERVAL,
     DEBUG_ENABLED,
     FLOW_DURATION_SAMPLES,
     FLOW_DURATION_THRESHOLD,
@@ -165,6 +166,7 @@ class LightningApp:
 
         self._last_run_time: float = 0.0
         self._run_times: list = []
+        self._last_check_error_queue: float = 0.0
 
         # Path attributes can't get properly attached during the initialization, because the full name
         # is only available after all Flows and Works have been instantiated.
@@ -318,10 +320,12 @@ class LightningApp:
             return []
 
     def check_error_queue(self) -> None:
-        exception: Exception = self.get_state_changed_from_queue(self.error_queue)  # type: ignore[assignment,arg-type]
-        if isinstance(exception, Exception):
-            self.exception = exception
-            self.stage = AppStage.FAILED
+        if (time() - self._last_check_error_queue) > CHECK_ERROR_QUEUE_INTERVAL:
+            exception: Exception = self.get_state_changed_from_queue(self.error_queue)  # type: ignore[assignment,arg-type]
+            if isinstance(exception, Exception):
+                self.exception = exception
+                self.stage = AppStage.FAILED
+            self._last_check_error_queue = time()
 
     @property
     def flows(self) -> List[Union[LightningWork, "LightningFlow"]]:

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/app/core/constants.py

@@ -70,6 +70,7 @@ DOT_IGNORE_FILENAME = ".lightningignore"
 LIGHTNING_COMPONENT_PUBLIC_REGISTRY = "https://lightning.ai/v1/components"
 LIGHTNING_APPS_PUBLIC_REGISTRY = "https://lightning.ai/v1/apps"
 LIGHTNING_MODELS_PUBLIC_REGISTRY = "https://lightning.ai/v1/models"
+ENABLE_ORCHESTRATOR = bool(int(os.getenv("ENABLE_ORCHESTRATOR", "1")))
 
 LIGHTNING_CLOUDSPACE_HOST = os.getenv("LIGHTNING_CLOUDSPACE_HOST")
 LIGHTNING_CLOUDSPACE_EXPOSED_PORT_COUNT = int(os.getenv("LIGHTNING_CLOUDSPACE_EXPOSED_PORT_COUNT", "0"))
@@ -99,6 +100,7 @@ SYS_CUSTOMIZATIONS_SYNC_ROOT = "/tmp/sys-customizations-sync"  # todo
 SYS_CUSTOMIZATIONS_SYNC_PATH = ".sys-customizations-sync"
 
 BATCH_DELTA_COUNT = int(os.getenv("BATCH_DELTA_COUNT", "128"))
+CHECK_ERROR_QUEUE_INTERVAL = float(os.getenv("CHECK_ERROR_QUEUE_INTERVAL", "30"))
 
 
 def enable_multiple_works_in_default_container() -> bool:

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/app/runners/cloud.py

@@ -34,7 +34,7 @@ from lightning_cloud.openapi import (
     CloudspaceIdRunsBody,
     Externalv1LightningappInstance,
     Gridv1ImageSpec,
-    IdGetBody1,
+    IdGetBody,
     ProjectIdCloudspacesBody,
     V1BuildSpec,
     V1CloudSpace,
@@ -1027,7 +1027,7 @@ class CloudRuntime(Runtime):
             project_id=project_id,
             cloudspace_id=cloudspace_id,
             id=run_id,
-            body=IdGetBody1(
+            body=IdGetBody(
                 cluster_id=cluster_id,
                 name=run_name,
                 desired_state=desired_state,

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/app/runners/multiprocess.py

@@ -81,16 +81,17 @@ class MultiProcessRuntime(Runtime):
 
             _set_flow_context()
 
-            storage_orchestrator = StorageOrchestrator(
-                self.app,
-                self.app.request_queues,
-                self.app.response_queues,
-                self.app.copy_request_queues,
-                self.app.copy_response_queues,
-            )
-            self.threads.append(storage_orchestrator)
-            storage_orchestrator.setDaemon(True)
-            storage_orchestrator.start()
+            if constants.ENABLE_ORCHESTRATOR:
+                storage_orchestrator = StorageOrchestrator(
+                    self.app,
+                    self.app.request_queues,
+                    self.app.response_queues,
+                    self.app.copy_request_queues,
+                    self.app.copy_response_queues,
+                )
+                self.threads.append(storage_orchestrator)
+                storage_orchestrator.setDaemon(True)
+                storage_orchestrator.start()
 
             if self.start_server:
                 self.app.should_publish_changes_to_api = True

{lightning-2.3.0.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/app/utilities/network.py

@@ -89,18 +89,36 @@ _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,
+        # Any 4xx and 5xx statuses except
+        # 400 Bad Request
+        # 401 Unauthorized
+        # 403 Forbidden
+        # 404 Not Found
+        status_forcelist={
+            402,
+            *range(405, 600),
+        },
+        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 +175,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.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/fabric/CHANGELOG.md

@@ -9,11 +9,14 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Enabled consolidating distributed checkpoints through `fabric consolidate` in the new CLI [#19560](https://github.com/Lightning-AI/pytorch-lightning/pull/19560))
+- Enabled consolidating distributed checkpoints through `fabric consolidate` in the new CLI ([#19560](https://github.com/Lightning-AI/pytorch-lightning/pull/19560))
 
--
+- Added the ability to explicitly mark forward methods in Fabric via `_FabricModule.mark_forward_method()` ([#19690](https://github.com/Lightning-AI/pytorch-lightning/pull/19690))
+
+- Added support for PyTorch 2.3 ([#19708](https://github.com/Lightning-AI/pytorch-lightning/pull/19708))
+
+- Added `ModelParallelStrategy` to support 2D parallelism ([#19846](https://github.com/Lightning-AI/pytorch-lightning/pull/19846), [#19852](https://github.com/Lightning-AI/pytorch-lightning/pull/19852), [#19870](https://github.com/Lightning-AI/pytorch-lightning/pull/19870), [#19872](https://github.com/Lightning-AI/pytorch-lightning/pull/19872))
 
--
 
 ### Changed
 
@@ -39,7 +42,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 +50,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))
+- Fixed a matrix shape mismatch issue when running a model loaded from a quantized checkpoint (bitsandbytes) ([#19886](https://github.com/Lightning-AI/lightning/pull/19886))
 
-- Fixed issue where some model methods couldn't be monkeypatched after being Fabric wrapped ([#19705](https://github.com/Lightning-AI/pytorch-lightning/pull/19705))
+-
 
 -
 
 
+## [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.dev20240526/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.dev20240407 → lightning-2.3.0.dev20240526}/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.dev20240407 → lightning-2.3.0.dev20240526}/src/lightning/fabric/connector.py

@@ -62,6 +62,7 @@ from lightning.fabric.strategies import (
 )
 from lightning.fabric.strategies.ddp import _DDP_FORK_ALIASES
 from lightning.fabric.strategies.fsdp import _FSDP_ALIASES, FSDPStrategy
+from lightning.fabric.strategies.model_parallel import ModelParallelStrategy
 from lightning.fabric.utilities import rank_zero_info, rank_zero_warn
 from lightning.fabric.utilities.device_parser import _determine_root_gpu_device
 from lightning.fabric.utilities.imports import _IS_INTERACTIVE
@@ -429,7 +430,7 @@ class _Connector:
                 f" platform. We recommed `Fabric(strategy='ddp_spawn')` instead."
             )
         if (
-            strategy_flag in _FSDP_ALIASES or isinstance(self._strategy_flag, FSDPStrategy)
+            strategy_flag in _FSDP_ALIASES or type(self._strategy_flag) is FSDPStrategy
         ) and self._accelerator_flag not in ("cuda", "gpu"):
             raise ValueError(
                 "You selected the FSDP strategy but FSDP is only available on GPU. Set `Fabric(accelerator='gpu', ...)`"
@@ -460,6 +461,12 @@ class _Connector:
             return DeepSpeedPrecision(self._precision_input)  # type: ignore
         if isinstance(self.strategy, FSDPStrategy):
             return FSDPPrecision(precision=self._precision_input)  # type: ignore[arg-type]
+        mp_precision_supported = ("32-true", "bf16-mixed", "bf16-true", "16-true")
+        if isinstance(self.strategy, ModelParallelStrategy) and self._precision_input not in mp_precision_supported:
+            raise ValueError(
+                f"The `ModelParallelStrategy` does not support `Fabric(..., precision={self._precision_input!r})`."
+                f" Choose a different precision among: {', '.join(mp_precision_supported)}."
+            )
         if self._precision_input in ("16-true", "bf16-true"):
             return HalfPrecision(self._precision_input)  # type: ignore
         if self._precision_input == "32-true":