update

Author: justusschock

src/lightning/fabric/utilities/imports.py

@@ -50,6 +50,7 @@
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
 _DEEPSPEED_GREATER_EQUAL_0_16 = RequirementCache("deepspeed>=0.16.0")
 _ENTERPRISE_AVAILABLE = RequirementCache("pytorch_lightning_enterprise")
+_TRANSFORMER_ENGINE_AVAILABLE = RequirementCache("transformer_engine>=0.11.0")
 
 
 def _raise_enterprise_not_available() -> None:

tests/tests_pytorch/deprecated_api/test_no_removal_version.py

@@ -2,6 +2,7 @@
 from unittest.mock import Mock
 
 import pytest
+import pytorch_lightning_enterprise.utils.imports
 import torch.nn
 
 import lightning.fabric
@@ -62,6 +63,7 @@ def test_fsdp_precision_plugin():
 
 def test_bitsandbytes_precision_plugin(monkeypatch):
     monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
+    monkeypatch.setattr(pytorch_lightning_enterprise.utils.imports, "_BITSANDBYTES_AVAILABLE", True)
     bitsandbytes_mock = Mock()
     monkeypatch.setitem(sys.modules, "bitsandbytes", bitsandbytes_mock)
 
@@ -107,7 +109,8 @@ def test_precision_plugin():
 
 
 def test_transformer_engine_precision_plugin(monkeypatch):
-    monkeypatch.setattr(lightning.fabric.plugins.precision.transformer_engine, "_TRANSFORMER_ENGINE_AVAILABLE", True)
+    monkeypatch.setattr(lightning.fabric.utilities.imports, "_TRANSFORMER_ENGINE_AVAILABLE", True)
+    monkeypatch.setattr(pytorch_lightning_enterprise.utils.imports, "_TRANSFORMER_ENGINE_AVAILABLE", True)
     transformer_engine_mock = Mock()
     monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock)
     monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", Mock())

tests/tests_pytorch/plugins/precision/test_transformer_engine.py

@@ -25,7 +25,7 @@
 
 
 def test_transformer_engine_precision_plugin(monkeypatch):
-    module = lightning.fabric.plugins.precision.transformer_engine
+    module = lightning.fabric.utilities.imports
     if module._TRANSFORMER_ENGINE_AVAILABLE:
         pytest.skip("Assumes transformer_engine is unavailable")
     monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True)