forward XLA precision

Author: justusschock

src/lightning/fabric/plugins/precision/bitsandbytes.py

@@ -57,7 +57,7 @@ def __init__(
             BitsandbytesPrecision as EnterpriseBitsandbytesPrecision,
         )
 
-        self.bitsandbytes_impl = EnterpriseBitsandbytesPrecision(mode, dtype, ignore_modules)
+        self.bitsandbytes_impl = EnterpriseBitsandbytesPrecision(mode=mode, dtype=dtype, ignore_modules=ignore_modules)
 
     @override
     def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:

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

@@ -45,10 +45,10 @@ def __init__(self, precision: _PRECISION_INPUT) -> None:
         super().__init__()
         _raise_enterprise_not_available()
         from pytorch_lightning_enterprise.fabric.plugins.precision.deepspeed import (
-            DeepSpeedPrecision as EnterpriseDeepSpeedPrecision,
+            DeepSpeedPrecisionFabric as EnterpriseDeepSpeedPrecision,
         )
 
-        self.deepspeed_impl = EnterpriseDeepSpeedPrecision(precision)
+        self.deepspeed_impl = EnterpriseDeepSpeedPrecision(precision=precision)
 
     @override
     def convert_module(self, module: Module) -> Module:

src/lightning/fabric/plugins/precision/xla.py

@@ -11,14 +11,12 @@
 # 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.
-import os
 from typing import Any, Literal
 
-import torch
-from typing_extensions import get_args, override
+from typing_extensions import override
 
-from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
 from lightning.fabric.plugins.precision.precision import Precision
+from lightning.fabric.utilities.imports import _raise_enterprise_not_available
 from lightning.fabric.utilities.types import Optimizable
 
 _PRECISION_INPUT = Literal["32-true", "16-true", "bf16-true"]
@@ -37,37 +35,20 @@ class XLAPrecision(Precision):
     """
 
     def __init__(self, precision: _PRECISION_INPUT) -> None:
-        if not _XLA_AVAILABLE:
-            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
-        supported_precision = get_args(_PRECISION_INPUT)
-        if precision not in supported_precision:
-            raise ValueError(
-                f"`precision={precision!r})` is not supported in XLA."
-                f" `precision` must be one of: {supported_precision}."
-            )
-        self.precision = precision
+        super().__init__()
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.fabric.plugins.precision.xla import XLAPrecision as EnterpriseXLAPrecision
 
-        if precision == "16-true":
-            os.environ["XLA_USE_F16"] = "1"
-            self._desired_dtype = torch.float16
-        elif precision == "bf16-true":
-            os.environ["XLA_USE_BF16"] = "1"
-            self._desired_dtype = torch.bfloat16
-        else:
-            self._desired_dtype = torch.float32
+        self.xla_impl = EnterpriseXLAPrecision(precision=precision)
 
     @override
     def optimizer_step(
         self,
         optimizer: Optimizable,
         **kwargs: Any,
     ) -> Any:
-        import torch_xla.core.xla_model as xm
-
-        # you always want to `xm.mark_step()` after `optimizer.step` for better performance, so we set `barrier=True`
-        return xm.optimizer_step(optimizer, optimizer_args=kwargs, barrier=True)
+        return self.xla_impl.optimizer_step(optimizer, **kwargs)
 
     @override
     def teardown(self) -> None:
-        os.environ.pop("XLA_USE_BF16", None)
-        os.environ.pop("XLA_USE_F16", None)
+        return self.xla_impl.teardown()