Merge branch 'master' into dependabot-pip-requirements-torch-2.9.0

Author: bhimrazy

.github/CODEOWNERS

@@ -5,19 +5,19 @@
 # the repo. Unless a later match takes precedence,
 # @global-owner1 and @global-owner2 will be requested for
 # review when someone opens a pull request.
-* @lantiga @borda @tchaton @justusschock @ethanwharris
+* @lantiga @tchaton @justusschock @ethanwharris
 
 # Docs
-/.github/*.md                               @williamfalcon @lantiga @borda
+/.github/*.md                               @williamfalcon @lantiga
 /docs/source-fabric/index.rst               @williamfalcon @lantiga
 /docs/source-pytorch/index.rst              @williamfalcon @lantiga
 /docs/source-pytorch/levels                 @williamfalcon @lantiga
 
 /.github/CODEOWNERS                  @williamfalcon
 /SECURITY.md                         @williamfalcon @lantiga
 /README.md                           @williamfalcon @lantiga
-/src/pytorch_lightning/__about__.py  @williamfalcon @lantiga @borda
-/src/lightning_fabric/__about__.py   @williamfalcon @lantiga @borda
+/src/pytorch_lightning/__about__.py  @williamfalcon @lantiga
+/src/lightning_fabric/__about__.py   @williamfalcon @lantiga
 
 /src/lightning/fabric/loggers        @williamfalcon
 /src/lightning/pytorch/loggers       @williamfalcon

.github/CONTRIBUTING.md

@@ -212,6 +212,8 @@ We welcome any useful contribution! For your convenience here's a recommended wo
    - [Test README](https://github.com/Lightning-AI/pytorch-lightning/blob/master/tests/README.md)
    - [CI/CD README](https://github.com/Lightning-AI/pytorch-lightning/tree/master/.github/workflows#readme)
 
+1. Once you have a PR opened (and thereby a PR number), please update the respective changelog for [fabric](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/CHANGELOG.md) or [pytorch](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/CHANGELOG.md) subpackage depending on where you made your changes.
+
 1. When you feel ready for integrating your work, mark your PR "Ready for review".
 
    - Your code should be readable and follow the project's design principles.

.github/workflows/docs-build.yml

@@ -125,7 +125,10 @@ jobs:
         working-directory: ./docs/source-${{ matrix.pkg-name }}
         # allow failing link check and doctest if you run with dispatch
         continue-on-error: ${{ (matrix.target == 'doctest' || matrix.target == 'linkcheck') && github.event_name == 'workflow_dispatch' }}
-        run: make ${{ matrix.target }} --debug --jobs $(nproc) SPHINXOPTS="$BUILD_SPHINX_OPTS"
+        run: |
+          # temp fix: https://github.com/Lightning-AI/pytorch-lightning/actions/runs/19440502586/job/55622388642?pr=21354#step:11:4596
+          uv pip install -U fastapi
+          make ${{ matrix.target }} --debug --jobs $(nproc) SPHINXOPTS="$BUILD_SPHINX_OPTS"
 
       - name: Keep artifact
         if: github.event_name == 'pull_request'

Makefile

@@ -10,7 +10,7 @@ export PACKAGE_NAME=pytorch
 
 # In Lightning Studio, the `lightning` package comes pre-installed.
 # Uninstall it first to ensure the editable install works correctly.
-setup:
+setup: update
 	uv pip uninstall lightning pytorch-lightning lightning-fabric || true
 	uv pip install -r requirements.txt \
 	    -r requirements/pytorch/base.txt \

docs/source-pytorch/community/governance.rst

@@ -19,15 +19,15 @@ Role: All final decisions related to Lightning.
 Maintainers
 -----------
 - Luca Antiga (`lantiga <https://github.com/lantiga>`_)
-- Jirka Borovec (`Borda <https://github.com/Borda>`_)
+- Ethan Harris (`ethanwharris <https://github.com/ethanwharris>`_) (Torchbearer founder)
 - Justus Schock (`justusschock <https://github.com/justusschock>`_)
 
 
 Emeritus Maintainers
 --------------------
-- Ethan Harris (`ethanwharris <https://github.com/ethanwharris>`_) (Torchbearer founder)
 - Nicki Skafte (`SkafteNicki <https://github.com/SkafteNicki>`_)
 - Thomas Chaton (`tchaton <https://github.com/tchaton>`_)
+- Jirka Borovec (`Borda <https://github.com/Borda>`_)
 
 
 Alumni

docs/source-pytorch/conf.py

@@ -604,10 +604,7 @@ def package_list_from_file(file):
 from lightning.pytorch.cli import _JSONARGPARSE_SIGNATURES_AVAILABLE as _JSONARGPARSE_AVAILABLE
 from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE
 from lightning.fabric.loggers.tensorboard import _TENSORBOARD_AVAILABLE, _TENSORBOARDX_AVAILABLE
-from lightning.pytorch.loggers.neptune import _NEPTUNE_AVAILABLE
-from lightning.pytorch.loggers.comet import _COMET_AVAILABLE
-from lightning.pytorch.loggers.mlflow import _MLFLOW_AVAILABLE
-from lightning.pytorch.loggers.wandb import _WANDB_AVAILABLE
+from lightning.fabric.utilities.imports import _COMET_AVAILABLE, _MLFLOW_AVAILABLE, _NEPTUNE_AVAILABLE, _WANDB_AVAILABLE
 """
 coverage_skip_undoc_in_source = True
 

requirements/fabric/base.txt

@@ -6,3 +6,4 @@ fsspec[http] >=2022.5.0, <2025.11.0
 packaging >=20.0, <=25.0
 typing-extensions >4.5.0, <4.16.0
 lightning-utilities >=0.10.0, <0.16.0
+pytorch-lightning-enterprise >=2.6.0

requirements/pytorch/base.txt

@@ -9,3 +9,4 @@ torchmetrics >0.7.0, <1.9.0
 packaging >=20.0, <=25.0
 typing-extensions >4.5.0, <4.16.0
 lightning-utilities >=0.10.0, <0.16.0
+pytorch-lightning-enterprise >=2.6.0

src/lightning/fabric/CHANGELOG.md

@@ -4,18 +4,7 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
-
-## [unreleased] - YYYY-MM-DD
-
-### Added
-
--
-
-
-### Removed
-
--
-
+## [2.6.0] - 2025-11-21
 
 ### Changed
 
@@ -25,7 +14,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed issue in detecting MPIEnvironment with partial mpi4py installation ([#21353](https://github.com/Lightning-AI/pytorch-lightning/pull/21353))
+
+- Learning rate scheduler is stepped at the end of epoch when `on_train_batch_start` returns -1 ([#21296](https://github.com/Lightning-AI/pytorch-lightning/issues/21296)).
+
+
+- Fixed FSDP mixed precision semantics and added user warning ([#21361](https://github.com/Lightning-AI/pytorch-lightning/pull/21361))
 
 
 ---

src/lightning/fabric/accelerators/xla.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools
+import warnings
 from typing import Any, Union
 
 import torch
@@ -20,7 +21,11 @@
 
 from lightning.fabric.accelerators.accelerator import Accelerator
 from lightning.fabric.accelerators.registry import _AcceleratorRegistry
-from lightning.fabric.utilities.device_parser import _check_data_type
+from lightning.fabric.utilities.imports import _raise_enterprise_not_available
+
+_XLA_AVAILABLE = RequirementCache("torch_xla>=1.13", "torch_xla")
+_XLA_GREATER_EQUAL_2_1 = RequirementCache("torch_xla>=2.1")
+_XLA_GREATER_EQUAL_2_5 = RequirementCache("torch_xla>=2.5")
 
 
 class XLAAccelerator(Accelerator):
@@ -31,38 +36,38 @@ class XLAAccelerator(Accelerator):
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        if not _XLA_AVAILABLE:
-            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
-        if not _using_pjrt():
-            raise RuntimeError("The XLA XRT runtime is not supported anymore.")
+        _raise_enterprise_not_available()
         super().__init__(*args, **kwargs)
 
+        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
+
+        self.accelerator_impl = EnterpriseXLAAccelerator(*args, **kwargs)
+
     @override
     def setup_device(self, device: torch.device) -> None:
-        pass
+        return self.accelerator_impl.setup_device(device)
 
     @override
     def teardown(self) -> None:
-        pass
+        return self.accelerator_impl.teardown()
 
     @staticmethod
     @override
     def parse_devices(devices: Union[int, str, list[int]]) -> Union[int, list[int]]:
         """Accelerator device parsing logic."""
-        return _parse_tpu_devices(devices)
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
+
+        return EnterpriseXLAAccelerator.parse_devices(devices)
 
     @staticmethod
     @override
     def get_parallel_devices(devices: Union[int, list[int]]) -> list[torch.device]:
         """Gets parallel devices for the Accelerator."""
-        devices = _parse_tpu_devices(devices)
-        if isinstance(devices, int):
-            return [torch.device("xla", i) for i in range(devices)]
-        # list of devices is not supported, just a specific index, fine to access [0]
-        return [torch.device("xla", devices[0])]
-        # we cannot create `xla_device` here because processes have not been spawned yet (this is called in the
-        # accelerator connector init). However, there doesn't seem to be a problem with instantiating `torch.device`.
-        # it will be replaced with `xla_device` (also a torch.device`, but with extra logic) in the strategy
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
+
+        return EnterpriseXLAAccelerator.get_parallel_devices(devices)
 
     @staticmethod
     @override
@@ -71,16 +76,10 @@ def get_parallel_devices(devices: Union[int, list[int]]) -> list[torch.device]:
     @functools.lru_cache(maxsize=1)
     def auto_device_count() -> int:
         """Get the devices when set to auto."""
-        if not _XLA_AVAILABLE:
-            return 0
-        if _XLA_GREATER_EQUAL_2_1:
-            from torch_xla._internal import tpu
-
-            return tpu.num_available_devices()
-        from torch_xla.experimental import tpu
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
 
-        device_count_on_version = {2: 8, 3: 8, 4: 4}
-        return device_count_on_version.get(tpu.version(), 8)
+        return EnterpriseXLAAccelerator.auto_device_count()
 
     @staticmethod
     @override
@@ -92,6 +91,9 @@ def is_available() -> bool:
             # XLA may raise these exceptions if it's not properly configured. This needs to be avoided for the cases
             # when `torch_xla` is imported but not used
             return False
+        except ModuleNotFoundError as e:
+            warnings.warn(str(e))
+            return False
 
     @staticmethod
     @override
@@ -106,74 +108,3 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No
             cls,
             description=cls.__name__,
         )
-
-
-# PJRT support requires this minimum version
-_XLA_AVAILABLE = RequirementCache("torch_xla>=1.13", "torch_xla")
-_XLA_GREATER_EQUAL_2_1 = RequirementCache("torch_xla>=2.1")
-_XLA_GREATER_EQUAL_2_5 = RequirementCache("torch_xla>=2.5")
-
-
-def _using_pjrt() -> bool:
-    # `using_pjrt` is removed in torch_xla 2.5
-    if _XLA_GREATER_EQUAL_2_5:
-        from torch_xla import runtime as xr
-
-        return xr.device_type() is not None
-    # delete me when torch_xla 2.2 is the min supported version, where XRT support has been dropped.
-    if _XLA_GREATER_EQUAL_2_1:
-        from torch_xla import runtime as xr
-
-        return xr.using_pjrt()
-
-    from torch_xla.experimental import pjrt
-
-    return pjrt.using_pjrt()
-
-
-def _parse_tpu_devices(devices: Union[int, str, list[int]]) -> Union[int, list[int]]:
-    """Parses the TPU devices given in the format as accepted by the
-    :class:`~lightning.pytorch.trainer.trainer.Trainer` and :class:`~lightning.fabric.Fabric`.
-
-    Args:
-        devices: An int of 1 or string '1' indicates that 1 core with multi-processing should be used
-            An int 8 or string '8' indicates that all 8 cores with multi-processing should be used
-            A single element list of int or string can be used to indicate the specific TPU core to use.
-
-    Returns:
-        A list of tpu cores to be used.
-
-    """
-    _check_data_type(devices)
-    if isinstance(devices, str):
-        devices = _parse_tpu_devices_str(devices)
-    _check_tpu_devices_valid(devices)
-    return devices
-
-
-def _check_tpu_devices_valid(devices: object) -> None:
-    device_count = XLAAccelerator.auto_device_count()
-    if (
-        # support number of devices
-        isinstance(devices, int)
-        and devices in {1, device_count}
-        # support picking a specific device
-        or isinstance(devices, (list, tuple))
-        and len(devices) == 1
-        and 0 <= devices[0] <= device_count - 1
-    ):
-        return
-    raise ValueError(
-        f"`devices` can only be 'auto', 1, {device_count} or [<0-{device_count - 1}>] for TPUs. Got {devices!r}"
-    )
-
-
-def _parse_tpu_devices_str(devices: str) -> Union[int, list[int]]:
-    devices = devices.strip()
-    try:
-        return int(devices)
-    except ValueError:
-        try:
-            return [int(x.strip()) for x in devices.split(",") if len(x) > 0]
-        except ValueError:
-            raise ValueError(f"Could not parse the selected TPU devices: {devices!r}")