Repository: huggingface/accelerate

Issue:
invalid `deepspeed_plugin.zero_stage` when using ds_config file
### System Info

```Shell
`accelerate@main`
```

### Reproduction

when using a stage 3 config with ds_config file, this returns `2` instead of expected `3`:
```
    print(f"STAGE {AcceleratorState().deepspeed_plugin.zero_stage}")
```

I think this issue comes from:

https://github.com/huggingface/accelerate/blob/6e4e8702030222979d7ff54a2d8b0406cdf96393/src/accelerate/utils/dataclasses.py#L489-L490

why does it use the default of 2? this would be incorrect in 3 out of 4 cases as there are 4 stages.

If `zero_stage` isn't explicitly passed there is the config file to derive this value from. Bur perhaps you'd want to use this API to explicitly express the value that was passed via `zero_stage` - so the other possibility is to return `None` if it wasn't set by the API or the env var.

But the default should definitely not be `2`. 

@pacman100 


Discussion:

Resolving patch:
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -515,17 +515,25 @@ def __post_init__(self):
                 raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.")
 
             self._deepspeed_config_checks()
-            kwargs = {
-                "gradient_accumulation_steps": self.gradient_accumulation_steps,
-                "gradient_clipping": self.gradient_clipping if self.gradient_clipping else 1.0,
-                "zero_optimization.stage": self.zero_stage,
-                "zero_optimization.offload_optimizer.device": self.offload_optimizer_device,
-                "zero_optimization.offload_param.device": self.offload_param_device,
-                "zero_optimization.stage3_gather_16bit_weights_on_model_save": self.zero3_save_16bit_model,
+            plugin_to_config_mapping = {
+                "gradient_accumulation_steps": "gradient_accumulation_steps",
+                "gradient_clipping": "gradient_clipping",
+                "zero_stage": "zero_optimization.stage",
+                "offload_optimizer_device": "zero_optimization.offload_optimizer.device",
+                "offload_param_device": "zero_optimization.offload_param.device",
+                "zero3_save_16bit_model": "zero_optimization.stage3_gather_16bit_weights_on_model_save",
             }
+            kwargs = {v: getattr(self, k) for k, v in plugin_to_config_mapping.items() if getattr(self, k) is not None}
             for key in kwargs.keys():
                 self.fill_match(key, **kwargs, must_match=False)
             self.hf_ds_config.set_stage_and_offload()
+
+            # filling the missing values in the class attributes from the DeepSpeed config
+            # when using the DeepSpeed config file.
+            for key, value in plugin_to_config_mapping.items():
+                config_value = self.hf_ds_config.get_value(value)
+                if config_value is not None and config_value != "auto":
+                    setattr(self, key, config_value)
         else:
             config = {
                 "train_batch_size": "auto",


Test patch:
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -656,6 +656,14 @@ def test_ds_config_assertions(self):
                 in str(cm.exception)
             )
 
+    @parameterized.expand(stages, name_func=parameterized_custom_name_func)
+    def test_ds_config(self, stage):
+        deepspeed_plugin = DeepSpeedPlugin(
+            hf_ds_config=self.ds_config_file[stage],
+            zero3_init_flag=True,
+        )
+        self.assertEqual(deepspeed_plugin.zero_stage, int(stage.replace("zero", "")))
+
     def test_basic_run(self):
         mod_file = inspect.getfile(accelerate.test_utils)
         test_file_path = os.path.sep.join(