Merge branch 'master' into deepspeed_exclude_frozen

Borda · web-flow · commit 43090f81bbf0 · 2025-08-21T18:11:01.000+02:00
diff --git a/.actions/assistant.py b/.actions/assistant.py
@@ -341,47 +341,6 @@ def create_mirror_package(source_dir: str, package_mapping: dict[str, str]) -> N
 
 
 class AssistantCLI:
-    @staticmethod
-    def requirements_prune_pkgs(packages: Sequence[str], req_files: Sequence[str] = REQUIREMENT_FILES_ALL) -> None:
-        """Remove some packages from given requirement files."""
-        if isinstance(req_files, str):
-            req_files = [req_files]
-        for req in req_files:
-            AssistantCLI._prune_packages(req, packages)
-
-    @staticmethod
-    def _prune_packages(req_file: str, packages: Sequence[str]) -> None:
-        """Remove some packages from given requirement files."""
-        path = Path(req_file)
-        assert path.exists()
-        text = path.read_text()
-        lines = text.splitlines()
-        final = []
-        for line in lines:
-            ln_ = line.strip()
-            if not ln_ or ln_.startswith("#"):
-                final.append(line)
-                continue
-            req = list(_parse_requirements([ln_]))[0]
-            if req.name not in packages:
-                final.append(line)
-        print(final)
-        path.write_text("\n".join(final) + "\n")
-
-    @staticmethod
-    def _replace_min(fname: str) -> None:
-        with open(fname, encoding="utf-8") as fopen:
-            req = fopen.read().replace(">=", "==")
-        with open(fname, "w", encoding="utf-8") as fwrite:
-            fwrite.write(req)
-
-    @staticmethod
-    def replace_oldest_ver(requirement_fnames: Sequence[str] = REQUIREMENT_FILES_ALL) -> None:
-        """Replace the min package version by fixed one."""
-        for fname in requirement_fnames:
-            print(fname)
-            AssistantCLI._replace_min(fname)
-
     @staticmethod
     def copy_replace_imports(
         source_dir: str,
diff --git a/.azure/gpu-tests-fabric.yml b/.azure/gpu-tests-fabric.yml
@@ -99,8 +99,16 @@ jobs:
         displayName: "Image info & NVIDIA"
 
       - bash: |
-          python .actions/assistant.py replace_oldest_ver
+          set -ex
           pip install "cython<3.0" wheel  # for compatibility
+          pip install -U "lightning-utilities[cli]"
+          cd requirements/fabric
+          # replace range by pin minimal requirements
+          python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'strategies.txt']"
+          # drop deepspeed since it is not supported by our minimal Torch requirements
+          python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files strategies.txt
+          # uninstall deepspeed since some older docker images have it pre-installed
+          pip uninstall -y deepspeed
         condition: contains(variables['Agent.JobName'], 'oldest')
         displayName: "setting oldest dependencies"
 
diff --git a/.azure/gpu-tests-pytorch.yml b/.azure/gpu-tests-pytorch.yml
@@ -103,8 +103,16 @@ jobs:
         displayName: "Image info & NVIDIA"
 
       - bash: |
-          python .actions/assistant.py replace_oldest_ver
+          set -ex
           pip install "cython<3.0" wheel  # for compatibility
+          pip install -U "lightning-utilities[cli]"
+          cd requirements/pytorch
+          # replace range by pin minimal requirements
+          python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'extra.txt', 'strategies.txt', 'examples.txt']"
+          # drop deepspeed since it is not supported by our minimal Torch requirements
+          python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files strategies.txt
+          # uninstall deepspeed since some older docker images have it pre-installed
+          pip uninstall -y deepspeed
         condition: contains(variables['Agent.JobName'], 'oldest')
         displayName: "setting oldest dependencies"
 
diff --git a/.github/workflows/ci-tests-fabric.yml b/.github/workflows/ci-tests-fabric.yml
@@ -94,7 +94,9 @@ jobs:
       - name: Set min. dependencies
         if: ${{ matrix.requires == 'oldest' }}
         run: |
-          python .actions/assistant.py replace_oldest_ver
+          cd requirements/fabric
+          pip install -U "lightning-utilities[cli]"
+          python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'strategies.txt', 'test.txt']"
           pip install "cython<3.0" wheel
           pip install "pyyaml==5.4" --no-build-isolation
 
@@ -140,7 +142,8 @@ jobs:
         run: |
           pip install -e ".[${EXTRA_PREFIX}test,${EXTRA_PREFIX}strategies]" \
             -U --upgrade-strategy=eager --prefer-binary \
-            --extra-index-url="${TORCH_URL}" --find-links="${PYPI_CACHE_DIR}"
+            --extra-index-url="${TORCH_URL}" \
+            --find-links="${PYPI_CACHE_DIR}"
           pip list
       - name: Dump handy wheels
         if: github.event_name == 'push' && github.ref == 'refs/heads/master'
diff --git a/.github/workflows/ci-tests-pytorch.yml b/.github/workflows/ci-tests-pytorch.yml
@@ -99,7 +99,9 @@ jobs:
       - name: Set min. dependencies
         if: ${{ matrix.requires == 'oldest' }}
         run: |
-          python .actions/assistant.py replace_oldest_ver
+          cd requirements/pytorch
+          pip install -U "lightning-utilities[cli]"
+          python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'extra.txt', 'strategies.txt', 'examples.txt', 'test.txt']"
           pip install "cython<3.0" wheel
           pip install "pyyaml==5.4" --no-build-isolation
 
@@ -139,7 +141,9 @@ jobs:
           pip install ".[${EXTRA_PREFIX}extra,${EXTRA_PREFIX}test,${EXTRA_PREFIX}strategies]" \
             -U --upgrade-strategy=eager --prefer-binary \
             -r requirements/_integrations/accelerators.txt \
-            --extra-index-url="${TORCH_URL}" --find-links="${PYPI_CACHE_DIR}" --find-links="https://download.pytorch.org/whl/torch-tensorrt"
+            --extra-index-url="${TORCH_URL}" \
+            --find-links="${PYPI_CACHE_DIR}" \
+            --find-links="https://download.pytorch.org/whl/torch-tensorrt"
           pip list
       - name: Drop LAI from extensions
         if: ${{ matrix.pkg-name != 'lightning' }}
diff --git a/.lightning/workflows/fabric.yml b/.lightning/workflows/fabric.yml
@@ -43,7 +43,10 @@ run: |
 
   if [ "${TORCH_VER}" == "2.1" ]; then
     echo "Set oldest versions"
-    python .actions/assistant.py replace_oldest_ver
+    cd requirements/fabric
+    pip install -U "lightning-utilities[cli]"
+    python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'strategies.txt']"
+    cd ../..
     pip install "cython<3.0" wheel  # for compatibility
   fi
 
diff --git a/.lightning/workflows/pytorch.yml b/.lightning/workflows/pytorch.yml
@@ -43,7 +43,10 @@ run: |
 
   if [ "${TORCH_VER}" == "2.1" ]; then
     recho "Set oldest versions"
-    python .actions/assistant.py replace_oldest_ver
+    cd requirements/pytorch
+    pip install -U "lightning-utilities[cli]"
+    python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'extra.txt', 'strategies.txt', 'examples.txt']"
+    cd ../..
     pip install "cython<3.0" wheel  # for compatibility
   fi
 
diff --git a/dockers/release/Dockerfile b/dockers/release/Dockerfile
@@ -21,6 +21,7 @@ FROM pytorchlightning/pytorch_lightning:base-cuda${CUDA_VERSION}-py${PYTHON_VERS
 LABEL maintainer="Lightning-AI <https://github.com/Lightning-AI>"
 
 ARG LIGHTNING_VERSION=""
+ARG PYTORCH_VERSION
 
 COPY ./ /home/pytorch-lightning/
 
@@ -39,7 +40,14 @@ RUN \
     fi && \
     # otherwise there is collision with folder name and pkg name on Pypi
     cd pytorch-lightning && \
-    pip install setuptools==75.6.0 && \
+    # pip install setuptools==75.6.0 && \
+    pip install -U "lightning-utilities[cli]" && \
+    # drop deepspeed since it is not supported by our minimal Torch requirements \
+    echo "PYTORCH_VERSION is: '$PYTORCH_VERSION'" && \
+    if [[ "$PYTORCH_VERSION" =~ ^(2\.1|2\.2|2\.3|2\.4)$ ]]; then \
+        python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files requirements/fabric/strategies.txt ; \
+        python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files requirements/pytorch/strategies.txt ; \
+    fi && \
     PACKAGE_NAME=lightning pip install '.[extra,loggers,strategies]' --no-cache-dir && \
     PACKAGE_NAME=pytorch pip install '.[extra,loggers,strategies]' --no-cache-dir && \
     cd .. && \
diff --git a/docs/source-pytorch/conf.py b/docs/source-pytorch/conf.py
@@ -127,7 +127,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:
         "https://pytorch.org/docs/stable/", "https://pytorch.org/docs/{torch.__version__}/", _PATH_ROOT
     )
     adjust_linked_external_docs(
-        "https://lightning.ai/docs/torchmetrics", "https://lightning.ai/docs/torchmetrics/v{torchmetrics.__version__}/", _PATH_ROOT, version_digits=3
+        "https://lightning.ai/docs/torchmetrics/stable/", "https://lightning.ai/docs/torchmetrics/v{torchmetrics.__version__}/", _PATH_ROOT, version_digits=3
     )
     adjust_linked_external_docs(
         "https://lightning.ai/docs/fabric/stable/", "https://lightning.ai/docs/fabric/{lightning_fabric.__version__}/", _PATH_ROOT, version_digits=3
diff --git a/docs/source-pytorch/extensions/logging.rst b/docs/source-pytorch/extensions/logging.rst
@@ -120,6 +120,10 @@ methods to log from anywhere in a :doc:`LightningModule <../common/lightning_mod
 .. note::
     Everything explained below applies to both :meth:`~lightning.pytorch.core.LightningModule.log` or :meth:`~lightning.pytorch.core.LightningModule.log_dict` methods.
 
+.. note::
+
+    When using TorchMetrics with Lightning, we recommend referring to the `TorchMetrics Lightning integration documentation <https://lightning.ai/docs/torchmetrics/stable/pages/lightning.html>`_ for logging best practices, common pitfalls, and proper usage patterns.
+
 Depending on where the :meth:`~lightning.pytorch.core.LightningModule.log` method is called, Lightning auto-determines
 the correct logging mode for you. Of course you can override the default behavior by manually setting the
 :meth:`~lightning.pytorch.core.LightningModule.log` parameters.
diff --git a/pyproject.toml b/pyproject.toml
@@ -179,6 +179,7 @@ markers = [
     "cloud: Run the cloud tests for example",
 ]
 filterwarnings = [
+    # "error::DeprecationWarning",
     "error::FutureWarning",
     "ignore::FutureWarning:onnxscript", # Temporary ignore until onnxscript is updated
 ]
diff --git a/requirements/fabric/strategies.txt b/requirements/fabric/strategies.txt
@@ -5,5 +5,5 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.9.3, <=0.9.3; platform_system != "Windows" and platform_system != "Darwin"  # strict
+deepspeed >=0.14.1,<=0.15.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
 bitsandbytes >=0.45.2,<0.47.0; platform_system != "Darwin"
diff --git a/requirements/fabric/test.txt b/requirements/fabric/test.txt
@@ -1,9 +1,10 @@
 coverage ==7.10.4
-numpy >=1.17.2, <1.27.0
+numpy >=1.21.0, <1.27.0
 pytest ==8.4.1
 pytest-cov ==6.2.1
 pytest-timeout ==2.4.0
 pytest-rerunfailures ==15.1
 pytest-random-order ==1.2.0
-click ==8.1.8
-tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
+click ==8.1.8; python_version < "3.11"
+click ==8.2.1; python_version > "3.10"
+tensorboardX >=2.6, <2.7.0  # todo: relax it back to `>=2.2` after fixing tests
diff --git a/requirements/pytorch/loggers.info b/requirements/pytorch/loggers.info
@@ -4,4 +4,4 @@ neptune >=1.0.0
 comet-ml >=3.31.0
 mlflow >=1.0.0
 wandb >=0.12.10
-tensorboard >=2.9.1
+tensorboard >=2.11
diff --git a/requirements/pytorch/strategies.txt b/requirements/pytorch/strategies.txt
@@ -3,4 +3,4 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.9.3, <=0.9.3; platform_system != "Windows" and platform_system != "Darwin"  # strict
+deepspeed >=0.14.1,<=0.15.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -8,16 +8,16 @@ pytest-random-order ==1.2.0
 # needed in tests
 cloudpickle >=1.3, <3.2.0
 scikit-learn >0.22.1, <1.8.0
-numpy >=1.17.2, <1.27.0
-onnx >=1.12.0, <1.19.0
+numpy >1.20.0, <1.27.0
+onnx >1.12.0, <1.19.0
 onnxruntime >=1.12.0, <1.23.0
-onnxscript >= 0.2.2, <0.4.0
+onnxscript >= 0.1.0, <0.4.0
 psutil <7.0.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
 
-tensorboard >=2.9.1, <2.21.0  # for `TensorBoardLogger`
+tensorboard >=2.11, <2.21.0  # for `TensorBoardLogger`
 
 --find-links https://download.pytorch.org/whl/torch-tensorrt
 torch-tensorrt; platform_system == "Linux" and python_version >= "3.12"
diff --git a/src/lightning/fabric/strategies/deepspeed.py b/src/lightning/fabric/strategies/deepspeed.py
@@ -47,7 +47,6 @@
     from torch.optim.lr_scheduler import _LRScheduler
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
-_DEEPSPEED_GREATER_EQUAL_0_14_1 = RequirementCache("deepspeed>=0.14.1")
 
 
 # TODO(fabric): Links in the docstrings to PL-specific deepspeed user docs need to be replaced.
@@ -509,10 +508,7 @@ def load_checkpoint(
             )
         engine = engines[0]
 
-        if _DEEPSPEED_GREATER_EQUAL_0_14_1:
-            from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
-        else:
-            from deepspeed.runtime import DeepSpeedOptimizer
+        from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
 
         optimzer_state_requested = any(isinstance(item, (Optimizer, DeepSpeedOptimizer)) for item in state.values())
 
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -31,11 +31,17 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed `AsyncCheckpointIO` snapshots tensors to avoid race with parameter mutation ([#21079](https://github.com/Lightning-AI/pytorch-lightning/pull/21079))
+
+
+- Fixed `AsyncCheckpointIO` threadpool exception if calling fit or validate more than one ([#20952](https://github.com/Lightning-AI/pytorch-lightning/pull/20952))
 
 
 - Fixed learning rate not being correctly set after using `LearningRateFinder` callback ([#21068](https://github.com/Lightning-AI/pytorch-lightning/pull/21068))
 
+
+- Fixed misalignment column while using rich model summary in `DeepSpeedstrategy` ([#21100](https://github.com/Lightning-AI/pytorch-lightning/pull/21100))
+
 ---
 
 ## [2.5.3] - 2025-08-13
diff --git a/src/lightning/pytorch/callbacks/rich_model_summary.py b/src/lightning/pytorch/callbacks/rich_model_summary.py
@@ -79,18 +79,21 @@ def summarize(
         from rich.table import Table
 
         console = get_console()
+        column_names = list(zip(*summary_data))[0]
 
         header_style: str = summarize_kwargs.get("header_style", "bold magenta")
         table = Table(header_style=header_style)
         table.add_column(" ", style="dim")
         table.add_column("Name", justify="left", no_wrap=True)
         table.add_column("Type")
         table.add_column("Params", justify="right")
+
+        if "Params per Device" in column_names:
+            table.add_column("Params per Device", justify="right")
+
         table.add_column("Mode")
         table.add_column("FLOPs", justify="right")
 
-        column_names = list(zip(*summary_data))[0]
-
         for column_name in ["In sizes", "Out sizes"]:
             if column_name in column_names:
                 table.add_column(column_name, justify="right", style="white")
diff --git a/src/lightning/pytorch/plugins/io/async_plugin.py b/src/lightning/pytorch/plugins/io/async_plugin.py
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py b/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py
diff --git a/tests/tests_fabric/strategies/launchers/test_multiprocessing_integration.py b/tests/tests_fabric/strategies/launchers/test_multiprocessing_integration.py
diff --git a/tests/tests_pytorch/plugins/test_async_checkpoint.py b/tests/tests_pytorch/plugins/test_async_checkpoint.py
diff --git a/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py b/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py
diff --git a/tests/tests_pytorch/utilities/test_compile.py b/tests/tests_pytorch/utilities/test_compile.py
diff --git a/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py b/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py

Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:`
`127`	`127`	`"https://pytorch.org/docs/stable/", "https://pytorch.org/docs/{torch.__version__}/", _PATH_ROOT`
`128`	`128`	`)`
`129`	`129`	`adjust_linked_external_docs(`
`130`		`- "https://lightning.ai/docs/torchmetrics", "https://lightning.ai/docs/torchmetrics/v{torchmetrics.__version__}/", _PATH_ROOT, version_digits=3`
	`130`	`+ "https://lightning.ai/docs/torchmetrics/stable/", "https://lightning.ai/docs/torchmetrics/v{torchmetrics.__version__}/", _PATH_ROOT, version_digits=3`
`131`	`131`	`)`
`132`	`132`	`adjust_linked_external_docs(`
`133`	`133`	`"https://lightning.ai/docs/fabric/stable/", "https://lightning.ai/docs/fabric/{lightning_fabric.__version__}/", _PATH_ROOT, version_digits=3`