Merge branch 'master' into 21296-fix-lr-update-on-train-epoch-skip

Author: LTMeyer

docs/source-pytorch/cli/lightning_cli_intermediate_2.rst

@@ -201,9 +201,10 @@ If the scheduler you want needs other arguments, add them via the CLI (no need t
 
 .. code:: bash
 
-    python main.py fit --optimizer=Adam --lr_scheduler=ReduceLROnPlateau --lr_scheduler.monitor=epoch
+    python main.py fit --optimizer=Adam --lr_scheduler=ReduceLROnPlateau --lr_scheduler.monitor=train_loss
 
-Furthermore, any custom subclass of ``torch.optim.lr_scheduler.LRScheduler`` can be used as learning rate scheduler:
+(assuming you have a ``train_loss`` metric logged). Furthermore, any custom subclass of
+``torch.optim.lr_scheduler.LRScheduler`` can be used as learning rate scheduler:
 
 .. code:: python
 
@@ -212,7 +213,6 @@ Furthermore, any custom subclass of ``torch.optim.lr_scheduler.LRScheduler`` can
     from lightning.pytorch.cli import LightningCLI
     from lightning.pytorch.demos.boring_classes import DemoModel, BoringDataModule
 
-
     class LitLRScheduler(torch.optim.lr_scheduler.CosineAnnealingLR):
         def step(self):
             print("⚡", "using LitLRScheduler", "⚡")

docs/source-pytorch/common/precision_basic.rst

@@ -39,6 +39,14 @@ However, this setting can sometimes lead to unstable training.
 
     Trainer(precision="16-true")
 
+.. warning::
+
+    Float16 cannot represent values smaller than ~6e-5. Values like Adam's default ``eps=1e-8`` become zero, which can cause
+    NaN during training. Increase ``eps`` to 1e-4 or higher, and avoid extremely small values in your model weights and data.
+
+.. note::
+
+    BFloat16 (``"bf16-mixed"`` or ``"bf16-true"``) has better numerical stability with a wider dynamic range.
 
 ----
 

requirements/pytorch/test.txt

@@ -13,7 +13,7 @@ numpy >1.20.0, <1.27.0
 onnx >1.12.0, <1.20.0
 onnxruntime >=1.12.0, <1.24.0
 onnxscript >= 0.1.0, < 0.5.0
-psutil <7.1.1 # for `DeviceStatsMonitor`
+psutil <7.1.2 # 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

src/lightning/pytorch/cli.py

@@ -66,6 +66,13 @@
 
 
 class ReduceLROnPlateau(torch.optim.lr_scheduler.ReduceLROnPlateau):
+    """Custom ReduceLROnPlateau scheduler that extends PyTorch's ReduceLROnPlateau.
+
+    This class adds a `monitor` attribute to the standard PyTorch ReduceLROnPlateau to specify which metric should be
+    tracked for learning rate adjustment.
+
+    """
+
     def __init__(self, optimizer: Optimizer, monitor: str, *args: Any, **kwargs: Any) -> None:
         super().__init__(optimizer, *args, **kwargs)
         self.monitor = monitor

tests/tests_pytorch/utilities/test_model_summary.py

@@ -319,13 +319,16 @@ def test_empty_model_size(max_depth):
 
 
 @pytest.mark.parametrize(
-    "accelerator",
+    ("accelerator", "precision"),
     [
-        pytest.param("gpu", marks=RunIf(min_cuda_gpus=1)),
-        pytest.param("mps", marks=RunIf(mps=True)),
+        pytest.param("gpu", "16-true", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("gpu", "32-true", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("gpu", "64-true", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("mps", "16-true", marks=RunIf(mps=True)),
+        pytest.param("mps", "32-true", marks=RunIf(mps=True)),
+        # Note: "64-true" with "mps" is skipped because MPS does not support float64
     ],
 )
-@pytest.mark.parametrize("precision", ["16-true", "32-true", "64-true"])
 def test_model_size_precision(tmp_path, accelerator, precision):
     """Test model size for different precision types."""
     model = PreCalculatedModel(precision=int(precision.split("-")[0]))