Skip to content

[UC] Execute a shard-to-shard transfer within a staging budget - #1

Open
0z5a wants to merge 1 commit into
uc/v02-b1-transfer-planfrom
uc/v02-b2-transfer-executor
Open

0z5a wants to merge 1 commit into
uc/v02-b1-transfer-planfrom
uc/v02-b2-transfer-executor

Conversation

@0z5a

@0z5a 0z5a commented Sep 22, 2026 •

Copy link
Copy Markdown
Owner

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, its deepspeed.comm entry 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 through deepspeed.comm. execute_transfers agrees 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.

Workload Rebuild→extract Per-tensor direct Batched direct Batched vs rebuild Correctness
GPTBigCode tiny checkpoint, 64 tensors / 83,161 elements 15.512 ms 62.728 ms 11.440 ms 1.36×; 5.48× vs per-tensor All target shards exact; model output max absolute difference 0
4096×4096 FP32 parameter, 67.1 MB 3.242 ms 2.300 ms Not measured Per-tensor direct: 1.41× All target shards exact

The model is optimum-internal-testing/tiny-random-gpt_bigcode-multi_query-False at revision 7e0d5dc36afe26cb54a250c39cde5b517b17a7be; the downloaded weight file SHA-256 is 9c0f7e2ae40319c9fd8d39a1d671cd68e50fc753dcc24a4e5c6b56d123732e9c. 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.

BigCode transfer policy Segments Written elements Two-GPU copy
Default lowest holder 7 768 Exact, including guard regions
Holder with the target's private Q rows 6 768 Exact, including guard regions
Highest-ranked KV holder 7 768 Exact target shards

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 extract oracle.

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.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a
0z5a force-pushed the uc/v02-b2-transfer-executor branch from 9cecb11 to e7ff688 Compare September 26, 2026 23:35
@0z5a
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant