[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

examples/app_multi_node/train_fabric.py

@@ -15,7 +15,7 @@ def run(self):
         )
 
         # 2. Create Fabric.
-        fabric = Fabric(strategy="ddp", precision='16-mixed')
+        fabric = Fabric(strategy="ddp", precision="16-mixed")
         model, optimizer = fabric.setup(model, torch.optim.SGD(model.parameters(), lr=0.01))
         criterion = torch.nn.MSELoss()
 

examples/pl_hpu/mnist_sample.py

@@ -63,7 +63,7 @@ def configure_optimizers(self):
             "accelerator": "hpu",
             "devices": 1,
             "max_epochs": 1,
-            "plugins": lazy_instance(HPUPrecisionPlugin, precision='16-mixed'),
+            "plugins": lazy_instance(HPUPrecisionPlugin, precision="16-mixed"),
         },
         run=False,
         save_config_kwargs={"overwrite": True},

src/lightning/pytorch/plugins/precision/amp.py

@@ -34,7 +34,10 @@ class MixedPrecisionPlugin(PrecisionPlugin):
     """
 
     def __init__(
-        self, precision: Literal["16-mixed", "bf16-mixed"], device: str, scaler: Optional[torch.cuda.amp.GradScaler] = None
+        self,
+        precision: Literal["16-mixed", "bf16-mixed"],
+        device: str,
+        scaler: Optional[torch.cuda.amp.GradScaler] = None,
     ) -> None:
         self.precision = cast(Literal["16-mixed", "bf16-mixed"], str(precision))
         if scaler is None and self.precision == "16-mixed":

src/lightning/pytorch/plugins/precision/deepspeed.py

@@ -34,7 +34,6 @@
 _PRECISION_INPUT = Literal["32-true", "16-mixed", "bf16-mixed"]
 
 
-
 class DeepSpeedPrecisionPlugin(PrecisionPlugin):
     """Precision plugin for DeepSpeed integration.
 

src/lightning/pytorch/plugins/precision/hpu.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import cast, Literal, Optional, Union
+from typing import cast, Literal, Optional
 
 from typing_extensions import get_args
 

src/lightning/pytorch/plugins/precision/ipu.py

@@ -29,6 +29,7 @@
 
 _PRECISION_INPUT = Literal["32-true", "16-mixed"]
 
+
 class IPUPrecisionPlugin(PrecisionPlugin):
     """Precision plugin for IPU integration.
 

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -20,6 +20,13 @@
 import torch
 from typing_extensions import get_args
 
+from lightning.fabric.connector import (
+    _PRECISION_INPUT,
+    _PRECISION_INPUT_INT,
+    _PRECISION_INPUT_STR,
+    _PRECISION_INPUT_STR_LEGACY,
+    _PRECISION_INPUT_STR_LEGACY_CONVERSION,
+)
 from lightning.fabric.plugins.environments import (
     ClusterEnvironment,
     KubeflowEnvironment,
@@ -71,7 +78,6 @@
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.imports import _LIGHTNING_COLOSSALAI_AVAILABLE
 from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
-from lightning.fabric.connector import _PRECISION_INPUT, _PRECISION_INPUT_STR_LEGACY, _PRECISION_INPUT_STR_LEGACY_CONVERSION, _PRECISION_INPUT_STR, _PRECISION_INPUT_INT
 
 log = logging.getLogger(__name__)
 
@@ -241,7 +247,9 @@ def _check_config_and_set_final_flags(
 
         self._accelerator_flag = accelerator
 
-        supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + get_args(_PRECISION_INPUT_STR_LEGACY)
+        supported_precision = (
+            get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + get_args(_PRECISION_INPUT_STR_LEGACY)
+        )
         if precision not in supported_precision:
             raise MisconfigurationException(
                 f"Precision {repr(precision)} is invalid. Allowed precision values: {supported_precision}"
@@ -250,7 +258,7 @@ def _check_config_and_set_final_flags(
         precision = str(precision)  # convert int flags to str here to enable the legacy-conversion below
 
         if precision in get_args(_PRECISION_INPUT_STR_LEGACY):
-            if not str(precision)[:2] in ('32', '64'):
+            if str(precision)[:2] not in ("32", "64"):
                 rank_zero_warn(
                     f"{precision} is supported for historical reasons but its usage is discouraged. "
                     f"Please set your precision to {_PRECISION_INPUT_STR_LEGACY_CONVERSION[precision]} instead!"

src/lightning/pytorch/trainer/trainer.py

@@ -122,7 +122,7 @@ def __init__(
         accelerator: Optional[Union[str, Accelerator]] = None,
         strategy: Optional[Union[str, Strategy]] = None,
         sync_batchnorm: bool = False,
-        precision: _PRECISION_INPUT = '32-true',
+        precision: _PRECISION_INPUT = "32-true",
         enable_model_summary: bool = True,
         num_sanity_val_steps: int = 2,
         profiler: Optional[Union[Profiler, str]] = None,

tests/tests_pytorch/accelerators/test_hpu.py

@@ -61,7 +61,7 @@ def test_all_stages(tmpdir, hpus):
         fast_dev_run=True,
         accelerator="hpu",
         devices=hpus,
-        precision='16-mixed',
+        precision="16-mixed",
     )
     trainer.fit(model)
     trainer.validate(model)

tests/tests_pytorch/accelerators/test_ipu.py

@@ -183,7 +183,12 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> Non
 
     model = IPUModel()
     trainer = Trainer(
-        default_root_dir=tmpdir, fast_dev_run=True, accelerator="ipu", devices=1, precision='16-mixed', callbacks=TestCallback()
+        default_root_dir=tmpdir,
+        fast_dev_run=True,
+        accelerator="ipu",
+        devices=1,
+        precision="16-mixed",
+        callbacks=TestCallback(),
     )
     assert isinstance(trainer.strategy.precision_plugin, IPUPrecisionPlugin)
     assert trainer.strategy.precision_plugin.precision == "16-mixed"
@@ -203,7 +208,12 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
     model = IPUModel()
     model = model.half()
     trainer = Trainer(
-        default_root_dir=tmpdir, fast_dev_run=True, accelerator="ipu", devices=1, precision='16-mixed', callbacks=TestCallback()
+        default_root_dir=tmpdir,
+        fast_dev_run=True,
+        accelerator="ipu",
+        devices=1,
+        precision="16-mixed",
+        callbacks=TestCallback(),
     )
 
     assert isinstance(trainer.strategy, IPUStrategy)
@@ -534,7 +544,7 @@ def configure_optimizers(self):
 def test_precision_plugin():
     """Ensure precision plugin value is set correctly."""
 
-    plugin = IPUPrecisionPlugin(precision='16-mixed')
+    plugin = IPUPrecisionPlugin(precision="16-mixed")
     assert plugin.precision == "16-mixed"
 
 

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -103,7 +103,7 @@ def test_resume_legacy_checkpoints(tmpdir, pl_version: str):
             default_root_dir=str(tmpdir),
             accelerator="auto",
             devices=1,
-            precision=('16-mixed' if torch.cuda.is_available() else '32-true'),
+            precision=("16-mixed" if torch.cuda.is_available() else "32-true"),
             callbacks=[stop],
             max_epochs=21,
             accumulate_grad_batches=2,

tests/tests_pytorch/models/test_amp.py

@@ -52,9 +52,9 @@ def _assert_autocast_enabled(self):
 @pytest.mark.parametrize(
     ("strategy", "precision", "devices"),
     (
-        ("single_device", '16-mixed', 1),
+        ("single_device", "16-mixed", 1),
         ("single_device", "bf16-mixed'", 1),
-        ("ddp_spawn", '16-mixed', 2),
+        ("ddp_spawn", "16-mixed", 2),
         ("ddp_spawn", "bf16-mixed", 2),
     ),
 )
@@ -83,7 +83,7 @@ def test_amp_cpus(tmpdir, strategy, precision, devices):
 
 
 @pytest.mark.parametrize("strategy", [None, "ddp_spawn"])
-@pytest.mark.parametrize("precision", ['16-mixed', pytest.param("bf16-mixed", marks=RunIf(bf16_cuda=True))])
+@pytest.mark.parametrize("precision", ["16-mixed", pytest.param("bf16-mixed", marks=RunIf(bf16_cuda=True))])
 @pytest.mark.parametrize(
     "devices", (pytest.param(1, marks=RunIf(min_cuda_gpus=1)), pytest.param(2, marks=RunIf(min_cuda_gpus=2)))
 )
@@ -135,7 +135,7 @@ def test_amp_gpu_ddp_slurm_managed(tmpdir):
         accelerator="gpu",
         devices=[0],
         strategy="ddp_spawn",
-        precision='16-mixed',
+        precision="16-mixed",
         callbacks=[checkpoint],
         logger=logger,
     )
@@ -153,7 +153,7 @@ def test_precision_16_clip_gradients(mock_clip_grad_norm, clip_val, tmpdir):
         enable_progress_bar=False,
         max_epochs=1,
         devices=1,
-        precision='16-mixed',
+        precision="16-mixed",
         limit_train_batches=4,
         limit_val_batches=0,
         gradient_clip_val=clip_val,

tests/tests_pytorch/models/test_ddp_fork_amp.py

@@ -24,7 +24,7 @@
 def test_amp_gpus_ddp_fork():
     """Ensure the use of AMP with `ddp_fork` (or associated alias strategies) does not generate CUDA initialization
     errors."""
-    _ = MixedPrecisionPlugin(precision='16-mixed', device="cuda")
+    _ = MixedPrecisionPlugin(precision="16-mixed", device="cuda")
     with multiprocessing.get_context("fork").Pool(1) as pool:
         in_bad_fork = pool.apply(torch.cuda._is_in_bad_fork)
     assert not in_bad_fork

tests/tests_pytorch/models/test_hooks.py

@@ -401,9 +401,9 @@ def _predict_batch(trainer, model, batches):
     [
         {},
         # these precision plugins modify the optimization flow, so testing them explicitly
-        pytest.param(dict(accelerator="gpu", devices=1, precision='16-mixed'), marks=RunIf(min_cuda_gpus=1)),
+        pytest.param(dict(accelerator="gpu", devices=1, precision="16-mixed"), marks=RunIf(min_cuda_gpus=1)),
         pytest.param(
-            dict(accelerator="gpu", devices=1, precision='16-mixed', strategy="deepspeed"),
+            dict(accelerator="gpu", devices=1, precision="16-mixed", strategy="deepspeed"),
             marks=RunIf(min_cuda_gpus=1, standalone=True, deepspeed=True),
         ),
     ],
@@ -453,7 +453,7 @@ def training_step(self, batch, batch_idx):
         "loops": ANY,
     }
     using_deepspeed = kwargs.get("strategy") == "deepspeed"
-    if kwargs.get("precision") == '16-mixed' and not using_deepspeed:
+    if kwargs.get("precision") == "16-mixed" and not using_deepspeed:
         saved_ckpt[trainer.precision_plugin.__class__.__qualname__] = ANY
     device = torch.device("cuda:0" if "accelerator" in kwargs and kwargs["accelerator"] == "gpu" else "cpu")
     expected = [

tests/tests_pytorch/models/test_tpu.py

@@ -104,7 +104,7 @@ def test_model_16bit_tpu_devices_1(tmpdir):
     """Make sure model trains on TPU."""
     trainer_options = dict(
         default_root_dir=tmpdir,
-        precision='16-mixed',
+        precision="16-mixed",
         enable_progress_bar=False,
         max_epochs=2,
         accelerator="tpu",
@@ -124,7 +124,7 @@ def test_model_16bit_tpu_index(tmpdir, tpu_core):
     """Make sure model trains on TPU."""
     trainer_options = dict(
         default_root_dir=tmpdir,
-        precision='16-mixed',
+        precision="16-mixed",
         enable_progress_bar=False,
         max_epochs=2,
         accelerator="tpu",
@@ -146,7 +146,7 @@ def test_model_16bit_tpu_devices_8(tmpdir):
     """Make sure model trains on TPU."""
     trainer_options = dict(
         default_root_dir=tmpdir,
-        precision='16-mixed',
+        precision="16-mixed",
         enable_progress_bar=False,
         max_epochs=1,
         accelerator="tpu",

tests/tests_pytorch/plugins/precision/hpu/test_hpu.py

@@ -77,7 +77,7 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
         fast_dev_run=True,
         accelerator="hpu",
         devices=1,
-        plugins=[HPUPrecisionPlugin(precision='16-mixed', **hmp_params)],
+        plugins=[HPUPrecisionPlugin(precision="16-mixed", **hmp_params)],
         callbacks=TestCallback(),
     )
 

tests/tests_pytorch/plugins/precision/test_amp.py

@@ -23,7 +23,7 @@
 def test_clip_gradients():
     """Test that `.clip_gradients()` is a no-op when clipping is disabled."""
     optimizer = Mock(spec=Optimizer)
-    precision = MixedPrecisionPlugin(precision='16-mixed', device="cuda:0", scaler=Mock())
+    precision = MixedPrecisionPlugin(precision="16-mixed", device="cuda:0", scaler=Mock())
     precision.clip_grad_by_value = Mock()
     precision.clip_grad_by_norm = Mock()
     precision.clip_gradients(optimizer)
@@ -47,7 +47,7 @@ def test_optimizer_amp_scaling_support_in_step_method():
     gradient clipping (example: fused Adam)."""
 
     optimizer = Mock(_step_supports_amp_scaling=True)
-    precision = MixedPrecisionPlugin(precision='16-mixed', device="cuda:0", scaler=Mock())
+    precision = MixedPrecisionPlugin(precision="16-mixed", device="cuda:0", scaler=Mock())
 
     with pytest.raises(RuntimeError, match="The current optimizer.*does not allow for gradient clipping"):
         precision.clip_gradients(optimizer, clip_val=1.0)

tests/tests_pytorch/plugins/precision/test_amp_integration.py

@@ -38,7 +38,7 @@ def run(fused=False):
             default_root_dir=tmpdir,
             accelerator="cuda",
             devices=1,
-            precision='16-mixed',
+            precision="16-mixed",
             max_steps=5,
             logger=False,
             enable_checkpointing=False,

tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py

@@ -19,4 +19,4 @@
 
 def test_invalid_precision_with_deepspeed_precision():
     with pytest.raises(ValueError, match="is not supported. `precision` must be one of"):
-        DeepSpeedPrecisionPlugin(precision='64-true')
+        DeepSpeedPrecisionPlugin(precision="64-true")

tests/tests_pytorch/plugins/test_amp_plugins.py

@@ -54,10 +54,10 @@ class MyAMP(MixedPrecisionPlugin):
 def test_amp_ddp(cuda_count_2, strategy, devices, custom_plugin, plugin_cls):
     plugin = None
     if custom_plugin:
-        plugin = plugin_cls('16-mixed', "cpu")
+        plugin = plugin_cls("16-mixed", "cpu")
     trainer = Trainer(
         fast_dev_run=True,
-        precision='16-mixed',
+        precision="16-mixed",
         accelerator="gpu",
         devices=devices,
         strategy=strategy,
@@ -137,7 +137,7 @@ def test_amp_gradient_unscale(tmpdir, accum: int):
         strategy="ddp_spawn",
         accelerator="gpu",
         devices=2,
-        precision='16-mixed',
+        precision="16-mixed",
         # use a tiny value to make sure it works
         gradient_clip_val=1e-3,
         gradient_clip_algorithm="value",
@@ -179,7 +179,7 @@ def configure_optimizers(self):
                 torch.optim.SGD(self.layer2.parameters(), lr=0.1),
             ]
 
-    trainer = Trainer(default_root_dir=tmpdir, accelerator="gpu", devices=1, fast_dev_run=1, precision='16-mixed')
+    trainer = Trainer(default_root_dir=tmpdir, accelerator="gpu", devices=1, fast_dev_run=1, precision="16-mixed")
     model = CustomBoringModel()
     trainer.fit(model)
 

tests/tests_pytorch/plugins/test_double_plugin.py

@@ -135,7 +135,7 @@ def on_fit_start(self):
 def test_double_precision(tmpdir, boring_model):
     model = boring_model()
 
-    trainer = Trainer(max_epochs=2, default_root_dir=tmpdir, fast_dev_run=2, precision='64-true', log_every_n_steps=1)
+    trainer = Trainer(max_epochs=2, default_root_dir=tmpdir, fast_dev_run=2, precision="64-true", log_every_n_steps=1)
     trainer.fit(model)
     trainer.test(model)
     trainer.predict(model)
@@ -152,7 +152,7 @@ def test_double_precision_ddp(tmpdir):
         accelerator="gpu",
         devices=2,
         fast_dev_run=2,
-        precision='64-true',
+        precision="64-true",
         log_every_n_steps=1,
     )
     trainer.fit(model)

tests/tests_pytorch/strategies/test_ddp.py

@@ -96,7 +96,7 @@ def setup(self, stage: str) -> None:
 
 
 @RunIf(min_cuda_gpus=2, standalone=True)
-@pytest.mark.parametrize("precision", ('16-mixed', '32-true'))
+@pytest.mark.parametrize("precision", ("16-mixed", "32-true"))
 def test_ddp_wrapper(tmpdir, precision):
     """Test parameters to ignore are carried over for DDP."""