Merge branch 'master' into weights-only-compatibility

Author: Borda

.github/workflows/probot-check-group.yml

@@ -12,14 +12,14 @@ jobs:
   required-jobs:
     runs-on: ubuntu-latest
     if: github.event.pull_request.draft == false
-    timeout-minutes: 61 # in case something is wrong with the internal timeout
+    timeout-minutes: 71 # in case something is wrong with the internal timeout
     steps:
       - uses: Lightning-AI/[email protected]
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:
           job: check-group
           interval: 180 # seconds
-          timeout: 60 # minutes
+          timeout: 70 # minutes
           maintainers: "Lightning-AI/lai-frameworks"
           owner: "carmocca"

.lightning/workflows/fabric.yml

@@ -4,7 +4,7 @@ trigger:
   pull_request:
     branches: ["master", "release/stable"]
 
-timeout: "55" # minutes
+timeout: "60" # minutes
 parametrize:
   matrix: {}
   include:

.lightning/workflows/pytorch.yml

@@ -4,7 +4,7 @@ trigger:
   pull_request:
     branches: ["master", "release/stable"]
 
-timeout: "55" # minutes
+timeout: "60" # minutes
 parametrize:
   matrix: {}
   include:

docs/source-pytorch/common/checkpointing_intermediate.rst

@@ -21,7 +21,13 @@ For fine-grained control over checkpointing behavior, use the :class:`~lightning
         checkpoint_callback = ModelCheckpoint(dirpath="my/path/", save_top_k=2, monitor="val_loss")
         trainer = Trainer(callbacks=[checkpoint_callback])
         trainer.fit(model)
-        checkpoint_callback.best_model_path
+
+        # Access best and last model checkpoint directly from the callback
+        print(checkpoint_callback.best_model_path)
+        print(checkpoint_callback.last_model_path)
+        # Or via the trainer
+        print(trainer.checkpoint_callback.best_model_path)
+        print(trainer.checkpoint_callback.last_model_path)
 
 Any value that has been logged via *self.log* in the LightningModule can be monitored.
 

requirements/docs.txt

@@ -3,7 +3,7 @@ myst-parser >=0.18.1, <5.0.0
 nbsphinx >=0.8.5, <=0.9.7
 nbconvert >7.14, <7.17
 pandoc >=1.0, <=2.4
-docutils>=0.18.1,<=0.22
+docutils>=0.18.1,<=0.22.2
 sphinxcontrib-fulltoc >=1.0, <=1.2.0
 sphinxcontrib-mockautodoc
 sphinx-autobuild

requirements/pytorch/test.txt

@@ -12,7 +12,7 @@ numpy >1.20.0, <1.27.0
 onnx >1.12.0, <1.20.0
 onnxruntime >=1.12.0, <1.23.0
 onnxscript >= 0.1.0, < 0.5.0
-psutil <7.0.1 # for `DeviceStatsMonitor`
+psutil <7.1.1 # for `DeviceStatsMonitor`
 pandas >2.0, <2.4.0  # needed in benchmarks
 fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in App
 uvicorn  # for `ServableModuleValidator`  # not setting version as re-defined in App

requirements/typing.txt

@@ -1,4 +1,4 @@
-mypy==1.18.1
+mypy==1.18.2
 torch==2.8.0
 
 types-Markdown

src/lightning/pytorch/CHANGELOG.md

@@ -33,6 +33,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Default to `RichProgressBar` and `RichModelSummary` if the rich package is available. Fallback to TQDMProgressBar and ModelSummary otherwise ([#20896](https://github.com/Lightning-AI/pytorch-lightning/pull/20896))
 
 
+- Add MPS accelerator support for mixed precision ([#21209](https://github.com/Lightning-AI/pytorch-lightning/pull/21209))
+
+
 ### Removed
 
 -

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -204,11 +204,11 @@ class ModelCheckpoint(Checkpoint):
         ... )
 
         # retrieve the best checkpoint after training
-        checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
-        trainer = Trainer(callbacks=[checkpoint_callback])
-        model = ...
-        trainer.fit(model)
-        checkpoint_callback.best_model_path
+        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
+        >>> trainer = Trainer(callbacks=[checkpoint_callback])
+        >>> model = ...  # doctest: +SKIP
+        >>> trainer.fit(model)  # doctest: +SKIP
+        >>> print(checkpoint_callback.best_model_path)  # doctest: +SKIP
 
     .. tip:: Saving and restoring multiple checkpoint callbacks at the same time is supported under variation in the
         following arguments:

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -515,7 +515,7 @@ def _check_and_init_precision(self) -> Precision:
             rank_zero_info(
                 f"Using {'16bit' if self._precision_flag == '16-mixed' else 'bfloat16'} Automatic Mixed Precision (AMP)"
             )
-            device = "cpu" if self._accelerator_flag == "cpu" else "cuda"
+            device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"
             return MixedPrecision(self._precision_flag, device)  # type: ignore[arg-type]
 
         raise RuntimeError("No precision set")