PyPI - lightning - Versions diffs - 2.3.0.dev20240602__tar.gz → 2.3.0.dev20240616__tar.gz - Mend

lightning 2.3.0.dev20240602tar.gz → 2.3.0.dev20240616tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (572) hide show

{lightning-2.3.0.dev20240602/src/lightning.egg-info → lightning-2.3.0.dev20240616}/PKG-INFO RENAMED Viewed

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.3.0.dev20240602
+Version: 2.3.0.dev20240616
 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.
@@ -74,11 +74,11 @@ License-File: LICENSE
 ______________________________________________________________________
 <p align="center">
-  <a href="https://lightning.ai/">Lightning.ai</a> •
-  <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.dev20240602">Docs</a> •
+  <a href="https://lightning.ai/">Lightning AI</a> •
+  <a href="#examples">Examples</a> •
+  <a href="#pytorch-lightning-train-and-deploy-pytorch-at-scale">PyTorch Lightning</a> •
+  <a href="#lightning-fabric-expert-control">Fabric</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240616">Docs</a> •
   <a href="#community">Community</a> •
   <a href="https://lightning.ai/docs/pytorch/stable/generated/CONTRIBUTING.html">Contribute</a> •
 </p>
@@ -101,6 +101,20 @@ ______________________________________________________________________
 </div>
+<div align="center">
+<p align="center">
+&nbsp;
+<a target="_blank" href="https://lightning.ai/docs/pytorch/latest/starter/introduction.html#define-a-lightningmodule">
+  <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/get-started-badge.svg" height="36px" alt="Get started"/>
+</a>
+</p>
+</div>
 ## Install Lightning
 Simple installation from PyPI
@@ -113,15 +127,11 @@ pip install lightning
 ______________________________________________________________________
-## Lightning has 4 core packages
+## Lightning has 2 core packages
 [PyTorch Lightning: Train and deploy PyTorch at scale](#pytorch-lightning-train-and-deploy-pytorch-at-scale).
 <br/>
 [Lightning Fabric: Expert control](#lightning-fabric-expert-control).
-<br/>
-[Lightning Data: Blazing fast, distributed streaming of training data from cloud storage](https://github.com/Lightning-AI/pytorch-lightning/tree/master/src/lightning/data).
-<br/>
-[Lightning Apps: Build AI products and ML workflows](#lightning-apps-build-ai-products-and-ml-workflows).
 Lightning gives you granular control over how much abstraction you want to add over PyTorch.
@@ -129,7 +139,9 @@ Lightning gives you granular control over how much abstraction you want to add o
     <img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/continuum.png" width="80%">
 </div>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 # PyTorch Lightning: Train and Deploy PyTorch at Scale
@@ -139,6 +151,15 @@ PyTorch Lightning is just organized PyTorch - Lightning disentangles PyTorch cod
 ______________________________________________________________________
+### Examples
+Explore various types of training possible with PyTorch Lightning. Pretrain and finetune ANY kind of model to perform ANY task like classification, segmentation, summarization and more:
+| Task  | Description | Run |
+|---|---|---|
+| [Hello world](#hello-simple-model)  | Pretrain - Hello world example | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/pytorch-lightning-hello-world"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
+| [Image segmentation](https://lightning.ai/lightning-ai/studios/image-segmentation-with-pytorch-lightning)  | Finetune - ResNet-50 model to segment images | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/image-segmentation-with-pytorch-lightning"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
+| [Text classification](https://lightning.ai/lightning-ai/studios/text-classification-with-pytorch-lightning)  | Finetune - text classifier (BERT model) | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/text-classification-with-pytorch-lightning"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
 ### Hello simple model
 ```python
@@ -340,7 +361,10 @@ ______________________________________________________________________
 ______________________________________________________________________
-# Lightning Fabric: Expert control.
+&nbsp;
+&nbsp;
+# Lightning Fabric: Expert control
 Run on any device at any scale with expert-level control over PyTorch training loop and scaling strategy. You can even write your own Trainer.
@@ -522,62 +546,8 @@ ______________________________________________________________________
 ______________________________________________________________________
-# Lightning Apps: Build AI products and ML workflows
-Lightning Apps remove the cloud infrastructure boilerplate so you can focus on solving the research or business problems. Lightning Apps can run on the Lightning Cloud, your own cluster or a private cloud.
-<div align="center">
-    <img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/lightning-apps-teaser.png" width="80%">
-</div>
-## Hello Lightning app world
-```python
-# app.py
-import lightning as L
-class TrainComponent(L.LightningWork):
-    def run(self, x):
-        print(f"train a model on {x}")
-class AnalyzeComponent(L.LightningWork):
-    def run(self, x):
-        print(f"analyze model on {x}")
-class WorkflowOrchestrator(L.LightningFlow):
-    def __init__(self) -> None:
-        super().__init__()
-        self.train = TrainComponent(cloud_compute=L.CloudCompute("cpu"))
-        self.analyze = AnalyzeComponent(cloud_compute=L.CloudCompute("gpu"))
-    def run(self):
-        self.train.run("CPU machine 1")
-        self.analyze.run("GPU machine 2")
-app = L.LightningApp(WorkflowOrchestrator())
-```
-Run on the cloud or locally
-```bash
-# run on the cloud
-lightning run app app.py --setup --cloud
-# run locally
-lightning run app app.py
-```
-______________________________________________________________________
-<div align="center">
-    <a href="https://lightning.ai/docs/app/stable/">Read the Lightning Apps docs</a>
-</div>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Examples
@@ -608,7 +578,8 @@ ______________________________________________________________________
 - [Logistic Regression](https://lightning-bolts.readthedocs.io/en/stable/models/classic_ml.html#logistic-regression)
 - [Linear Regression](https://lightning-bolts.readthedocs.io/en/stable/models/classic_ml.html#linear-regression)
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Continuous Integration
@@ -623,7 +594,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.dev20240602)](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.dev20240616)](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)                 |
@@ -632,7 +603,8 @@ Lightning is rigorously tested across multiple CPUs, GPUs and TPUs and against m
 </center>
 </details>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Community

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/README.md RENAMED Viewed

@@ -12,10 +12,10 @@
 ______________________________________________________________________
 <p align="center">
-  <a href="https://lightning.ai/">Lightning.ai</a> •
-  <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://lightning.ai/">Lightning AI</a> •
+  <a href="#examples">Examples</a> •
+  <a href="#pytorch-lightning-train-and-deploy-pytorch-at-scale">PyTorch Lightning</a> •
+  <a href="#lightning-fabric-expert-control">Fabric</a> •
   <a href="https://pytorch-lightning.readthedocs.io/en/stable/">Docs</a> •
   <a href="#community">Community</a> •
   <a href="https://lightning.ai/docs/pytorch/stable/generated/CONTRIBUTING.html">Contribute</a> •
@@ -39,6 +39,20 @@ ______________________________________________________________________
 </div>
+<div align="center">
+<p align="center">
+&nbsp;
+<a target="_blank" href="https://lightning.ai/docs/pytorch/latest/starter/introduction.html#define-a-lightningmodule">
+  <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/get-started-badge.svg" height="36px" alt="Get started"/>
+</a>
+</p>
+</div>
 ## Install Lightning
 Simple installation from PyPI
@@ -92,15 +106,11 @@ pip install -iU https://test.pypi.org/simple/ pytorch-lightning
 ______________________________________________________________________
-## Lightning has 4 core packages
+## Lightning has 2 core packages
 [PyTorch Lightning: Train and deploy PyTorch at scale](#pytorch-lightning-train-and-deploy-pytorch-at-scale).
 <br/>
 [Lightning Fabric: Expert control](#lightning-fabric-expert-control).
-<br/>
-[Lightning Data: Blazing fast, distributed streaming of training data from cloud storage](https://github.com/Lightning-AI/pytorch-lightning/tree/master/src/lightning/data).
-<br/>
-[Lightning Apps: Build AI products and ML workflows](#lightning-apps-build-ai-products-and-ml-workflows).
 Lightning gives you granular control over how much abstraction you want to add over PyTorch.
@@ -108,7 +118,9 @@ Lightning gives you granular control over how much abstraction you want to add o
     <img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/continuum.png" width="80%">
 </div>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 # PyTorch Lightning: Train and Deploy PyTorch at Scale
@@ -118,6 +130,15 @@ PyTorch Lightning is just organized PyTorch - Lightning disentangles PyTorch cod
 ______________________________________________________________________
+### Examples
+Explore various types of training possible with PyTorch Lightning. Pretrain and finetune ANY kind of model to perform ANY task like classification, segmentation, summarization and more:
+| Task  | Description | Run |
+|---|---|---|
+| [Hello world](#hello-simple-model)  | Pretrain - Hello world example | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/pytorch-lightning-hello-world"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
+| [Image segmentation](https://lightning.ai/lightning-ai/studios/image-segmentation-with-pytorch-lightning)  | Finetune - ResNet-50 model to segment images | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/image-segmentation-with-pytorch-lightning"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
+| [Text classification](https://lightning.ai/lightning-ai/studios/text-classification-with-pytorch-lightning)  | Finetune - text classifier (BERT model) | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/text-classification-with-pytorch-lightning"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
 ### Hello simple model
 ```python
@@ -319,7 +340,10 @@ ______________________________________________________________________
 ______________________________________________________________________
-# Lightning Fabric: Expert control.
+&nbsp;
+&nbsp;
+# Lightning Fabric: Expert control
 Run on any device at any scale with expert-level control over PyTorch training loop and scaling strategy. You can even write your own Trainer.
@@ -501,62 +525,8 @@ ______________________________________________________________________
 ______________________________________________________________________
-# Lightning Apps: Build AI products and ML workflows
-Lightning Apps remove the cloud infrastructure boilerplate so you can focus on solving the research or business problems. Lightning Apps can run on the Lightning Cloud, your own cluster or a private cloud.
-<div align="center">
-    <img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/lightning-apps-teaser.png" width="80%">
-</div>
-## Hello Lightning app world
-```python
-# app.py
-import lightning as L
-class TrainComponent(L.LightningWork):
-    def run(self, x):
-        print(f"train a model on {x}")
-class AnalyzeComponent(L.LightningWork):
-    def run(self, x):
-        print(f"analyze model on {x}")
-class WorkflowOrchestrator(L.LightningFlow):
-    def __init__(self) -> None:
-        super().__init__()
-        self.train = TrainComponent(cloud_compute=L.CloudCompute("cpu"))
-        self.analyze = AnalyzeComponent(cloud_compute=L.CloudCompute("gpu"))
-    def run(self):
-        self.train.run("CPU machine 1")
-        self.analyze.run("GPU machine 2")
-app = L.LightningApp(WorkflowOrchestrator())
-```
-Run on the cloud or locally
-```bash
-# run on the cloud
-lightning run app app.py --setup --cloud
-# run locally
-lightning run app app.py
-```
-______________________________________________________________________
-<div align="center">
-    <a href="https://lightning.ai/docs/app/stable/">Read the Lightning Apps docs</a>
-</div>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Examples
@@ -587,7 +557,8 @@ ______________________________________________________________________
 - [Logistic Regression](https://lightning-bolts.readthedocs.io/en/stable/models/classic_ml.html#logistic-regression)
 - [Linear Regression](https://lightning-bolts.readthedocs.io/en/stable/models/classic_ml.html#linear-regression)
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Continuous Integration
@@ -611,7 +582,8 @@ Lightning is rigorously tested across multiple CPUs, GPUs and TPUs and against m
 </center>
 </details>
-______________________________________________________________________
+&nbsp;
+&nbsp;
 ## Community

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/pyproject.toml RENAMED Viewed

@@ -12,11 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
-[metadata]
-name = "lightning"
-author = "Lightning-AI et al."
-url = "https://github.com/Lightning-AI/lightning"
 [build-system]
 requires = [
     "setuptools",

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/requirements/app/app.txt RENAMED Viewed

@@ -23,7 +23,7 @@ Jinja2 <3.2.0
 PyYAML <=6.0.1
 requests <2.32.0
 rich >=12.3.0, <13.6.0
-urllib3 <2.1.0
+urllib3 <2.0.0
 uvicorn <0.24.0
 websocket-client <1.7.0
 websockets <11.1.0

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/CHANGELOG.md RENAMED Viewed

@@ -5,57 +5,33 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
-## [unReleased] - 2024-MM-DD
+## [2.3.0] - 2024-06-13
 ### Added
+- Added sanitization for classes before logging them as hyperparameters ([#19771](https://github.com/Lightning-AI/pytorch-lightning/pull/19771))
 - 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))
+- Added a call to `torch.distributed.destroy_process_group` in atexit handler if process group needs destruction ([#19931](https://github.com/Lightning-AI/pytorch-lightning/pull/19931))
+- Added support for configuring hybrid-sharding by passing a tuple for the `FSDPStrategy(device_mesh=...)` argument ([#19504](https://github.com/Lightning-AI/pytorch-lightning/pull/19504))
 ### Changed
 - Renamed `lightning run model` to `fabric run` ([#19442](https://github.com/Lightning-AI/pytorch-lightning/pull/19442), [#19527](https://github.com/Lightning-AI/pytorch-lightning/pull/19527))
 - The `Fabric.rank_zero_first` context manager now uses a barrier without timeout to avoid long-running tasks to be interrupted ([#19448](https://github.com/Lightning-AI/lightning/pull/19448))
 - Fabric now raises an error if you forget to call `fabric.backward()` when it is needed by the strategy or precision selection ([#19447](https://github.com/Lightning-AI/lightning/pull/19447), [#19493](https://github.com/Lightning-AI/lightning/pull/19493))
 - `_BackwardSyncControl` can now control what to do when gradient accumulation is disabled ([#19577](https://github.com/Lightning-AI/lightning/pull/19577))
-### Deprecated
--
--
--
 ### Removed
 - Removed support for PyTorch 1.13 ([#19706](https://github.com/Lightning-AI/lightning/pull/19706))
--
--
 ### Fixed
 - 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))
--
--
 ## [2.2.2] - 2024-04-11

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/connector.py RENAMED Viewed

@@ -313,7 +313,8 @@ class _Connector:
                 f" using {accelerator_name} accelerator."
             )
-    def _choose_auto_accelerator(self) -> str:
+    @staticmethod
+    def _choose_auto_accelerator() -> str:
         """Choose the accelerator type (str) based on availability when ``accelerator='auto'``."""
         if XLAAccelerator.is_available():
             return "tpu"

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/strategies/fsdp.py RENAMED Viewed

@@ -74,12 +74,14 @@ from lightning.fabric.utilities.seed import reset_seed
 from lightning.fabric.utilities.types import _PATH, _Stateful
 if TYPE_CHECKING:
+    from torch.distributed.device_mesh import DeviceMesh
     from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, MixedPrecision, ShardingStrategy
     from torch.distributed.fsdp.wrap import ModuleWrapPolicy
     _POLICY = Union[Set[Type[Module]], Callable[[Module, bool, int], bool], ModuleWrapPolicy]
     _SHARDING_STRATEGY = Union[ShardingStrategy, Literal["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD", "HYBRID_SHARD"]]
 _FSDP_ALIASES = ("fsdp", "fsdp_cpu_offload")
@@ -117,10 +119,14 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
             - ``"SHARD_GRAD_OP"``: Shards gradients and optimizer states only. Model parameters get replicated.
             - ``"NO_SHARD"``: No sharding (identical to regular DDP).
             - ``"HYBRID_SHARD"``: Shards model parameters, gradients, and optimizer states within a single machine, but
-              replicates across machines.
+              replicates across machines. See also the `device_mesh` parameter below.
             Also accepts a :class:`torch.distributed.fsdp.ShardingStrategy` enum value.
+        device_mesh: A tuple `(replication size, sharding size)` that defines over how many devices to shard and
+            replicate the model. The product of the two numbers must equal the world size. Only valid in combination
+            with the `HYBRID_SHARD` sharding strategy.
         state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
             - ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
@@ -146,6 +152,7 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
         activation_checkpointing_policy: Optional["_POLICY"] = None,
         sharding_strategy: "_SHARDING_STRATEGY" = "FULL_SHARD",
         state_dict_type: Literal["full", "sharded"] = "sharded",
+        device_mesh: Optional[Union[Tuple[int], "DeviceMesh"]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -163,6 +170,11 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
         # Enables joint setup of model and optimizer, multiple optimizer param groups, and `torch.compile()`
         self._fsdp_kwargs.setdefault("use_orig_params", True)
+        if device_mesh is not None:
+            if not _TORCH_GREATER_EQUAL_2_2:
+                raise ValueError("The `device_mesh` argument is only supported in torch >= 2.2.")
+            self._fsdp_kwargs["device_mesh"] = device_mesh
         self._activation_checkpointing_kwargs = _activation_checkpointing_kwargs(
             activation_checkpointing, activation_checkpointing_policy
         )
@@ -244,6 +256,12 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
         super().setup_environment()
         self._setup_distributed()
+        # if 'device_mesh' in the `_fsdp_kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
+        if isinstance(self._fsdp_kwargs.get("device_mesh"), tuple):
+            from torch.distributed.device_mesh import init_device_mesh
+            self._fsdp_kwargs["device_mesh"] = init_device_mesh("cuda", self._fsdp_kwargs["device_mesh"])
     @override
     def setup_module_and_optimizers(
         self, module: Module, optimizers: List[Optimizer]

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/strategies/model_parallel.py RENAMED Viewed

@@ -275,12 +275,7 @@ class ModelParallelStrategy(ParallelStrategy):
         state: Optional[Union[Module, Optimizer, Dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
     ) -> Dict[str, Any]:
-        """Load the contents from a checkpoint and restore the state of the given objects.
-        Currently does not support loading the optimizer state if the model is distributed but the checkpoint is a full,
-        non-distributed checkpoint.
-        """
+        """Load the contents from a checkpoint and restore the state of the given objects."""
         if not state:
             raise ValueError(
                 f"Got {type(self).__name__}.load_checkpoint(..., state={state!r}) but a state with at least "
@@ -559,14 +554,14 @@ def _load_raw_module_state(
         state_dict_options = StateDictOptions(
             broadcast_from_rank0=True,  # type: ignore[call-arg]
             full_state_dict=True,
-            strict=strict,  # gets ignored at the moment
+            # must be set False to allow loading each param separately below
+            strict=False,
         )
         for submodule_name, submodule in module.named_modules():
             for param_name, _ in _named_parameters_and_buffers_to_load(submodule):
                 full_param_name = f"{submodule_name}{'.' if submodule_name else ''}{param_name}"
                 if full_param_name not in state_dict:
-                    # Note: PyTorch does not currently respect the `strict` setting in state_dict_options!
                     if not strict:
                         continue
                     raise KeyError(

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/utilities/distributed.py RENAMED Viewed

@@ -1,3 +1,4 @@
+import atexit
 import contextlib
 import logging
 import os
@@ -291,6 +292,10 @@ def _init_dist_connection(
     log.info(f"Initializing distributed: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
     torch.distributed.init_process_group(torch_distributed_backend, rank=global_rank, world_size=world_size, **kwargs)
+    if torch_distributed_backend == "nccl":
+        # PyTorch >= 2.4 warns about undestroyed NCCL process group, so we need to do it at program exit
+        atexit.register(_destroy_dist_connection)
     # On rank=0 let everyone know training is starting
     rank_zero_info(
         f"{'-' * 100}\n"
@@ -300,6 +305,11 @@ def _init_dist_connection(
     )
+def _destroy_dist_connection() -> None:
+    if _distributed_is_initialized():
+        torch.distributed.destroy_process_group()
 def _get_default_process_group_backend_for_device(device: torch.device) -> str:
     return "nccl" if device.type == "cuda" else "gloo"

{lightning-2.3.0.dev20240602 → lightning-2.3.0.dev20240616}/src/lightning/fabric/utilities/logger.py RENAMED Viewed

@@ -11,6 +11,8 @@
 # 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 inspect
 import json
 from argparse import Namespace
 from dataclasses import asdict, is_dataclass
@@ -52,8 +54,11 @@ def _sanitize_callable_params(params: Dict[str, Any]) -> Dict[str, Any]:
     """
     def _sanitize_callable(val: Any) -> Any:
-        # Give them one chance to return a value. Don't go rabbit hole of recursive call
+        if inspect.isclass(val):
+            # If it's a class, don't try to instantiate it, just return the name
+            return val.__name__
         if callable(val):
+            # Callables get a chance to return a name
             try:
                 _val = val()
                 if callable(_val):

lightning 2.3.0.dev20240602__tar.gz → 2.3.0.dev20240616__tar.gz

lightning 2.3.0.dev20240602tar.gz → 2.3.0.dev20240616tar.gz