add support for precision="16-mixed"

src/lightning_colossalai/precision.py

@@ -31,13 +31,13 @@ class ColossalAIPrecisionPlugin(PrecisionPlugin):
             If precison is not 16.
     """
 
-    def __init__(self, precision: Literal["16", 16] = 16) -> None:
+    def __init__(self, precision: Literal["16", 16, "16-mixed"] = 16) -> None:
         if precision not in ("16", 16):
             raise ValueError(
                 f"`Trainer(strategy='colossalai', precision={precision!r})` is not supported."
                 " Consider setting `precision=16`."
             )
-        self.precision = cast(Literal["16"], str(precision))
+        self.precision = cast(Literal["16", "16-mixed"], str(precision))
 
     def backward(  # type: ignore[override]
         self,

src/lightning_colossalai/strategy.py

@@ -300,7 +300,7 @@ def setup_precision_plugin(self) -> None:
 
     def setup(self, trainer: Trainer) -> None:
         precision = self.precision_plugin.precision
-        if precision != "16":
+        if precision not in ("16", "16-mixed"):
             raise ValueError(
                 f"`Trainer(strategy='colossalai', precision={precision!r})` is not supported."
                 " Consider setting `precision=16`."