[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

src/lightning/pytorch/strategies/deepspeed.py

@@ -521,12 +521,16 @@ def model_sharded_context(self) -> Generator[None, None, None]:
         import deepspeed
 
         self._init_config_if_needed()
-        
+
         # If detect 'mics_shard_size'>0 in config['zero_optimization'], alter to use deepspeed.zero.MiCS_Init()
         # https://deepspeed.readthedocs.io/en/latest/zero3.html#mics-configurations
         #! default deepspeed 0.9.0 is not compatible
-        if 'zero_optimization' in self.config and 'mics_shard_size' in self.config['zero_optimization']\
-            and self.config['zero_optimization']['mics_shard_size'] > 0 and self.zero_stage_3:
+        if (
+            "zero_optimization" in self.config
+            and "mics_shard_size" in self.config["zero_optimization"]
+            and self.config["zero_optimization"]["mics_shard_size"] > 0
+            and self.zero_stage_3
+        ):
             with deepspeed.zero.MiCS_Init(
                 enabled=self.zero_stage_3,
                 remote_device=self.remote_device,

tests/tests_pytorch/strategies/test_deepspeed.py

@@ -1283,7 +1283,7 @@ def test_deepspeed_load_checkpoint_validate_path(tmp_path):
 
 @RunIf(min_cuda_gpus=2, standalone=True, deepspeed=True)
 def test_deepspeed_multigpu_stage_3_MiCS_support(tmp_path):
-    """Test to ensure we can use DeepSpeed with basic ZeRO Stage 3 MiCS Support"""
+    """Test to ensure we can use DeepSpeed with basic ZeRO Stage 3 MiCS Support."""
     model = ModelParallelBoringModel()
     strategy = DeepSpeedStrategy(stage=3)
     strategy.config["zero_optimization"]["stage"] = 3
@@ -1302,11 +1302,11 @@ def test_deepspeed_multigpu_stage_3_MiCS_support(tmp_path):
     )
     trainer.test(model)
     trainer.fit(model)
-    
+
     _assert_save_model_is_equal(model, tmp_path, trainer)
     assert isinstance(trainer.strategy, DeepSpeedStrategy)
-    assert 'zero_optimization' in trainer.strategy.config
-    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] == False
+    assert "zero_optimization" in trainer.strategy.config
+    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] is False
     assert trainer.strategy.config["zero_optimization"]["mics_shard_size"] == 1
     assert trainer.strategy.config["zero_optimization"]["stage"] == 3
 
@@ -1317,9 +1317,9 @@ def test_deepspeed_multigpu_stage_3_MiCS_offload_param_support(tmp_path):
         However, in some past pratice, offload param + mics + torchrun will cause inner exception in multi-node environment. \
         Probably this exception is caused by torchrun, not deepspeed. """
     model = ModelParallelBoringModel()
-    strategy = DeepSpeedStrategy(stage=3,offload_params_device="cpu")
+    strategy = DeepSpeedStrategy(stage=3, offload_params_device="cpu")
     strategy.config["zero_optimization"]["stage"] = 3
-    strategy.config["zero_optimization"]["mics_shard_size"] = 1 
+    strategy.config["zero_optimization"]["mics_shard_size"] = 1
     strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] = False
     trainer = Trainer(
         default_root_dir=tmp_path,
@@ -1336,18 +1336,19 @@ def test_deepspeed_multigpu_stage_3_MiCS_offload_param_support(tmp_path):
 
     _assert_save_model_is_equal(model, tmp_path, trainer)
     assert isinstance(trainer.strategy, DeepSpeedStrategy)
-    assert 'zero_optimization' in trainer.strategy.config
-    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] == False
+    assert "zero_optimization" in trainer.strategy.config
+    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] is False
     assert trainer.strategy.config["zero_optimization"]["mics_shard_size"] == 1
     assert trainer.strategy.config["zero_optimization"]["stage"] == 3
 
+
 @RunIf(min_cuda_gpus=2, standalone=True, deepspeed=True)
 def test_deepspeed_multigpu_stage_3_MiCS_offload_param_optimizer_support(tmp_path):
-    """Test to ensure we can use DeepSpeed with ZeRO Stage param & optimizer offload 3 MiCS Support"""
+    """Test to ensure we can use DeepSpeed with ZeRO Stage param & optimizer offload 3 MiCS Support."""
     model = ModelParallelBoringModel()
-    strategy = DeepSpeedStrategy(stage=3,offload_params_device="cpu", offload_optimizer_device="cpu")
+    strategy = DeepSpeedStrategy(stage=3, offload_params_device="cpu", offload_optimizer_device="cpu")
     strategy.config["zero_optimization"]["stage"] = 3
-    strategy.config["zero_optimization"]["mics_shard_size"] = 1 
+    strategy.config["zero_optimization"]["mics_shard_size"] = 1
     strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] = False
     trainer = Trainer(
         default_root_dir=tmp_path,
@@ -1364,23 +1365,24 @@ def test_deepspeed_multigpu_stage_3_MiCS_offload_param_optimizer_support(tmp_pat
 
     _assert_save_model_is_equal(model, tmp_path, trainer)
     assert isinstance(trainer.strategy, DeepSpeedStrategy)
-    assert 'zero_optimization' in trainer.strategy.config
-    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] == False
+    assert "zero_optimization" in trainer.strategy.config
+    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] is False
     assert trainer.strategy.config["zero_optimization"]["mics_shard_size"] == 1
     assert trainer.strategy.config["zero_optimization"]["stage"] == 3
 
 
 @RunIf(min_cuda_gpus=4, standalone=True, deepspeed=True)
 def test_deepspeed_multigpu_stage_3_hierarchical_MiCS_support(tmp_path):
-    """Test to ensure we can use DeepSpeed with ZeRO Stage 3 MiCS Support ('mics_hierarchical_params_gather' = True)."""
+    """Test to ensure we can use DeepSpeed with ZeRO Stage 3 MiCS Support ('mics_hierarchical_params_gather' =
+    True)."""
     model = ModelParallelBoringModel()
     strategy = DeepSpeedStrategy(stage=3)
     strategy.config["zero_optimization"]["stage"] = 3
     strategy.config["zero_optimization"]["mics_shard_size"] = 2
     strategy.config["zero_optimization"]["offload_param"] = {}
     strategy.config["zero_optimization"]["offload_optimizer"] = {}
     strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] = True
-    #Forming a 2 x 2 hierarchy
+    # Forming a 2 x 2 hierarchy
     trainer = Trainer(
         default_root_dir=tmp_path,
         strategy=strategy,
@@ -1396,8 +1398,7 @@ def test_deepspeed_multigpu_stage_3_hierarchical_MiCS_support(tmp_path):
 
     _assert_save_model_is_equal(model, tmp_path, trainer)
     assert isinstance(trainer.strategy, DeepSpeedStrategy)
-    assert 'zero_optimization' in trainer.strategy.config
-    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] == True
+    assert "zero_optimization" in trainer.strategy.config
+    assert trainer.strategy.config["zero_optimization"]["mics_hierarchical_params_gather"] is True
     assert trainer.strategy.config["zero_optimization"]["mics_shard_size"] == 2
     assert trainer.strategy.config["zero_optimization"]["stage"] == 3
-