Merge branch 'bugfix/20350_grad_acc_fix' of https://github.com/Sohaib-Ahmed21/pytorch-lightning into bugfix/20350_grad_acc_fix

Author: Sohaib-Ahmed21

.github/workflows/docs-build.yml

@@ -125,10 +125,7 @@ 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: |
-          # 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"
+        run: make ${{ matrix.target }} --debug --jobs $(nproc) SPHINXOPTS="$BUILD_SPHINX_OPTS"
 
       - name: Keep artifact
         if: github.event_name == 'pull_request'

docs/source-fabric/guide/callbacks.rst

@@ -83,6 +83,30 @@ The :meth:`~lightning.fabric.fabric.Fabric.call` calls the callback objects in t
 Not all objects registered via ``Fabric(callbacks=...)`` must implement a method with the given name.
 The ones that have a matching method name will get called.
 
+The different callbacks can have different method signatures. Fabric automatically filters keyword arguments based on
+each callback's function signature, allowing callbacks with different signatures to work together seamlessly.
+
+.. code-block:: python
+
+    class TrainingMetricsCallback:
+        def on_train_epoch_end(self, train_loss):
+            print(f"Training loss: {train_loss:.4f}")
+
+    class ValidationMetricsCallback:
+        def on_train_epoch_end(self, val_accuracy):
+            print(f"Validation accuracy: {val_accuracy:.4f}")
+
+    class ComprehensiveCallback:
+        def on_train_epoch_end(self, epoch, **kwargs):
+            print(f"Epoch {epoch} complete with metrics: {kwargs}")
+
+    fabric = Fabric(
+        callbacks=[TrainingMetricsCallback(), ValidationMetricsCallback(), ComprehensiveCallback()]
+    )
+
+    # Each callback receives only the arguments it can handle
+    fabric.call("on_train_epoch_end", epoch=5, train_loss=0.1, val_accuracy=0.95, learning_rate=0.001)
+
 
 ----
 

docs/source-pytorch/conf.py

@@ -604,7 +604,10 @@ 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.fabric.utilities.imports import _COMET_AVAILABLE, _MLFLOW_AVAILABLE, _NEPTUNE_AVAILABLE, _WANDB_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
 """
 coverage_skip_undoc_in_source = True
 

requirements/fabric/base.txt

@@ -6,4 +6,3 @@ 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,4 +9,3 @@ 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

requirements/pytorch/extra.txt

@@ -5,7 +5,7 @@
 matplotlib>3.1, <3.11.0
 omegaconf >=2.2.3, <2.4.0
 hydra-core >=1.2.0, <1.4.0
-jsonargparse[signatures,jsonnet] >=4.39.0, <4.43.0
+jsonargparse[signatures,jsonnet] >=4.39.0, <4.44.0
 rich >=12.3.0, <14.3.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.45.2,<0.47.0; platform_system != "Darwin"

src/lightning/fabric/CHANGELOG.md

@@ -4,6 +4,21 @@ 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
+
+- Added kwargs-filtering for `Fabric.call` to support different callback method signatures ([#21258](https://github.com/Lightning-AI/pytorch-lightning/pull/21258))
+
+
+### Removed
+
+-
+
+
+---
+
 ## [2.6.0] - 2025-11-21
 
 ### Changed

src/lightning/fabric/accelerators/xla.py

@@ -12,7 +12,6 @@
 # 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
@@ -21,11 +20,7 @@
 
 from lightning.fabric.accelerators.accelerator import Accelerator
 from lightning.fabric.accelerators.registry import _AcceleratorRegistry
-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")
+from lightning.fabric.utilities.device_parser import _check_data_type
 
 
 class XLAAccelerator(Accelerator):
@@ -36,38 +31,38 @@ class XLAAccelerator(Accelerator):
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        _raise_enterprise_not_available()
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
+        if not _using_pjrt():
+            raise RuntimeError("The XLA XRT runtime is not supported anymore.")
         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:
-        return self.accelerator_impl.setup_device(device)
+        pass
 
     @override
     def teardown(self) -> None:
-        return self.accelerator_impl.teardown()
+        pass
 
     @staticmethod
     @override
     def parse_devices(devices: Union[int, str, list[int]]) -> Union[int, list[int]]:
         """Accelerator device parsing logic."""
-        _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
-
-        return EnterpriseXLAAccelerator.parse_devices(devices)
+        return _parse_tpu_devices(devices)
 
     @staticmethod
     @override
     def get_parallel_devices(devices: Union[int, list[int]]) -> list[torch.device]:
         """Gets parallel devices for the Accelerator."""
-        _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
-
-        return EnterpriseXLAAccelerator.get_parallel_devices(devices)
+        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
 
     @staticmethod
     @override
@@ -76,10 +71,16 @@ 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."""
-        _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.accelerators.xla import XLAAccelerator as EnterpriseXLAAccelerator
+        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
 
-        return EnterpriseXLAAccelerator.auto_device_count()
+        device_count_on_version = {2: 8, 3: 8, 4: 4}
+        return device_count_on_version.get(tpu.version(), 8)
 
     @staticmethod
     @override
@@ -91,9 +92,6 @@ 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
@@ -108,3 +106,74 @@ 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}")

src/lightning/fabric/fabric.py

@@ -985,6 +985,34 @@ def train_function(fabric):
             )
         return self._wrap_and_launch(function, self, *args, **kwargs)
 
+    def _filter_kwargs_for_callback(self, method: Callable, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Filter keyword arguments to only include those that match the callback method's signature.
+
+        Args:
+            method: The callback method to inspect
+            kwargs: The keyword arguments to filter
+
+        Returns:
+            A filtered dictionary of keyword arguments that match the method's signature
+
+        """
+        try:
+            sig = inspect.signature(method)
+        except (ValueError, TypeError):
+            # If we can't inspect the signature, pass all kwargs to maintain backward compatibility
+            return kwargs
+
+        filtered_kwargs = {}
+        for name, param in sig.parameters.items():
+            # If the method accepts **kwargs, pass all original kwargs directly
+            if param.kind == inspect.Parameter.VAR_KEYWORD:
+                return kwargs
+            # If the parameter exists in the incoming kwargs, add it to filtered_kwargs
+            if name in kwargs:
+                filtered_kwargs[name] = kwargs[name]
+
+        return filtered_kwargs
+
     def call(self, hook_name: str, *args: Any, **kwargs: Any) -> None:
         r"""Trigger the callback methods with the given name and arguments.
 
@@ -994,7 +1022,9 @@ def call(self, hook_name: str, *args: Any, **kwargs: Any) -> None:
         Args:
             hook_name: The name of the callback method.
             *args: Optional positional arguments that get passed down to the callback method.
-            **kwargs: Optional keyword arguments that get passed down to the callback method.
+            **kwargs: Optional keyword arguments that get passed down to the callback method. Keyword arguments
+                that are not present in the callback's signature will be filtered out automatically, allowing
+                callbacks to have different signatures for the same hook.
 
         Example::
 
@@ -1016,13 +1046,8 @@ def on_train_epoch_end(self, results):
                 )
                 continue
 
-            method(*args, **kwargs)
-
-            # TODO(fabric): handle the following signatures
-            # method(self, fabric|trainer, x, y=1)
-            # method(self, fabric|trainer, *args, x, y=1)
-            # method(self, *args, y=1)
-            # method(self, *args, **kwargs)
+            filtered_kwargs = self._filter_kwargs_for_callback(method, kwargs)
+            method(*args, **filtered_kwargs)
 
     def log(self, name: str, value: Any, step: Optional[int] = None) -> None:
         """Log a scalar to all loggers that were added to Fabric.