lightning 2.4.0.dev20240721__tar.gz → 2.4.0.dev20240804__tar.gz

{lightning-2.4.0.dev20240721/src/lightning.egg-info → lightning-2.4.0.dev20240804}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning
-Version: 2.4.0.dev20240721
+Version: 2.4.0.dev20240804
 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,22 +27,22 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-Provides-Extra: fabric-strategies
-Provides-Extra: fabric-test
-Provides-Extra: fabric-examples
-Provides-Extra: pytorch-strategies
+Provides-Extra: pytorch-examples
 Provides-Extra: pytorch-extra
 Provides-Extra: pytorch-test
-Provides-Extra: pytorch-examples
+Provides-Extra: pytorch-strategies
+Provides-Extra: fabric-examples
+Provides-Extra: fabric-test
+Provides-Extra: fabric-strategies
 Provides-Extra: fabric-all
 Provides-Extra: fabric-dev
 Provides-Extra: pytorch-all
 Provides-Extra: pytorch-dev
-Provides-Extra: strategies
-Provides-Extra: test
 Provides-Extra: examples
-Provides-Extra: data
 Provides-Extra: extra
+Provides-Extra: test
+Provides-Extra: strategies
+Provides-Extra: data
 Provides-Extra: all
 Provides-Extra: dev
 License-File: LICENSE
@@ -65,7 +65,7 @@ ______________________________________________________________________
   <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.4.0.dev20240721">Docs</a> •
+  <a href="https://pytorch-lightning.readthedocs.io/en/2.4.0.dev20240804">Docs</a> •
   <a href="#community">Community</a> •
   <a href="https://lightning.ai/docs/pytorch/stable/generated/CONTRIBUTING.html">Contribute</a> •
 </p>
@@ -141,13 +141,16 @@ ______________________________________________________________________
 ### 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> |   
+| 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 classification](https://lightning.ai/lightning-ai/studios/image-classification-with-pytorch-lightning) | Finetune - ResNet-34 model to classify images of cars          | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/image-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> |   
+| [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> |   
+| [Object detection](https://lightning.ai/lightning-ai/studios/object-detection-with-pytorch-lightning)       | Finetune - Faster R-CNN model to detect objects                   | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/object-detection-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> |   
 | [Text summarization](https://lightning.ai/lightning-ai/studios/text-summarization-with-pytorch-lightning)   | Finetune - text summarization (Hugging Face transformer model) | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/text-summarization-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> |   
-| [Audio generation](https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator)        | Finetune - audio generator (transformer model)                | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |   
+| [Audio generation](https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator)        | Finetune - audio generator (transformer model)                 | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |   
+| [LLM finetuning](https://lightning.ai/lightning-ai/studios/finetune-an-llm-with-pytorch-lightning)          | Finetune - LLM (Meta Llama 3.1 8B)                | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-an-llm-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
 
@@ -583,7 +586,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.4.0.dev20240721)](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.4.0.dev20240804)](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.4.0.dev20240721 → lightning-2.4.0.dev20240804}/README.md

@@ -133,13 +133,16 @@ ______________________________________________________________________
 ### 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> |   
+| 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 classification](https://lightning.ai/lightning-ai/studios/image-classification-with-pytorch-lightning) | Finetune - ResNet-34 model to classify images of cars          | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/image-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> |   
+| [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> |   
+| [Object detection](https://lightning.ai/lightning-ai/studios/object-detection-with-pytorch-lightning)       | Finetune - Faster R-CNN model to detect objects                   | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/object-detection-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> |   
 | [Text summarization](https://lightning.ai/lightning-ai/studios/text-summarization-with-pytorch-lightning)   | Finetune - text summarization (Hugging Face transformer model) | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/text-summarization-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> |   
-| [Audio generation](https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator)        | Finetune - audio generator (transformer model)                | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |   
+| [Audio generation](https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator)        | Finetune - audio generator (transformer model)                 | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |   
+| [LLM finetuning](https://lightning.ai/lightning-ai/studios/finetune-an-llm-with-pytorch-lightning)          | Finetune - LLM (Meta Llama 3.1 8B)                | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/finetune-an-llm-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
 

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/CHANGELOG.md

@@ -35,7 +35,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed an attribute error when loading a checkpoint into a quantized model using the `_lazy_load()` function ([#20121](https://github.com/Lightning-AI/lightning/pull/20121))
+
 
 -
 

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/plugins/collectives/torch_collective.py

@@ -95,11 +95,11 @@ class TorchCollective(Collective):
 
     @override
     def send(self, tensor: Tensor, dst: int, tag: int = 0) -> None:
-        dist.send(tensor, dst, tag=tag, group=self.group)
+        dist.send(tensor, dst, tag=tag, group=self.group)  # type: ignore[arg-type]
 
     @override
     def recv(self, tensor: Tensor, src: Optional[int] = None, tag: int = 0) -> Tensor:
-        dist.recv(tensor, src, tag=tag, group=self.group)
+        dist.recv(tensor, src, tag=tag, group=self.group)  # type: ignore[arg-type]
         return tensor
 
     def all_gather_object(self, object_list: List[Any], obj: Any) -> List[Any]:

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/plugins/precision/amp.py

@@ -40,7 +40,7 @@ class MixedPrecision(Precision):
         self,
         precision: Literal["16-mixed", "bf16-mixed"],
         device: str,
-        scaler: Optional["torch.cuda.amp.GradScaler"] = None,
+        scaler: Optional["torch.amp.GradScaler"] = None,
     ) -> None:
         if precision not in ("16-mixed", "bf16-mixed"):
             raise ValueError(

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/strategies/fsdp.py

@@ -864,7 +864,7 @@ def _distributed_checkpoint_save(converted_state: Dict[str, Any], path: Path) ->
 
         # let torch automatically infer the writer to use. This might also support fsspec paths in the future
         # https://github.com/pytorch/pytorch/issues/118036
-        save(converted_state, checkpoint_id=path)  # type: ignore[call-arg]
+        save(converted_state, checkpoint_id=path)
     else:  # deprecated
         from torch.distributed.checkpoint import FileSystemWriter
 
@@ -883,7 +883,7 @@ def _distributed_checkpoint_load(module_state: Dict[str, Any], path: Path) -> No
 
         # let torch automatically infer the reader to use. This might also support fsspec paths in the future
         # https://github.com/pytorch/pytorch/issues/118036
-        load(module_state, checkpoint_id=path)  # type: ignore[call-arg]
+        load(module_state, checkpoint_id=path)
     else:  # deprecated
         from torch.distributed.checkpoint import FileSystemReader
 

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/strategies/model_parallel.py

@@ -172,7 +172,7 @@ class ModelParallelStrategy(ParallelStrategy):
         )
 
     @override
-    def setup_module(self, module: TModel) -> TModel:
+    def setup_module(self, module: Module) -> Module:
         from torch.distributed.fsdp import FullyShardedDataParallel
 
         if any(isinstance(mod, FullyShardedDataParallel) for mod in module.modules()):
@@ -181,7 +181,7 @@ class ModelParallelStrategy(ParallelStrategy):
                 f" The `{self.__class__.__name__}` only supports the new FSDP2 APIs in PyTorch >= 2.4."
             )
 
-        module = self._parallelize_fn(module, self.device_mesh)
+        module = self._parallelize_fn(module, self.device_mesh)  # type: ignore[arg-type]
         if not isinstance(module, Module):
             raise TypeError(
                 f"The `parallelize_fn` must return a `nn.Module` instance, but got: {type(module).__name__}"
@@ -462,7 +462,7 @@ def _load_checkpoint(
         _load_raw_module_state(checkpoint.pop(module_key), module, strict=strict)
 
         state_dict_options = StateDictOptions(
-            broadcast_from_rank0=True,  # type: ignore[call-arg]
+            broadcast_from_rank0=True,
             full_state_dict=True,
             strict=strict,
         )
@@ -546,7 +546,7 @@ def _load_raw_module_state(
         from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
 
         state_dict_options = StateDictOptions(
-            broadcast_from_rank0=True,  # type: ignore[call-arg]
+            broadcast_from_rank0=True,
             full_state_dict=True,
             # must be set False to allow loading each param separately below
             strict=False,

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/utilities/load.py

@@ -160,7 +160,7 @@ class _NotYetLoadedTensor:
             return getattr(self.metatensor, name)
 
         # materializing these is needed for quantization (see lit-gpt)
-        if name in {"contiguous", "cuda", "half"}:
+        if name in {"contiguous", "cuda", "half", "data"}:
             return getattr(self._load_tensor(), name)
 
         raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/fabric/utilities/logger.py

@@ -95,7 +95,7 @@ def _flatten_dict(params: MutableMapping[Any, Any], delimiter: str = "/", parent
     result: Dict[str, Any] = {}
     for k, v in params.items():
         new_key = parent_key + delimiter + str(k) if parent_key else str(k)
-        if is_dataclass(v):
+        if is_dataclass(v) and not isinstance(v, type):
             v = asdict(v)
         elif isinstance(v, Namespace):
             v = vars(v)

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added a flag `verbose` to the `seed_everything()` function ([#20108](https://github.com/Lightning-AI/pytorch-lightning/pull/20108))
 
+-
 
 ### Changed
 
@@ -44,6 +45,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed an issue that would cause too many printouts of the seed info when using `seed_everything()` ([#20108](https://github.com/Lightning-AI/pytorch-lightning/pull/20108))
 
+- Fixed `_LoggerConnector`'s `_ResultMetric` to move all registered keys to the device of the logged value if needed ([#19814](https://github.com/Lightning-AI/pytorch-lightning/issues/19814))
+
 
 
 ## [2.3.0] - 2024-06-13

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/callbacks/progress/rich_progress.py

@@ -187,9 +187,6 @@ if _RICH_AVAILABLE:
                     value = f"{value:{self._metrics_format}}"
                 yield f"{name}: {value}"
 
-else:
-    Task, Style = Any, Any  # type: ignore[assignment, misc]
-
 
 @dataclass
 class RichProgressBarTheme:
@@ -209,14 +206,14 @@ class RichProgressBarTheme:
 
     """
 
-    description: Union[str, Style] = "white"
-    progress_bar: Union[str, Style] = "#6206E0"
-    progress_bar_finished: Union[str, Style] = "#6206E0"
-    progress_bar_pulse: Union[str, Style] = "#6206E0"
-    batch_progress: Union[str, Style] = "white"
-    time: Union[str, Style] = "grey54"
-    processing_speed: Union[str, Style] = "grey70"
-    metrics: Union[str, Style] = "white"
+    description: Union[str, "Style"] = "white"
+    progress_bar: Union[str, "Style"] = "#6206E0"
+    progress_bar_finished: Union[str, "Style"] = "#6206E0"
+    progress_bar_pulse: Union[str, "Style"] = "#6206E0"
+    batch_progress: Union[str, "Style"] = "white"
+    time: Union[str, "Style"] = "grey54"
+    processing_speed: Union[str, "Style"] = "grey70"
+    metrics: Union[str, "Style"] = "white"
     metrics_text_delimiter: str = " "
     metrics_format: str = ".3f"
 
@@ -298,25 +295,25 @@ class RichProgressBar(ProgressBar):
         return not self.is_enabled
 
     @property
-    def train_progress_bar(self) -> Task:
+    def train_progress_bar(self) -> "Task":
         assert self.progress is not None
         assert self.train_progress_bar_id is not None
         return self.progress.tasks[self.train_progress_bar_id]
 
     @property
-    def val_sanity_check_bar(self) -> Task:
+    def val_sanity_check_bar(self) -> "Task":
         assert self.progress is not None
         assert self.val_sanity_progress_bar_id is not None
         return self.progress.tasks[self.val_sanity_progress_bar_id]
 
     @property
-    def val_progress_bar(self) -> Task:
+    def val_progress_bar(self) -> "Task":
         assert self.progress is not None
         assert self.val_progress_bar_id is not None
         return self.progress.tasks[self.val_progress_bar_id]
 
     @property
-    def test_progress_bar(self) -> Task:
+    def test_progress_bar(self) -> "Task":
         assert self.progress is not None
         assert self.test_progress_bar_id is not None
         return self.progress.tasks[self.test_progress_bar_id]

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/callbacks/pruning.py

@@ -211,8 +211,8 @@ class ModelPruning(Callback):
         # need to ignore typing here since pytorch base class does not define the PRUNING_TYPE attribute
         if use_global_unstructured and pruning_fn.PRUNING_TYPE != "unstructured":  # type: ignore
             raise MisconfigurationException(
-                'Only the "unstructured" PRUNING_TYPE is supported with `use_global_unstructured=True`.'  # type: ignore
-                f" Found method {pruning_fn} of type {pruning_fn.PRUNING_TYPE}. "
+                'Only the "unstructured" PRUNING_TYPE is supported with `use_global_unstructured=True`.'
+                f" Found method {pruning_fn} of type {pruning_fn.PRUNING_TYPE}. "  # type: ignore[union-attr]
             )
 
         self.pruning_fn = pruning_fn
@@ -308,7 +308,7 @@ class ModelPruning(Callback):
 
     def _apply_local_pruning(self, amount: float) -> None:
         for module, name in self._parameters_to_prune:
-            self.pruning_fn(module, name=name, amount=amount)
+            self.pruning_fn(module, name=name, amount=amount)  # type: ignore[call-arg]
 
     def _resolve_global_kwargs(self, amount: float) -> Dict[str, Any]:
         self._global_kwargs["amount"] = amount

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/callbacks/stochastic_weight_avg.py

@@ -17,7 +17,7 @@ Stochastic Weight Averaging Callback
 """
 
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Union, cast
+from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast
 
 import torch
 from torch import Tensor, nn
@@ -42,7 +42,7 @@ class StochasticWeightAveraging(Callback):
         swa_lrs: Union[float, List[float]],
         swa_epoch_start: Union[int, float] = 0.8,
         annealing_epochs: int = 10,
-        annealing_strategy: str = "cos",
+        annealing_strategy: Literal["cos", "linear"] = "cos",
         avg_fn: Optional[_AVG_FN] = None,
         device: Optional[Union[torch.device, str]] = torch.device("cpu"),
     ):
@@ -303,14 +303,14 @@ class StochasticWeightAveraging(Callback):
                 dtype=module.running_var.dtype,
             )
             self.momenta[module] = module.momentum
-            module.momentum = None  # type: ignore[assignment]
+            module.momentum = None
             assert module.num_batches_tracked is not None
             module.num_batches_tracked *= 0
 
     def reset_momenta(self) -> None:
         """Adapted from https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L164-L165."""
         for bn_module in self.momenta:
-            bn_module.momentum = self.momenta[bn_module]  # type: ignore[assignment]
+            bn_module.momentum = self.momenta[bn_module]
 
     @staticmethod
     def update_parameters(

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/core/datamodule.py

@@ -235,7 +235,7 @@ class LightningDataModule(DataHooks, HyperparametersMixin):
 
         """
         loaded = _load_from_checkpoint(
-            cls,  # type: ignore[arg-type]
+            cls,
             checkpoint_path,
             map_location=map_location,
             hparams_file=hparams_file,

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/core/module.py

@@ -1580,7 +1580,7 @@ class LightningModule(
 
         """
         loaded = _load_from_checkpoint(
-            cls,  # type: ignore[arg-type]
+            cls,
             checkpoint_path,
             map_location,
             hparams_file,

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/demos/__init__.py

@@ -1 +1,2 @@
+from lightning.pytorch.demos.lstm import LightningLSTM, SequenceSampler, SimpleLSTM  # noqa: F401
 from lightning.pytorch.demos.transformer import LightningTransformer, Transformer, WikiText2  # noqa: F401

lightning-2.4.0.dev20240804/src/lightning/pytorch/demos/lstm.py

@@ -0,0 +1,98 @@
+"""Demo of a simple LSTM language model.
+
+Code is adapted from the PyTorch examples at
+https://github.com/pytorch/examples/blob/main/word_language_model
+
+"""
+
+from typing import Iterator, List, Optional, Sized, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader, Sampler
+
+from lightning.pytorch.core import LightningModule
+from lightning.pytorch.demos.transformer import WikiText2
+
+
+class SimpleLSTM(nn.Module):
+    def __init__(
+        self, vocab_size: int = 33278, ninp: int = 512, nhid: int = 512, nlayers: int = 4, dropout: float = 0.2
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.drop = nn.Dropout(dropout)
+        self.encoder = nn.Embedding(vocab_size, ninp)
+        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout, batch_first=True)
+        self.decoder = nn.Linear(nhid, vocab_size)
+        self.nlayers = nlayers
+        self.nhid = nhid
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
+        nn.init.zeros_(self.decoder.bias)
+        nn.init.uniform_(self.decoder.weight, -0.1, 0.1)
+
+    def forward(self, input: Tensor, hidden: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+        emb = self.drop(self.encoder(input))
+        output, hidden = self.rnn(emb, hidden)
+        output = self.drop(output)
+        decoded = self.decoder(output).view(-1, self.vocab_size)
+        return F.log_softmax(decoded, dim=1), hidden
+
+    def init_hidden(self, batch_size: int) -> Tuple[Tensor, Tensor]:
+        weight = next(self.parameters())
+        return (
+            weight.new_zeros(self.nlayers, batch_size, self.nhid),
+            weight.new_zeros(self.nlayers, batch_size, self.nhid),
+        )
+
+
+class SequenceSampler(Sampler[List[int]]):
+    def __init__(self, dataset: Sized, batch_size: int) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.chunk_size = len(self.dataset) // self.batch_size
+
+    def __iter__(self) -> Iterator[List[int]]:
+        n = len(self.dataset)
+        for i in range(self.chunk_size):
+            yield list(range(i, n - (n % self.batch_size), self.chunk_size))
+
+    def __len__(self) -> int:
+        return self.chunk_size
+
+
+class LightningLSTM(LightningModule):
+    def __init__(self, vocab_size: int = 33278):
+        super().__init__()
+        self.model = SimpleLSTM(vocab_size=vocab_size)
+        self.hidden: Optional[Tuple[Tensor, Tensor]] = None
+
+    def on_train_epoch_end(self) -> None:
+        self.hidden = None
+
+    def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
+        input, target = batch
+        if self.hidden is None:
+            self.hidden = self.model.init_hidden(input.size(0))
+        self.hidden = (self.hidden[0].detach(), self.hidden[1].detach())
+        output, self.hidden = self.model(input, self.hidden)
+        loss = F.nll_loss(output, target.view(-1))
+        self.log("train_loss", loss, prog_bar=True)
+        return loss
+
+    def prepare_data(self) -> None:
+        WikiText2(download=True)
+
+    def train_dataloader(self) -> DataLoader:
+        dataset = WikiText2()
+        return DataLoader(dataset, batch_sampler=SequenceSampler(dataset, batch_size=20))
+
+    def configure_optimizers(self) -> Optimizer:
+        return torch.optim.SGD(self.parameters(), lr=20.0)

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/plugins/precision/amp.py

@@ -40,7 +40,7 @@ class MixedPrecision(Precision):
         self,
         precision: Literal["16-mixed", "bf16-mixed"],
         device: str,
-        scaler: Optional["torch.cuda.amp.GradScaler"] = None,
+        scaler: Optional["torch.amp.GradScaler"] = None,
     ) -> None:
         if precision not in ("16-mixed", "bf16-mixed"):
             raise ValueError(

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/plugins/precision/precision.py

@@ -95,7 +95,7 @@ class Precision(FabricPrecision, CheckpointHooks):
     def _wrap_closure(
         self,
         model: "pl.LightningModule",
-        optimizer: Optimizer,
+        optimizer: Steppable,
         closure: Callable[[], Any],
     ) -> Any:
         """This double-closure allows makes sure the ``closure`` is executed before the ``on_before_optimizer_step``

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/trainer/connectors/logger_connector/result.py

@@ -91,7 +91,7 @@ class _Sync:
         fn = self.no_op if self.fn is None or not self.should or self.rank_zero_only else self.fn
         # save the function as `_fn` as the meta are being re-created and the object references need to match.
         # ignore typing, bad support for `partial`: mypy/issues/1484
-        self._fn: Callable = partial(fn, reduce_op=self.op, group=self.group)  # type: ignore [arg-type]
+        self._fn: Callable = partial(fn, reduce_op=self.op, group=self.group)  # type: ignore[arg-type,operator,misc]
 
     @property
     def __call__(self) -> Any:
@@ -400,26 +400,19 @@ class _ResultCollection(dict):
 
         # register logged value if it doesn't exist
         if key not in self:
-            self.register_key(key, meta, value)
+            metric = _ResultMetric(meta, isinstance(value, Tensor))
+            self[key] = metric
 
         # check the stored metadata and the current one match
         elif meta != self[key].meta:
             raise MisconfigurationException(
                 f"You called `self.log({name}, ...)` twice in `{fx}` with different arguments. This is not allowed"
             )
+        self[key].to(value.device)
 
         batch_size = self._extract_batch_size(self[key], batch_size, meta)
         self.update_metrics(key, value, batch_size)
 
-    def register_key(self, key: str, meta: _Metadata, value: _VALUE) -> None:
-        """Create one _ResultMetric object per value.
-
-        Value can be provided as a nested collection
-
-        """
-        metric = _ResultMetric(meta, isinstance(value, Tensor)).to(value.device)
-        self[key] = metric
-
     def update_metrics(self, key: str, value: _VALUE, batch_size: int) -> None:
         result_metric = self[key]
         # performance: avoid calling `__call__` to avoid the checks in `torch.nn.Module._call_impl`

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/tuner/lr_finder.py

@@ -132,7 +132,9 @@ class _LRFinder:
         trainer.strategy.lr_scheduler_configs = [LRSchedulerConfig(scheduler, interval="step")]
         _validate_optimizers_attached(trainer.optimizers, trainer.lr_scheduler_configs)
 
-    def plot(self, suggest: bool = False, show: bool = False, ax: Optional["Axes"] = None) -> Optional["plt.Figure"]:
+    def plot(
+        self, suggest: bool = False, show: bool = False, ax: Optional["Axes"] = None
+    ) -> Optional[Union["plt.Figure", "plt.SubFigure"]]:
         """Plot results from lr_find run
         Args:
             suggest: if True, will mark suggested lr to use with a red point
@@ -151,10 +153,11 @@ class _LRFinder:
         lrs = self.results["lr"]
         losses = self.results["loss"]
 
+        fig: Optional[Union[plt.Figure, plt.SubFigure]]
         if ax is None:
             fig, ax = plt.subplots()
         else:
-            fig = ax.figure  # type: ignore[assignment]
+            fig = ax.figure
 
         # Plot loss as a function of the learning rate
         ax.plot(lrs, losses)
@@ -460,7 +463,7 @@ class _LinearLR(LRScheduler):
         super().__init__(optimizer, last_epoch)
 
     @override
-    def get_lr(self) -> List[float]:  # type: ignore[override]
+    def get_lr(self) -> List[float]:
         curr_iter = self.last_epoch + 1
         r = curr_iter / self.num_iter
 
@@ -497,7 +500,7 @@ class _ExponentialLR(LRScheduler):
         super().__init__(optimizer, last_epoch)
 
     @override
-    def get_lr(self) -> List[float]:  # type: ignore[override]
+    def get_lr(self) -> List[float]:
         curr_iter = self.last_epoch + 1
         r = curr_iter / self.num_iter
 

{lightning-2.4.0.dev20240721 → lightning-2.4.0.dev20240804}/src/lightning/pytorch/utilities/model_helpers.py

@@ -127,6 +127,9 @@ class _restricted_classmethod_impl(Generic[_T, _P, _R_co]):
         return wrapper
 
 
-# trick static type checkers into thinking it's a @classmethod
-# https://github.com/microsoft/pyright/issues/5865
-_restricted_classmethod = classmethod if TYPE_CHECKING else _restricted_classmethod_impl
+if TYPE_CHECKING:
+    # trick static type checkers into thinking it's a @classmethod
+    # https://github.com/microsoft/pyright/issues/5865
+    _restricted_classmethod = classmethod
+else:
+    _restricted_classmethod = _restricted_classmethod_impl

lightning-2.4.0.dev20240804/src/lightning/version.info

@@ -0,0 +1 @@
+2.4.0.dev20240804