diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 194dc3555387..9d89ff053c23 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -628,10 +628,11 @@ 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" + 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()) diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 01dd50053319..fb4294d1e72b 100755 --- a/docs/_pages/config-json.md +++ b/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 diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index 7f8d2cbd52df..ec701c05119b 100644 --- a/docs/_tutorials/autotp-training.md +++ b/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. diff --git a/docs/code-docs/source/training.rst b/docs/code-docs/source/training.rst index aae26007c15e..0e4c43a5309b 100644 --- a/docs/code-docs/source/training.rst +++ b/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: diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/model_parallelism/test_autotp_training.py index dcd90076d48f..32d3b1740e42 100644 --- a/tests/unit/model_parallelism/test_autotp_training.py +++ b/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