Repository navigation
Conversation
0z5a
force-pushed
the
uc/v02-b1-transfer-plan
branch
from
September 24, 2026 06:03
f4e1e16 to
b8f0f3b
Compare
0z5a
force-pushed
the
uc/v02-b2-transfer-executor
branch
from
September 24, 2026 06:05
8fcd090 to
9cecb11
Compare
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
0z5a
force-pushed
the
uc/v02-b2-transfer-executor
branch
from
September 26, 2026 23:35
9cecb11 to
e7ff688
Compare
0z5a
marked this pull request as ready for review
September 27, 2026 04:10
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 10, 2026
… the parameter (deepspeedai#8623) ### Summary Compose source and target affine maps into shard-to-shard copy descriptors without rebuilding the logical parameter. This PR remains metadata only and based on `master`. Resolved copies now fold when both physical shard address mappings continue along an axis. Their descriptor no longer claims a single `logical_origin`: a copy can span two distinct logical regions. The planner keeps every validated replica candidate until a target intersection is resolved. It uses the lowest holder by default and accepts a caller-supplied `replica_selector` for other choices. Invalid selections fail explicitly. The spec distinguishes homogeneous input map pieces from resolved copies. The regression suite is committed in `tests/unit/checkpoint/test_affine_transfer_plan.py`. ### Review result | BigCode TP2→TP4, 48×8 | Previous plan | Updated plan | Target elements written | Result | |---|---:|---:|---:|---| | Lowest holder | 8 segments | 7 segments | 768 | Exact shard contents and guarded storage offsets | | Holder with the target's private Q rows | Unavailable | 6 segments | 768 | Exact shard contents and guarded storage offsets | | Highest-ranked KV holder | Unavailable | 7 segments | 768 | Exact shard contents and guarded storage offsets | The 7/6/7 counts apply to this layout and these three holder policies. For targets 0 and 1, the private Q rows are on source 0; for targets 2 and 3, they are on source 1. Selecting that same source for the replicated KV block makes the query and KV reads adjacent, allowing the sixth segment. Selecting the highest-ranked KV holder for every target does not. The previous 8-segment plan was valid; this change compresses its copy schedule. The selector does not change the input map's `locations` or relax its homogeneity rule. On the 8× RTX 5090 host, Python 3.12.13 and torch 2.13.0+cu130, the initial B1 CPU suite passed: **98 passed** (`PYTHONPATH=$PWD CUDA_VISIBLE_DEVICES="" DS_ACCELERATOR=cpu python -m pytest tests/unit/checkpoint/test_affine_transfer_plan.py -q`). Tests cover exact address oracles, nonzero storage offsets and guard regions, replica offsets, invalid selectors, serialization, budgets, determinism and the no-materialization contract. Follow-up at `c2ac3a4`: the targeted suite passed **119 tests** in the existing 0z5a Python environment using task-local PR modules. New cases cover stride rank, raw-piece and merge budgets, missing shard pieces, declared holders, complete target coverage, replica boundary mismatch, and serialized plan size. No environment packages were changed. The stacked [B2 executor PR](0z5a#1) has two-GPU copy and model E2E results, including the direct-versus-rebuild speed table. B1 itself does not move data, so it makes no transport speed claim. ### Scope Scale 1 and matching logical shapes only. No checkpoint format change, topology search, CUDA graph, online weight update or optimizer coordinate transform. Refs deepspeedai#8230 and deepspeedai#8252. ### CI follow-up (October 10) The failed `cpu-torch-latest / unit tests` run 36671037164 had three rollout test failures because its Transformers main revision requires the configuration's `per_layer_config` contract. The three test models now use the real `PretrainedConfig` rather than a partial `SimpleNamespace` subclass. Merged current master `967803b0fe1a2b4a2738fac676ded4a85d7ce491` and resolved conflicts with its concurrent cache-config fix, retaining all upstream rollout tests and behavior. The affine transfer planner and its regressions remain intact. Local CPU validation on macOS arm64, Python 3.12.13 and PyTorch 2.14.1: - Transformers 5.19.0.dev0 at the failing run's exact Git revision `5c9f079ca0651631cd8c9c10ebf8d6ac2d760ea1`: reproduced all three original failures before the fix. After the master merge, both affected rollout test files plus the affine transfer-plan suite report **157 passed, 1 skipped**. The skip requires native CUDA shared-prefill support. - Transformers 5.14.0: the same three files report **157 passed, 1 skipped** after the merge. - All applicable scoped pre-commit hooks and diff checks passed. The refreshed PR is mergeable, and DCO and test collection have passed. The full Linux CPU workflow is currently `action_required` with no jobs started ([run 38042624340](https://github.com/deepspeedai/DeepSpeed/actions/runs/38042624340)); maintainer approval is needed before execution. These local results do not claim a full CI or GPU run, or Transformers 4.x rollout compatibility. CI fixture fix: `dc765cac01e67a14cc501f00b9d6fef52b796533`; conflict refresh: `a2e60af442dd2207e92a8b6a0f90fc6d6f7af9a4`. --------- Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Execute the transfer plan from B1 #8623 with bounded staging and batched P2P. This PR remains stacked on
uc/v02-b1-transfer-plan.Tested parent:
75fe58c09c390e19a948a25aed2410a31616f538. The fork PR diff contains only the executor, itsdeepspeed.commentry point and its tests.The executor uses the physical-copy descriptor without
logical_origin, includes a schema version in its schedule digest, and posts P2P throughdeepspeed.comm.execute_transfersagrees on a batch schedule once and packs small copies into one message per peer and direction within the staging budget. Each request carries a stable parameter key, so ranks that order otherwise identical parameters differently reject the batch before data moves. The committed tests cover FP32/FP16/BF16 copies, strided source packing, multiple rounds, replicated sources, digest disagreement, collective refusal, multi-tensor packing and key-order mismatch.E2E and speed comparison
Two RTX 5090 GPUs on one host, FP32 and NCCL. Routes were synchronized on each GPU; each sample is the slower rank, and the table reports medians. Direct means plan once and execute shard-to-shard; rebuild means gather source shards, reconstruct the full parameter, then extract target shards. The same source and target layouts and exact-value oracle were used for each comparison.
The model is
optimum-internal-testing/tiny-random-gpt_bigcode-multi_query-Falseat revision7e0d5dc36afe26cb54a250c39cde5b517b17a7be; the downloaded weight file SHA-256 is9c0f7e2ae40319c9fd8d39a1d671cd68e50fc753dcc24a4e5c6b56d123732e9c. The model route generated 344 resolved copies. The tiny-model row is five paired samples per route on the same two GPUs; the 67.1 MB row is the earlier four-sample single-parameter run. The weight file was removed after validation.The six-segment policy selects source 0 for targets 0 and 1, and source 1 for targets 2 and 3, matching each target's private Q rows. Picking the highest KV holder for every target still needs seven segments. The committed two-GPU test executes all three plans against the target-map
extractoracle.Validation at the final branch code: B1 CPU tests 98 passed; B2 manual two-rank suite 12 cases per rank passed. The full 64-tensor model and 67.1 MB E2E timing rows were measured on the unchanged production transfer path before the review-only test and spec updates. The 8-rank disjoint-endpoint case remains unrun because those GPUs were in use by other processes.
Scope
Single node, one device per process, synchronous on return. No CUDA graph, overlap with compute, online rollout integration or topology-aware replica selection. Refs deepspeedai#8230 and deepspeedai#8252.