Conversation
Signed-off-by: Mingyang Wu <mingyangwu@cuhk.edu.hk>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![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: f2ea14b93d
ℹ️ 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".
| i: transformer_layers[i * layers_per_pp:(i + 1) * layers_per_pp] | ||
| for i in range(0, self.pp_degree) | ||
| } | ||
| layers_per_pp, remainder = divmod(len(transformer_layers), self.pp_degree) |
There was a problem hiding this comment.
This is a non-merge commit, but its message has no Signed-off-by trailer, so it violates the repository's commit requirements and may be rejected by the DCO/CI checks. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
Fixes #3588.
For 70 transformer layers and 12 pipeline stages,
DeepSpeedCheckpointcurrently exposes only 60 layers throughget_pp_transformer_map, while its separate file mapping creates 14 stages. This can omit layers when callers iterate over the configured pipeline ranks. With fewer layers than stages, file mapping also divides by zero.Partition all transformer layers into the requested number of contiguous, balanced stages, then use that single partition to build the checkpoint file map. Divisible layouts retain their existing assignment, and empty stages return an empty state list. Tensor-parallel shard partitioning and merging are unchanged.
Validation
Added a parametrized regression using small real checkpoint files and public checkpoint loading methods. It covers the reported 70-layer/12-stage layout, uneven pipeline reduction, divisible partitions, empty stages, zero transformer layers, a single stage, a custom final-norm position, and both retained and merged tensor-parallel shards.
git diff --check: passed.Executed on macOS arm64 CPU with Python 3.10.18 and PyTorch 2.8.0:
PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ DS_ACCELERATOR=cpu OMP_NUM_THREADS=1 PYTHONPATH="$PWD:$PWD/tests" \ python -m pytest tests/unit/checkpoint/test_reshape_checkpoint.py \ --torch_ver=2.8.0 -q --tb=shortNo GPU training or compiled-op tests were run. The tests exercise checkpoint discovery, layer assignment and loading, rather than training convergence.
This contribution was prepared with AI assistance. The reported regression and formatting checks were executed locally.