Merge branch 'Lightning-AI:master' into master

Author: KAVYANSHTYAGI

.github/workflows/docker-build.yml

@@ -97,7 +97,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.0"]
+        pytorch_version: ["2.6.0", "2.7.1"]
         cuda_version: ["12.6.3"]
         include:
           # These are the base images for PL release docker images.
@@ -108,7 +108,7 @@ jobs:
           - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.1" }
           - { 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.0", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.7.1", cuda_version: "12.6.3" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3

requirements/typing.txt

@@ -1,5 +1,5 @@
 mypy==1.15.0
-torch==2.7.0
+torch==2.7.1
 
 types-Markdown
 types-PyYAML

src/lightning/fabric/connector.py

@@ -492,7 +492,7 @@ def _check_and_init_precision(self) -> Precision:
                 if self._precision_input == "16-mixed"
                 else "Using bfloat16 Automatic Mixed Precision (AMP)"
             )
-            device = "cpu" if self._accelerator_flag == "cpu" else "cuda"
+            device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"
             return MixedPrecision(precision=self._precision_input, device=device)  # type: ignore[arg-type]
 
         raise RuntimeError("No precision set")

src/lightning/pytorch/core/module.py

@@ -1472,6 +1472,10 @@ def forward(self, x):
                     )
                 example_inputs = self.example_input_array
 
+            if kwargs.get("check_inputs") is not None:
+                kwargs["check_inputs"] = self._on_before_batch_transfer(kwargs["check_inputs"])
+                kwargs["check_inputs"] = self._apply_batch_transfer_handler(kwargs["check_inputs"])
+
             # automatically send example inputs to the right device and use trace
             example_inputs = self._on_before_batch_transfer(example_inputs)
             example_inputs = self._apply_batch_transfer_handler(example_inputs)

tests/tests_fabric/test_connector.py

@@ -405,6 +405,13 @@ def test_unsupported_strategy_types_on_cpu_and_fallback():
     assert isinstance(connector.strategy, DDPStrategy)
 
 
+@RunIf(mps=True)
+@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
+def test_mps_enabled_with_float16_or_bfloat16_precision(precision):
+    connector = _Connector(accelerator="mps", precision=precision)
+    assert connector.precision.device == "mps"
+
+
 def test_invalid_accelerator_choice():
     with pytest.raises(ValueError, match="You selected an invalid accelerator name: `accelerator='cocofruit'`"):
         _Connector(accelerator="cocofruit")

tests/tests_fabric/test_fabric.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+import warnings
 from contextlib import nullcontext
 from re import escape
 from unittest import mock
@@ -735,6 +736,22 @@ def test_autocast():
     fabric._precision.forward_context().__exit__.assert_called()
 
 
+@RunIf(mps=True)
+@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
+def test_autocast_does_not_use_cuda_on_mps(precision):
+    """Ensure Fabric.autocast on MPS does not fall back to CUDA when using (bf)16-mixed precision."""
+    fabric = Fabric(accelerator="mps", precision=precision)
+    fabric.launch()
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        with fabric.autocast():
+            pass
+
+    for warning in w:
+        assert "device_type of 'cuda'" not in str(warning.message)
+
+
 def test_no_backward_sync():
     """Test that `Fabric.no_backward_sync()` validates the strategy and model is compatible."""
     fabric = Fabric(devices=1)

tests/tests_pytorch/models/test_torchscript.py

@@ -105,6 +105,26 @@ def test_torchscript_device(device_str):
     assert script_output.device == device
 
 
+@pytest.mark.parametrize(
+    "device_str",
+    [
+        "cpu",
+        pytest.param("cuda:0", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("mps:0", marks=RunIf(mps=True)),
+    ],
+)
+def test_torchscript_device_with_check_inputs(device_str):
+    """Test that scripted module is on the correct device."""
+    device = torch.device(device_str)
+    model = BoringModel().to(device)
+    model.example_input_array = torch.randn(5, 32)
+
+    check_inputs = torch.rand(5, 32)
+
+    script = model.to_torchscript(method="trace", check_inputs=check_inputs)
+    assert isinstance(script, torch.jit.ScriptModule)
+
+
 def test_torchscript_retain_training_state():
     """Test that torchscript export does not alter the training mode of original model."""
     model = BoringModel()