DeepSpeed

ae1b183f - [UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter (#8623)

[UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter (#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](https://github.com/0z5a/DeepSpeed/pull/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 #8230 and #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>