Merge branch 'master' into ci/lit

Borda · web-flow · commit 925eee6c6a64 · 2025-08-08T23:52:29.000+03:00
diff --git a/.github/checkgroup.yml b/.github/checkgroup.yml
@@ -135,7 +135,8 @@ subprojects:
       - "build-pl (3.11, 2.4, 12.1.1)"
       - "build-pl (3.12, 2.5, 12.1.1)"
       - "build-pl (3.12, 2.6, 12.4.1)"
-      - "build-pl (3.12, 2.7, 12.6.3, true)"
+      - "build-pl (3.12, 2.7, 12.6.3)"
+      - "build-pl (3.12, 2.8, 12.6.3, true)"
 
   # SECTION: lightning_fabric
 
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -49,7 +49,8 @@ jobs:
           - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.6", cuda_version: "12.4.1" }
-          - { python_version: "3.12", pytorch_version: "2.7", cuda_version: "12.6.3", latest: "true" }
+          - { python_version: "3.12", pytorch_version: "2.7", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.8", cuda_version: "12.6.3", latest: "true" }
     steps:
       - uses: actions/checkout@v4
         with:
@@ -97,7 +98,7 @@ jobs:
         # adding dome more images as Thunder mainly using python 3.10,
         # and we need to support integrations as for example LitGPT
         python_version: ["3.10"]
-        pytorch_version: ["2.6.0", "2.7.1"]
+        pytorch_version: ["2.7.1", "2.8.0"]
         cuda_version: ["12.6.3"]
         include:
           # These are the base images for PL release docker images.
@@ -109,6 +110,7 @@ jobs:
           - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.6.0", cuda_version: "12.4.1" }
           - { python_version: "3.12", pytorch_version: "2.7.1", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.8.0", cuda_version: "12.6.3" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3
@@ -129,6 +131,7 @@ jobs:
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch_version }}
             CUDA_VERSION=${{ matrix.cuda_version }}
+            MAKE_FLAGS="-j2"
           file: dockers/base-cuda/Dockerfile
           push: ${{ env.PUSH_NIGHTLY }}
           tags: "pytorchlightning/pytorch_lightning:base-cuda-py${{ matrix.python_version }}-torch${{ env.PT_VERSION }}-cuda${{ matrix.cuda_version }}"
@@ -157,6 +160,8 @@ jobs:
         continue-on-error: true
         uses: docker/build-push-action@v6
         with:
+          build-args: |
+            PYTORCH_VERSION="25.04"
           file: dockers/nvidia/Dockerfile
           push: false
         timeout-minutes: 55
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -58,7 +58,7 @@ repos:
         #args: ["--write-changes"] # uncomment if you want to get automatic fixing
 
   - repo: https://github.com/PyCQA/docformatter
-    rev: 06907d0267368b49b9180eed423fae5697c1e909 # todo: fix for docformatter after last 1.7.5
+    rev: v1.7.7
     hooks:
       - id: docformatter
         additional_dependencies: [tomli]
@@ -70,7 +70,7 @@ repos:
       - id: sphinx-lint
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.4
+    rev: v0.12.2
     hooks:
       # try to fix what is possible
       - id: ruff
diff --git a/dockers/base-cuda/Dockerfile b/dockers/base-cuda/Dockerfile
@@ -19,8 +19,9 @@ ARG CUDA_VERSION=11.7.1
 FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}
 
 ARG PYTHON_VERSION=3.10
-ARG PYTORCH_VERSION=2.1
+ARG PYTORCH_VERSION=2.8
 ARG MAX_ALLOWED_NCCL=2.22.3
+ARG MAKE_FLAGS="-j$(nproc)"
 
 SHELL ["/bin/bash", "-c"]
 # https://techoverflow.net/2019/05/18/how-to-fix-configuring-tzdata-interactive-input-when-building-docker-images/
@@ -30,8 +31,7 @@ ENV \
     PATH="$PATH:/root/.local/bin" \
     CUDA_TOOLKIT_ROOT_DIR="/usr/local/cuda" \
     MKL_THREADING_LAYER="GNU" \
-    # MAKEFLAGS="-j$(nproc)"
-    MAKEFLAGS="-j2"
+    MAKEFLAGS=${MAKE_FLAGS}
 
 RUN \
     CUDA_VERSION_MM=${CUDA_VERSION%.*} && \
diff --git a/dockers/nvidia/Dockerfile b/dockers/nvidia/Dockerfile
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG PYTORCH_VERSION=22.09
+ARG PYTORCH_VERSION=24.05
 
 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes
 FROM nvcr.io/nvidia/pytorch:${PYTORCH_VERSION}-py3
diff --git a/dockers/release/Dockerfile b/dockers/release/Dockerfile
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 ARG PYTHON_VERSION=3.10
-ARG PYTORCH_VERSION=2.0
-ARG CUDA_VERSION=11.8.0
+ARG PYTORCH_VERSION=2.8
+ARG CUDA_VERSION=12.6.3
 
 FROM pytorchlightning/pytorch_lightning:base-cuda-py${PYTHON_VERSION}-torch${PYTORCH_VERSION}-cuda${CUDA_VERSION}
 
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -27,11 +27,20 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Fixed
 
 - fix progress bar console clearing for Rich `14.1+` ([#21016](https://github.com/Lightning-AI/pytorch-lightning/pull/21016))
+
+
 - fix `AdvancedProfiler` to handle nested profiling actions for Python 3.12+ ([#20809](https://github.com/Lightning-AI/pytorch-lightning/pull/20809))
 
 
+- Fix support for more dtypes in `ModelSummary` ([#21034](https://github.com/Lightning-AI/pytorch-lightning/pull/21034))
+
+
 - Fixed metrics in `RichProgressBar` being updated according to user provided `refresh_rate` ([#21032](https://github.com/Lightning-AI/pytorch-lightning/pull/21032))
 
+
+- Fix `save_last` behavior in the absence of validation ([#20960](https://github.com/Lightning-AI/pytorch-lightning/pull/20960))
+
+
 ---
 
 ## [2.5.2] - 2025-06-20
diff --git a/src/lightning/pytorch/callbacks/model_checkpoint.py b/src/lightning/pytorch/callbacks/model_checkpoint.py
@@ -344,6 +344,13 @@ def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModul
                 self._save_topk_checkpoint(trainer, monitor_candidates)
             self._save_last_checkpoint(trainer, monitor_candidates)
 
+    @override
+    def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Ensure save_last=True is applied when training ends."""
+        if self.save_last and not self._last_checkpoint_saved:
+            monitor_candidates = self._monitor_candidates(trainer)
+            self._save_last_checkpoint(trainer, monitor_candidates)
+
     @override
     def state_dict(self) -> dict[str, Any]:
         return {
diff --git a/src/lightning/pytorch/callbacks/progress/rich_progress.py b/src/lightning/pytorch/callbacks/progress/rich_progress.py
@@ -184,7 +184,7 @@ def render(self, task: "Task") -> Text:
 
         def _generate_metrics_texts(self) -> Generator[str, None, None]:
             for name, value in self._metrics.items():
-                if not isinstance(value, str):
+                if not isinstance(value, (str, int)):
                     value = f"{value:{self._metrics_format}}"
                 yield f"{name}: {value}"
 
diff --git a/src/lightning/pytorch/core/optimizer.py b/src/lightning/pytorch/core/optimizer.py
@@ -274,7 +274,7 @@ def _configure_schedulers_automatic_opt(schedulers: list, monitor: Optional[str]
             scheduler["reduce_on_plateau"] = scheduler.get(
                 "reduce_on_plateau", isinstance(scheduler["scheduler"], optim.lr_scheduler.ReduceLROnPlateau)
             )
-            if scheduler["reduce_on_plateau"] and scheduler.get("monitor", None) is None:
+            if scheduler["reduce_on_plateau"] and scheduler.get("monitor") is None:
                 raise MisconfigurationException(
                     "The lr scheduler dict must include a monitor when a `ReduceLROnPlateau` scheduler is used."
                     ' For example: {"optimizer": optimizer, "lr_scheduler":'
diff --git a/src/lightning/pytorch/demos/transformer.py b/src/lightning/pytorch/demos/transformer.py
@@ -54,15 +54,24 @@ def __init__(
 
         self.ninp = ninp
         self.vocab_size = vocab_size
-        self.src_mask = None
+        self.src_mask: Optional[Tensor] = None
+
+    def generate_square_subsequent_mask(self, size: int) -> Tensor:
+        """Generate a square mask for the sequence to prevent future tokens from being seen."""
+        mask = torch.triu(torch.ones(size, size), diagonal=1)
+        mask = mask.float().masked_fill(mask == 1, float("-inf")).masked_fill(mask == 0, 0.0)
+        return mask
 
     def forward(self, inputs: Tensor, target: Tensor, mask: Optional[Tensor] = None) -> Tensor:
         _, t = inputs.shape
 
-        # we assume target is already shifted w.r.t. inputs
+        # Generate source mask to prevent future token leakage
+        if self.src_mask is None or self.src_mask.size(0) != t:
+            self.src_mask = self.generate_square_subsequent_mask(t).to(inputs.device)
+
+        # Generate target mask if not provided
         if mask is None:
-            mask = torch.tril(torch.ones(t, t, device=inputs.device)) == 1
-            mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, 0.0)
+            mask = self.generate_square_subsequent_mask(t).to(inputs.device)
 
         src = self.pos_encoder(self.embedding(inputs) * math.sqrt(self.ninp))
         target = self.pos_encoder(self.embedding(target) * math.sqrt(self.ninp))
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py
@@ -26,6 +26,7 @@
 from torch.utils.hooks import RemovableHandle
 
 import lightning.pytorch as pl
+from lightning.fabric.utilities import rank_zero_warn
 from lightning.fabric.utilities.distributed import _is_dtensor
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_4
 from lightning.pytorch.utilities.model_helpers import _ModuleMode
@@ -227,7 +228,22 @@ def __init__(self, model: "pl.LightningModule", max_depth: int = 1) -> None:
         self._layer_summary = self.summarize()
         # 1 byte -> 8 bits
         # TODO: how do we compute precision_megabytes in case of mixed precision?
-        precision_to_bits = {"64": 64, "32": 32, "16": 16, "bf16": 16}
+        precision_to_bits = {
+            "64": 64,
+            "32": 32,
+            "16": 16,
+            "bf16": 16,
+            "16-true": 16,
+            "bf16-true": 16,
+            "32-true": 32,
+            "64-true": 64,
+        }
+        if self._model._trainer and self._model.trainer.precision not in precision_to_bits:
+            rank_zero_warn(
+                f"Precision {self._model.trainer.precision} is not supported by the model summary. "
+                " Estimated model size in MB will not be accurate. Using 32 bits instead.",
+                category=UserWarning,
+            )
         precision = precision_to_bits.get(self._model.trainer.precision, 32) if self._model._trainer else 32
         self._precision_megabytes = (precision / 8.0) * 1e-6
 
diff --git a/tests/tests_pytorch/callbacks/test_finetuning_callback.py b/tests/tests_pytorch/callbacks/test_finetuning_callback.py
@@ -109,8 +109,8 @@ def configure_optimizers(self):
     model.validation_step = None
     callback = TestBackboneFinetuningWarningCallback(unfreeze_backbone_at_epoch=3, verbose=False)
 
+    trainer = Trainer(limit_train_batches=1, default_root_dir=tmp_path, callbacks=[callback, chk], max_epochs=2)
     with pytest.warns(UserWarning, match="Did you init your optimizer in"):
-        trainer = Trainer(limit_train_batches=1, default_root_dir=tmp_path, callbacks=[callback, chk], max_epochs=2)
         trainer.fit(model)
 
     assert model.backbone.has_been_used
diff --git a/tests/tests_pytorch/checkpointing/test_model_checkpoint.py b/tests/tests_pytorch/checkpointing/test_model_checkpoint.py
@@ -1666,3 +1666,30 @@ def val_dataloader(self) -> DataLoader:
     trainer_kwargs["max_epochs"] = 4
     trainer = Trainer(**trainer_kwargs, callbacks=ModelCheckpoint(**mc_kwargs))
     trainer.fit(model, ckpt_path=checkpoint_path)
+
+
+def test_save_last_without_save_on_train_epoch_and_without_val(tmp_path):
+    """Test that save_last=True works correctly when save_on_train_epoch_end=False in a model without validation."""
+
+    # Remove validation methods to test the edge case
+    model = BoringModel()
+    model.validation_step = None
+    model.val_dataloader = None
+
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=tmp_path,
+        save_last=True,
+        save_on_train_epoch_end=False,
+    )
+
+    trainer = Trainer(
+        max_epochs=2,
+        callbacks=[checkpoint_callback],
+        logger=False,
+        enable_progress_bar=False,
+    )
+
+    trainer.fit(model)
+
+    # save_last=True should always save last.ckpt
+    assert (tmp_path / "last.ckpt").exists()
diff --git a/tests/tests_pytorch/profilers/test_profiler.py b/tests/tests_pytorch/profilers/test_profiler.py
@@ -73,9 +73,9 @@ def test_simple_profiler_durations(simple_profiler, action: str, expected: list)
     np.testing.assert_allclose(simple_profiler.recorded_durations[action], expected, rtol=0.2)
 
 
-def test_simple_profiler_overhead(simple_profiler, n_iter=5):
+def test_simple_profiler_overhead(simple_profiler):
     """Ensure that the profiler doesn't introduce too much overhead during training."""
-    for _ in range(n_iter):
+    for _ in range(5):
         with simple_profiler.profile("no-op"):
             pass
 
@@ -284,8 +284,9 @@ def test_advanced_profiler_durations(advanced_profiler, action: str, expected: l
 
 
 @pytest.mark.flaky(reruns=3)
-def test_advanced_profiler_overhead(advanced_profiler, n_iter=5):
+def test_advanced_profiler_overhead(advanced_profiler):
     """Ensure that the profiler doesn't introduce too much overhead during training."""
+    n_iter = 5
     for _ in range(n_iter):
         with advanced_profiler.profile("no-op"):
             pass
@@ -620,8 +621,8 @@ def test_pytorch_profiler_raises_warning_for_limited_steps(tmp_path, trainer_con
     warning_cache.clear()
     with pytest.warns(UserWarning, match="not enough steps to properly record traces"):
         getattr(trainer, trainer_fn)(model)
-        assert trainer.profiler._schedule is None
-        warning_cache.clear()
+    assert trainer.profiler._schedule is None
+    warning_cache.clear()
 
 
 def test_profile_callbacks(tmp_path):
diff --git a/tests/tests_pytorch/trainer/logging_/test_eval_loop_logging.py b/tests/tests_pytorch/trainer/logging_/test_eval_loop_logging.py
@@ -234,9 +234,9 @@ def on_test_epoch_end(self):
 
 
 @pytest.mark.parametrize("suffix", [False, True])
-def test_multi_dataloaders_add_suffix_properly(tmp_path, suffix):
+def test_multi_dataloaders_add_suffix_properly(suffix, tmp_path):
     class TestModel(BoringModel):
-        def test_step(self, batch, batch_idx, dataloader_idx=0):
+        def test_step(self, batch, batch_idx, dataloader_idx=0):  # noqa: PT028
             out = super().test_step(batch, batch_idx)
             self.log("test_loss", out["y"], on_step=True, on_epoch=True)
             return out
@@ -441,7 +441,7 @@ def on_test_epoch_end(self, _, pl_module):
     class TestModel(BoringModel):
         seen_losses = {i: [] for i in range(num_dataloaders)}
 
-        def test_step(self, batch, batch_idx, dataloader_idx=0):
+        def test_step(self, batch, batch_idx, dataloader_idx=0):  # noqa: PT028
             loss = super().test_step(batch, batch_idx)["y"]
             self.log("test_loss", loss)
             self.seen_losses[dataloader_idx].append(loss)
diff --git a/tests/tests_pytorch/trainer/test_config_validator.py b/tests/tests_pytorch/trainer/test_config_validator.py
@@ -16,7 +16,6 @@
 import pytest
 import torch
 
-from lightning.fabric.utilities.warnings import PossibleUserWarning
 from lightning.pytorch import LightningDataModule, LightningModule, Trainer
 from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
 from lightning.pytorch.trainer.configuration_validator import (
@@ -46,20 +45,19 @@ def test_wrong_configure_optimizers(tmp_path):
         trainer.fit(model)
 
 
-def test_fit_val_loop_config(tmp_path):
+@pytest.mark.parametrize("model_attrib", ["validation_step", "val_dataloader"])
+def test_fit_val_loop_config(model_attrib, tmp_path):
     """When either val loop or val data are missing raise warning."""
     trainer = Trainer(default_root_dir=tmp_path, max_epochs=1)
 
-    # no val data has val loop
-    with pytest.warns(UserWarning, match=r"You passed in a `val_dataloader` but have no `validation_step`"):
-        model = BoringModel()
-        model.validation_step = None
-        trainer.fit(model)
-
-    # has val loop but no val data
-    with pytest.warns(PossibleUserWarning, match=r"You defined a `validation_step` but have no `val_dataloader`"):
-        model = BoringModel()
-        model.val_dataloader = None
+    model = BoringModel()
+    setattr(model, model_attrib, None)
+    match_msg = (
+        r"You passed in a `val_dataloader` but have no `validation_step`"
+        if model_attrib == "validation_step"
+        else "You defined a `validation_step` but have no `val_dataloader`"
+    )
+    with pytest.warns(UserWarning, match=match_msg):
         trainer.fit(model)
 
 
diff --git a/tests/tests_pytorch/trainer/test_dataloaders.py b/tests/tests_pytorch/trainer/test_dataloaders.py
diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py

Original file line number	Diff line number	Diff line change
`@@ -274,7 +274,7 @@ def _configure_schedulers_automatic_opt(schedulers: list, monitor: Optional[str]`
`274`	`274`	`scheduler["reduce_on_plateau"] = scheduler.get(`
`275`	`275`	`"reduce_on_plateau", isinstance(scheduler["scheduler"], optim.lr_scheduler.ReduceLROnPlateau)`
`276`	`276`	`)`
`277`		`- if scheduler["reduce_on_plateau"] and scheduler.get("monitor", None) is None:`
	`277`	`+ if scheduler["reduce_on_plateau"] and scheduler.get("monitor") is None:`
`278`	`278`	`raise MisconfigurationException(`
`279`	`279`	"The lr scheduler dict must include a monitor when a `ReduceLROnPlateau` scheduler is used."
`280`	`280`	`' For example: {"optimizer": optimizer, "lr_scheduler":'`