Add configurable dtype for ZeRO checkpoint export - #8318
Conversation
Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>
![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: f38e9a01b9
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -0,0 +1,49 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
When this script is copied into a checkpoint, _change_recovery_script_permissions makes it executable and the documentation tells users to invoke it as ./zero_to_torch.py, but the file has no shebang. Executing such a copied script directly causes the shell to interpret the Python source and fail immediately; add a Python shebang as used by zero_to_fp32.py.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7ffebd1. The recovery entry point now starts with a Python shebang. The standalone regression also verifies the copied script retains the shebang.
| @@ -0,0 +1,49 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
This is a single-parent, non-merge commit, but its commit message contains no Signed-off-by: trailer, despite stating that the sign-off is present. Add the author sign-off so the commit satisfies the repository's DCO requirement.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The original commit f38e9a0 already contains Signed-off-by: gaoxiaomo 165135449+gaoxiaomo@users.noreply.github.com, and the repository DCO check is passing. The follow-up commit 7ffebd1 is signed off as well.
|
|
||
| import argparse | ||
|
|
||
| import deepspeed.utils.zero_to_fp32 as zero_to_fp32 |
There was a problem hiding this comment.
When the checkpoint is recovered under a different installed DeepSpeed version, this qualified import resolves deepspeed.utils.zero_to_fp32 from that installation rather than the zero_to_fp32.py copied beside this script. If the installed version predates this change, argument-parser construction fails because OUTPUT_DTYPE_NAMES is absent; other version mismatches silently run converter code unrelated to the saved checkpoint. The recovery entry point should load its accompanying script instead.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7ffebd1. Package imports now use the relative sibling module, while direct execution imports zero_to_fp32.py from the copied script directory. The new regression runs a copied zero_to_torch.py beside a stub converter and verifies that the local converter is invoked.
Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>
1 participant
What this changes
Closes #4032.
This adds lower-precision ZeRO checkpoint export without changing the existing
zero_to_fp32.pybehavior:zero_to_torch.pywith a required--dtypeoption for float32, float16, or bfloat16 output.convert_zero_checkpoint_to_state_dict(...)for the same functionality from Python.convert_zero_checkpoint_to_fp32_state_dict(...)as a backward-compatible fp32 wrapper.Example:
./zero_to_torch.py . checkpoint-bf16 --dtype bfloat16Validation
Focused unit tests cover dtype validation, FP16/BF16 conversion, shared tensors, saved checkpoint files, and CLI argument forwarding:
I also ran an end-to-end ZeRO-3 conversion with world size 2 on two RTX 4090 GPUs. The test saved a real distributed checkpoint, reconstructed both formats, loaded both output files, and checked every tensor.
Repository checks:
git diff --checkpassed.AI assisted with implementation suggestions and test preparation. I reviewed and understood every change and accept responsibility for maintenance and reviewer follow-up.