lightning 2.3.0.dev20240526__tar.gz → 2.3.0.dev20240609__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.3.0.dev20240526
+Version: 2.3.0.dev20240609
 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.
@@ -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.dev20240526">Docs</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240609">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.dev20240526)](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.dev20240609)](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.dev20240526 → lightning-2.3.0.dev20240609}/pyproject.toml

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/CHANGELOG.md

@@ -9,6 +9,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### 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))
@@ -17,6 +19,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - 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
 

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/connector.py

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/strategies/fsdp.py

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/strategies/model_parallel.py

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/utilities/distributed.py

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/utilities/logger.py

@@ -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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/fabric/utilities/throughput.py

@@ -296,7 +296,7 @@ def measure_flops(
         raise ImportError("`measure_flops` requires PyTorch >= 2.1.")
     from torch.utils.flop_counter import FlopCounterMode
 
-    flop_counter = FlopCounterMode(model, display=False)
+    flop_counter = FlopCounterMode(display=False)
     with flop_counter:
         if loss_fn is None:
             forward_fn()

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/CHANGELOG.md

@@ -20,6 +20,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `ModelParallelStrategy` to support 2D parallelism ([#19878](https://github.com/Lightning-AI/pytorch-lightning/pull/19878), [#19888](https://github.com/Lightning-AI/pytorch-lightning/pull/19888))
 
+- 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
@@ -28,7 +31,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Relaxed the requirement for custom batch samplers to expose `drop_last` for prediction ([#19678](https://github.com/Lightning-AI/pytorch-lightning/pull/19678))
 
--
+- It is no longer allowed to skip `training_step()` by returning `None` in distributed training ([#19918](https://github.com/Lightning-AI/pytorch-lightning/pull/19918))
+
 
 ### Deprecated
 
@@ -53,7 +57,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed `WandbLogger.log_hyperparameters()` raising an error if hyperparameters are not JSON serializable ([#19769](https://github.com/Lightning-AI/pytorch-lightning/pull/19769))
 
--
+
+- Fixed an issue with the LightningCLI not being able to set the `ModelCheckpoint(save_last=...)` argument ([#19808](https://github.com/Lightning-AI/pytorch-lightning/pull/19808))
+
+- Fixed an issue causing ValueError for certain object such as TorchMetrics when dumping hyperparameters to YAML ([#19804](https://github.com/Lightning-AI/pytorch-lightning/pull/19804))
+
+
+- Fixed resetting `epoch_loop.restarting` to avoid full validation run after `LearningRateFinder` ([#19818](https://github.com/Lightning-AI/pytorch-lightning/issues/19818))
 
 
 ## [2.2.2] - 2024-04-11

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -27,7 +27,7 @@ import warnings
 from copy import deepcopy
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Set
+from typing import Any, Dict, Literal, Optional, Set, Union
 from weakref import proxy
 
 import torch
@@ -216,7 +216,7 @@ class ModelCheckpoint(Checkpoint):
         filename: Optional[str] = None,
         monitor: Optional[str] = None,
         verbose: bool = False,
-        save_last: Optional[Literal[True, False, "link"]] = None,
+        save_last: Optional[Union[bool, Literal["link"]]] = None,
         save_top_k: int = 1,
         save_weights_only: bool = False,
         mode: str = "min",

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/core/module.py

@@ -405,7 +405,7 @@ class LightningModule(
         The default behavior per hook is documented here: :ref:`extensions/logging:Automatic Logging`.
 
         Args:
-            name: key to log.
+            name: key to log. Must be identical across all processes if using DDP or any other distributed strategy.
             value: value to log. Can be a ``float``, ``Tensor``, or a ``Metric``.
             prog_bar: if ``True`` logs to the progress bar.
             logger: if ``True`` logs to the logger.
@@ -569,6 +569,7 @@ class LightningModule(
 
         Args:
             dictionary: key value pairs.
+                Keys must be identical across all processes if using DDP or any other distributed strategy.
                 The values can be a ``float``, ``Tensor``, ``Metric``, or ``MetricCollection``.
             prog_bar: if ``True`` logs to the progress base.
             logger: if ``True`` logs to the logger.

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/core/saving.py

@@ -359,7 +359,7 @@ def save_hparams_to_yaml(config_yaml: _PATH, hparams: Union[dict, Namespace], us
         try:
             v = v.name if isinstance(v, Enum) else v
             yaml.dump(v)
-        except TypeError:
+        except (TypeError, ValueError):
             warn(f"Skipping '{k}' parameter because it is not possible to safely dump to YAML.")
             hparams[k] = type(v).__name__
         else:

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/loggers/logger.py

@@ -15,11 +15,11 @@
 
 import functools
 import operator
+import statistics
 from abc import ABC
 from collections import defaultdict
 from typing import Any, Callable, Dict, Mapping, Optional, Sequence
 
-import numpy as np
 from typing_extensions import override
 
 from lightning.fabric.loggers import Logger as FabricLogger
@@ -100,7 +100,7 @@ class DummyLogger(Logger):
 def merge_dicts(  # pragma: no cover
     dicts: Sequence[Mapping],
     agg_key_funcs: Optional[Mapping] = None,
-    default_func: Callable[[Sequence[float]], float] = np.mean,
+    default_func: Callable[[Sequence[float]], float] = statistics.mean,
 ) -> Dict:
     """Merge a sequence with dictionaries into one dictionary by aggregating the same keys with some given function.
 
@@ -126,7 +126,7 @@ def merge_dicts(  # pragma: no cover
         >>> d2 = {'a': 1.1, 'b': 2.2, 'v': 1, 'd': {'d1': 2, 'd2': 3}}
         >>> d3 = {'a': 1.1, 'v': 2.3, 'd': {'d3': 3, 'd4': {'d5': 1}}}
         >>> dflt_func = min
-        >>> agg_funcs = {'a': np.mean, 'v': max, 'd': {'d1': sum}}
+        >>> agg_funcs = {'a': statistics.mean, 'v': max, 'd': {'d1': sum}}
         >>> pprint.pprint(merge_dicts([d1, d2, d3], agg_funcs, dflt_func))
         {'a': 1.3,
          'b': 2.0,

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/loops/optimization/automatic.py

@@ -314,8 +314,14 @@ class _AutomaticOptimization(_Loop):
         """
         trainer = self.trainer
 
-        # manually capture logged metrics
         training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
         self.trainer.strategy.post_training_step()  # unused hook - call anyway for backward compatibility
 
+        if training_step_output is None and trainer.world_size > 1:
+            raise RuntimeError(
+                "Skipping the `training_step` by returning None in distributed training is not supported."
+                " It is recommended that you rewrite your training logic to avoid having to skip the step in the first"
+                " place."
+            )
+
         return self.output_result_cls.from_training_step_output(training_step_output, trainer.accumulate_grad_batches)

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/strategies/fsdp.py

@@ -16,7 +16,21 @@ import shutil
 from contextlib import contextmanager, nullcontext
 from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Literal, Mapping, Optional, Set, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 import torch
 from lightning_utilities.core.rank_zero import rank_zero_only as utils_rank_zero_only
@@ -53,7 +67,10 @@ from lightning.fabric.utilities.distributed import (
     _sync_ddp_if_available,
 )
 from lightning.fabric.utilities.distributed import group as _group
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
+from lightning.fabric.utilities.imports import (
+    _TORCH_GREATER_EQUAL_2_1,
+    _TORCH_GREATER_EQUAL_2_2,
+)
 from lightning.fabric.utilities.init import _EmptyInit, _has_meta_device_parameters_or_buffers
 from lightning.fabric.utilities.load import _lazy_load, _materialize_tensors
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
@@ -70,6 +87,7 @@ from lightning.pytorch.utilities.model_helpers import is_overridden
 from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_only, rank_zero_warn
 
 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
 
@@ -114,10 +132,14 @@ class FSDPStrategy(ParallelStrategy):
             - ``"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.
@@ -147,6 +169,7 @@ class FSDPStrategy(ParallelStrategy):
         activation_checkpointing_policy: Optional["_POLICY"] = None,
         sharding_strategy: "_SHARDING_STRATEGY" = "FULL_SHARD",
         state_dict_type: Literal["full", "sharded"] = "full",
+        device_mesh: Optional[Union[Tuple[int], "DeviceMesh"]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -162,6 +185,12 @@ class FSDPStrategy(ParallelStrategy):
         self.cpu_offload = _init_cpu_offload(cpu_offload)
         self.mixed_precision = mixed_precision
         self.kwargs = _auto_wrap_policy_kwargs(auto_wrap_policy, kwargs)
+
+        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.kwargs["device_mesh"] = device_mesh
+
         self.sharding_strategy = _init_sharding_strategy(sharding_strategy, self.kwargs)
 
         # Avoids the need for user to reference params in `configure_optimizers` via
@@ -242,6 +271,12 @@ class FSDPStrategy(ParallelStrategy):
         assert self.cluster_environment is not None
         _init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
 
+        # if 'device_mesh' in the `kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
+        if isinstance(self.kwargs.get("device_mesh"), tuple):
+            from torch.distributed.device_mesh import init_device_mesh
+
+            self.kwargs["device_mesh"] = init_device_mesh("cuda", self.kwargs["device_mesh"])
+
     def _get_process_group_backend(self) -> str:
         return self._process_group_backend or _get_default_process_group_backend_for_device(self.root_device)
 

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -328,7 +328,8 @@ class _AcceleratorConnector:
                 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."""
         if XLAAccelerator.is_available():
             return "tpu"

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/trainer/trainer.py

@@ -1654,8 +1654,8 @@ class Trainer:
     def estimated_stepping_batches(self) -> Union[int, float]:
         r"""The estimated number of batches that will ``optimizer.step()`` during training.
 
-        This accounts for gradient accumulation and the current trainer configuration. This might sets up your training
-        dataloader if hadn't been set up already.
+        This accounts for gradient accumulation and the current trainer configuration. This might be used when setting
+        up your training dataloader, if it hasn't been set up already.
 
         .. code-block:: python
 

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/tuner/lr_finder.py

@@ -190,7 +190,7 @@ class _LRFinder:
         losses = losses[torch.isfinite(losses)]
 
         if len(losses) < 2:
-            # computing np.gradient requires at least 2 points
+            # computing torch.gradient requires at least 2 points
             log.error(
                 "Failed to compute suggestion for learning rate because there are not enough points. Increase the loop"
                 " iteration limits or the size of your dataset/dataloader."
@@ -301,6 +301,7 @@ def _lr_find(
     trainer._checkpoint_connector.restore(ckpt_path)
     trainer.strategy.remove_checkpoint(ckpt_path)
     trainer.fit_loop.restarting = False  # reset restarting flag as checkpoint restoring sets it to True
+    trainer.fit_loop.epoch_loop.restarting = False  # reset restarting flag as checkpoint restoring sets it to True
     trainer.fit_loop.epoch_loop.val_loop._combined_loader = None
 
     return lr_finder

{lightning-2.3.0.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning/pytorch/utilities/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """General utilities."""
 
-import numpy
+import torch
 
 from lightning.fabric.utilities import (
     LightningEnum,
@@ -55,6 +55,6 @@ __all__ = [
     "suggested_max_num_workers",
 ]
 
-FLOAT16_EPSILON = numpy.finfo(numpy.float16).eps
-FLOAT32_EPSILON = numpy.finfo(numpy.float32).eps
-FLOAT64_EPSILON = numpy.finfo(numpy.float64).eps
+FLOAT16_EPSILON = torch.finfo(torch.float16).eps
+FLOAT32_EPSILON = torch.finfo(torch.float32).eps
+FLOAT64_EPSILON = torch.finfo(torch.float64).eps

lightning-2.3.0.dev20240609/src/lightning/version.info

@@ -0,0 +1 @@
+2.3.0.dev20240609

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.3.0.dev20240526
+Version: 2.3.0.dev20240609
 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.
@@ -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.dev20240526">Docs</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.3.0.dev20240609">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.dev20240526)](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.dev20240609)](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.dev20240526 → lightning-2.3.0.dev20240609}/src/lightning.egg-info/SOURCES.txt

@@ -23,7 +23,7 @@ requirements/pytorch/extra.txt
 requirements/pytorch/strategies.txt
 requirements/pytorch/test.txt
 src/version.info
-src/lightning/9da8d09d39dabf8f876897b29a49edf1b4d01405126f6faa
+src/lightning/4a4ff9257060ce7cb0debe4b69e27ff0d444d2aaba3d9f61
 src/lightning/__about__.py
 src/lightning/__init__.py
 src/lightning/__main__.py

lightning-2.3.0.dev20240609/src/version.info

@@ -0,0 +1 @@
+2.3.0.dev20240609

lightning-2.3.0.dev20240526/src/lightning/version.info

@@ -1 +0,0 @@
-2.3.0.dev20240526

lightning-2.3.0.dev20240526/src/version.info

@@ -1 +0,0 @@
-2.3.0.dev20240526