Remove the Bagua integration (#19445)

Author: awaelchli

.azure/gpu-tests-pytorch.yml

@@ -108,7 +108,7 @@ jobs:
       - bash: |
           pip install -q -r .actions/requirements.txt
           python .actions/assistant.py requirements_prune_pkgs \
-            --packages="[lightning-colossalai,lightning-bagua]" \
+            --packages="[lightning-colossalai]" \
             --req_files="[requirements/_integrations/strategies.txt]"
         displayName: "Prune packages" # these have installation issues
 

requirements/_integrations/strategies.txt

@@ -2,4 +2,3 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 lightning-colossalai >=0.1.0
-lightning-bagua >=0.1.0

src/lightning/pytorch/CHANGELOG.md

@@ -32,7 +32,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Removed
 
--
+- Removed the Bagua integration (`Trainer(strategy="bagua")`) ([#19445](https://github.com/Lightning-AI/lightning/pull/19445))
 
 -
 

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -63,7 +63,6 @@
 from lightning.pytorch.strategies.ddp import _DDP_FORK_ALIASES
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.imports import (
-    _LIGHTNING_BAGUA_AVAILABLE,
     _LIGHTNING_COLOSSALAI_AVAILABLE,
     _graphcore_available_and_importable,
     _habana_available_and_importable,
@@ -196,9 +195,6 @@ def _check_config_and_set_final_flags(
         if strategy == "colossalai" and not _LIGHTNING_COLOSSALAI_AVAILABLE:
             raise ModuleNotFoundError(str(_LIGHTNING_COLOSSALAI_AVAILABLE))
 
-        if strategy == "bagua" and not _LIGHTNING_BAGUA_AVAILABLE:
-            raise ModuleNotFoundError(str(_LIGHTNING_BAGUA_AVAILABLE))
-
         if strategy != "auto" and strategy not in self._registered_strategies and not isinstance(strategy, Strategy):
             raise ValueError(
                 f"You selected an invalid strategy name: `strategy={strategy!r}`."
@@ -421,11 +417,6 @@ def _choose_and_init_cluster_environment(self) -> ClusterEnvironment:
         ):
             if env_type.detect():
                 return env_type()
-        if _LIGHTNING_BAGUA_AVAILABLE:
-            from lightning_bagua import BaguaEnvironment
-
-            if BaguaEnvironment.detect():
-                return BaguaEnvironment()
         return LightningEnvironment()
 
     def _choose_strategy(self) -> Union[Strategy, str]:
@@ -690,13 +681,6 @@ def _register_external_accelerators_and_strategies() -> None:
         if "colossalai" not in StrategyRegistry:
             ColossalAIStrategy.register_strategies(StrategyRegistry)
 
-    if _LIGHTNING_BAGUA_AVAILABLE:
-        from lightning_bagua import BaguaStrategy
-
-        # TODO: Prevent registering multiple times
-        if "bagua" not in StrategyRegistry:
-            BaguaStrategy.register_strategies(StrategyRegistry)
-
     if _habana_available_and_importable():
         from lightning_habana import HPUAccelerator, HPUParallelStrategy, SingleHPUStrategy
 

src/lightning/pytorch/utilities/imports.py

@@ -28,7 +28,6 @@
 _OMEGACONF_AVAILABLE = package_available("omegaconf")
 _TORCHVISION_AVAILABLE = RequirementCache("torchvision")
 _LIGHTNING_COLOSSALAI_AVAILABLE = RequirementCache("lightning-colossalai")
-_LIGHTNING_BAGUA_AVAILABLE = RequirementCache("lightning-bagua")
 
 
 @functools.lru_cache(maxsize=128)

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -872,21 +872,6 @@ def test_colossalai_external_strategy(monkeypatch):
     assert isinstance(trainer.strategy, ColossalAIStrategy)
 
 
-@RunIf(min_cuda_gpus=1)  # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite
-@pytest.mark.xfail(raises=ImportError, reason="Not updated to latest API")
-@pytest.mark.skipif(not package_available("lightning_bagua"), reason="Requires Bagua Strategy")
-def test_bagua_external_strategy(monkeypatch):
-    with mock.patch(
-        "lightning.pytorch.trainer.connectors.accelerator_connector._LIGHTNING_BAGUA_AVAILABLE", False
-    ), pytest.raises(ModuleNotFoundError):
-        Trainer(strategy="bagua")
-
-    from lightning_bagua import BaguaStrategy
-
-    trainer = Trainer(strategy="bagua")
-    assert isinstance(trainer.strategy, BaguaStrategy)
-
-
 class DeviceMock(Mock):
     def __instancecheck__(self, instance):
         return True