Validate the precision input earlier (Lightning-AI#9763)

Author: carmocca

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -133,6 +133,10 @@ def __init__(
         self.sync_batchnorm = sync_batchnorm
         self.benchmark = benchmark
         self.replace_sampler_ddp = replace_sampler_ddp
+        if not PrecisionType.supported_type(precision):
+            raise MisconfigurationException(
+                f"Precision {repr(precision)} is invalid. Allowed precision values: {PrecisionType.supported_types()}"
+            )
         self.precision = precision
         self.amp_type = amp_type.lower() if isinstance(amp_type, str) else None
         self.amp_level = amp_level
@@ -662,10 +666,6 @@ def select_precision_plugin(self) -> PrecisionPlugin:
 
                 return ApexMixedPrecisionPlugin(self.amp_level)
 
-        raise MisconfigurationException(
-            f"Precision {self.precision} is invalid. Allowed precision values: {PrecisionType.supported_types()}"
-        )
-
     def select_training_type_plugin(self) -> TrainingTypePlugin:
         if isinstance(self.accelerator, Accelerator) and self.accelerator.training_type_plugin is not None:
             plugin = self.accelerator.training_type_plugin

pytorch_lightning/utilities/enums.py

@@ -62,6 +62,7 @@ class PrecisionType(LightningEnum):
     FLOAT = "32"
     FULL = "64"
     BFLOAT = "bf16"
+    MIXED = "mixed"
 
     @staticmethod
     def supported_type(precision: Union[str, int]) -> bool:

tests/accelerators/test_accelerator_connector.py

@@ -709,7 +709,7 @@ def test_device_type_when_training_plugin_gpu_passed(tmpdir, plugin):
 @pytest.mark.parametrize("precision", [1, 12, "invalid"])
 def test_validate_precision_type(tmpdir, precision):
 
-    with pytest.raises(MisconfigurationException, match=f"Precision {precision} is invalid"):
+    with pytest.raises(MisconfigurationException, match=f"Precision {repr(precision)} is invalid"):
         Trainer(precision=precision)
 
 

tests/utilities/test_enums.py

@@ -27,7 +27,7 @@ def test_consistency():
 
 
 def test_precision_supported_types():
-    assert PrecisionType.supported_types() == ["16", "32", "64", "bf16"]
+    assert PrecisionType.supported_types() == ["16", "32", "64", "bf16", "mixed"]
     assert PrecisionType.supported_type(16)
     assert PrecisionType.supported_type("16")
     assert not PrecisionType.supported_type(1)