add support for precision="16-mixed"

Lightning-Universe · Feb 17, 2023 · c87e31e · c87e31e
1 parent cd1d0db
commit c87e31e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/lightning_colossalai/precision.py b/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,

diff --git a/src/lightning_colossalai/strategy.py b/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`."