From f93e79ee9f3221aac9eb4d862d00a8ae02f11434 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Wed, 15 Feb 2023 16:55:12 +0100 Subject: [PATCH 01/26] adapt precision in fabric --- src/lightning/fabric/cli.py | 8 +-- src/lightning/fabric/connector.py | 56 ++++++++++++------- src/lightning/fabric/fabric.py | 6 +- src/lightning/fabric/plugins/precision/amp.py | 17 +++--- .../fabric/plugins/precision/deepspeed.py | 14 ++--- .../fabric/plugins/precision/double.py | 2 +- .../fabric/plugins/precision/fsdp.py | 8 +-- .../fabric/plugins/precision/precision.py | 8 ++- .../fabric/plugins/precision/tpu_bf16.py | 2 +- src/lightning/fabric/strategies/deepspeed.py | 12 ++-- src/lightning/fabric/strategies/fsdp.py | 5 +- 11 files changed, 80 insertions(+), 58 deletions(-) diff --git a/src/lightning/fabric/cli.py b/src/lightning/fabric/cli.py index 6ade6d5ce1039..1410a1b3deb17 100644 --- a/src/lightning/fabric/cli.py +++ b/src/lightning/fabric/cli.py @@ -20,6 +20,7 @@ from lightning_utilities.core.imports import RequirementCache from lightning.fabric.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator +from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT_STR from lightning.fabric.strategies import STRATEGY_REGISTRY from lightning.fabric.utilities.device_parser import _parse_gpu_ids @@ -28,7 +29,6 @@ _CLICK_AVAILABLE = RequirementCache("click") _SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu") -_SUPPORTED_PRECISION = ("64", "32", "16", "bf16") def _get_supported_strategies() -> List[str]: @@ -106,11 +106,11 @@ def _get_supported_strategies() -> List[str]: ) @click.option( "--precision", - type=click.Choice(_SUPPORTED_PRECISION), + type=click.Choice(_PRECISION_INPUT_STR), default="32", help=( - "Double precision (``64``), full precision (``32``), half precision (``16``) or bfloat16 precision" - " (``'bf16'``)" + "Double precision (``64-true``), full precision (``32-true``), half precision (``16-mixed``) or " + "bfloat16 precision (``'bf16-mixed'``)" ), ) @click.argument("script_args", nargs=-1, type=click.UNPROCESSED) diff --git a/src/lightning/fabric/connector.py b/src/lightning/fabric/connector.py index b0e997dcf83a2..e1d5b7dabd527 100644 --- a/src/lightning/fabric/connector.py +++ b/src/lightning/fabric/connector.py @@ -42,7 +42,13 @@ ) from lightning.fabric.plugins.precision.double import DoublePrecision from lightning.fabric.plugins.precision.fsdp import FSDPPrecision -from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT, _PRECISION_INPUT_INT, _PRECISION_INPUT_STR +from lightning.fabric.plugins.precision.precision import ( + _PRECISION_INPUT, + _PRECISION_INPUT_INT, + _PRECISION_INPUT_STR, + _PRECISION_INPUT_STR_LEGACY, + _PRECISION_INPUT_STR_LEGACY_CONVERSION, +) from lightning.fabric.strategies import ( DeepSpeedStrategy, ParallelStrategy, @@ -107,7 +113,7 @@ def __init__( strategy = self._argument_from_env("strategy", strategy, default=None) devices = self._argument_from_env("devices", devices, default=None) num_nodes = self._argument_from_env("num_nodes", num_nodes, default=1) - precision = self._argument_from_env("precision", precision, default=32) + precision = self._argument_from_env("precision", precision, default="32-true") # 1. Parsing flags # Get registered strategies, built-in accelerators and precision plugins @@ -119,7 +125,7 @@ def __init__( # For devices: Assign gpus, etc. to the accelerator flag and devices flag self._strategy_flag: Optional[Union[Strategy, str]] = None self._accelerator_flag: Optional[Union[Accelerator, str]] = None - self._precision_input: _PRECISION_INPUT_STR = "32" + self._precision_input: _PRECISION_INPUT_STR = "32-true" self._precision_instance: Optional[Precision] = None self._cluster_environment_flag: Optional[Union[ClusterEnvironment, str]] = None self._parallel_devices: List[Union[int, torch.device, str]] = [] @@ -220,10 +226,22 @@ def _check_config_and_set_final_flags( self._accelerator_flag = accelerator - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + 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 ValueError(f"Precision {repr(precision)} is invalid. Allowed precision values: {supported_precision}") - self._precision_input = cast(_PRECISION_INPUT_STR, str(precision)) + + precision = str(precision) # convert int flags to str here to enable the legacy-conversion below + + if precision in get_args(_PRECISION_INPUT_STR_LEGACY): + rank_zero_warn( + f"{precision} is supported for historical reasons but its usage is discouraged. " + "Please set your precision to {_PRECISION_INPUT_STR_LEGACY_CONVERSION[precision]} instead!" + ) + precision = _PRECISION_INPUT_STR_LEGACY_CONVERSION[precision] + + self._precision_input = cast(_PRECISION_INPUT_STR, precision) if plugins: plugins_flags_types: Dict[str, int] = Counter() @@ -453,34 +471,34 @@ def _check_and_init_precision(self) -> Precision: return self._precision_instance if isinstance(self.accelerator, TPUAccelerator): - if self._precision_input == "32": + if self._precision_input == "32-true": return TPUPrecision() - elif self._precision_input in ("16", "bf16"): - if self._precision_input == "16": + elif self._precision_input in ("16-mixed", "bf16-mixed"): + if self._precision_input == "16-mixed": rank_zero_warn( - "You passed `Fabric(accelerator='tpu', precision=16)` but AMP" - " is not supported with TPUs. Using `precision='bf16'` instead." + "You passed `Fabric(accelerator='tpu', precision='16-mixed')` but AMP with fp16" + " is not supported with TPUs. Using `precision='bf16-mixed'` instead." ) return TPUBf16Precision() if isinstance(self.strategy, DeepSpeedStrategy): return DeepSpeedPrecision(self._precision_input) # type: ignore - if self._precision_input == "32": + if self._precision_input == "32-true": return Precision() - if self._precision_input == "64": + if self._precision_input == "64-true": return DoublePrecision() - if self._precision_input == "16" and self._accelerator_flag == "cpu": + if self._precision_input == "16-mixed" and self._accelerator_flag == "cpu": rank_zero_warn( - "You passed `Fabric(accelerator='cpu', precision=16)` but AMP is not supported on CPU." - " Using `precision='bf16'` instead." + "You passed `Fabric(accelerator='cpu', precision='16-mixed')` but AMP is not supported on CPU with " + "fp16. Using `precision='bf16-mixed'` instead." ) - self._precision_input = "bf16" + self._precision_input = "bf16-mixed" - if self._precision_input in ("16", "bf16"): + if self._precision_input in ("16-mixed", "bf16-mixed"): rank_zero_info( "Using 16-bit Automatic Mixed Precision (AMP)" - if self._precision_input == "16" + if self._precision_input == "16-mixed" else "Using bfloat16 Automatic Mixed Precision (AMP)" ) device = "cpu" if self._accelerator_flag == "cpu" else "cuda" @@ -494,7 +512,7 @@ def _check_and_init_precision(self) -> Precision: def _validate_precision_choice(self) -> None: """Validate the combination of choices for precision, and accelerator.""" if isinstance(self.accelerator, TPUAccelerator): - if self._precision_input == "64": + if self._precision_input == "64-true": raise NotImplementedError( "`Fabric(accelerator='tpu', precision=64)` is not implemented." " Please, open an issue in `https://github.com/Lightning-AI/lightning/issues`" diff --git a/src/lightning/fabric/fabric.py b/src/lightning/fabric/fabric.py index 7e41844243673..47a3097244258 100644 --- a/src/lightning/fabric/fabric.py +++ b/src/lightning/fabric/fabric.py @@ -67,8 +67,8 @@ class Fabric: devices: Number of devices to train on (``int``), which GPUs to train on (``list`` or ``str``), or ``"auto"``. The value applies per node. num_nodes: Number of GPU nodes for distributed training. - precision: Double precision (``64``), full precision (``32``), half precision (``16``), - or bfloat16 precision (``"bf16"``). + precision: Double precision (``"64-true"``), full precision (``"32"``), half precision AMP (``"16-mixed"``), + or bfloat16 precision AMP (``"bf16-mixed"``). plugins: One or several custom plugins callbacks: A single callback or a list of callbacks. A callback can contain any arbitrary methods that can be invoked through :meth:`~lightning.fabric.fabric.Fabric.call` by the user. @@ -82,7 +82,7 @@ def __init__( strategy: Optional[Union[str, Strategy]] = None, devices: Optional[Union[List[int], str, int]] = None, num_nodes: int = 1, - precision: _PRECISION_INPUT = 32, + precision: _PRECISION_INPUT = "32-true", plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]] = None, callbacks: Optional[Union[List[Any], Any]] = None, loggers: Optional[Union[Logger, List[Logger]]] = None, diff --git a/src/lightning/fabric/plugins/precision/amp.py b/src/lightning/fabric/plugins/precision/amp.py index 5fa6752ae4e0c..65b5dfb9161ef 100644 --- a/src/lightning/fabric/plugins/precision/amp.py +++ b/src/lightning/fabric/plugins/precision/amp.py @@ -29,21 +29,24 @@ class MixedPrecision(Precision): """Plugin for Automatic Mixed Precision (AMP) training with ``torch.autocast``. Args: - precision: Whether to use ``torch.float16`` (``16``) or ``torch.bfloat16`` (``'bf16'``). + precision: Whether to use ``torch.float16`` (``'16-mixed'``) or ``torch.bfloat16`` (``'bf16-mixed'``). device: The device for ``torch.autocast``. scaler: An optional :class:`torch.cuda.amp.GradScaler` to use. """ def __init__( - self, precision: Literal["16", 16, "bf16"], 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", "bf16"], str(precision)) - if scaler is None and self.precision == "16": + self.precision = cast(Literal["16-mixed", "bf16-mixed"], str(precision)) + if scaler is None and self.precision == "16-mixed": with _patch_cuda_is_available(): # if possible, we defer CUDA initialization to support strategies that will attempt forks scaler = torch.cuda.amp.GradScaler() - if scaler is not None and self.precision == "bf16": - raise ValueError(f"`precision='bf16'` does not use a scaler, found {scaler}.") + if scaler is not None and self.precision == "bf16-mixed": + raise ValueError(f"`precision='bf16-mixed'` does not use a scaler, found {scaler}.") self.device = device self.scaler = scaler @@ -53,7 +56,7 @@ def forward_context(self) -> Generator[None, None, None]: yield def convert_input(self, data: Tensor) -> Tensor: - precision_to_type = {"bf16": torch.bfloat16, "16": torch.float16} + precision_to_type = {"bf16-mixed": torch.bfloat16, "16-mixed": torch.float16} dst_type = precision_to_type[self.precision] return _convert_fp_tensor(data, dst_type) diff --git a/src/lightning/fabric/plugins/precision/deepspeed.py b/src/lightning/fabric/plugins/precision/deepspeed.py index c06a43d8c03c0..2d813da1ec657 100644 --- a/src/lightning/fabric/plugins/precision/deepspeed.py +++ b/src/lightning/fabric/plugins/precision/deepspeed.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 Any, cast, Literal, TYPE_CHECKING, Union +from typing import Any, cast, Literal, TYPE_CHECKING import torch from torch import Tensor @@ -27,16 +27,14 @@ if _DEEPSPEED_AVAILABLE: # type: ignore[has-type] import deepspeed -_PRECISION_INPUT_INT = Literal[32, 16] -_PRECISION_INPUT_STR = Literal["32", "16", "bf16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] +_PRECISION_INPUT = Literal["32-true", "16-mixed", "bf16-mixed"] class DeepSpeedPrecision(Precision): """Precision plugin for DeepSpeed integration. Args: - precision: Full precision (32), half precision (16) or bfloat16 precision (bf16). + precision: Full precision (32-true), half precision (16-mixed) or bfloat16 precision (bf16-mixed). Raises: ValueError: @@ -44,16 +42,16 @@ class DeepSpeedPrecision(Precision): """ def __init__(self, precision: _PRECISION_INPUT) -> None: - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + supported_precision = get_args(_PRECISION_INPUT) if precision not in supported_precision: raise ValueError( f"`precision={precision!r})` is not supported in DeepSpeed." f" `precision` must be one of: {supported_precision}." ) - self.precision = cast(_PRECISION_INPUT_STR, str(precision)) + self.precision = cast(_PRECISION_INPUT, precision) def convert_input(self, data: Tensor) -> Tensor: - precision_to_type = {"bf16": torch.bfloat16, "16": torch.float16, "32": torch.float32} + precision_to_type = {"bf16-mixed": torch.bfloat16, "16-mixed": torch.float16, "32-true": torch.float32} dst_type = precision_to_type[self.precision] return _convert_fp_tensor(data, dst_type) diff --git a/src/lightning/fabric/plugins/precision/double.py b/src/lightning/fabric/plugins/precision/double.py index 652fe8ede5076..687dec35ed568 100644 --- a/src/lightning/fabric/plugins/precision/double.py +++ b/src/lightning/fabric/plugins/precision/double.py @@ -25,7 +25,7 @@ class DoublePrecision(Precision): """Plugin for training with double (``torch.float64``) precision.""" - precision: Literal["64"] = "64" + precision: Literal["64-true"] = "64-true" def convert_module(self, module: Module) -> Module: return module.double() diff --git a/src/lightning/fabric/plugins/precision/fsdp.py b/src/lightning/fabric/plugins/precision/fsdp.py index dc471049df2d0..10de06ab2707a 100644 --- a/src/lightning/fabric/plugins/precision/fsdp.py +++ b/src/lightning/fabric/plugins/precision/fsdp.py @@ -27,7 +27,7 @@ class FSDPPrecision(MixedPrecision): """AMP for Fully Sharded Data Parallel training.""" def __init__( - self, precision: Literal["16", 16, "bf16"], device: str, scaler: Optional["ShardedGradScaler"] = None + self, precision: Literal["16-mixed", "bf16-mixed"], device: str, scaler: Optional["ShardedGradScaler"] = None ) -> None: if not _TORCH_GREATER_EQUAL_1_12: raise NotImplementedError("`FSDPPrecision` is supported from PyTorch v1.12.0 onwards.") @@ -37,16 +37,16 @@ def __init__( super().__init__( precision=precision, device=device, - scaler=(ShardedGradScaler() if scaler is None and str(precision) == "16" else None), + scaler=(ShardedGradScaler() if scaler is None and precision == "16-mixed" else None), ) @property def mixed_precision_config(self) -> "TorchMixedPrecision": from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision as TorchMixedPrecision - if self.precision == "16": + if self.precision == "16-mixed": dtype = torch.float16 - elif self.precision == "bf16": + elif self.precision == "bf16-mixed": dtype = torch.bfloat16 else: raise ValueError(f"Was unable to infer precision type, received {self.precision!r}.") diff --git a/src/lightning/fabric/plugins/precision/precision.py b/src/lightning/fabric/plugins/precision/precision.py index b609ee78a21e9..2c8e3625cdc86 100644 --- a/src/lightning/fabric/plugins/precision/precision.py +++ b/src/lightning/fabric/plugins/precision/precision.py @@ -23,8 +23,10 @@ from lightning.fabric.utilities.types import _PARAMETERS, Optimizable _PRECISION_INPUT_INT = Literal[64, 32, 16] -_PRECISION_INPUT_STR = Literal["64", "32", "16", "bf16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] +_PRECISION_INPUT_STR_LEGACY_CONVERSION = {"64": "64-true", "32": "32-true", "16": "16-mixed", "bf16": "bf16-mixed"} +_PRECISION_INPUT_STR_LEGACY = Literal["64", "32", "16", "bf16"] +_PRECISION_INPUT_STR = Literal["16-mixed", "bf16-mixed", "32-true", "64-true"] +_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR, _PRECISION_INPUT_STR_LEGACY] class Precision: @@ -33,7 +35,7 @@ class Precision: The class attribute precision must be overwritten in child classes. The default value reflects fp32 training. """ - precision: _PRECISION_INPUT_STR = "32" + precision: _PRECISION_INPUT_STR = "32-true" def convert_module(self, module: Module) -> Module: """Convert the module parameters to the precision type this plugin handles. diff --git a/src/lightning/fabric/plugins/precision/tpu_bf16.py b/src/lightning/fabric/plugins/precision/tpu_bf16.py index 79654a9c041a3..d0b9dc5b686e6 100644 --- a/src/lightning/fabric/plugins/precision/tpu_bf16.py +++ b/src/lightning/fabric/plugins/precision/tpu_bf16.py @@ -24,7 +24,7 @@ class TPUBf16Precision(TPUPrecision): """Plugin that enables bfloats on TPUs.""" - precision: Literal["bf16"] = "bf16" + precision: Literal["bf16-mixed"] = "bf16-mixed" def __init__(self) -> None: super().__init__() diff --git a/src/lightning/fabric/strategies/deepspeed.py b/src/lightning/fabric/strategies/deepspeed.py index dab9b7ec5886e..290b2d05c6334 100644 --- a/src/lightning/fabric/strategies/deepspeed.py +++ b/src/lightning/fabric/strategies/deepspeed.py @@ -105,8 +105,8 @@ def __init__( Arguments: - zero_optimization: Enable ZeRO optimization. This is compatible with either ``precision=16`` or - ``precision="bf16"``. + zero_optimization: Enable ZeRO optimization. This is compatible with either ``precision="16-mixed"`` or + ``precision="bf16-mixed"``. stage: Different stages of the ZeRO Optimizer. 0 is disabled, 1 is optimizer state partitioning, 2 is optimizer+gradient state partitioning, @@ -350,9 +350,9 @@ def module_sharded_context(self) -> Generator[None, None, None]: if self.zero_stage_3: assert self._config_initialized - if self.precision.precision == "16": + if self.precision.precision == "16-mixed": dtype = torch.float16 - elif self.precision.precision == "bf16": + elif self.precision.precision == "bf16-mixed": dtype = torch.bfloat16 else: dtype = torch.float32 @@ -604,7 +604,7 @@ def _format_config(self) -> None: def _format_precision_config(self) -> None: assert isinstance(self.config, dict) - if self.precision.precision == "16": + if self.precision.precision == "16-mixed": if "fp16" not in self.config: # FP16 is a DeepSpeed standalone AMP implementation rank_zero_info("Enabling DeepSpeed FP16.") @@ -616,7 +616,7 @@ def _format_precision_config(self) -> None: "hysteresis": self.hysteresis, "min_loss_scale": self.min_loss_scale, } - elif "bf16" not in self.config and self.precision.precision == "bf16": + elif "bf16" not in self.config and self.precision.precision == "bf16-mixed": rank_zero_info("Enabling DeepSpeed BF16.") self.config["bf16"] = {"enabled": True} diff --git a/src/lightning/fabric/strategies/fsdp.py b/src/lightning/fabric/strategies/fsdp.py index ebb27dfb15b49..b465bc9a94126 100644 --- a/src/lightning/fabric/strategies/fsdp.py +++ b/src/lightning/fabric/strategies/fsdp.py @@ -76,8 +76,9 @@ class FSDPStrategy(ParallelStrategy, _Sharded): backward_prefetch: This is an experimental feature that is subject to change in the near future. It allows users to enable two different backward prefetching algorithms to help backward communication and computation overlapping. The pros and cons of each algorithm is explained in the class ``BackwardPrefetch``. - mixed_precision: Mixed Precision config. By default, Lightning will enable FP16 if ``precision=16`` or BF16 - if ``precision=bf16`` unless a config is passed in. This is only available in PyTorch 1.12 and later. + mixed_precision: Mixed Precision config. By default, Lightning will enable FP16 if ``precision="16-mixed"`` or + BF16 if ``precision="bf16-mixed"`` unless a config is passed in. + This is only available in PyTorch 1.12 and later. activation_checkpointing: A single layer or a list of layer classes for which you want to enable activation checkpointing. This is typically your transformer block (including attention + feed-forward). Enabling this can free up a significant amount of memory at the cost of speed since activations in From 682131c6b3996fc317af97619d0191f1fd135d0b Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Wed, 15 Feb 2023 16:55:21 +0100 Subject: [PATCH 02/26] adapt fabric tests to new precision --- .../plugins/precision/test_amp.py | 18 +++--- .../plugins/precision/test_amp_integration.py | 8 +-- .../plugins/precision/test_deepspeed.py | 6 +- .../precision/test_deepspeed_integration.py | 2 +- .../precision/test_double_integration.py | 2 +- .../plugins/precision/test_fsdp.py | 4 +- .../strategies/test_deepspeed_integration.py | 14 ++--- .../strategies/test_fsdp_integration.py | 2 +- tests/tests_fabric/test_connector.py | 61 +++++++++++++++---- 9 files changed, 77 insertions(+), 40 deletions(-) diff --git a/tests/tests_fabric/plugins/precision/test_amp.py b/tests/tests_fabric/plugins/precision/test_amp.py index fca1bd9140968..ae96107a552f4 100644 --- a/tests/tests_fabric/plugins/precision/test_amp.py +++ b/tests/tests_fabric/plugins/precision/test_amp.py @@ -20,21 +20,21 @@ def test_amp_precision_default_scaler(): - precision = MixedPrecision(precision=16, device=Mock()) + precision = MixedPrecision(precision="16-mixed", device=Mock()) assert isinstance(precision.scaler, torch.cuda.amp.GradScaler) def test_amp_precision_scaler_with_bf16(): - with pytest.raises(ValueError, match="`precision='bf16'` does not use a scaler"): - MixedPrecision(precision="bf16", device=Mock(), scaler=Mock()) + with pytest.raises(ValueError, match="`precision='bf16-mixed'` does not use a scaler"): + MixedPrecision(precision="bf16-mixed", device=Mock(), scaler=Mock()) - precision = MixedPrecision(precision="bf16", device=Mock()) + precision = MixedPrecision(precision="bf16-mixed", device=Mock()) assert precision.scaler is None def test_amp_precision_forward_context(): """Test to ensure that the context manager correctly is set to bfloat16 on CPU and CUDA.""" - precision = MixedPrecision(precision=16, device="cuda") + precision = MixedPrecision(precision="16-mixed", device="cuda") assert precision.device == "cuda" assert isinstance(precision.scaler, torch.cuda.amp.GradScaler) assert torch.get_default_dtype() == torch.float32 @@ -42,7 +42,7 @@ def test_amp_precision_forward_context(): # check with str due to a bug upstream: https://github.com/pytorch/pytorch/issues/65786 assert str(torch.get_autocast_gpu_dtype()) in ("torch.float16", "torch.half") - precision = MixedPrecision(precision="bf16", device="cpu") + precision = MixedPrecision(precision="bf16-mixed", device="cpu") assert precision.device == "cpu" assert precision.scaler is None with precision.forward_context(): @@ -56,7 +56,7 @@ def test_amp_precision_forward_context(): def test_amp_precision_backward(): - precision = MixedPrecision(precision="mixed", device="cuda") + precision = MixedPrecision(precision="16-mixed", device="cuda") precision.scaler = Mock() precision.scaler.scale = Mock(side_effect=(lambda x: x)) tensor = Mock() @@ -67,7 +67,7 @@ def test_amp_precision_backward(): def test_amp_precision_optimizer_step_with_scaler(): - precision = MixedPrecision(precision="mixed", device="cuda") + precision = MixedPrecision(precision="16-mixed", device="cuda") precision.scaler = Mock() optimizer = Mock() @@ -77,7 +77,7 @@ def test_amp_precision_optimizer_step_with_scaler(): def test_amp_precision_optimizer_step_without_scaler(): - precision = MixedPrecision(precision="bf16", device="cuda") + precision = MixedPrecision(precision="bf16-mixed", device="cuda") assert precision.scaler is None optimizer = Mock() diff --git a/tests/tests_fabric/plugins/precision/test_amp_integration.py b/tests/tests_fabric/plugins/precision/test_amp_integration.py index 29c809d7c7b7b..133060b45a687 100644 --- a/tests/tests_fabric/plugins/precision/test_amp_integration.py +++ b/tests/tests_fabric/plugins/precision/test_amp_integration.py @@ -61,10 +61,10 @@ def after_backward(self, model): @pytest.mark.parametrize( "accelerator, precision, expected_dtype", [ - ("cpu", 16, torch.bfloat16), - ("cpu", "bf16", torch.bfloat16), - pytest.param("cuda", 16, torch.float16, marks=RunIf(min_cuda_gpus=1)), - pytest.param("cuda", "bf16", torch.bfloat16, marks=RunIf(min_cuda_gpus=1, bf16_cuda=True)), + ("cpu", "16-mixed", torch.bfloat16), + ("cpu", "bf16-mixed", torch.bfloat16), + pytest.param("cuda", "16-mixed", torch.float16, marks=RunIf(min_cuda_gpus=1)), + pytest.param("cuda", "bf16-mixed", torch.bfloat16, marks=RunIf(min_cuda_gpus=1, bf16_cuda=True)), ], ) def test_amp(accelerator, precision, expected_dtype): diff --git a/tests/tests_fabric/plugins/precision/test_deepspeed.py b/tests/tests_fabric/plugins/precision/test_deepspeed.py index 5eed7c0a4d933..6f8316e0d8f9e 100644 --- a/tests/tests_fabric/plugins/precision/test_deepspeed.py +++ b/tests/tests_fabric/plugins/precision/test_deepspeed.py @@ -23,11 +23,11 @@ def test_invalid_precision_with_deepspeed_precision(): with pytest.raises(ValueError, match="is not supported in DeepSpeed. `precision` must be one of"): - DeepSpeedPrecision(precision=64) + DeepSpeedPrecision(precision="64-true") def test_deepspeed_precision_backward(): - precision = DeepSpeedPrecision(precision=32) + precision = DeepSpeedPrecision(precision="32-true") tensor = Mock() model = Mock() precision.backward(tensor, model, "positional-arg", keyword="arg") @@ -45,7 +45,7 @@ def test_deepspeed_engine_is_steppable(engine): def test_deepspeed_precision_optimizer_step(): - precision = DeepSpeedPrecision(precision=32) + precision = DeepSpeedPrecision(precision="32-true") optimizer = model = Mock() precision.optimizer_step(optimizer, lr_kwargs=dict()) model.step.assert_called_once_with(lr_kwargs=dict()) diff --git a/tests/tests_fabric/plugins/precision/test_deepspeed_integration.py b/tests/tests_fabric/plugins/precision/test_deepspeed_integration.py index 5dfc670603db0..544649c163166 100644 --- a/tests/tests_fabric/plugins/precision/test_deepspeed_integration.py +++ b/tests/tests_fabric/plugins/precision/test_deepspeed_integration.py @@ -22,7 +22,7 @@ @RunIf(deepspeed=True) -@pytest.mark.parametrize("precision", ["bf16", 16, 32]) +@pytest.mark.parametrize("precision", ["bf16-mixed", "16-mixed", "32-true"]) @mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False) def test_deepspeed_precision_choice(_, precision): """Test to ensure precision plugin is correctly chosen. diff --git a/tests/tests_fabric/plugins/precision/test_double_integration.py b/tests/tests_fabric/plugins/precision/test_double_integration.py index 7af8ac6e28416..012d1f39623b0 100644 --- a/tests/tests_fabric/plugins/precision/test_double_integration.py +++ b/tests/tests_fabric/plugins/precision/test_double_integration.py @@ -50,5 +50,5 @@ def after_backward(self, model): def test_double_precision(): - fabric = DoublePrecisionBoringFabric(precision=64) + fabric = DoublePrecisionBoringFabric(precision="64-true") fabric.run() diff --git a/tests/tests_fabric/plugins/precision/test_fsdp.py b/tests/tests_fabric/plugins/precision/test_fsdp.py index cee1afe35941e..fb121dc191b7b 100644 --- a/tests/tests_fabric/plugins/precision/test_fsdp.py +++ b/tests/tests_fabric/plugins/precision/test_fsdp.py @@ -23,11 +23,11 @@ @mock.patch("lightning.fabric.plugins.precision.fsdp._TORCH_GREATER_EQUAL_1_12", False) def test_fsdp_precision_support(*_): with pytest.raises(NotImplementedError, match="`FSDPPrecision` is supported from PyTorch v1.12.0"): - FSDPPrecision(precision=16, device="cuda") + FSDPPrecision(precision="16-mixed", device="cuda") @RunIf(min_torch="1.12", min_cuda_gpus=1) -@pytest.mark.parametrize("precision, expected", [(16, torch.float16), ("bf16", torch.bfloat16)]) +@pytest.mark.parametrize("precision, expected", [("16-mixed", torch.float16), ("bf16-mixed", torch.bfloat16)]) def test_fsdp_precision_config(precision, expected): plugin = FSDPPrecision(precision=precision, device="cuda") config = plugin.mixed_precision_config diff --git a/tests/tests_fabric/strategies/test_deepspeed_integration.py b/tests/tests_fabric/strategies/test_deepspeed_integration.py index a068d19cbc31d..6e4232459ac3b 100644 --- a/tests/tests_fabric/strategies/test_deepspeed_integration.py +++ b/tests/tests_fabric/strategies/test_deepspeed_integration.py @@ -151,7 +151,7 @@ def run(self): strategy=DeepSpeedStrategy(), accelerator="cuda", devices=1, - precision=16, + precision="16-mixed", ) fabric.run() @@ -175,7 +175,7 @@ def run(self): ) fabric = RunFabric( strategy=strategy, - precision=16, + precision="16-mixed", accelerator="cuda", devices=1, ) @@ -212,7 +212,7 @@ def run(self): ) fabric = RunFabric( strategy=strategy, - precision=16, + precision="16-mixed", accelerator="cuda", devices=1, ) @@ -247,7 +247,7 @@ def test_deepspeed_multigpu_stage_3(): strategy=DeepSpeedStrategy(stage=3), accelerator="cuda", devices=2, - precision=16, + precision="16-mixed", ) fabric.run() @@ -318,9 +318,9 @@ def step(self, model, batch): assert model.layer.weight.dtype == torch.bfloat16 return super().step(model, batch) - fabric = RunFabric(accelerator="cuda", devices=2, strategy="deepspeed_stage_3", precision="bf16") + fabric = RunFabric(accelerator="cuda", devices=2, strategy="deepspeed_stage_3", precision="bf16-mixed") assert isinstance(fabric._strategy.precision, DeepSpeedPrecision) - assert fabric._strategy.precision.precision == "bf16" + assert fabric._strategy.precision.precision == "bf16-mixed" assert fabric._strategy.config["zero_optimization"]["stage"] == 3 fabric.run() @@ -361,7 +361,7 @@ def test_deepspeed_save_load_checkpoint_zero_3(stage, tmp_path): """Test that DeepSpeed stage 1, 2, and 3 model checkpoints can be saved and loaded successfully.""" from deepspeed import DeepSpeedEngine - fabric = Fabric(accelerator="cuda", devices=2, strategy=DeepSpeedStrategy(stage=stage), precision="bf16") + fabric = Fabric(accelerator="cuda", devices=2, strategy=DeepSpeedStrategy(stage=stage), precision="bf16-mixed") fabric.launch() checkpoint_path = fabric.broadcast(tmp_path / "deepspeed-checkpoint") diff --git a/tests/tests_fabric/strategies/test_fsdp_integration.py b/tests/tests_fabric/strategies/test_fsdp_integration.py index c80daee9723cc..c2040a5c4eab2 100644 --- a/tests/tests_fabric/strategies/test_fsdp_integration.py +++ b/tests/tests_fabric/strategies/test_fsdp_integration.py @@ -86,7 +86,7 @@ def _custom_auto_wrap_policy(module, recurse, unwrapped_params: int, min_num_par @RunIf(min_cuda_gpus=2, skip_windows=True, standalone=True, min_torch="1.13") -@pytest.mark.parametrize("precision", (16, pytest.param("bf16", marks=RunIf(bf16_cuda=True)))) +@pytest.mark.parametrize("precision", ("16-mixed", pytest.param("bf16-mixed", marks=RunIf(bf16_cuda=True)))) @pytest.mark.parametrize("manual_wrapping", [True, False]) def test_fsdp_train_save_load(manual_wrapping, precision): """Test FSDP training, saving and loading with different wrapping and precision settings.""" diff --git a/tests/tests_fabric/test_connector.py b/tests/tests_fabric/test_connector.py index 3d9f56b5076b7..f41ddc177bc24 100644 --- a/tests/tests_fabric/test_connector.py +++ b/tests/tests_fabric/test_connector.py @@ -19,6 +19,7 @@ import pytest import torch import torch.distributed +from lightning_utilities.test.warning import no_warning_call from tests_fabric.helpers.runif import RunIf import lightning.fabric @@ -438,6 +439,39 @@ def test_validate_precision_type(precision): _Connector(precision=precision) +@pytest.mark.parametrize( + "precision,expected_precision,should_warn", + [ + (16, "16-mixed", True), + ("16", "16-mixed", True), + ("16-mixed", "16-mixed", False), + ("bf16", "bf16-mixed", True), + ("bf16-mixed", "bf16-mixed", False), + (32, "32-true", True), + ("32", "32-true", True), + ("32-true", "32-true", False), + (64, "64-true", True), + ("64", "64-true", True), + ("64-true", "64-true", False), + ], +) +# mock cuda as available to not be limited by dtype and accelerator compatibility - this is tested elsewhere +@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}) +@mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=1) +@mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False) +def test_precision_conversion(patch1, patch2, precision, expected_precision, should_warn): + warn_context = pytest.warns if should_warn else no_warning_call + with warn_context( + UserWarning, + match=( + f"{precision} is supported for historical reasons but its usage is discouraged. " + f"Please set your precision to {expected_precision} instead!" + ), + ): + connector = _Connector(precision=precision, accelerator="cuda") + assert connector._precision_input == expected_precision + + def test_multi_device_default_strategy(): """The default strategy when multiple devices are selected is "ddp" with the subprocess launcher.""" connector = _Connector(strategy=None, accelerator="cpu", devices=2) @@ -632,14 +666,14 @@ def test_strategy_choice_ddp_cpu_slurm(strategy): @mock.patch.dict(os.environ, {}, clear=True) @mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False) def test_unsupported_tpu_choice(_, tpu_available): - with pytest.raises(NotImplementedError, match=r"accelerator='tpu', precision=64\)` is not implemented"): - _Connector(accelerator="tpu", precision=64) + with pytest.raises(NotImplementedError, match=r"accelerator='tpu', precision=64-true\)` is not implemented"): + _Connector(accelerator="tpu", precision="64-true") # if user didn't set strategy, _Connector will choose the TPUSingleStrategy or XLAStrategy with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"), pytest.warns( - UserWarning, match=r"accelerator='tpu', precision=16\)` but AMP is not supported" + UserWarning, match=r"accelerator='tpu', precision='16-mixed'\)` but AMP with fp16 is not supported" ): - _Connector(accelerator="tpu", precision=16, strategy="ddp") + _Connector(accelerator="tpu", precision="16-mixed", strategy="ddp") # wrong precision plugin type strategy = XLAStrategy(accelerator=TPUAccelerator(), precision=Precision()) @@ -760,8 +794,11 @@ def test_ddp_fork_on_unsupported_platform(_, __, strategy): def test_precision_selection_16_on_cpu_warns(): - with pytest.warns(UserWarning, match=r"precision=16\)` but AMP is not supported on CPU. Using `precision='bf16"): - _Connector(precision=16) + with pytest.warns( + UserWarning, + match=r"precision='16-mixed'\)` but AMP with fp16 is not supported on CPU. Using `precision='bf16-mixed'", + ): + _Connector(precision="16-mixed") class MyAMP(MixedPrecision): @@ -777,9 +814,9 @@ class MyAMP(MixedPrecision): def test_precision_selection_amp_ddp(strategy, devices, is_custom_plugin, plugin_cls): plugin = None if is_custom_plugin: - plugin = plugin_cls(16, "cpu") + plugin = plugin_cls("16-mixed", "cpu") connector = _Connector( - precision=16, + precision="16-mixed", devices=devices, strategy=strategy, plugins=plugin, @@ -794,7 +831,7 @@ def test_strategy_str_passed_being_case_insensitive(_, strategy, strategy_cls): assert isinstance(connector.strategy, strategy_cls) -@pytest.mark.parametrize("precision", ["64", "32", "16", "bf16"]) +@pytest.mark.parametrize("precision", ["64-true", "32-true", "16-mixed", "bf16-mixed"]) @mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=1) def test_precision_from_environment(_, precision): """Test that the precision input can be set through the environment variable.""" @@ -856,9 +893,9 @@ def test_arguments_from_environment_collision(): with pytest.raises(ValueError, match="`Fabric\\(num_nodes=2, ...\\)` but .* `--num_nodes=3`"): _Connector(num_nodes=2) - with mock.patch.dict(os.environ, {"LT_PRECISION": "16"}): - with pytest.raises(ValueError, match="`Fabric\\(precision=64, ...\\)` but .* `--precision=16`"): - _Connector(precision=64) + with mock.patch.dict(os.environ, {"LT_PRECISION": "16-mixed"}): + with pytest.raises(ValueError, match="`Fabric\\(precision=64-true, ...\\)` but .* `--precision=16-mixed`"): + _Connector(precision="64-true") @RunIf(min_torch="1.12") From ff870f2fb8c7862896512ede5f1c32e3bfc223e3 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 10:45:32 +0100 Subject: [PATCH 03/26] update docs --- .../source-pytorch/fabric/api/fabric_args.rst | 11 +++++---- .../fabric/fundamentals/launch.rst | 8 ++++--- .../fabric/fundamentals/precision.rst | 24 +++++++------------ .../fabric/guide/multi_node/cloud.rst | 2 +- 4 files changed, 21 insertions(+), 24 deletions(-) diff --git a/docs/source-pytorch/fabric/api/fabric_args.rst b/docs/source-pytorch/fabric/api/fabric_args.rst index 244129549a110..2006a3d1670d9 100644 --- a/docs/source-pytorch/fabric/api/fabric_args.rst +++ b/docs/source-pytorch/fabric/api/fabric_args.rst @@ -112,23 +112,24 @@ Learn more about :ref:`distributed multi-node training on clusters `_). +Fabric supports double precision (64 bit), full precision (32 bit), or half-precision (16 bit) floating point operation (including `bfloat16 `_). Half precision, or mixed precision, combines 32 and 16-bit floating points to reduce the memory footprint during model training. +Automatic mixed precision settings are denoted by a ``"-mixed"`` suffix, while settings that only work in the specified precision have a ``"-true"`` suffix. This can result in improved performance, achieving significant speedups on modern GPUs. .. code-block:: python # Default used by the Fabric - fabric = Fabric(precision=32, devices=1) + fabric = Fabric(precision="32-true", devices=1) # 16-bit (mixed) precision - fabric = Fabric(precision=16, devices=1) + fabric = Fabric(precision="16-mixed", devices=1) # 16-bit bfloat precision - fabric = Fabric(precision="bf16", devices=1) + fabric = Fabric(precision="bf16-mixed", devices=1) # 64-bit (double) precision - fabric = Fabric(precision=64, devices=1) + fabric = Fabric(precision="64-true", devices=1) See also: :doc:`../fundamentals/precision` diff --git a/docs/source-pytorch/fabric/fundamentals/launch.rst b/docs/source-pytorch/fabric/fundamentals/launch.rst index 9a49d9b050f4f..95959afd83084 100644 --- a/docs/source-pytorch/fabric/fundamentals/launch.rst +++ b/docs/source-pytorch/fabric/fundamentals/launch.rst @@ -68,9 +68,11 @@ This is essentially the same as running ``python path/to/your/script.py``, but i --main-port, --main_port INTEGER The main port to connect to the main machine. - --precision [64|32|16|bf16] Double precision (``64``), full precision - (``32``), half precision (``16``) or - bfloat16 precision (``'bf16'``) + --precision [16-mixed|bf16-mixed|32-true|64-true] + Double precision (``64-true``), full + precision (``32-true``), half precision + (``16-mixed``) or bfloat16 precision + (``'bf16-mixed'``) --help Show this message and exit. diff --git a/docs/source-pytorch/fabric/fundamentals/precision.rst b/docs/source-pytorch/fabric/fundamentals/precision.rst index d3a5b2ab726df..7a9a0c4692881 100644 --- a/docs/source-pytorch/fabric/fundamentals/precision.rst +++ b/docs/source-pytorch/fabric/fundamentals/precision.rst @@ -24,26 +24,23 @@ This is how you select the precision in Fabric: from lightning.fabric import Fabric # This is the default - fabric = Fabric(precision=32) + fabric = Fabric(precision="32-true") # FP16 mixed precision - fabric = Fabric(precision=16) - - # Precision values can also be set as a string - fabric = Fabric(precision="16") + fabric = Fabric(precision="16-mixed) # BFloat16 precision (Volta GPUs and later) - fabric = Fabric(precision="bf16") + fabric = Fabric(precision="bf16-mixed") # Double precision - fabric = Fabric(precision=64) + fabric = Fabric(precision="64-true") The same values can also be set through the :doc:`command line interface `: .. code-block:: bash - lightning run model train.py --precision=bf16 + lightning run model train.py --precision=bf16-mixed .. note:: @@ -70,14 +67,11 @@ This is how you enable FP16 in Fabric: .. code-block:: python # Select FP16 mixed precision - fabric = Fabric(precision=16) - - # Or as a string - fabric = Fabric(precision="16") + fabric = Fabric(precision="16-mixed") .. note:: - When using TPUs, setting ``precision=16`` will enable bfloat16, the only supported half-precision type on TPUs. + When using TPUs, setting ``precision="16-mixed"`` will enable bfloat16 based mixed precision, the only supported half-precision type on TPUs. ---- @@ -94,7 +88,7 @@ For more information, see `this TPU performance blog post `__ with the dtype set to ``bfloat16``, with no gradient scaling. @@ -117,7 +111,7 @@ Fabric automatically casts the data type and operations in the ``forward`` of yo .. code-block:: python - fabric = Fabric(precision="bf16") + fabric = Fabric(precision="bf16-mixed") model = ... optimizer = ... diff --git a/docs/source-pytorch/fabric/guide/multi_node/cloud.rst b/docs/source-pytorch/fabric/guide/multi_node/cloud.rst index 117e64c6e3592..f5a7d9b352b0e 100644 --- a/docs/source-pytorch/fabric/guide/multi_node/cloud.rst +++ b/docs/source-pytorch/fabric/guide/multi_node/cloud.rst @@ -50,7 +50,7 @@ Launch multi-node training in the cloud def run(self): # Set up Fabric # The `devices` and `num_nodes` gets set by Lightning automatically - fabric = L.Fabric(strategy="ddp", precision=16) + fabric = L.Fabric(strategy="ddp", precision="16-mixed") # Your training code model = ... From 02591b07d6b3193307a67889b979db5a14bef334 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 10:45:36 +0100 Subject: [PATCH 04/26] fix cli --- src/lightning/fabric/cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/lightning/fabric/cli.py b/src/lightning/fabric/cli.py index 1410a1b3deb17..40615ff7ca55d 100644 --- a/src/lightning/fabric/cli.py +++ b/src/lightning/fabric/cli.py @@ -16,6 +16,7 @@ import re from argparse import Namespace from typing import Any, List, Optional +from typing_extensions import get_args from lightning_utilities.core.imports import RequirementCache @@ -106,7 +107,7 @@ def _get_supported_strategies() -> List[str]: ) @click.option( "--precision", - type=click.Choice(_PRECISION_INPUT_STR), + type=click.Choice(get_args(_PRECISION_INPUT_STR)), default="32", help=( "Double precision (``64-true``), full precision (``32-true``), half precision (``16-mixed``) or " From a5e48483a2187422e62b54c680a11d25339185f3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Feb 2023 09:47:54 +0000 Subject: [PATCH 05/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning/fabric/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/fabric/cli.py b/src/lightning/fabric/cli.py index 40615ff7ca55d..c974f831296b9 100644 --- a/src/lightning/fabric/cli.py +++ b/src/lightning/fabric/cli.py @@ -16,9 +16,9 @@ import re from argparse import Namespace from typing import Any, List, Optional -from typing_extensions import get_args from lightning_utilities.core.imports import RequirementCache +from typing_extensions import get_args from lightning.fabric.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT_STR From 457b6d17e16a9417529500579cae2d6ad7ab895c Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 10:49:46 +0100 Subject: [PATCH 06/26] changelog --- src/lightning/fabric/CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md index 77d93955929ee..3e826b52d57f4 100644 --- a/src/lightning/fabric/CHANGELOG.md +++ b/src/lightning/fabric/CHANGELOG.md @@ -37,6 +37,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Enabled all shorthand strategy names that can be supported in the CLI ([#16485](https://github.com/Lightning-AI/lightning/pull/16485)) +- Changed arguments for precision settings (from [64|32|16|bf16] to ["64-true"|"32-true"|"16-mixed"|"bf16-mixed"]) ([#16767](https://github.com/Lightning-AI/lightning/pull/16767)) + ### Deprecated - From 0299a29f777287473ae32064cc23ca48722c41a9 Mon Sep 17 00:00:00 2001 From: Justus Schock <12886177+justusschock@users.noreply.github.com> Date: Thu, 16 Feb 2023 11:38:03 +0100 Subject: [PATCH 07/26] cli fixes --- src/lightning/fabric/cli.py | 2 +- tests/tests_fabric/test_cli.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lightning/fabric/cli.py b/src/lightning/fabric/cli.py index c974f831296b9..f6b5bf7112d6e 100644 --- a/src/lightning/fabric/cli.py +++ b/src/lightning/fabric/cli.py @@ -108,7 +108,7 @@ def _get_supported_strategies() -> List[str]: @click.option( "--precision", type=click.Choice(get_args(_PRECISION_INPUT_STR)), - default="32", + default="32-true", help=( "Double precision (``64-true``), full precision (``32-true``), half precision (``16-mixed``) or " "bfloat16 precision (``'bf16-mixed'``)" diff --git a/tests/tests_fabric/test_cli.py b/tests/tests_fabric/test_cli.py index a10fd11c41eeb..df7832c12dd1c 100644 --- a/tests/tests_fabric/test_cli.py +++ b/tests/tests_fabric/test_cli.py @@ -43,7 +43,7 @@ def test_cli_env_vars_defaults(monkeypatch, fake_script): assert "LT_STRATEGY" not in os.environ assert os.environ["LT_DEVICES"] == "1" assert os.environ["LT_NUM_NODES"] == "1" - assert os.environ["LT_PRECISION"] == "32" + assert os.environ["LT_PRECISION"] == "32-true" @pytest.mark.parametrize("accelerator", ["cpu", "gpu", "cuda", pytest.param("mps", marks=RunIf(mps=True))]) @@ -120,7 +120,7 @@ def test_cli_env_vars_num_nodes(num_nodes, monkeypatch, fake_script): assert os.environ["LT_NUM_NODES"] == num_nodes -@pytest.mark.parametrize("precision", ["64", "32", "16", "bf16"]) +@pytest.mark.parametrize("precision", ["64-true", "32-true", "16-mixed", "bf16-mixed"]) @mock.patch.dict(os.environ, os.environ.copy(), clear=True) def test_cli_env_vars_precision(precision, monkeypatch, fake_script): monkeypatch.setattr(torch.distributed, "run", Mock()) From 2e66197f5f7866066193f5fe0bd81d1d5b10374e Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 13:03:29 +0100 Subject: [PATCH 08/26] fix tests and warnings --- src/lightning/fabric/connector.py | 10 +++++----- src/lightning/fabric/plugins/precision/amp.py | 2 +- tests/tests_fabric/test_connector.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/lightning/fabric/connector.py b/src/lightning/fabric/connector.py index e1d5b7dabd527..05e76f926d097 100644 --- a/src/lightning/fabric/connector.py +++ b/src/lightning/fabric/connector.py @@ -104,7 +104,7 @@ def __init__( strategy: Optional[Union[str, Strategy]] = None, devices: Optional[Union[List[int], str, int]] = None, num_nodes: int = 1, - precision: _PRECISION_INPUT = 32, + precision: _PRECISION_INPUT = "32-true", plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]] = None, ) -> None: @@ -237,7 +237,7 @@ def _check_config_and_set_final_flags( if precision in get_args(_PRECISION_INPUT_STR_LEGACY): rank_zero_warn( f"{precision} is supported for historical reasons but its usage is discouraged. " - "Please set your precision to {_PRECISION_INPUT_STR_LEGACY_CONVERSION[precision]} instead!" + f"Please set your precision to {_PRECISION_INPUT_STR_LEGACY_CONVERSION[precision]} instead!" ) precision = _PRECISION_INPUT_STR_LEGACY_CONVERSION[precision] @@ -490,8 +490,8 @@ def _check_and_init_precision(self) -> Precision: if self._precision_input == "16-mixed" and self._accelerator_flag == "cpu": rank_zero_warn( - "You passed `Fabric(accelerator='cpu', precision='16-mixed')` but AMP is not supported on CPU with " - "fp16. Using `precision='bf16-mixed'` instead." + "You passed `Fabric(accelerator='cpu', precision='16-mixed')` but AMP with fp16 is not supported on " + "CPU. Using `precision='bf16-mixed'` instead." ) self._precision_input = "bf16-mixed" @@ -514,7 +514,7 @@ def _validate_precision_choice(self) -> None: if isinstance(self.accelerator, TPUAccelerator): if self._precision_input == "64-true": raise NotImplementedError( - "`Fabric(accelerator='tpu', precision=64)` is not implemented." + "`Fabric(accelerator='tpu', precision='64-true')` is not implemented." " Please, open an issue in `https://github.com/Lightning-AI/lightning/issues`" " requesting this feature." ) diff --git a/src/lightning/fabric/plugins/precision/amp.py b/src/lightning/fabric/plugins/precision/amp.py index 65b5dfb9161ef..e7ff4858220d2 100644 --- a/src/lightning/fabric/plugins/precision/amp.py +++ b/src/lightning/fabric/plugins/precision/amp.py @@ -92,4 +92,4 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None: def _autocast_context_manager(self) -> torch.autocast: # the dtype could be automatically inferred but we need to manually set it due to a bug upstream # https://github.com/pytorch/pytorch/issues/67233 - return torch.autocast(self.device, dtype=torch.bfloat16 if self.precision == "bf16" else torch.half) + return torch.autocast(self.device, dtype=torch.bfloat16 if self.precision == "bf16-mixed" else torch.half) diff --git a/tests/tests_fabric/test_connector.py b/tests/tests_fabric/test_connector.py index f41ddc177bc24..f13d461312488 100644 --- a/tests/tests_fabric/test_connector.py +++ b/tests/tests_fabric/test_connector.py @@ -666,7 +666,7 @@ def test_strategy_choice_ddp_cpu_slurm(strategy): @mock.patch.dict(os.environ, {}, clear=True) @mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False) def test_unsupported_tpu_choice(_, tpu_available): - with pytest.raises(NotImplementedError, match=r"accelerator='tpu', precision=64-true\)` is not implemented"): + with pytest.raises(NotImplementedError, match=r"accelerator='tpu', precision='64-true'\)` is not implemented"): _Connector(accelerator="tpu", precision="64-true") # if user didn't set strategy, _Connector will choose the TPUSingleStrategy or XLAStrategy @@ -894,7 +894,7 @@ def test_arguments_from_environment_collision(): _Connector(num_nodes=2) with mock.patch.dict(os.environ, {"LT_PRECISION": "16-mixed"}): - with pytest.raises(ValueError, match="`Fabric\\(precision=64-true, ...\\)` but .* `--precision=16-mixed`"): + with pytest.raises(ValueError, match="`Fabric\\(precision='64-true', ...\\)` but .* `--precision=16-mixed`"): _Connector(precision="64-true") From af09b84860ceb274811cf4153eafd64780510995 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 15:05:57 +0100 Subject: [PATCH 09/26] update PL docs --- .../source-pytorch/common/precision_basic.rst | 24 ++++++++++++++----- .../common/precision_expert.rst | 2 +- .../common/precision_intermediate.rst | 2 +- docs/source-pytorch/common/trainer.rst | 4 ++-- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/docs/source-pytorch/common/precision_basic.rst b/docs/source-pytorch/common/precision_basic.rst index 3cc0b3a9677be..0b8706a194b68 100644 --- a/docs/source-pytorch/common/precision_basic.rst +++ b/docs/source-pytorch/common/precision_basic.rst @@ -20,11 +20,11 @@ Higher precision, such as the 64-bit floating-point, can be used for highly sens 16-bit Precision **************** -Use 16-bit precision to cut your memory consumption in half so that you can train and deploy larger models. If your GPUs are [`Tensor Core `_] GPUs, you can also get a ~3x speed improvement. Half precision can sometimes lead to unstable training. +Use 16-bit mixed precision to lower your memory consumption by up to half so that you can train and deploy larger models. If your GPUs are [`Tensor Core `_] GPUs, you can also get a ~3x speed improvement. Half precision can sometimes lead to unstable training. .. code:: - Trainer(precision=16) + Trainer(precision='16-mixed') ---- @@ -36,6 +36,12 @@ Use 16-bit precision to cut your memory consumption in half so that you can trai .. testcode:: + Trainer(precision='32-true') + + # or + Trainer(precision='32') + + # or Trainer(precision=32) ---- @@ -48,6 +54,12 @@ For certain scientific computations, 64-bit precision enables more accurate mode .. testcode:: + Trainer(precision='64-true') + + # or + Trainer(precision='64') + + # or Trainer(precision=64) .. note:: @@ -70,22 +82,22 @@ Precision support by accelerator - GPU - TPU - IPU - * - 16 + * - 16 Mixed - No - Yes - No - Yes - * - BFloat16 + * - BFloat16 Mixed - Yes - Yes - Yes - No - * - 32 + * - 32 True - Yes - Yes - Yes - Yes - * - 64 + * - 64 True - Yes - Yes - No diff --git a/docs/source-pytorch/common/precision_expert.rst b/docs/source-pytorch/common/precision_expert.rst index 34bc95568c962..7a6c2dada1c17 100644 --- a/docs/source-pytorch/common/precision_expert.rst +++ b/docs/source-pytorch/common/precision_expert.rst @@ -20,7 +20,7 @@ You can also customize and pass your own Precision Plugin by subclassing the :cl .. code-block:: python class CustomPrecisionPlugin(PrecisionPlugin): - precision = 16 + precision = '16-mixed' ... diff --git a/docs/source-pytorch/common/precision_intermediate.rst b/docs/source-pytorch/common/precision_intermediate.rst index 52ad86d004e0b..7cdd929ad0e4b 100644 --- a/docs/source-pytorch/common/precision_intermediate.rst +++ b/docs/source-pytorch/common/precision_intermediate.rst @@ -63,7 +63,7 @@ Since computation happens in FP16, there is a chance of numerical instability du .. note:: - When using TPUs, setting ``precision=16`` will enable bfloat16, the only supported half precision type on TPUs. + When using TPUs, setting ``precision='16-mixed'`` will enable bfloat16, the only supported half precision type on TPUs. .. testcode:: :skipif: not torch.cuda.is_available() diff --git a/docs/source-pytorch/common/trainer.rst b/docs/source-pytorch/common/trainer.rst index fd8b3af0cf982..3ea6436f7c537 100644 --- a/docs/source-pytorch/common/trainer.rst +++ b/docs/source-pytorch/common/trainer.rst @@ -926,10 +926,10 @@ Half precision, or mixed precision, is the combined use of 32 and 16 bit floatin trainer = Trainer(precision=32) # 16-bit precision - trainer = Trainer(precision=16, accelerator="gpu", devices=1) # works only on CUDA + trainer = Trainer(precision="16-mixed", accelerator="gpu", devices=1) # works only on CUDA # bfloat16 precision - trainer = Trainer(precision="bf16") + trainer = Trainer(precision="bf16-mixed") # 64-bit precision trainer = Trainer(precision=64) From 747a08c8aa5ea26ff04d59fdee134d3ebd491ed8 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 15:06:06 +0100 Subject: [PATCH 10/26] update examples --- examples/app_multi_node/train_fabric.py | 2 +- examples/pl_hpu/mnist_sample.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/app_multi_node/train_fabric.py b/examples/app_multi_node/train_fabric.py index 1bb2ecd313202..4fd7bee094697 100644 --- a/examples/app_multi_node/train_fabric.py +++ b/examples/app_multi_node/train_fabric.py @@ -15,7 +15,7 @@ def run(self): ) # 2. Create Fabric. - fabric = Fabric(strategy="ddp", precision=16) + fabric = Fabric(strategy="ddp", precision='16-mixed') model, optimizer = fabric.setup(model, torch.optim.SGD(model.parameters(), lr=0.01)) criterion = torch.nn.MSELoss() diff --git a/examples/pl_hpu/mnist_sample.py b/examples/pl_hpu/mnist_sample.py index ccb60e7c9de14..09613e295f47a 100644 --- a/examples/pl_hpu/mnist_sample.py +++ b/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), + "plugins": lazy_instance(HPUPrecisionPlugin, precision='16-mixed'), }, run=False, save_config_kwargs={"overwrite": True}, From de2281dfbb97c16eecb81a8b149c9b9e250e79eb Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 15:06:19 +0100 Subject: [PATCH 11/26] update src code --- .../pytorch/plugins/precision/amp.py | 12 ++-- .../pytorch/plugins/precision/deepspeed.py | 11 ++-- .../pytorch/plugins/precision/double.py | 2 +- .../pytorch/plugins/precision/fsdp.py | 8 +-- .../pytorch/plugins/precision/hpu.py | 10 ++-- .../pytorch/plugins/precision/ipu.py | 13 ++--- .../pytorch/plugins/precision/tpu_bf16.py | 2 +- src/lightning/pytorch/strategies/deepspeed.py | 12 ++-- src/lightning/pytorch/strategies/fsdp.py | 4 +- .../connectors/accelerator_connector.py | 55 +++++++++++-------- src/lightning/pytorch/trainer/trainer.py | 7 ++- 11 files changed, 70 insertions(+), 66 deletions(-) diff --git a/src/lightning/pytorch/plugins/precision/amp.py b/src/lightning/pytorch/plugins/precision/amp.py index ce984070ae7a5..bddcf2b480a75 100644 --- a/src/lightning/pytorch/plugins/precision/amp.py +++ b/src/lightning/pytorch/plugins/precision/amp.py @@ -34,15 +34,15 @@ class MixedPrecisionPlugin(PrecisionPlugin): """ def __init__( - self, precision: Literal["16", 16, "bf16"], 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", "bf16"], str(precision)) - if scaler is None and self.precision == "16": + self.precision = cast(Literal["16-mixed", "bf16-mixed"], str(precision)) + if scaler is None and self.precision == "16-mixed": with _patch_cuda_is_available(): # if possible, we defer CUDA initialization to support strategies that will attempt forks scaler = torch.cuda.amp.GradScaler() - if scaler is not None and self.precision == "bf16": - raise MisconfigurationException(f"`precision='bf16'` does not use a scaler, found {scaler}.") + if scaler is not None and self.precision == "bf16-mixed": + raise MisconfigurationException(f"`precision='bf16-mixed'` does not use a scaler, found {scaler}.") self.device = device self.scaler = scaler @@ -97,7 +97,7 @@ def clip_gradients( def autocast_context_manager(self) -> torch.autocast: # the dtype could be automatically inferred but we need to manually set it due to a bug upstream # https://github.com/pytorch/pytorch/issues/67233 - return torch.autocast(self.device, dtype=torch.bfloat16 if self.precision == "bf16" else torch.half) + return torch.autocast(self.device, dtype=torch.bfloat16 if self.precision == "bf16-mixed" else torch.half) @contextmanager def forward_context(self) -> Generator[None, None, None]: diff --git a/src/lightning/pytorch/plugins/precision/deepspeed.py b/src/lightning/pytorch/plugins/precision/deepspeed.py index 20f5748aa444e..99f8d06173fa8 100644 --- a/src/lightning/pytorch/plugins/precision/deepspeed.py +++ b/src/lightning/pytorch/plugins/precision/deepspeed.py @@ -31,9 +31,8 @@ warning_cache = WarningCache() -_PRECISION_INPUT_INT = Literal[32, 16] -_PRECISION_INPUT_STR = Literal["32", "16", "bf16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] +_PRECISION_INPUT = Literal["32-true", "16-mixed", "bf16-mixed"] + class DeepSpeedPrecisionPlugin(PrecisionPlugin): @@ -46,14 +45,14 @@ class DeepSpeedPrecisionPlugin(PrecisionPlugin): If unsupported ``precision`` is provided. """ - def __init__(self, precision: Literal["32", 32, "16", 16, "bf16"]) -> None: - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + def __init__(self, precision: Literal["32-true", "16-mixed", "bf16-mixed"]) -> None: + supported_precision = get_args(_PRECISION_INPUT) if precision not in supported_precision: raise ValueError( f"`Trainer(strategy='deepspeed', precision={precision!r})` is not supported." f" `precision` must be one of: {supported_precision}." ) - self.precision = cast(_PRECISION_INPUT_STR, str(precision)) + self.precision = cast(_PRECISION_INPUT, str(precision)) def backward( # type: ignore[override] self, diff --git a/src/lightning/pytorch/plugins/precision/double.py b/src/lightning/pytorch/plugins/precision/double.py index e008097046637..77fa9c4171a2b 100644 --- a/src/lightning/pytorch/plugins/precision/double.py +++ b/src/lightning/pytorch/plugins/precision/double.py @@ -72,7 +72,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: class DoublePrecisionPlugin(PrecisionPlugin): """Plugin for training with double (``torch.float64``) precision.""" - precision: Literal["64"] = "64" + precision: Literal["64-true"] = "64-true" def connect( self, model: nn.Module, optimizers: List[Optimizer], lr_schedulers: List[Any] diff --git a/src/lightning/pytorch/plugins/precision/fsdp.py b/src/lightning/pytorch/plugins/precision/fsdp.py index 7e1d6a5250294..1561bd693f037 100644 --- a/src/lightning/pytorch/plugins/precision/fsdp.py +++ b/src/lightning/pytorch/plugins/precision/fsdp.py @@ -31,12 +31,12 @@ class FSDPMixedPrecisionPlugin(MixedPrecisionPlugin): """AMP for Fully Sharded Data Parallel (FSDP) Training.""" def __init__( - self, precision: Literal["16", 16, "bf16"], device: str, scaler: Optional[ShardedGradScaler] = None + self, precision: Literal["16-mixed", "bf16-mixed"], device: str, scaler: Optional[ShardedGradScaler] = None ) -> None: if not _TORCH_GREATER_EQUAL_1_12: raise MisconfigurationException("`FSDPMixedPrecisionPlugin` is supported from PyTorch v1.12.0 onwards.") super().__init__( - precision, device, scaler=(ShardedGradScaler() if scaler is None and str(precision) == "16" else None) + precision, device, scaler=(ShardedGradScaler() if scaler is None and str(precision) == "16-mixed" else None) ) def clip_grad_by_norm(self, *_: Any, **__: Any) -> None: @@ -52,9 +52,9 @@ def clip_grad_by_norm(self, *_: Any, **__: Any) -> None: @property def mixed_precision_config(self) -> Optional[MixedPrecision]: assert MixedPrecision is not None - if self.precision == "16": + if self.precision == "16-mixed": dtype = torch.float16 - elif self.precision == "bf16": + elif self.precision == "bf16-mixed": dtype = torch.bfloat16 else: raise MisconfigurationException(f"Was unable to infer precision type, received {self.precision!r}.") diff --git a/src/lightning/pytorch/plugins/precision/hpu.py b/src/lightning/pytorch/plugins/precision/hpu.py index 8d805deae1da8..65863e9f77d22 100644 --- a/src/lightning/pytorch/plugins/precision/hpu.py +++ b/src/lightning/pytorch/plugins/precision/hpu.py @@ -22,9 +22,7 @@ if _HPU_AVAILABLE: from habana_frameworks.torch.hpex import hmp -_PRECISION_INPUT_INT = Literal[32, 16] -_PRECISION_INPUT_STR = Literal["32", "16", "bf16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] +_PRECISION_INPUT = Literal["32-true", "16-mixed", "bf16-mixed"] class HPUPrecisionPlugin(PrecisionPlugin): @@ -48,14 +46,14 @@ def __init__( ) -> None: if not _HPU_AVAILABLE: raise MisconfigurationException("HPU precision plugin requires HPU devices.") - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + supported_precision = get_args(_PRECISION_INPUT) if precision not in supported_precision: raise ValueError( f"`Trainer(accelerator='hpu', precision={precision!r})` is not supported." f" `precision` must be one of: {supported_precision}." ) - self.precision = cast(_PRECISION_INPUT_STR, str(precision)) - if self.precision in ("16", "bf16"): + self.precision = cast(_PRECISION_INPUT, str(precision)) + if self.precision in ("16-mixed", "bf16-mixed"): hmp.convert( opt_level=opt_level, bf16_file_path=bf16_file_path, fp32_file_path=fp32_file_path, isVerbose=verbose ) diff --git a/src/lightning/pytorch/plugins/precision/ipu.py b/src/lightning/pytorch/plugins/precision/ipu.py index f82bc07ac2119..95737fb996d49 100644 --- a/src/lightning/pytorch/plugins/precision/ipu.py +++ b/src/lightning/pytorch/plugins/precision/ipu.py @@ -27,27 +27,24 @@ warning_cache = WarningCache() -_PRECISION_INPUT_INT = Literal[32, 16] -_PRECISION_INPUT_STR = Literal["32", "16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] - +_PRECISION_INPUT = Literal["32-true", "16-mixed"] class IPUPrecisionPlugin(PrecisionPlugin): """Precision plugin for IPU integration. Raises: ValueError: - If the precision is neither 16 nor 32. + If the precision is neither 16-mixed nor 32-true. """ - def __init__(self, precision: Literal["32", 32, "16", 16]) -> None: - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + def __init__(self, precision: Literal["32-true", "16-mixed"]) -> None: + supported_precision = get_args(_PRECISION_INPUT) if precision not in supported_precision: raise ValueError( f"`Trainer(accelerator='ipu', precision={precision!r})` is not supported." f" `precision` must be one of: {supported_precision}." ) - self.precision = cast(_PRECISION_INPUT_STR, str(precision)) + self.precision = cast(_PRECISION_INPUT, str(precision)) def backward( # type: ignore[override] self, diff --git a/src/lightning/pytorch/plugins/precision/tpu_bf16.py b/src/lightning/pytorch/plugins/precision/tpu_bf16.py index 814af173d0464..0ff320d1ea807 100644 --- a/src/lightning/pytorch/plugins/precision/tpu_bf16.py +++ b/src/lightning/pytorch/plugins/precision/tpu_bf16.py @@ -23,7 +23,7 @@ class TPUBf16PrecisionPlugin(TPUPrecisionPlugin): """Plugin that enables bfloats on TPUs.""" - precision: Literal["bf16"] = "bf16" + precision: Literal["bf16"] = "bf16-mixed" def connect( self, model: nn.Module, optimizers: List[Optimizer], lr_schedulers: List[Any] diff --git a/src/lightning/pytorch/strategies/deepspeed.py b/src/lightning/pytorch/strategies/deepspeed.py index 4b7a806cbf95e..7778cab07c47f 100644 --- a/src/lightning/pytorch/strategies/deepspeed.py +++ b/src/lightning/pytorch/strategies/deepspeed.py @@ -127,8 +127,8 @@ def __init__( Arguments: - zero_optimization: Enable ZeRO optimization. This is compatible with either `precision=16` or - `precision="bf16"`. + zero_optimization: Enable ZeRO optimization. This is compatible with either `precision="16-mixed"` or + `precision="bf16-mixed"`. stage: Different stages of the ZeRO Optimizer. 0 is disabled, 1 is optimizer state partitioning, 2 is optimizer+gradient state partitioning, @@ -505,9 +505,9 @@ def model_sharded_context(self) -> Generator[None, None, None]: if self.zero_stage_3: assert self._config_initialized - if self.precision_plugin.precision == "16": + if self.precision_plugin.precision == "16-mixed": dtype = torch.float16 - elif self.precision_plugin.precision == "bf16": + elif self.precision_plugin.precision == "bf16-mixed": dtype = torch.bfloat16 else: dtype = torch.float32 @@ -641,7 +641,7 @@ def _auto_select_batch_size(self) -> int: def _format_precision_config(self) -> None: assert isinstance(self.config, dict) - if self.precision_plugin.precision == "16": + if self.precision_plugin.precision == "16-mixed": if "fp16" not in self.config: # FP16 is a DeepSpeed standalone AMP implementation rank_zero_info("Enabling DeepSpeed FP16.") @@ -653,7 +653,7 @@ def _format_precision_config(self) -> None: "hysteresis": self.hysteresis, "min_loss_scale": self.min_loss_scale, } - elif "bf16" not in self.config and self.precision_plugin.precision == "bf16": + elif "bf16" not in self.config and self.precision_plugin.precision == "bf16-mixed": rank_zero_info("Enabling DeepSpeed BF16.") self.config["bf16"] = {"enabled": True} diff --git a/src/lightning/pytorch/strategies/fsdp.py b/src/lightning/pytorch/strategies/fsdp.py index 0ac1709ad3680..f58dfc1db90f8 100644 --- a/src/lightning/pytorch/strategies/fsdp.py +++ b/src/lightning/pytorch/strategies/fsdp.py @@ -99,8 +99,8 @@ class FSDPStrategy(ParallelStrategy): algorithms to help backward communication and computation overlapping. The pros and cons of each algorithm is explained in the class ``BackwardPrefetch``. mixed_precision: - Mixed Precision config. By default, Lightning will enable FP16 if ``precision=16`` - or BF16 if ``precision=bf16`` unless a config is passed in. + Mixed Precision config. By default, Lightning will enable FP16 if ``precision="16-mixed"`` + or BF16 if ``precision="bf16-mixed"`` unless a config is passed in. This is only available in PyTorch 1.12 and later. activation_checkpointing: A single layer or a list of layer classes for which you want to enable activation checkpointing. This is typically your transformer block (including attention + feed-forward). diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py index d803ec58c25e5..4185efe87407c 100644 --- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py +++ b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py @@ -71,13 +71,11 @@ 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__) _LITERAL_WARN = Literal["warn"] -_PRECISION_INPUT_INT = Literal[64, 32, 16] -_PRECISION_INPUT_STR = Literal["64", "32", "16", "bf16"] -_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR] class AcceleratorConnector: @@ -88,7 +86,7 @@ def __init__( accelerator: Optional[Union[str, Accelerator]] = None, strategy: Optional[Union[str, Strategy]] = None, plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]] = None, - precision: _PRECISION_INPUT = 32, + precision: _PRECISION_INPUT = "32-true", sync_batchnorm: bool = False, benchmark: Optional[bool] = None, replace_sampler_ddp: bool = True, @@ -136,7 +134,7 @@ def __init__( # Set each valid flag to `self._x_flag` after validation self._strategy_flag: Optional[Union[Strategy, str]] = None self._accelerator_flag: Optional[Union[Accelerator, str]] = None - self._precision_flag: _PRECISION_INPUT_STR = "32" + self._precision_flag: _PRECISION_INPUT_STR = "32-true" self._precision_plugin_flag: Optional[PrecisionPlugin] = None self._cluster_environment_flag: Optional[Union[ClusterEnvironment, str]] = None self._parallel_devices: List[Union[int, torch.device, str]] = [] @@ -243,12 +241,23 @@ def _check_config_and_set_final_flags( self._accelerator_flag = accelerator - supported_precision = get_args(_PRECISION_INPUT_STR) + get_args(_PRECISION_INPUT_INT) + 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}" ) - self._precision_flag = cast(_PRECISION_INPUT_STR, str(precision)) + + 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'): + 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!" + ) + precision = _PRECISION_INPUT_STR_LEGACY_CONVERSION[precision] + + self._precision_flag = cast(_PRECISION_INPUT_STR, precision) if plugins: plugins_flags_types: Dict[str, int] = Counter() @@ -518,13 +527,13 @@ def _check_and_init_precision(self) -> PrecisionPlugin: if isinstance(self.accelerator, HPUAccelerator): return HPUPrecisionPlugin(self._precision_flag) # type: ignore if isinstance(self.accelerator, TPUAccelerator): - if self._precision_flag == "32": + if self._precision_flag == "32-true": return TPUPrecisionPlugin() - elif self._precision_flag in ("16", "bf16"): - if self._precision_flag == "16": + elif self._precision_flag in ("16-mixed", "bf16-mixed"): + if self._precision_flag == "16-mixed": rank_zero_warn( - "You passed `Trainer(accelerator='tpu', precision=16)` but AMP" - " is not supported with TPUs. Using `precision='bf16'` instead." + "You passed `Trainer(accelerator='tpu', precision='16-mixed')` but AMP with fp16" + " is not supported on TPUs. Using `precision='bf16-mixed'` instead." ) return TPUBf16PrecisionPlugin() @@ -537,21 +546,21 @@ def _check_and_init_precision(self) -> PrecisionPlugin: if isinstance(self.strategy, DeepSpeedStrategy): return DeepSpeedPrecisionPlugin(self._precision_flag) - if self._precision_flag == "32": + if self._precision_flag == "32-true": return PrecisionPlugin() - if self._precision_flag == "64": + if self._precision_flag == "64-true": return DoublePrecisionPlugin() - if self._precision_flag == "16" and self._accelerator_flag == "cpu": + if self._precision_flag == "16-mixed" and self._accelerator_flag == "cpu": rank_zero_warn( - "You passed `Trainer(accelerator='cpu', precision=16)` but AMP is not supported on CPU." - " Using `precision='bf16'` instead." + "You passed `Trainer(accelerator='cpu', precision='16-mixed')` but AMP with fp16 is not supported on " + "CPU. Using `precision='bf16-mixed'` instead." ) - self._precision_flag = "bf16" + self._precision_flag = "bf16-mixed" - if self._precision_flag in ("16", "bf16"): + if self._precision_flag in ("16-mixed", "bf16-mixed"): rank_zero_info( - f"Using {'16bit' if self._precision_flag == 16 else 'bfloat16'} Automatic Mixed Precision (AMP)" + f"Using {'16bit' if self._precision_flag == '16-mixed' else 'bfloat16'} Automatic Mixed Precision (AMP)" ) device = "cpu" if self._accelerator_flag == "cpu" else "cuda" @@ -564,9 +573,9 @@ def _check_and_init_precision(self) -> PrecisionPlugin: def _validate_precision_choice(self) -> None: """Validate the combination of choices for precision, AMP type, and accelerator.""" if isinstance(self.accelerator, TPUAccelerator): - if self._precision_flag == "64": + if self._precision_flag == "64-true": raise MisconfigurationException( - "`Trainer(accelerator='tpu', precision=64)` is not implemented." + "`Trainer(accelerator='tpu', precision='64-true')` is not implemented." " Please, open an issue in `https://github.com/Lightning-AI/lightning/issues`" " requesting this feature." ) @@ -578,7 +587,7 @@ def _validate_precision_choice(self) -> None: f" found: {self._precision_plugin_flag}." ) if isinstance(self.accelerator, HPUAccelerator): - if self._precision_flag not in ("16", "bf16", "32"): + if self._precision_flag not in ("16-mixed", "bf16-mixed", "32-true"): raise MisconfigurationException( f"`Trainer(accelerator='hpu', precision={self._precision_flag!r})` is not supported." ) diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index 2926982ecc94c..cf5f33c0d89d3 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/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, + precision: _PRECISION_INPUT = '32-true', enable_model_summary: bool = True, num_sanity_val_steps: int = 2, profiler: Optional[Union[Profiler, str]] = None, @@ -226,9 +226,10 @@ def __init__( plugins: Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins. Default: ``None``. - precision: Double precision (64), full precision (32), half precision (16) or bfloat16 precision (bf16). + precision: Double precision (64, '64' or '64-true'), full precision (32, '32' or '32-true'), + 16bit mixed precision (16, '16', '16-mixed') or bfloat16 mixed precision ('bf16', 'bf16-mixed'). Can be used on CPU, GPU, TPUs, HPUs or IPUs. - Default: ``32``. + Default: ``'32-true'``. max_epochs: Stop training once this number of epochs is reached. Disabled by default (None). If both max_epochs and max_steps are not specified, defaults to ``max_epochs = 1000``. From b4e2e0181ea35dc8c4224360ae0d913da0dc6cf2 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Thu, 16 Feb 2023 15:06:24 +0100 Subject: [PATCH 12/26] update tests --- tests/tests_pytorch/accelerators/test_hpu.py | 2 +- tests/tests_pytorch/accelerators/test_ipu.py | 16 +++--- .../checkpointing/test_legacy_checkpoints.py | 2 +- .../helpers/deterministic_model.py | 2 +- tests/tests_pytorch/models/test_amp.py | 18 +++--- .../tests_pytorch/models/test_ddp_fork_amp.py | 2 +- tests/tests_pytorch/models/test_hooks.py | 6 +- tests/tests_pytorch/models/test_tpu.py | 6 +- .../plugins/precision/hpu/test_hpu.py | 16 +++--- .../plugins/precision/test_amp.py | 4 +- .../plugins/precision/test_amp_integration.py | 2 +- .../precision/test_deepspeed_precision.py | 2 +- .../tests_pytorch/plugins/test_amp_plugins.py | 10 ++-- .../plugins/test_double_plugin.py | 4 +- tests/tests_pytorch/strategies/test_ddp.py | 2 +- .../strategies/test_deepspeed_strategy.py | 56 +++++++++---------- tests/tests_pytorch/strategies/test_fsdp.py | 12 ++-- .../tests_pytorch/strategies/test_registry.py | 2 +- .../connectors/test_accelerator_connector.py | 18 +++--- .../optimization/test_manual_optimization.py | 16 +++--- tests/tests_pytorch/trainer/test_trainer.py | 8 +-- .../tuner/test_scale_batch_size.py | 2 +- .../test_deepspeed_collate_checkpoint.py | 2 +- .../utilities/test_deepspeed_model_summary.py | 2 +- .../utilities/test_torchdistx.py | 2 +- 25 files changed, 107 insertions(+), 107 deletions(-) diff --git a/tests/tests_pytorch/accelerators/test_hpu.py b/tests/tests_pytorch/accelerators/test_hpu.py index 6307a78b1c815..fc08b1ee069fa 100644 --- a/tests/tests_pytorch/accelerators/test_hpu.py +++ b/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, + precision='16-mixed', ) trainer.fit(model) trainer.validate(model) diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index 01978eaeb9c8e..ab4d44f579491 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/tests/tests_pytorch/accelerators/test_ipu.py @@ -178,15 +178,15 @@ def test_optimization(tmpdir): def test_half_precision(tmpdir): class TestCallback(Callback): def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None: - assert trainer.precision == "16" + assert trainer.precision == "16-mixed" raise SystemExit model = IPUModel() trainer = Trainer( - default_root_dir=tmpdir, fast_dev_run=True, accelerator="ipu", devices=1, precision=16, 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" + assert trainer.strategy.precision_plugin.precision == "16-mixed" with pytest.raises(SystemExit): trainer.fit(model) @@ -195,7 +195,7 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> Non def test_pure_half_precision(tmpdir): class TestCallback(Callback): def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: - assert trainer.strategy.precision_plugin.precision == "16" + assert trainer.strategy.precision_plugin.precision == "16-mixed" for param in trainer.strategy.model.parameters(): assert param.dtype == torch.float16 raise SystemExit @@ -203,12 +203,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, callbacks=TestCallback() + default_root_dir=tmpdir, fast_dev_run=True, accelerator="ipu", devices=1, precision='16-mixed', callbacks=TestCallback() ) assert isinstance(trainer.strategy, IPUStrategy) assert isinstance(trainer.strategy.precision_plugin, IPUPrecisionPlugin) - assert trainer.strategy.precision_plugin.precision == "16" + assert trainer.strategy.precision_plugin.precision == "16-mixed" changed_dtypes = [torch.float, torch.float64] data = [torch.zeros((1), dtype=dtype) for dtype in changed_dtypes] @@ -534,8 +534,8 @@ def configure_optimizers(self): def test_precision_plugin(): """Ensure precision plugin value is set correctly.""" - plugin = IPUPrecisionPlugin(precision=16) - assert plugin.precision == "16" + plugin = IPUPrecisionPlugin(precision='16-mixed') + assert plugin.precision == "16-mixed" @RunIf(ipu=True) diff --git a/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py b/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py index 829c498e1e7c5..4ede1b8baacd6 100644 --- a/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py +++ b/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 if torch.cuda.is_available() else 32), + precision=('16-mixed' if torch.cuda.is_available() else '32-true'), callbacks=[stop], max_epochs=21, accumulate_grad_batches=2, diff --git a/tests/tests_pytorch/helpers/deterministic_model.py b/tests/tests_pytorch/helpers/deterministic_model.py index b5a4b588881c2..158406b4b7435 100644 --- a/tests/tests_pytorch/helpers/deterministic_model.py +++ b/tests/tests_pytorch/helpers/deterministic_model.py @@ -98,7 +98,7 @@ def configure_optimizers__lr_on_plateau_step(self): def backward(self, loss, *args, **kwargs): if self.assert_backward: - if self.trainer.precision == "16": + if self.trainer.precision == "16-mixed": assert loss > 171 * 1000 else: assert loss == 171.0 diff --git a/tests/tests_pytorch/models/test_amp.py b/tests/tests_pytorch/models/test_amp.py index d7c6922362141..0bcffcd49c28e 100644 --- a/tests/tests_pytorch/models/test_amp.py +++ b/tests/tests_pytorch/models/test_amp.py @@ -29,7 +29,7 @@ class AMPTestModel(BoringModel): def step(self, batch): self._assert_autocast_enabled() output = self(batch) - is_bfloat16 = self.trainer.precision_plugin.precision == "bf16" + is_bfloat16 = self.trainer.precision_plugin.precision == "bf16-mixed" assert output.dtype == torch.float16 if not is_bfloat16 else torch.bfloat16 loss = self.loss(output) return loss @@ -37,7 +37,7 @@ def step(self, batch): def predict_step(self, batch, batch_idx, dataloader_idx=0): self._assert_autocast_enabled() output = self(batch) - is_bfloat16 = self.trainer.precision_plugin.precision == "bf16" + is_bfloat16 = self.trainer.precision_plugin.precision == "bf16-mixed" assert output.dtype == torch.float16 if not is_bfloat16 else torch.bfloat16 return output @@ -52,10 +52,10 @@ def _assert_autocast_enabled(self): @pytest.mark.parametrize( ("strategy", "precision", "devices"), ( - ("single_device", 16, 1), - ("single_device", "bf16", 1), - ("ddp_spawn", 16, 2), - ("ddp_spawn", "bf16", 2), + ("single_device", '16-mixed', 1), + ("single_device", "bf16-mixed'", 1), + ("ddp_spawn", '16-mixed', 2), + ("ddp_spawn", "bf16-mixed", 2), ), ) def test_amp_cpus(tmpdir, strategy, precision, devices): @@ -83,7 +83,7 @@ def test_amp_cpus(tmpdir, strategy, precision, devices): @pytest.mark.parametrize("strategy", [None, "ddp_spawn"]) -@pytest.mark.parametrize("precision", [16, pytest.param("bf16", 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, + 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, + precision='16-mixed', limit_train_batches=4, limit_val_batches=0, gradient_clip_val=clip_val, diff --git a/tests/tests_pytorch/models/test_ddp_fork_amp.py b/tests/tests_pytorch/models/test_ddp_fork_amp.py index ae873ccad6eb0..7fba705e507bf 100644 --- a/tests/tests_pytorch/models/test_ddp_fork_amp.py +++ b/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, 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 diff --git a/tests/tests_pytorch/models/test_hooks.py b/tests/tests_pytorch/models/test_hooks.py index 44bbedc3a819d..41e715db3fcf7 100644 --- a/tests/tests_pytorch/models/test_hooks.py +++ b/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), 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, 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 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 = [ diff --git a/tests/tests_pytorch/models/test_tpu.py b/tests/tests_pytorch/models/test_tpu.py index 790f100fe3f58..ce6db19248f3d 100644 --- a/tests/tests_pytorch/models/test_tpu.py +++ b/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, + 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, + 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, + precision='16-mixed', enable_progress_bar=False, max_epochs=1, accelerator="tpu", diff --git a/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py b/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py index 718ef030eb507..5ddb295aa548b 100644 --- a/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py +++ b/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py @@ -34,15 +34,15 @@ def hmp_params(request): @RunIf(hpu=True) def test_precision_plugin(hmp_params): - plugin = HPUPrecisionPlugin(precision="bf16", **hmp_params) - assert plugin.precision == "bf16" + plugin = HPUPrecisionPlugin(precision="bf16-mixed", **hmp_params) + assert plugin.precision == "bf16-mixed" @RunIf(hpu=True) def test_mixed_precision(tmpdir, hmp_params: dict): class TestCallback(Callback): def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None: - assert trainer.precision == "bf16" + assert trainer.precision == "bf16-mixed" raise SystemExit model = BoringModel() @@ -51,12 +51,12 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> Non fast_dev_run=True, accelerator="hpu", devices=1, - plugins=[HPUPrecisionPlugin(precision="bf16", **hmp_params)], + plugins=[HPUPrecisionPlugin(precision="bf16-mixed", **hmp_params)], callbacks=TestCallback(), ) assert isinstance(trainer.strategy, SingleHPUStrategy) assert isinstance(trainer.strategy.precision_plugin, HPUPrecisionPlugin) - assert trainer.strategy.precision_plugin.precision == "bf16" + assert trainer.strategy.precision_plugin.precision == "bf16-mixed" with pytest.raises(SystemExit): trainer.fit(model) @@ -65,7 +65,7 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> Non def test_pure_half_precision(tmpdir, hmp_params: dict): class TestCallback(Callback): def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: - assert trainer.precision == "16" + assert trainer.precision == "16-mixed" for param in trainer.strategy.model.parameters(): assert param.dtype == torch.float16 raise SystemExit @@ -77,13 +77,13 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: fast_dev_run=True, accelerator="hpu", devices=1, - plugins=[HPUPrecisionPlugin(precision=16, **hmp_params)], + plugins=[HPUPrecisionPlugin(precision='16-mixed', **hmp_params)], callbacks=TestCallback(), ) assert isinstance(trainer.strategy, SingleHPUStrategy) assert isinstance(trainer.strategy.precision_plugin, HPUPrecisionPlugin) - assert trainer.strategy.precision_plugin.precision == "16" + assert trainer.strategy.precision_plugin.precision == "16-mixed" with pytest.raises(RuntimeError, match=r"float16/half is not supported on Gaudi."): trainer.fit(model) diff --git a/tests/tests_pytorch/plugins/precision/test_amp.py b/tests/tests_pytorch/plugins/precision/test_amp.py index 189386cb90502..fed890e471175 100644 --- a/tests/tests_pytorch/plugins/precision/test_amp.py +++ b/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, 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, 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) diff --git a/tests/tests_pytorch/plugins/precision/test_amp_integration.py b/tests/tests_pytorch/plugins/precision/test_amp_integration.py index 0d7fb3f8e2bc0..895a67c6fecc2 100644 --- a/tests/tests_pytorch/plugins/precision/test_amp_integration.py +++ b/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, + precision='16-mixed', max_steps=5, logger=False, enable_checkpointing=False, diff --git a/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py b/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py index 8420c5c793aec..2b31972cc88d4 100644 --- a/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py +++ b/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) + DeepSpeedPrecisionPlugin(precision='64-true') diff --git a/tests/tests_pytorch/plugins/test_amp_plugins.py b/tests/tests_pytorch/plugins/test_amp_plugins.py index e542c01967cf7..c20656f8c3a96 100644 --- a/tests/tests_pytorch/plugins/test_amp_plugins.py +++ b/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, "cpu") + plugin = plugin_cls('16-mixed', "cpu") trainer = Trainer( fast_dev_run=True, - precision=16, + 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, + precision='16-mixed', # use a tiny value to make sure it works gradient_clip_val=1e-3, gradient_clip_algorithm="value", @@ -179,14 +179,14 @@ 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) + trainer = Trainer(default_root_dir=tmpdir, accelerator="gpu", devices=1, fast_dev_run=1, precision='16-mixed') model = CustomBoringModel() trainer.fit(model) def test_cpu_amp_precision_context_manager(tmpdir): """Test to ensure that the context manager correctly is set to CPU + bfloat16.""" - plugin = MixedPrecisionPlugin("bf16", "cpu") + plugin = MixedPrecisionPlugin("bf16-mixed", "cpu") assert plugin.device == "cpu" assert plugin.scaler is None context_manager = plugin.autocast_context_manager() diff --git a/tests/tests_pytorch/plugins/test_double_plugin.py b/tests/tests_pytorch/plugins/test_double_plugin.py index 9c93f09cad221..b03a660b962b0 100644 --- a/tests/tests_pytorch/plugins/test_double_plugin.py +++ b/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, 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, + precision='64-true', log_every_n_steps=1, ) trainer.fit(model) diff --git a/tests/tests_pytorch/strategies/test_ddp.py b/tests/tests_pytorch/strategies/test_ddp.py index 248e42bd7e69d..833a53f8e799d 100644 --- a/tests/tests_pytorch/strategies/test_ddp.py +++ b/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, 32)) +@pytest.mark.parametrize("precision", ('16-mixed', '32-true')) def test_ddp_wrapper(tmpdir, precision): """Test parameters to ignore are carried over for DDP.""" diff --git a/tests/tests_pytorch/strategies/test_deepspeed_strategy.py b/tests/tests_pytorch/strategies/test_deepspeed_strategy.py index e6eeff8c36f5f..c5b67c0bcb6bb 100644 --- a/tests/tests_pytorch/strategies/test_deepspeed_strategy.py +++ b/tests/tests_pytorch/strategies/test_deepspeed_strategy.py @@ -139,12 +139,12 @@ def test_deepspeed_precision_choice(cuda_count_1, tmpdir): default_root_dir=tmpdir, accelerator="gpu", strategy="deepspeed", - precision=16, + precision='16-mixed', ) assert isinstance(trainer.strategy, DeepSpeedStrategy) assert isinstance(trainer.strategy.precision_plugin, DeepSpeedPrecisionPlugin) - assert trainer.strategy.precision_plugin.precision == "16" + assert trainer.strategy.precision_plugin.precision == "16-mixed" @RunIf(deepspeed=True) @@ -189,7 +189,7 @@ def backward(self, loss: Tensor, *args, **kwargs) -> None: strategy=DeepSpeedStrategy(), accelerator="gpu", devices=1, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -264,7 +264,7 @@ def configure_optimizers(self): accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', callbacks=[TestCB(), lr_monitor], logger=CSVLogger(tmpdir), enable_progress_bar=False, @@ -303,7 +303,7 @@ def on_train_start(self, trainer, pl_module) -> None: limit_val_batches=4, limit_test_batches=4, max_epochs=2, - precision=16, + precision='16-mixed', callbacks=[TestCB(), lr_monitor], logger=CSVLogger(tmpdir), enable_progress_bar=False, @@ -337,7 +337,7 @@ def on_train_start(self, trainer, pl_module) -> None: trainer = Trainer( default_root_dir=tmpdir, strategy=ds, - precision=16, + precision='16-mixed', accelerator="gpu", devices=1, callbacks=[TestCB()], @@ -380,7 +380,7 @@ def test_deepspeed_custom_activation_checkpointing_params_forwarded(tmpdir): default_root_dir=tmpdir, fast_dev_run=1, strategy=ds, - precision=16, + precision='16-mixed', accelerator="gpu", devices=1, enable_progress_bar=False, @@ -413,7 +413,7 @@ def setup(self, trainer, pl_module, stage=None) -> None: enable_progress_bar=False, max_epochs=1, strategy=DeepSpeedStrategy(config=deepspeed_zero_config), - precision=16, + precision='16-mixed', accelerator="gpu", devices=1, callbacks=[TestCallback()], @@ -433,7 +433,7 @@ def test_deepspeed_multigpu(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -476,7 +476,7 @@ def test_deepspeed_stage_3_save_warning(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -508,7 +508,7 @@ def test_deepspeed_multigpu_single_file(tmpdir): accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -524,7 +524,7 @@ def test_deepspeed_multigpu_single_file(tmpdir): accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -626,7 +626,7 @@ def test_deepspeed_multigpu_stage_3(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -646,7 +646,7 @@ def test_deepspeed_multigpu_stage_3_manual_optimization(tmpdir, deepspeed_config accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -672,7 +672,7 @@ def test_deepspeed_multigpu_stage_3_checkpointing(tmpdir, automatic_optimization strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=2, - precision=16, + precision='16-mixed', accumulate_grad_batches=accumulate_grad_batches, callbacks=[ck], enable_progress_bar=False, @@ -693,7 +693,7 @@ def test_deepspeed_multigpu_stage_3_checkpointing(tmpdir, automatic_optimization accelerator="gpu", devices=2, strategy=DeepSpeedStrategy(stage=3), - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -722,7 +722,7 @@ def test_deepspeed_multigpu_stage_3_warns_resume_training(tmpdir): strategy=DeepSpeedStrategy(stage=3, load_full_weights=True), accelerator="gpu", devices=1, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -751,7 +751,7 @@ def test_deepspeed_multigpu_stage_3_resume_training(tmpdir): strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=1, - precision=16, + precision='16-mixed', callbacks=[ck], enable_progress_bar=False, enable_model_summary=False, @@ -792,7 +792,7 @@ def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> max_epochs=2, limit_train_batches=1, limit_val_batches=0, - precision=16, + precision='16-mixed', callbacks=TestCallback(), enable_progress_bar=False, enable_model_summary=False, @@ -828,7 +828,7 @@ def on_train_batch_start(self, trainer, pl_module: LightningModule, batch: Any, devices=2, limit_train_batches=5, limit_val_batches=2, - precision=16, + precision='16-mixed', accumulate_grad_batches=2, callbacks=[verification_callback], enable_progress_bar=False, @@ -849,7 +849,7 @@ def test_deepspeed_multigpu_test(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -885,7 +885,7 @@ def on_train_epoch_start(self) -> None: accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -912,7 +912,7 @@ def on_train_epoch_start(self) -> None: accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -976,7 +976,7 @@ def test_deepspeed_multigpu_no_schedulers(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -998,7 +998,7 @@ def training_step(self, batch, batch_idx): accelerator="gpu", devices=1, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) @@ -1212,7 +1212,7 @@ def test_deepspeed_with_bfloat16_precision(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision="bf16", + precision="bf16-mixed", num_sanity_val_steps=0, enable_progress_bar=False, enable_model_summary=False, @@ -1220,7 +1220,7 @@ def test_deepspeed_with_bfloat16_precision(tmpdir): trainer.fit(model) assert isinstance(trainer.strategy.precision_plugin, DeepSpeedPrecisionPlugin) - assert trainer.strategy.precision_plugin.precision == "bf16" + assert trainer.strategy.precision_plugin.precision == "bf16-mixed" assert trainer.strategy.config["zero_optimization"]["stage"] == 3 assert trainer.strategy.config["bf16"]["enabled"] assert model.layer.weight.dtype == torch.bfloat16 @@ -1271,7 +1271,7 @@ def transfer_batch_to_device(self, batch, *args, **kwargs): return super().transfer_batch_to_device(batch, *args, **kwargs) model = CustomBoringModel() - trainer = Trainer(strategy="deepspeed", devices=1, accelerator="cuda", precision=16) + trainer = Trainer(strategy="deepspeed", devices=1, accelerator="cuda", precision='16-mixed') trainer.strategy.connect(model) batch = torch.zeros((1), dtype=torch.float32) batch = trainer.strategy.batch_to_device(batch) diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 42425f581765f..9cef1fe13f6f6 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -64,7 +64,7 @@ def on_predict_batch_end(self, *_) -> None: def _assert_layer_fsdp_instance(self) -> None: assert isinstance(self.layer, FullyShardedDataParallel) assert isinstance(self.trainer.strategy.precision_plugin, FSDPMixedPrecisionPlugin) - precision = torch.float16 if self.trainer.precision == "16" else torch.bfloat16 + precision = torch.float16 if self.trainer.precision == "16-mixed" else torch.bfloat16 assert self.layer.mixed_precision.param_dtype == precision assert self.layer.mixed_precision.reduce_dtype == precision assert self.layer.mixed_precision.buffer_dtype == precision @@ -100,7 +100,7 @@ def _assert_layer_fsdp_instance(self) -> None: assert isinstance(self.layer, torch.nn.Sequential) assert isinstance(self.trainer.strategy.precision_plugin, FSDPMixedPrecisionPlugin) - precision = torch.float16 if self.trainer.precision == "16" else torch.bfloat16 + precision = torch.float16 if self.trainer.precision == "16-mixed" else torch.bfloat16 for layer_num in [0, 2]: assert isinstance(self.layer[layer_num], FullyShardedDataParallel) assert self.layer[layer_num].mixed_precision.param_dtype == precision @@ -164,7 +164,7 @@ def test_invalid_on_cpu(tmpdir): @RunIf(min_torch="1.12", min_cuda_gpus=1) -@pytest.mark.parametrize("precision, expected", [(16, torch.float16), ("bf16", torch.bfloat16)]) +@pytest.mark.parametrize("precision, expected", [('16-mixed', torch.float16), ("bf16-mixed", torch.bfloat16)]) def test_precision_plugin_config(precision, expected): plugin = FSDPMixedPrecisionPlugin(precision=precision, device="cuda") config = plugin.mixed_precision_config @@ -191,7 +191,7 @@ def test_fsdp_strategy_sync_batchnorm(tmpdir): accelerator="gpu", devices=2, strategy="fsdp", - precision=16, + precision='16-mixed', max_epochs=1, sync_batchnorm=True, ) @@ -199,7 +199,7 @@ def test_fsdp_strategy_sync_batchnorm(tmpdir): @RunIf(min_cuda_gpus=1, skip_windows=True, standalone=True, min_torch="1.12") -@pytest.mark.parametrize("precision", (16, pytest.param("bf16", marks=RunIf(bf16_cuda=True)))) +@pytest.mark.parametrize("precision", ('16-mixed', pytest.param("bf16-mixed", marks=RunIf(bf16_cuda=True)))) def test_fsdp_strategy_checkpoint(tmpdir, precision): """Test to ensure that checkpoint is saved correctly when using a single GPU, and all stages can be run.""" model = TestFSDPModel() @@ -230,7 +230,7 @@ def test_fsdp_checkpoint_multi_gpus(tmpdir, model, strategy): accelerator="gpu", devices=2, strategy=strategy, - precision=16, + precision='16-mixed', max_epochs=1, limit_train_batches=2, limit_val_batches=2, diff --git a/tests/tests_pytorch/strategies/test_registry.py b/tests/tests_pytorch/strategies/test_registry.py index 270fb028fad7f..9fa2e26f95008 100644 --- a/tests/tests_pytorch/strategies/test_registry.py +++ b/tests/tests_pytorch/strategies/test_registry.py @@ -48,7 +48,7 @@ def test_strategy_registry_with_deepspeed_strategies(strategy_name, init_params) @pytest.mark.parametrize("strategy", ["deepspeed", "deepspeed_stage_2_offload", "deepspeed_stage_3"]) def test_deepspeed_strategy_registry_with_trainer(tmpdir, strategy): - trainer = Trainer(default_root_dir=tmpdir, strategy=strategy, precision=16) + trainer = Trainer(default_root_dir=tmpdir, strategy=strategy, precision='16-mixed') assert isinstance(trainer.strategy, DeepSpeedStrategy) diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 4621648e7c201..35f89bc1d882b 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -596,14 +596,14 @@ def test_check_fsdp_strategy_and_fallback(): def test_unsupported_tpu_choice(tpu_available): - with pytest.raises(MisconfigurationException, match=r"accelerator='tpu', precision=64\)` is not implemented"): - Trainer(accelerator="tpu", precision=64) + with pytest.raises(MisconfigurationException, match=r"accelerator='tpu', precision='64-true'\)` is not implemented"): + Trainer(accelerator="tpu", precision='64-true') # if user didn't set strategy, AcceleratorConnector will choose the TPUSingleStrategy or TPUSpawnStrategy with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"), pytest.warns( - UserWarning, match=r"accelerator='tpu', precision=16\)` but AMP is not supported" + UserWarning, match=r"accelerator='tpu', precision=16-mixed\)` but AMP with fp16 is not supported" ): - Trainer(accelerator="tpu", precision=16, strategy="ddp") + Trainer(accelerator="tpu", precision='16-mixed', strategy="ddp") @mock.patch("lightning.pytorch.accelerators.ipu.IPUAccelerator.is_available", return_value=True) @@ -613,10 +613,10 @@ def test_unsupported_ipu_choice(mock_ipu_acc_avail, monkeypatch): monkeypatch.setattr(ipu_, "_IPU_AVAILABLE", True) monkeypatch.setattr(ipu, "_IPU_AVAILABLE", True) - with pytest.raises(ValueError, match=r"accelerator='ipu', precision='bf16'\)` is not supported"): - Trainer(accelerator="ipu", precision="bf16") - with pytest.raises(ValueError, match=r"accelerator='ipu', precision='64'\)` is not supported"): - Trainer(accelerator="ipu", precision=64) + with pytest.raises(ValueError, match=r"accelerator='ipu', precision='bf16-mixed'\)` is not supported"): + Trainer(accelerator="ipu", precision="bf16-mixed") + with pytest.raises(ValueError, match=r"accelerator='ipu', precision='64-true'\)` is not supported"): + Trainer(accelerator="ipu", precision="64-true") @mock.patch("lightning.pytorch.accelerators.tpu._XLA_AVAILABLE", return_value=False) @@ -846,5 +846,5 @@ def test_colossalai_external_strategy(monkeypatch): from lightning_colossalai import ColossalAIStrategy - trainer = Trainer(strategy="colossalai", precision=16) + trainer = Trainer(strategy="colossalai", precision='16-mixed') assert isinstance(trainer.strategy, ColossalAIStrategy) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index 8a5bedf8efabe..d181dfa76067e 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -72,7 +72,7 @@ def configure_optimizers(self): @pytest.mark.parametrize( - "kwargs", [{}, pytest.param({"accelerator": "gpu", "devices": 1, "precision": 16}, marks=RunIf(min_cuda_gpus=1))] + "kwargs", [{}, pytest.param({"accelerator": "gpu", "devices": 1, "precision": '16-mixed'}, marks=RunIf(min_cuda_gpus=1))] ) def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): model = ManualOptModel() @@ -87,7 +87,7 @@ def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): **kwargs, ) - if kwargs.get("precision") == 16: + if kwargs.get("precision") == '16-mixed': # mock the scaler instead of the optimizer step because it can be skipped with NaNs scaler_step_patch = mock.patch.object( trainer.precision_plugin.scaler, "step", wraps=trainer.precision_plugin.scaler.step @@ -99,7 +99,7 @@ def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): assert bwd_mock.call_count == limit_train_batches * 3 assert trainer.global_step == limit_train_batches * 2 - if kwargs.get("precision") == 16: + if kwargs.get("precision") == '16-mixed': scaler_step_patch.stop() assert scaler_step.call_count == len(model.optimizers()) * limit_train_batches @@ -141,7 +141,7 @@ def test_multiple_optimizers_manual_amp(tmpdir, accelerator): max_epochs=1, log_every_n_steps=1, enable_model_summary=False, - precision=16, + precision='16-mixed', accelerator=accelerator, devices=1, ) @@ -224,7 +224,7 @@ def test_manual_optimization_and_return_tensor(tmpdir): limit_train_batches=10, limit_test_batches=0, limit_val_batches=0, - precision=16, + precision='16-mixed', strategy="ddp_spawn", accelerator="gpu", devices=2, @@ -309,7 +309,7 @@ def on_train_epoch_end(self, *_, **__): limit_train_batches=20, limit_test_batches=0, limit_val_batches=0, - precision=16, + precision='16-mixed', accelerator="gpu", devices=1, ) @@ -383,7 +383,7 @@ def on_before_optimizer_step(self, optimizer, *_): max_epochs=1, log_every_n_steps=1, enable_model_summary=False, - precision=16, + precision='16-mixed', accelerator="gpu", devices=1, ) @@ -848,7 +848,7 @@ def test_lr_scheduler_step_not_called(tmpdir): @RunIf(min_cuda_gpus=1) -@pytest.mark.parametrize("precision", [16, 32]) +@pytest.mark.parametrize("precision", ['16-mixed', '32-true']) def test_multiple_optimizers_logging(precision, tmpdir): """Tests that metrics are properly being logged.""" diff --git a/tests/tests_pytorch/trainer/test_trainer.py b/tests/tests_pytorch/trainer/test_trainer.py index 8d7b044d24a12..d1c42fa3499d0 100644 --- a/tests/tests_pytorch/trainer/test_trainer.py +++ b/tests/tests_pytorch/trainer/test_trainer.py @@ -1019,7 +1019,7 @@ def on_exception(self, trainer, pl_module, exception): assert isinstance(handle_interrupt_callback.exception, MisconfigurationException) -@pytest.mark.parametrize("precision", [32, pytest.param(16, marks=RunIf(min_cuda_gpus=1))]) +@pytest.mark.parametrize("precision", ['32-true', pytest.param('16-mixed', marks=RunIf(min_cuda_gpus=1))]) @RunIf(sklearn=True) def test_gradient_clipping_by_norm(tmpdir, precision): """Test gradient clipping by norm.""" @@ -1048,7 +1048,7 @@ def configure_gradient_clipping(self, *args, **kwargs): assert model.assertion_called -@pytest.mark.parametrize("precision", [32, pytest.param(16, marks=RunIf(min_cuda_gpus=1))]) +@pytest.mark.parametrize("precision", ['32-true', pytest.param('16-mixed', marks=RunIf(min_cuda_gpus=1))]) def test_gradient_clipping_by_value(tmpdir, precision): """Test gradient clipping by value.""" trainer = Trainer( @@ -1444,7 +1444,7 @@ def test_spawn_predict_return_predictions(tmpdir): @pytest.mark.parametrize("return_predictions", [None, False, True]) -@pytest.mark.parametrize("precision", [32, 64]) +@pytest.mark.parametrize("precision", ['32-true', '64-true']) def test_predict_return_predictions_cpu(return_predictions, precision, tmpdir): """Test that `return_predictions=True`.""" seed_everything(42) @@ -1455,7 +1455,7 @@ def test_predict_return_predictions_cpu(return_predictions, precision, tmpdir): if return_predictions or return_predictions is None: assert len(preds) == 1 assert preds[0].shape == torch.Size([1, 2]) - assert preds[0].dtype == (torch.float64 if precision == 64 else torch.float32) + assert preds[0].dtype == (torch.float64 if precision == '64-true' else torch.float32) @pytest.mark.parametrize(["max_steps", "max_epochs", "global_step"], [(10, 5, 10), (20, None, 20)]) diff --git a/tests/tests_pytorch/tuner/test_scale_batch_size.py b/tests/tests_pytorch/tuner/test_scale_batch_size.py index 08b94a4763a8f..ea55b513d1ab4 100644 --- a/tests/tests_pytorch/tuner/test_scale_batch_size.py +++ b/tests/tests_pytorch/tuner/test_scale_batch_size.py @@ -254,7 +254,7 @@ def test_error_on_dataloader_passed_to_fit(tmpdir): def test_auto_scale_batch_size_with_amp(tmpdir): before_batch_size = 2 model = BatchSizeModel(batch_size=before_batch_size) - trainer = Trainer(default_root_dir=tmpdir, max_steps=1, accelerator="gpu", devices=1, precision=16) + trainer = Trainer(default_root_dir=tmpdir, max_steps=1, accelerator="gpu", devices=1, precision='16-mixed') tuner = Tuner(trainer) tuner.scale_batch_size(model) after_batch_size = model.batch_size diff --git a/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py b/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py index 5c0cb588ebe5b..314501ef0e8ee 100644 --- a/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py +++ b/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py @@ -32,7 +32,7 @@ def test_deepspeed_collate_checkpoint(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision=16, + precision='16-mixed', enable_progress_bar=False, enable_model_summary=False, ) diff --git a/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py b/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py index f06c101db8246..8d734ad74649f 100644 --- a/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py +++ b/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py @@ -45,7 +45,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") - accelerator="gpu", fast_dev_run=True, devices=2, - precision=16, + precision='16-mixed', enable_model_summary=True, callbacks=[TestCallback()], ) diff --git a/tests/tests_pytorch/utilities/test_torchdistx.py b/tests/tests_pytorch/utilities/test_torchdistx.py index 9fee068cee9ab..1491dc5106443 100644 --- a/tests/tests_pytorch/utilities/test_torchdistx.py +++ b/tests/tests_pytorch/utilities/test_torchdistx.py @@ -55,7 +55,7 @@ def test_deferred_init_with_lightning_module(): ( {"accelerator": "auto", "devices": 1}, pytest.param( - {"strategy": "deepspeed_stage_3", "accelerator": "gpu", "devices": 2, "precision": 16}, + {"strategy": "deepspeed_stage_3", "accelerator": "gpu", "devices": 2, "precision": '16-mixed'}, marks=RunIf(min_cuda_gpus=2, deepspeed=True), ), ), From b55382401bb83dff30aa9e7dce85f7b77b7b2d71 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Feb 2023 15:07:18 +0000 Subject: [PATCH 13/26] [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 | 2 +- examples/pl_hpu/mnist_sample.py | 2 +- .../pytorch/plugins/precision/amp.py | 5 +- .../pytorch/plugins/precision/deepspeed.py | 1 - .../pytorch/plugins/precision/hpu.py | 2 +- .../pytorch/plugins/precision/ipu.py | 1 + .../connectors/accelerator_connector.py | 14 ++++-- src/lightning/pytorch/trainer/trainer.py | 2 +- tests/tests_pytorch/accelerators/test_hpu.py | 2 +- tests/tests_pytorch/accelerators/test_ipu.py | 16 ++++-- .../checkpointing/test_legacy_checkpoints.py | 2 +- tests/tests_pytorch/models/test_amp.py | 10 ++-- .../tests_pytorch/models/test_ddp_fork_amp.py | 2 +- tests/tests_pytorch/models/test_hooks.py | 6 +-- tests/tests_pytorch/models/test_tpu.py | 6 +-- .../plugins/precision/hpu/test_hpu.py | 2 +- .../plugins/precision/test_amp.py | 4 +- .../plugins/precision/test_amp_integration.py | 2 +- .../precision/test_deepspeed_precision.py | 2 +- .../tests_pytorch/plugins/test_amp_plugins.py | 8 +-- .../plugins/test_double_plugin.py | 4 +- tests/tests_pytorch/strategies/test_ddp.py | 2 +- .../strategies/test_deepspeed_strategy.py | 50 +++++++++---------- tests/tests_pytorch/strategies/test_fsdp.py | 8 +-- .../tests_pytorch/strategies/test_registry.py | 2 +- .../connectors/test_accelerator_connector.py | 10 ++-- .../optimization/test_manual_optimization.py | 17 ++++--- tests/tests_pytorch/trainer/test_trainer.py | 8 +-- .../tuner/test_scale_batch_size.py | 2 +- .../test_deepspeed_collate_checkpoint.py | 2 +- .../utilities/test_deepspeed_model_summary.py | 2 +- .../utilities/test_torchdistx.py | 2 +- 32 files changed, 112 insertions(+), 88 deletions(-) diff --git a/examples/app_multi_node/train_fabric.py b/examples/app_multi_node/train_fabric.py index 4fd7bee094697..335e1e73db6e0 100644 --- a/examples/app_multi_node/train_fabric.py +++ b/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() diff --git a/examples/pl_hpu/mnist_sample.py b/examples/pl_hpu/mnist_sample.py index 09613e295f47a..0ed24ad75403b 100644 --- a/examples/pl_hpu/mnist_sample.py +++ b/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}, diff --git a/src/lightning/pytorch/plugins/precision/amp.py b/src/lightning/pytorch/plugins/precision/amp.py index bddcf2b480a75..1ac94415e33dc 100644 --- a/src/lightning/pytorch/plugins/precision/amp.py +++ b/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": diff --git a/src/lightning/pytorch/plugins/precision/deepspeed.py b/src/lightning/pytorch/plugins/precision/deepspeed.py index 99f8d06173fa8..627026214eaf4 100644 --- a/src/lightning/pytorch/plugins/precision/deepspeed.py +++ b/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. diff --git a/src/lightning/pytorch/plugins/precision/hpu.py b/src/lightning/pytorch/plugins/precision/hpu.py index 65863e9f77d22..47a145807bcff 100644 --- a/src/lightning/pytorch/plugins/precision/hpu.py +++ b/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 diff --git a/src/lightning/pytorch/plugins/precision/ipu.py b/src/lightning/pytorch/plugins/precision/ipu.py index 95737fb996d49..e414bc693163e 100644 --- a/src/lightning/pytorch/plugins/precision/ipu.py +++ b/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. diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py index 4185efe87407c..f98535050324a 100644 --- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py +++ b/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!" diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index cf5f33c0d89d3..0f033641fe07e 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/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, diff --git a/tests/tests_pytorch/accelerators/test_hpu.py b/tests/tests_pytorch/accelerators/test_hpu.py index fc08b1ee069fa..b8ba801e7ede4 100644 --- a/tests/tests_pytorch/accelerators/test_hpu.py +++ b/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) diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index ab4d44f579491..650b4c56ec1f3 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/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" diff --git a/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py b/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py index 4ede1b8baacd6..86dd5c6cfe9b7 100644 --- a/tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py +++ b/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, diff --git a/tests/tests_pytorch/models/test_amp.py b/tests/tests_pytorch/models/test_amp.py index 0bcffcd49c28e..372cb10db349a 100644 --- a/tests/tests_pytorch/models/test_amp.py +++ b/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, diff --git a/tests/tests_pytorch/models/test_ddp_fork_amp.py b/tests/tests_pytorch/models/test_ddp_fork_amp.py index 7fba705e507bf..13434dcab69bf 100644 --- a/tests/tests_pytorch/models/test_ddp_fork_amp.py +++ b/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 diff --git a/tests/tests_pytorch/models/test_hooks.py b/tests/tests_pytorch/models/test_hooks.py index 41e715db3fcf7..50f6f36a0811d 100644 --- a/tests/tests_pytorch/models/test_hooks.py +++ b/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 = [ diff --git a/tests/tests_pytorch/models/test_tpu.py b/tests/tests_pytorch/models/test_tpu.py index ce6db19248f3d..031fef8e72ce8 100644 --- a/tests/tests_pytorch/models/test_tpu.py +++ b/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", diff --git a/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py b/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py index 5ddb295aa548b..54599f58448c0 100644 --- a/tests/tests_pytorch/plugins/precision/hpu/test_hpu.py +++ b/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(), ) diff --git a/tests/tests_pytorch/plugins/precision/test_amp.py b/tests/tests_pytorch/plugins/precision/test_amp.py index fed890e471175..4c86f02986894 100644 --- a/tests/tests_pytorch/plugins/precision/test_amp.py +++ b/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) diff --git a/tests/tests_pytorch/plugins/precision/test_amp_integration.py b/tests/tests_pytorch/plugins/precision/test_amp_integration.py index 895a67c6fecc2..8a64169e9f1fe 100644 --- a/tests/tests_pytorch/plugins/precision/test_amp_integration.py +++ b/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, diff --git a/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py b/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py index 2b31972cc88d4..b0ef260309639 100644 --- a/tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py +++ b/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") diff --git a/tests/tests_pytorch/plugins/test_amp_plugins.py b/tests/tests_pytorch/plugins/test_amp_plugins.py index c20656f8c3a96..5cca3a93aa518 100644 --- a/tests/tests_pytorch/plugins/test_amp_plugins.py +++ b/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) diff --git a/tests/tests_pytorch/plugins/test_double_plugin.py b/tests/tests_pytorch/plugins/test_double_plugin.py index b03a660b962b0..8d801d6eaf7eb 100644 --- a/tests/tests_pytorch/plugins/test_double_plugin.py +++ b/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) diff --git a/tests/tests_pytorch/strategies/test_ddp.py b/tests/tests_pytorch/strategies/test_ddp.py index 833a53f8e799d..f6470764d9016 100644 --- a/tests/tests_pytorch/strategies/test_ddp.py +++ b/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.""" diff --git a/tests/tests_pytorch/strategies/test_deepspeed_strategy.py b/tests/tests_pytorch/strategies/test_deepspeed_strategy.py index c5b67c0bcb6bb..76f248bb5264e 100644 --- a/tests/tests_pytorch/strategies/test_deepspeed_strategy.py +++ b/tests/tests_pytorch/strategies/test_deepspeed_strategy.py @@ -139,7 +139,7 @@ def test_deepspeed_precision_choice(cuda_count_1, tmpdir): default_root_dir=tmpdir, accelerator="gpu", strategy="deepspeed", - precision='16-mixed', + precision="16-mixed", ) assert isinstance(trainer.strategy, DeepSpeedStrategy) @@ -189,7 +189,7 @@ def backward(self, loss: Tensor, *args, **kwargs) -> None: strategy=DeepSpeedStrategy(), accelerator="gpu", devices=1, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -264,7 +264,7 @@ def configure_optimizers(self): accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", callbacks=[TestCB(), lr_monitor], logger=CSVLogger(tmpdir), enable_progress_bar=False, @@ -303,7 +303,7 @@ def on_train_start(self, trainer, pl_module) -> None: limit_val_batches=4, limit_test_batches=4, max_epochs=2, - precision='16-mixed', + precision="16-mixed", callbacks=[TestCB(), lr_monitor], logger=CSVLogger(tmpdir), enable_progress_bar=False, @@ -337,7 +337,7 @@ def on_train_start(self, trainer, pl_module) -> None: trainer = Trainer( default_root_dir=tmpdir, strategy=ds, - precision='16-mixed', + precision="16-mixed", accelerator="gpu", devices=1, callbacks=[TestCB()], @@ -380,7 +380,7 @@ def test_deepspeed_custom_activation_checkpointing_params_forwarded(tmpdir): default_root_dir=tmpdir, fast_dev_run=1, strategy=ds, - precision='16-mixed', + precision="16-mixed", accelerator="gpu", devices=1, enable_progress_bar=False, @@ -413,7 +413,7 @@ def setup(self, trainer, pl_module, stage=None) -> None: enable_progress_bar=False, max_epochs=1, strategy=DeepSpeedStrategy(config=deepspeed_zero_config), - precision='16-mixed', + precision="16-mixed", accelerator="gpu", devices=1, callbacks=[TestCallback()], @@ -433,7 +433,7 @@ def test_deepspeed_multigpu(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -476,7 +476,7 @@ def test_deepspeed_stage_3_save_warning(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -508,7 +508,7 @@ def test_deepspeed_multigpu_single_file(tmpdir): accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -524,7 +524,7 @@ def test_deepspeed_multigpu_single_file(tmpdir): accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -626,7 +626,7 @@ def test_deepspeed_multigpu_stage_3(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -646,7 +646,7 @@ def test_deepspeed_multigpu_stage_3_manual_optimization(tmpdir, deepspeed_config accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -672,7 +672,7 @@ def test_deepspeed_multigpu_stage_3_checkpointing(tmpdir, automatic_optimization strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=2, - precision='16-mixed', + precision="16-mixed", accumulate_grad_batches=accumulate_grad_batches, callbacks=[ck], enable_progress_bar=False, @@ -693,7 +693,7 @@ def test_deepspeed_multigpu_stage_3_checkpointing(tmpdir, automatic_optimization accelerator="gpu", devices=2, strategy=DeepSpeedStrategy(stage=3), - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -722,7 +722,7 @@ def test_deepspeed_multigpu_stage_3_warns_resume_training(tmpdir): strategy=DeepSpeedStrategy(stage=3, load_full_weights=True), accelerator="gpu", devices=1, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -751,7 +751,7 @@ def test_deepspeed_multigpu_stage_3_resume_training(tmpdir): strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=1, - precision='16-mixed', + precision="16-mixed", callbacks=[ck], enable_progress_bar=False, enable_model_summary=False, @@ -792,7 +792,7 @@ def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> max_epochs=2, limit_train_batches=1, limit_val_batches=0, - precision='16-mixed', + precision="16-mixed", callbacks=TestCallback(), enable_progress_bar=False, enable_model_summary=False, @@ -828,7 +828,7 @@ def on_train_batch_start(self, trainer, pl_module: LightningModule, batch: Any, devices=2, limit_train_batches=5, limit_val_batches=2, - precision='16-mixed', + precision="16-mixed", accumulate_grad_batches=2, callbacks=[verification_callback], enable_progress_bar=False, @@ -849,7 +849,7 @@ def test_deepspeed_multigpu_test(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -885,7 +885,7 @@ def on_train_epoch_start(self) -> None: accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -912,7 +912,7 @@ def on_train_epoch_start(self) -> None: accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -976,7 +976,7 @@ def test_deepspeed_multigpu_no_schedulers(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -998,7 +998,7 @@ def training_step(self, batch, batch_idx): accelerator="gpu", devices=1, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) @@ -1271,7 +1271,7 @@ def transfer_batch_to_device(self, batch, *args, **kwargs): return super().transfer_batch_to_device(batch, *args, **kwargs) model = CustomBoringModel() - trainer = Trainer(strategy="deepspeed", devices=1, accelerator="cuda", precision='16-mixed') + trainer = Trainer(strategy="deepspeed", devices=1, accelerator="cuda", precision="16-mixed") trainer.strategy.connect(model) batch = torch.zeros((1), dtype=torch.float32) batch = trainer.strategy.batch_to_device(batch) diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 9cef1fe13f6f6..05aec225204a4 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -164,7 +164,7 @@ def test_invalid_on_cpu(tmpdir): @RunIf(min_torch="1.12", min_cuda_gpus=1) -@pytest.mark.parametrize("precision, expected", [('16-mixed', torch.float16), ("bf16-mixed", torch.bfloat16)]) +@pytest.mark.parametrize("precision, expected", [("16-mixed", torch.float16), ("bf16-mixed", torch.bfloat16)]) def test_precision_plugin_config(precision, expected): plugin = FSDPMixedPrecisionPlugin(precision=precision, device="cuda") config = plugin.mixed_precision_config @@ -191,7 +191,7 @@ def test_fsdp_strategy_sync_batchnorm(tmpdir): accelerator="gpu", devices=2, strategy="fsdp", - precision='16-mixed', + precision="16-mixed", max_epochs=1, sync_batchnorm=True, ) @@ -199,7 +199,7 @@ def test_fsdp_strategy_sync_batchnorm(tmpdir): @RunIf(min_cuda_gpus=1, skip_windows=True, standalone=True, min_torch="1.12") -@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)))) def test_fsdp_strategy_checkpoint(tmpdir, precision): """Test to ensure that checkpoint is saved correctly when using a single GPU, and all stages can be run.""" model = TestFSDPModel() @@ -230,7 +230,7 @@ def test_fsdp_checkpoint_multi_gpus(tmpdir, model, strategy): accelerator="gpu", devices=2, strategy=strategy, - precision='16-mixed', + precision="16-mixed", max_epochs=1, limit_train_batches=2, limit_val_batches=2, diff --git a/tests/tests_pytorch/strategies/test_registry.py b/tests/tests_pytorch/strategies/test_registry.py index 9fa2e26f95008..f5e7384ea1f2c 100644 --- a/tests/tests_pytorch/strategies/test_registry.py +++ b/tests/tests_pytorch/strategies/test_registry.py @@ -48,7 +48,7 @@ def test_strategy_registry_with_deepspeed_strategies(strategy_name, init_params) @pytest.mark.parametrize("strategy", ["deepspeed", "deepspeed_stage_2_offload", "deepspeed_stage_3"]) def test_deepspeed_strategy_registry_with_trainer(tmpdir, strategy): - trainer = Trainer(default_root_dir=tmpdir, strategy=strategy, precision='16-mixed') + trainer = Trainer(default_root_dir=tmpdir, strategy=strategy, precision="16-mixed") assert isinstance(trainer.strategy, DeepSpeedStrategy) diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 35f89bc1d882b..b97fa36d64bac 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -596,14 +596,16 @@ def test_check_fsdp_strategy_and_fallback(): def test_unsupported_tpu_choice(tpu_available): - with pytest.raises(MisconfigurationException, match=r"accelerator='tpu', precision='64-true'\)` is not implemented"): - Trainer(accelerator="tpu", precision='64-true') + with pytest.raises( + MisconfigurationException, match=r"accelerator='tpu', precision='64-true'\)` is not implemented" + ): + Trainer(accelerator="tpu", precision="64-true") # if user didn't set strategy, AcceleratorConnector will choose the TPUSingleStrategy or TPUSpawnStrategy with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"), pytest.warns( UserWarning, match=r"accelerator='tpu', precision=16-mixed\)` but AMP with fp16 is not supported" ): - Trainer(accelerator="tpu", precision='16-mixed', strategy="ddp") + Trainer(accelerator="tpu", precision="16-mixed", strategy="ddp") @mock.patch("lightning.pytorch.accelerators.ipu.IPUAccelerator.is_available", return_value=True) @@ -846,5 +848,5 @@ def test_colossalai_external_strategy(monkeypatch): from lightning_colossalai import ColossalAIStrategy - trainer = Trainer(strategy="colossalai", precision='16-mixed') + trainer = Trainer(strategy="colossalai", precision="16-mixed") assert isinstance(trainer.strategy, ColossalAIStrategy) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index d181dfa76067e..ad6e0c69908af 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -72,7 +72,8 @@ def configure_optimizers(self): @pytest.mark.parametrize( - "kwargs", [{}, pytest.param({"accelerator": "gpu", "devices": 1, "precision": '16-mixed'}, marks=RunIf(min_cuda_gpus=1))] + "kwargs", + [{}, pytest.param({"accelerator": "gpu", "devices": 1, "precision": "16-mixed"}, marks=RunIf(min_cuda_gpus=1))], ) def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): model = ManualOptModel() @@ -87,7 +88,7 @@ def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): **kwargs, ) - if kwargs.get("precision") == '16-mixed': + if kwargs.get("precision") == "16-mixed": # mock the scaler instead of the optimizer step because it can be skipped with NaNs scaler_step_patch = mock.patch.object( trainer.precision_plugin.scaler, "step", wraps=trainer.precision_plugin.scaler.step @@ -99,7 +100,7 @@ def test_multiple_optimizers_manual_call_counts(tmpdir, kwargs): assert bwd_mock.call_count == limit_train_batches * 3 assert trainer.global_step == limit_train_batches * 2 - if kwargs.get("precision") == '16-mixed': + if kwargs.get("precision") == "16-mixed": scaler_step_patch.stop() assert scaler_step.call_count == len(model.optimizers()) * limit_train_batches @@ -141,7 +142,7 @@ def test_multiple_optimizers_manual_amp(tmpdir, accelerator): max_epochs=1, log_every_n_steps=1, enable_model_summary=False, - precision='16-mixed', + precision="16-mixed", accelerator=accelerator, devices=1, ) @@ -224,7 +225,7 @@ def test_manual_optimization_and_return_tensor(tmpdir): limit_train_batches=10, limit_test_batches=0, limit_val_batches=0, - precision='16-mixed', + precision="16-mixed", strategy="ddp_spawn", accelerator="gpu", devices=2, @@ -309,7 +310,7 @@ def on_train_epoch_end(self, *_, **__): limit_train_batches=20, limit_test_batches=0, limit_val_batches=0, - precision='16-mixed', + precision="16-mixed", accelerator="gpu", devices=1, ) @@ -383,7 +384,7 @@ def on_before_optimizer_step(self, optimizer, *_): max_epochs=1, log_every_n_steps=1, enable_model_summary=False, - precision='16-mixed', + precision="16-mixed", accelerator="gpu", devices=1, ) @@ -848,7 +849,7 @@ def test_lr_scheduler_step_not_called(tmpdir): @RunIf(min_cuda_gpus=1) -@pytest.mark.parametrize("precision", ['16-mixed', '32-true']) +@pytest.mark.parametrize("precision", ["16-mixed", "32-true"]) def test_multiple_optimizers_logging(precision, tmpdir): """Tests that metrics are properly being logged.""" diff --git a/tests/tests_pytorch/trainer/test_trainer.py b/tests/tests_pytorch/trainer/test_trainer.py index d1c42fa3499d0..1e5546d5a288b 100644 --- a/tests/tests_pytorch/trainer/test_trainer.py +++ b/tests/tests_pytorch/trainer/test_trainer.py @@ -1019,7 +1019,7 @@ def on_exception(self, trainer, pl_module, exception): assert isinstance(handle_interrupt_callback.exception, MisconfigurationException) -@pytest.mark.parametrize("precision", ['32-true', pytest.param('16-mixed', marks=RunIf(min_cuda_gpus=1))]) +@pytest.mark.parametrize("precision", ["32-true", pytest.param("16-mixed", marks=RunIf(min_cuda_gpus=1))]) @RunIf(sklearn=True) def test_gradient_clipping_by_norm(tmpdir, precision): """Test gradient clipping by norm.""" @@ -1048,7 +1048,7 @@ def configure_gradient_clipping(self, *args, **kwargs): assert model.assertion_called -@pytest.mark.parametrize("precision", ['32-true', pytest.param('16-mixed', marks=RunIf(min_cuda_gpus=1))]) +@pytest.mark.parametrize("precision", ["32-true", pytest.param("16-mixed", marks=RunIf(min_cuda_gpus=1))]) def test_gradient_clipping_by_value(tmpdir, precision): """Test gradient clipping by value.""" trainer = Trainer( @@ -1444,7 +1444,7 @@ def test_spawn_predict_return_predictions(tmpdir): @pytest.mark.parametrize("return_predictions", [None, False, True]) -@pytest.mark.parametrize("precision", ['32-true', '64-true']) +@pytest.mark.parametrize("precision", ["32-true", "64-true"]) def test_predict_return_predictions_cpu(return_predictions, precision, tmpdir): """Test that `return_predictions=True`.""" seed_everything(42) @@ -1455,7 +1455,7 @@ def test_predict_return_predictions_cpu(return_predictions, precision, tmpdir): if return_predictions or return_predictions is None: assert len(preds) == 1 assert preds[0].shape == torch.Size([1, 2]) - assert preds[0].dtype == (torch.float64 if precision == '64-true' else torch.float32) + assert preds[0].dtype == (torch.float64 if precision == "64-true" else torch.float32) @pytest.mark.parametrize(["max_steps", "max_epochs", "global_step"], [(10, 5, 10), (20, None, 20)]) diff --git a/tests/tests_pytorch/tuner/test_scale_batch_size.py b/tests/tests_pytorch/tuner/test_scale_batch_size.py index ea55b513d1ab4..533669c30f94a 100644 --- a/tests/tests_pytorch/tuner/test_scale_batch_size.py +++ b/tests/tests_pytorch/tuner/test_scale_batch_size.py @@ -254,7 +254,7 @@ def test_error_on_dataloader_passed_to_fit(tmpdir): def test_auto_scale_batch_size_with_amp(tmpdir): before_batch_size = 2 model = BatchSizeModel(batch_size=before_batch_size) - trainer = Trainer(default_root_dir=tmpdir, max_steps=1, accelerator="gpu", devices=1, precision='16-mixed') + trainer = Trainer(default_root_dir=tmpdir, max_steps=1, accelerator="gpu", devices=1, precision="16-mixed") tuner = Tuner(trainer) tuner.scale_batch_size(model) after_batch_size = model.batch_size diff --git a/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py b/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py index 314501ef0e8ee..9f4bcef723434 100644 --- a/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py +++ b/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py @@ -32,7 +32,7 @@ def test_deepspeed_collate_checkpoint(tmpdir): accelerator="gpu", devices=2, fast_dev_run=True, - precision='16-mixed', + precision="16-mixed", enable_progress_bar=False, enable_model_summary=False, ) diff --git a/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py b/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py index 8d734ad74649f..146ab1aa6601b 100644 --- a/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py +++ b/tests/tests_pytorch/utilities/test_deepspeed_model_summary.py @@ -45,7 +45,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") - accelerator="gpu", fast_dev_run=True, devices=2, - precision='16-mixed', + precision="16-mixed", enable_model_summary=True, callbacks=[TestCallback()], ) diff --git a/tests/tests_pytorch/utilities/test_torchdistx.py b/tests/tests_pytorch/utilities/test_torchdistx.py index 1491dc5106443..187a9a5c56084 100644 --- a/tests/tests_pytorch/utilities/test_torchdistx.py +++ b/tests/tests_pytorch/utilities/test_torchdistx.py @@ -55,7 +55,7 @@ def test_deferred_init_with_lightning_module(): ( {"accelerator": "auto", "devices": 1}, pytest.param( - {"strategy": "deepspeed_stage_3", "accelerator": "gpu", "devices": 2, "precision": '16-mixed'}, + {"strategy": "deepspeed_stage_3", "accelerator": "gpu", "devices": 2, "precision": "16-mixed"}, marks=RunIf(min_cuda_gpus=2, deepspeed=True), ), ), From 6976db715151cfb7fc6c0665a59e188df46659b1 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 11:51:16 +0100 Subject: [PATCH 14/26] update with latest fabric changes --- .../connectors/accelerator_connector.py | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py index dabf162875597..2a7ba20562dbf 100644 --- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py +++ b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py @@ -22,10 +22,8 @@ from lightning.fabric.connector import ( _PRECISION_INPUT, - _PRECISION_INPUT_INT, _PRECISION_INPUT_STR, - _PRECISION_INPUT_STR_LEGACY, - _PRECISION_INPUT_STR_LEGACY_CONVERSION, + _convert_precision_to_unified_args ) from lightning.fabric.plugins.environments import ( ClusterEnvironment, @@ -247,25 +245,7 @@ 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) - ) - if precision not in supported_precision: - raise MisconfigurationException( - f"Precision {repr(precision)} is invalid. Allowed precision values: {supported_precision}" - ) - - 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 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!" - ) - precision = _PRECISION_INPUT_STR_LEGACY_CONVERSION[precision] - - self._precision_flag = cast(_PRECISION_INPUT_STR, precision) + self._precision_flag = _convert_precision_to_unified_args(precision) if plugins: plugins_flags_types: Dict[str, int] = Counter() From a730bccb33e06be61306ac52ae17c56be20508a8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Feb 2023 10:53:31 +0000 Subject: [PATCH 15/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/trainer/connectors/accelerator_connector.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py index 2a7ba20562dbf..ced2daefb1508 100644 --- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py +++ b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py @@ -15,16 +15,11 @@ import logging import os from collections import Counter -from typing import cast, Dict, List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union import torch -from typing_extensions import get_args -from lightning.fabric.connector import ( - _PRECISION_INPUT, - _PRECISION_INPUT_STR, - _convert_precision_to_unified_args -) +from lightning.fabric.connector import _convert_precision_to_unified_args, _PRECISION_INPUT, _PRECISION_INPUT_STR from lightning.fabric.plugins.environments import ( ClusterEnvironment, KubeflowEnvironment, From e4eb724667f906b92542df68c8798bb855ba4469 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 12:17:45 +0100 Subject: [PATCH 16/26] fix tests --- tests/tests_pytorch/models/test_amp.py | 2 +- .../trainer/connectors/test_accelerator_connector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/models/test_amp.py b/tests/tests_pytorch/models/test_amp.py index 372cb10db349a..01d16e1c64adb 100644 --- a/tests/tests_pytorch/models/test_amp.py +++ b/tests/tests_pytorch/models/test_amp.py @@ -53,7 +53,7 @@ def _assert_autocast_enabled(self): ("strategy", "precision", "devices"), ( ("single_device", "16-mixed", 1), - ("single_device", "bf16-mixed'", 1), + ("single_device", "bf16-mixed", 1), ("ddp_spawn", "16-mixed", 2), ("ddp_spawn", "bf16-mixed", 2), ), diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 009f57d8db3fc..1797895da951c 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -413,7 +413,7 @@ def test_device_type_when_strategy_instance_gpu_passed(strategy_class, cuda_coun @pytest.mark.parametrize("precision", [1, 12, "invalid"]) def test_validate_precision_type(precision): - with pytest.raises(MisconfigurationException, match=f"Precision {repr(precision)} is invalid"): + with pytest.raises(ValueError, match=f"Precision {repr(precision)} is invalid"): Trainer(precision=precision) From 84e24aa88e19945beff9130675292d083c7eb5e0 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 12:48:23 +0100 Subject: [PATCH 17/26] update test to see results --- tests/tests_pytorch/accelerators/test_ipu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index 650b4c56ec1f3..5c65e3b0cb0b5 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/tests/tests_pytorch/accelerators/test_ipu.py @@ -228,7 +228,7 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64] data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes] new_data = trainer.strategy.batch_to_device(data) - assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)) + assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)), ''.join([f'{dtype}: {val.dtype}' for dtype, val in zip(not_changed_dtypes, new_data)]) with pytest.raises(SystemExit): trainer.fit(model) From 3f8c319d08d9013af6a0176cfdd7df53868e04c7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Feb 2023 11:49:20 +0000 Subject: [PATCH 18/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/tests_pytorch/accelerators/test_ipu.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index 5c65e3b0cb0b5..ca46178719f7f 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/tests/tests_pytorch/accelerators/test_ipu.py @@ -228,7 +228,9 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64] data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes] new_data = trainer.strategy.batch_to_device(data) - assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)), ''.join([f'{dtype}: {val.dtype}' for dtype, val in zip(not_changed_dtypes, new_data)]) + assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)), "".join( + [f"{dtype}: {val.dtype}" for dtype, val in zip(not_changed_dtypes, new_data)] + ) with pytest.raises(SystemExit): trainer.fit(model) From cbb85ae4b61b024b21865d881ada831b33bf2de3 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 13:04:55 +0100 Subject: [PATCH 19/26] update tests --- src/lightning/pytorch/plugins/precision/tpu_bf16.py | 2 +- tests/tests_pytorch/accelerators/test_ipu.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lightning/pytorch/plugins/precision/tpu_bf16.py b/src/lightning/pytorch/plugins/precision/tpu_bf16.py index 0ff320d1ea807..bef5989736a18 100644 --- a/src/lightning/pytorch/plugins/precision/tpu_bf16.py +++ b/src/lightning/pytorch/plugins/precision/tpu_bf16.py @@ -23,7 +23,7 @@ class TPUBf16PrecisionPlugin(TPUPrecisionPlugin): """Plugin that enables bfloats on TPUs.""" - precision: Literal["bf16"] = "bf16-mixed" + precision: Literal["bf16-mixed"] = "bf16-mixed" def connect( self, model: nn.Module, optimizers: List[Optimizer], lr_schedulers: List[Any] diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index ca46178719f7f..93d82295be603 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/tests/tests_pytorch/accelerators/test_ipu.py @@ -223,7 +223,7 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: changed_dtypes = [torch.float, torch.float64] data = [torch.zeros((1), dtype=dtype) for dtype in changed_dtypes] new_data = trainer.strategy.batch_to_device(data) - assert all(val.dtype is torch.half for val in new_data) + assert all(val.dtype is torch.half for val in new_data), ''.join([f'{dtype}: {val.dtype}' for dtype, val in zip(changed_dtypes, new_data)]) not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64] data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes] @@ -235,7 +235,7 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: with pytest.raises(SystemExit): trainer.fit(model) - + @RunIf(ipu=True) def test_device_iterations_ipu_strategy(tmpdir): class TestCallback(Callback): From 03482ce64e4dd1418d54c4d51a28bc2a3d9cc6ef Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Feb 2023 12:05:54 +0000 Subject: [PATCH 20/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/tests_pytorch/accelerators/test_ipu.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/accelerators/test_ipu.py b/tests/tests_pytorch/accelerators/test_ipu.py index 93d82295be603..1e03c3cf617bb 100644 --- a/tests/tests_pytorch/accelerators/test_ipu.py +++ b/tests/tests_pytorch/accelerators/test_ipu.py @@ -223,7 +223,9 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: changed_dtypes = [torch.float, torch.float64] data = [torch.zeros((1), dtype=dtype) for dtype in changed_dtypes] new_data = trainer.strategy.batch_to_device(data) - assert all(val.dtype is torch.half for val in new_data), ''.join([f'{dtype}: {val.dtype}' for dtype, val in zip(changed_dtypes, new_data)]) + assert all(val.dtype is torch.half for val in new_data), "".join( + [f"{dtype}: {val.dtype}" for dtype, val in zip(changed_dtypes, new_data)] + ) not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64] data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes] @@ -235,7 +237,7 @@ def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: with pytest.raises(SystemExit): trainer.fit(model) - + @RunIf(ipu=True) def test_device_iterations_ipu_strategy(tmpdir): class TestCallback(Callback): From 4d19de49e6c1e5d6c0849dd28e4ef975f2b60bbe Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 13:34:50 +0100 Subject: [PATCH 21/26] update utils --- src/lightning/pytorch/strategies/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lightning/pytorch/strategies/utils.py b/src/lightning/pytorch/strategies/utils.py index 1c3d72337786d..1a4c8c00c9215 100644 --- a/src/lightning/pytorch/strategies/utils.py +++ b/src/lightning/pytorch/strategies/utils.py @@ -32,9 +32,9 @@ def _call_register_strategies(registry: _StrategyRegistry, base_module: str) -> mod.register_strategies(registry) -def _fp_to_half(tensor: Tensor, precision: Literal["64", 64, "32", 32, "16", 16, "bf16"]) -> Tensor: - if str(precision) == "16": +def _fp_to_half(tensor: Tensor, precision: Literal["64-true", "32-true", "16-mixed", "bf16-mixed", ]) -> Tensor: + if str(precision) == "16-mixed": return _convert_fp_tensor(tensor, torch.half) - if precision == "bf16": + if precision == "bf16-mixed": return _convert_fp_tensor(tensor, torch.bfloat16) return tensor From 14b06adefe637f6dd7f7e7fe143e5d59b598dfc2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Feb 2023 12:35:45 +0000 Subject: [PATCH 22/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning/pytorch/strategies/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/strategies/utils.py b/src/lightning/pytorch/strategies/utils.py index 1a4c8c00c9215..f67fb55823a51 100644 --- a/src/lightning/pytorch/strategies/utils.py +++ b/src/lightning/pytorch/strategies/utils.py @@ -32,7 +32,15 @@ def _call_register_strategies(registry: _StrategyRegistry, base_module: str) -> mod.register_strategies(registry) -def _fp_to_half(tensor: Tensor, precision: Literal["64-true", "32-true", "16-mixed", "bf16-mixed", ]) -> Tensor: +def _fp_to_half( + tensor: Tensor, + precision: Literal[ + "64-true", + "32-true", + "16-mixed", + "bf16-mixed", + ], +) -> Tensor: if str(precision) == "16-mixed": return _convert_fp_tensor(tensor, torch.half) if precision == "bf16-mixed": From 01ef70ef4bde7ec461296788c3163bdd90b45663 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 14:42:26 +0100 Subject: [PATCH 23/26] update --- src/lightning/fabric/CHANGELOG.md | 2 -- src/lightning/pytorch/CHANGELOG.md | 2 ++ .../connectors/test_accelerator_connector.py | 24 +++++++++---------- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md index 1b9543d63e9cc..f53eeea7081fe 100644 --- a/src/lightning/fabric/CHANGELOG.md +++ b/src/lightning/fabric/CHANGELOG.md @@ -38,8 +38,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Renamed `strategy='tpu_spawn'` to `strategy='xla'` and `strategy='tpu_spawn_debug'` to `strategy='xla_debug'` ([#16781](https://github.com/Lightning-AI/lightning/pull/16781)) -- Changed arguments for precision settings (from [64|32|16|bf16] to ["64-true"|"32-true"|"16-mixed"|"bf16-mixed"]) ([#16767](https://github.com/Lightning-AI/lightning/pull/16767)) - - Changed arguments for precision settings (from [64|32|16|bf16] to ["64-true"|"32-true"|"16-mixed"|"bf16-mixed"]) ([#16767](https://github.com/Lightning-AI/lightning/pull/16767)) ### Deprecated diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index f39ecdb2bdc25..41d3009ef87b1 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -95,6 +95,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Renamed `strategy='tpu_spawn'` to `strategy='xla'` and `strategy='tpu_spawn_debug'` to `strategy='xla_debug'` ([#16781](https://github.com/Lightning-AI/lightning/pull/16781)) +- Changed arguments for precision settings (from [64|32|16|bf16] to ["64-true"|"32-true"|"16-mixed"|"bf16-mixed"]) ([#16783](https://github.com/Lightning-AI/lightning/pull/16783)) + ### Deprecated - diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 1797895da951c..c26053dfb0ff9 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -839,15 +839,15 @@ def get_defaults(cls): assert connector_default == trainer_defaults[name] -@RunIf(min_cuda_gpus=1) # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite -@pytest.mark.skipif(not package_available("lightning_colossalai"), reason="Requires Colossal AI Strategy") -def test_colossalai_external_strategy(monkeypatch): - with mock.patch( - "lightning.pytorch.trainer.connectors.accelerator_connector._LIGHTNING_COLOSSALAI_AVAILABLE", False - ), pytest.raises(ModuleNotFoundError): - Trainer(strategy="colossalai") - - from lightning_colossalai import ColossalAIStrategy - - trainer = Trainer(strategy="colossalai", precision="16-mixed") - assert isinstance(trainer.strategy, ColossalAIStrategy) +# @RunIf(min_cuda_gpus=1) # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite +# @pytest.mark.skipif(not package_available("lightning_colossalai"), reason="Requires Colossal AI Strategy") +# def test_colossalai_external_strategy(monkeypatch): +# with mock.patch( +# "lightning.pytorch.trainer.connectors.accelerator_connector._LIGHTNING_COLOSSALAI_AVAILABLE", False +# ), pytest.raises(ModuleNotFoundError): +# Trainer(strategy="colossalai") +# +# from lightning_colossalai import ColossalAIStrategy +# +# trainer = Trainer(strategy="colossalai", precision="16-mixed") +# assert isinstance(trainer.strategy, ColossalAIStrategy) From 69a46e5950cdbbd77191c0659344bcc98d8e8282 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Feb 2023 13:43:33 +0000 Subject: [PATCH 24/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../trainer/connectors/test_accelerator_connector.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index c26053dfb0ff9..2a91284d4ac89 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -20,7 +20,6 @@ import pytest import torch import torch.distributed -from lightning_utilities.core.imports import package_available import lightning.pytorch from lightning.fabric.plugins.environments import ( From 4181d60bdce6960392ca00729e40b9b594d74a79 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 14:46:56 +0100 Subject: [PATCH 25/26] update --- .../connectors/test_accelerator_connector.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 2a91284d4ac89..9cc3754d23764 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -838,15 +838,16 @@ def get_defaults(cls): assert connector_default == trainer_defaults[name] -# @RunIf(min_cuda_gpus=1) # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite -# @pytest.mark.skipif(not package_available("lightning_colossalai"), reason="Requires Colossal AI Strategy") -# def test_colossalai_external_strategy(monkeypatch): -# with mock.patch( -# "lightning.pytorch.trainer.connectors.accelerator_connector._LIGHTNING_COLOSSALAI_AVAILABLE", False -# ), pytest.raises(ModuleNotFoundError): -# Trainer(strategy="colossalai") -# -# from lightning_colossalai import ColossalAIStrategy -# -# trainer = Trainer(strategy="colossalai", precision="16-mixed") -# assert isinstance(trainer.strategy, ColossalAIStrategy) +@RunIf(min_cuda_gpus=1) # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite +@pytest.mark.skipif(not package_available("lightning_colossalai"), reason="Requires Colossal AI Strategy") +@pytest.mark.skip +def test_colossalai_external_strategy(monkeypatch): + with mock.patch( + "lightning.pytorch.trainer.connectors.accelerator_connector._LIGHTNING_COLOSSALAI_AVAILABLE", False + ), pytest.raises(ModuleNotFoundError): + Trainer(strategy="colossalai") + + from lightning_colossalai import ColossalAIStrategy + + trainer = Trainer(strategy="colossalai", precision="16-mixed") + assert isinstance(trainer.strategy, ColossalAIStrategy) From 9cfbc24851869542a274d5a24d5b2b2c017f3894 Mon Sep 17 00:00:00 2001 From: Justus Schock Date: Fri, 17 Feb 2023 14:52:39 +0100 Subject: [PATCH 26/26] revert pre-commit yet again --- .../trainer/connectors/test_accelerator_connector.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 9cc3754d23764..e98d4df2a9c54 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py @@ -20,6 +20,7 @@ import pytest import torch import torch.distributed +from lightning_utilities.core.imports import package_available import lightning.pytorch from lightning.fabric.plugins.environments import (