[AutoTP] Allow ZeRO stage 3 inference with tensor parallelism#8167
[AutoTP] Allow ZeRO stage 3 inference with tensor parallelism#8167delock wants to merge 3 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cf7ae19e99
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.mpu = groups | ||
| self.mpu._init_tp_mesh_device(tensor_model_parallel_size=self.autotp_size()) |
There was a problem hiding this comment.
For autotp_size > 1 with ZeRO stage 3, users can now enable a configuration where 16-bit export still drops TP shards: _consolidated_16bit_state_dict() checks zero_optimization_stage() == ZeroStageEnum.weights before the AutoTP branch, so save_16bit_model() when stage3_gather_16bit_weights_on_model_save is enabled gathers only ZeRO data-parallel partitions and rank 0 writes its local TP shard rather than the full weight. The removed guard was the only thing preventing silently truncated exports; please either keep blocking this combination or make the ZeRO-3 consolidation path also gather the AutoTP shards.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This comments had been addressed by allow autotp only in zero3-inference mode. If there is an optimizer, the assertion will block it. The compatibility between autotp and checkpoint saving should be tracked with a seperate issue and fixed in a seperate PR.
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>
cf7ae19 to
2e4247e
Compare
delock
changed the title
| # 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). |
There was a problem hiding this comment.
I think the comment can be removed since it is a duplicate of the error message.
Signed-off-by: Guokai Ma <guokai.ma@intel.com>
| 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.") |
There was a problem hiding this comment.
So this means we cannot apply AutoTP + ZeRO stage 3 to student models?
There was a problem hiding this comment.
Yes, student model need to have checkpoint saved, and zero stage 3 checkpoint saving does not support AutoTP yet. This would be taken care of in another PR (still work in progress).
|
|
||
| ### 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. |
There was a problem hiding this comment.
I feel this description is a little confusing. Shall we say ZeRO stage 3 is supported only for inference in the context of on-policy distillation?
There was a problem hiding this comment.
Zero inference can also support other scenario where model is too large to fit in GPU memory, so not only for OPD.
3 participants
Replace the unconditional zero_optimization_stage() <= 2 assert in _configure_tensor_parallel_states with a guard that only blocks AutoTP + ZeRO-3 when an optimizer is present. The ZeRO-Inference path (no optimizer → DummyOptim → DeepSpeedZeRoOffload) is now permitted; the training path (optimizer → DeepSpeedZeroOptimizer_Stage3) raises NotImplementedError.
AutoTP with ZeRO-3 checkpoint saving should be tracked seperatly in a seperate PR.