[AutoTP] Allow ZeRO stage 3 inference with tensor parallelism

Replace the unconditional `zero_optimization_stage() <= 2` guard with one that
only blocks `autotp + zero_stage == 3` when an optimizer is present. The
inference path (no optimizer -> DummyOptim -> DeepSpeedZeRoOffload) is now
permitted; the training path (optimizer -> DeepSpeedZeroOptimizer_Stage3) is
blocked with a clear NotImplementedError.

Rationale: training under autotp+ZeRO-3 itself works (forward/backward/step and
gradient norms match the non-TP baseline), but the TP-aware checkpoint
consolidation path (`_consolidated_16bit_state_dict` / `save_16bit_model`)
gathers only the ZeRO data-parallel partition and would silently write rank 0's
TP shard instead of the full weight, producing incomplete checkpoints. Training
support will land alongside the checkpoint fix (tracked separately).

Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Author: delock

deepspeed/runtime/engine.py

@@ -628,10 +628,15 @@ def _configure_tensor_parallel_states(self, model):
         and registering a pre-hook to ensure that the Dataloader inputs are consistent across ranks.
         """
         self._set_client_model(model)
-        # sanity check
-        # currently, the compatibility between 'autotp' and 'zero > 1' has not been validated
-        assert self.zero_optimization_stage(
-        ) <= 2, "Currently, the compatibility between 'autotp' and 'zero_stage = 3' has not been validated"
+        # AutoTP + ZeRO-3 is supported only for ZeRO-Inference (no optimizer).
+        # With an optimizer the TP-aware checkpoint consolidation path is not yet
+        # implemented, so save_16bit_model()/save_checkpoint would silently drop
+        # TP shards. Block the training path until that is fixed (tracked separately).
+        if self.zero_optimization_stage() > 2 and (self.client_optimizer or self.optimizer_name()):
+            raise NotImplementedError("AutoTP together with ZeRO stage 3 is currently supported only for inference "
+                                      "(no optimizer). Running it with an optimizer is disabled because TP-aware "
+                                      "checkpoint consolidation is not yet implemented and would silently produce "
+                                      "incomplete checkpoints.")
 
         self.mpu = groups
         self.mpu._init_tp_mesh_device(tensor_model_parallel_size=self.autotp_size())

docs/_pages/config-json.md

@@ -763,7 +763,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s
 | Submit requests to storage device in an overlapped fashion without waiting for completion of earlier requests. | `true`  |
 
 ### Tensor Parallel (AutoTP)
-Configure AutoTP tensor parallelism for training via the DeepSpeed config and hybrid TP + ZeRO. AutoTP supports ZeRO stages 0, 1, and 2 (stage 3 is not supported). `deepspeed.tp_model_init()` remains supported for backward compatibility but is not required when `tensor_parallel` is set in the config.
+Configure AutoTP tensor parallelism for training via the DeepSpeed config and hybrid TP + ZeRO. AutoTP supports ZeRO stages 0, 1, and 2. ZeRO stage 3 is supported only for inference (no optimizer); training with stage 3 is not yet supported. `deepspeed.tp_model_init()` remains supported for backward compatibility but is not required when `tensor_parallel` is set in the config.
 
 When a HuggingFace model provides a built-in `tp_plan` (via `model.config.base_model_tp_plan`), DeepSpeed automatically detects and uses it. In this case, neither `preset_model` nor `partition_config` is required -- just set `autotp_size`. If `partition_config` is also provided, it takes precedence over the model's `tp_plan`.
 ```json

docs/_tutorials/autotp-training.md

@@ -212,7 +212,7 @@ For Grouped Query Attention with different Q/K/V sizes:
 
 ## Limitations
 
-1. **ZeRO Stage 3 not supported**: AutoTP currently only works with ZeRO stages 0, 1, and 2.
+1. **ZeRO Stage 3 training not supported**: AutoTP with ZeRO stage 3 is supported only for inference (no optimizer). Training with an optimizer is blocked until TP-aware checkpoint consolidation is implemented.
 
 2. **TP size must divide model dimensions**: The tensor parallel size must evenly divide the attention head count and hidden dimensions.
 

docs/code-docs/source/training.rst

@@ -392,7 +392,7 @@ See :ref:`autotp-training-init-details` for more details.
     )
 
 .. note::
-   AutoTP training supports ZeRO stages 0, 1, and 2. ZeRO Stage 3 is not supported.
+   AutoTP training supports ZeRO stages 0, 1, and 2. ZeRO stage 3 is supported only for inference (no optimizer).
 
 .. _autotp-training-init-details:
 

tests/unit/model_parallelism/test_autotp_training.py

@@ -185,6 +185,58 @@ def test(self, tp_size: int):
         assert groups.get_data_parallel_world_size() == dp_size
 
 
+@pytest.mark.parametrize("tp_size", [2, 4])
+class TestAutoTpZeroStage3(DistributedTest):
+    """AutoTP + ZeRO stage 3 is supported only for ZeRO-Inference (no optimizer).
+
+    The training path (optimizer present) is blocked because TP-aware checkpoint
+    consolidation is not yet implemented; the inference path (no optimizer ->
+    DummyOptim -> DeepSpeedZeRoOffload) must initialize.
+    """
+
+    world_size = 4
+    reuse_dist_env = False
+
+    def _build_config(self, tp_size: int, with_optimizer: bool) -> dict:
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "autotp_size": tp_size,
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [r".*\.weight$"],
+                        "partition_type": "column",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 3
+            }
+        }
+        if with_optimizer:
+            config_dict["optimizer"] = {"type": "Adam", "params": {"lr": 1e-6}}
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+        return config_dict
+
+    def test_autotp_zero3_training_blocked(self, tp_size: int):
+        skip_on_device()
+        model = SimpleModel(hidden_dim=64)
+        config_dict = self._build_config(tp_size, with_optimizer=True)
+        with pytest.raises(NotImplementedError, match="only for inference"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+    def test_autotp_zero3_inference(self, tp_size: int):
+        skip_on_device()
+        model = SimpleModel(hidden_dim=64)
+        config_dict = self._build_config(tp_size, with_optimizer=False)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert groups.get_tensor_model_parallel_world_size() == tp_size
+
+
 class TestTpModelInitCompatibility(DistributedTest):
     world_size = 4
     reuse_dist_env = False