[0z5a]

Implements the planner side of #8935. It does not close the issue: the cost-aware policies are the policy layer under #8252.

What it adds

• plan_transfers(pairs, limits, *, source_processes, target_processes, element_size, min_split_bytes, replica_policy) in deepspeed/checkpoint/affine_transfer.py. It plans every parameter of one sync and calls one replica policy with every request at once.
• ReplicaRequest and ReplicaPart.
  • A request names the parameter, the target location, the region and its validated holders. When the caller passes the location→process mappings, it also names the processes behind them.
  • Regions with a single holder are included, so a policy sees load it cannot move.
  • Requests arrive in a fixed order: parameter, then target location, then origin.
• An answer is either a holder location, or ReplicaParts that cut the region along axis 0.
• The planner checks every answer before it emits any segment, and raises InvalidReplicaPolicyError otherwise:
  • one answer per request;
  • only validated holders;
  • parts that tile the request exactly along axis 0;
  • part fields that are not iterable coordinates, rather than a raw TypeError;
  • with min_split_bytes, every part of a split is at least that large;
  • the parts of one parameter stay within PlanBudget.max_pair_checks, checked from the answer's length before the parts are expanded.
• deepspeed/checkpoint/affine_replica_policy.py holds three policies:
  • first_holder: what plan_transfer does by default.
  • local_first_least_loaded: the plan_transfers default. It takes a holder in the target's own process first, then the sending process with the fewest bytes assigned so far in this batch. Ties go to the lowest location, and it never splits.
  • per_request(selector): runs an existing replica_selector as a batch policy.

Review follow-up (commit 773ddfd3)

• Non-iterable ReplicaPart.logical_origin/shape fields are refused as InvalidReplicaPolicyError; the pair budget now uses that error too, so every refused answer uses the documented class.
• An over-budget answer is refused on its length before _resolve expands it.
• plan_transfers, the request/answer protocol, the process mappings, element_size (one entry per parameter) and the split constraints are documented in docs/code-docs/source/affine-transfer.rst, the page from [checkpoint] Execute affine transfer plans with bounded peer staging #8809.
• Policy authors are told that every rank must return the same answers, so cost inputs must come from data every rank shares (for example fit_links on the calibration from [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) #8966), never from a rank-local measurement; the executor's schedule digest refuses a mismatch.
• The Define affine IR for checkpointing #8252 correctness theorem is cited as "Any consistent plan is exact" instead of by number.

What it does not change
plan_transfer behaves exactly as before; it shares its validation and emission steps with plan_transfers. AutoEP's plan_expert_transfer and its docs are untouched.

Decisions from #8935

• One sync per batch.
• Only axis-0 splits for now. Other split shapes are refused, so widening later is an additive change.
• Plain location→process mappings rather than TransferEndpoints, so the planner does not import the executor.
• Policies live in their own module.
• The per-request adapter stays.

min_split_bytes comes from my last comment on #8935 and has not been agreed yet. It is optional and off by default. The process_node mapping proposed there is deferred.

Tests
tests/unit/checkpoint/test_affine_replica_policy.py, CPU only:

• Random policies over 24 seeds, including random splits across holders, on five parameters: whole replicas, a partly replicated fused block, a vector and a scalar. Every plan writes each target element exactly once and matches extract. This is the check that matches the "Any consistent plan is exact" theorem of the correctness draft for Define affine IR for checkpointing #8252.
• first_holder and per_request reproduce plan_transfer's plans exactly.
• The default policy:
  • copies locally when it can;
  • balances load, counting single-holder load;
  • never assumes a local copy without mappings.
• Refused answers:
  • holders that are not candidates;
  • bools, floats, None and strings;
  • gaps, overlaps and parts outside the region;
  • cuts along axis 1, empty parts and parts with the wrong number of axes;
  • non-iterable part fields;
  • the wrong number of answers;
  • splits under min_split_bytes, and more parts than the budget allows, including parts accumulated across requests.
• The maps are validated before the policy runs.
• The pair budget is decided from the answer's length before the planner reads its parts.

With DS_ACCELERATOR=cpu on macOS arm64 (Python 3.11, torch 2.15.0), the new file passes (65 tests) together with test_affine_transfer_plan.py, test_autoep_affine.py and test_affine_shard_map.py (306 passed), and tests/unit/v1/moe/test_autoep_unit.py (which covers plan_expert_transfer) together with test_autoep_affine_universal.py pass too (215 passed, 8 skipped). yapf 0.43.0, flake8 5.0.4, codespell, the license check and the other pre-commit hooks are clean. The five rejection/budget tests added here fail on the previous planner and pass now.

How it fits with other PRs

• The branch is rebased on current master (after [checkpoint] Execute affine transfer plans with bounded peer staging #8809, [AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885 and Check that converting through an affine map writes every element once #8975); plan_transfer still shares _prepare / _emit with plan_transfers.
• @jinyouzhi: the affine_geometry.py move can land before or after this; I'll rebase in whichever order suits you.
• [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) #8966 supplies the link calibration (fit_links) that cost-aware policies will read; the planner does not consume it.

Scope, compatibility and AI assistance
Only the CPU planning interface, its tests and its documentation change: no training-loop, GPU, ZeRO, offload or precision path is touched, so no compatibility-matrix cell or on-device serving check is affected, and none was run. AI assistance was used: the initial implementation was co-authored with Claude Opus 5.5, this review follow-up was prepared with an AI coding agent, and the submitter reviewed every changed line and ran the tests above.

cc @Achyuthan-S @delock

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-05T09:59:49.703718Z	24881a4	Draft marked ready

ℹ️ 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" or "@codex security review".

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]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 24881a444e

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Document the new batch planning API

For downstream users adopting plan_transfers, this patch exposes the entry point in __all__ plus a new callback/split protocol and default selection behavior, but adds no user documentation; the existing AutoEP docs only explain plan_expert_transfer and the old per-request replica_selector, so there is no supported reference for the process mappings, element sizes, split constraints, or policy answers. Add the required documentation alongside the feature.

AGENTS.md reference: AGENTS.md:L71-L71

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Validate part fields before tuple conversion

When a custom policy returns a ReplicaPart whose logical_origin or shape is non-iterable (for example None), these tuple(...) calls raise a raw TypeError. The function's new contract says every invalid answer raises InvalidReplicaPolicyError; callers that catch AffineTransferError to reject a policy or fall back will instead abort, so convert and validate these fields under the policy-error path.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Enforce split budgets before materializing parts

When a custom policy accidentally returns a very large split for one request, _resolve materializes every LogicalBox and sorts the full answer before this cumulative check runs. Consequently even a tiny max_pair_checks does not bound planner work or memory, so a malformed policy can cause a large allocation or OOM before it is rejected; enforce the remaining per-parameter allowance while or before resolving the answer.

Useful? React with 👍 / 👎.

[Achyuthan-S]

I reviewed this from the policy side, since the cost-aware policy will sit on this interface. It has what that policy needs:

• Each request carries its bytes, the target process and every candidate's process.
• The whole sync arrives in one call, in a fixed order, so a policy can balance load over all of it, including the single-holder load it can't move.

The topology and the calibration (#8966's fit_links) can live inside the policy, so I agree the process_node mapping can stay out of the planner. The refactor keeps plan_transfer's selection and emission as they were, and local_first_least_loaded does what its docstring says.

Checked on CPU:

• At 24881a4: test_affine_replica_policy.py, test_affine_transfer_plan.py and test_affine_shard_map.py pass (251).
• Merged with current master (after #8809, #8885 and #8975), the merge is clean, and the planner, executor, AutoEP and ZeRO-affine tests pass together (583 passed, 8 skipped).

A few small things:

1. Codex's first P2 is real: a ReplicaPart with logical_origin=None raises a raw TypeError (I checked), not InvalidReplicaPolicyError. The second is minor: the policy's answer is already in memory by then, so checking len(answer) against the remaining max_pair_checks allowance before _resolve only saves the planner's own work on a malformed split. It's still cheap to add.
2. Docs: #8809 has landed, so the plan_transfers section can go into affine-transfer.rst now.
3. A note for policy authors: every rank has to return the same answers. #8809's schedule digest already refuses mismatched plans ("ranks disagree about the transfer schedule"), so this fails loudly. Still, one line in the docstring saying that cost inputs should come from data every rank shares, such as fit_links on the gathered calibration, would help.
4. The test docstring cites "Theorem 3 of the proof draft". In the PDF I posted it's Theorem 4, "Any consistent plan is exact". Citing it by name would survive renumbering.

LGTM from the policy side. The first part of 1 is worth doing before merge, and the rest is up to you.

[0z5a]

@Achyuthan-S thanks — all four are in 773ddfd3 (branch rebased on current master, after #8809, #8885 and #8975):

1. A ReplicaPart whose logical_origin or shape is not iterable now raises InvalidReplicaPolicyError; the pair budget uses the same class, so every refused answer uses the documented one. The budget is decided from len(answer) before _resolve expands the parts, and it still accumulates across requests. The tests pin both (test_the_pair_budget_is_checked_before_the_parts_are_read, test_parts_count_against_the_budget_across_requests); the five rejection/budget cases fail on the previous planner.
2. plan_transfers is documented in docs/code-docs/source/affine-transfer.rst: the request/answer protocol, source_processes/target_processes, element_size (one entry per parameter), the axis-0 split and min_split_bytes constraints, the policy error, and the three built-in policies.
3. The plan_transfers and affine_replica_policy docstrings now say answers must agree on every rank, so cost inputs must come from data every rank shares (for example fit_links on the gathered calibration, [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) #8966), and quote the executor's "ranks disagree about the transfer schedule (batch schedule)".
4. The test docstring cites "Any consistent plan is exact" by name.

Validation on macOS arm64 CPU (Python 3.11, torch 2.15.0): 65 tests in the new file, 306 across the four affine checkpoint suites, and 215 passed / 8 skipped for the AutoEP suites; pre-commit (yapf, flake8, codespell, license) is clean. No GPU path is touched.

[Achyuthan-S]

Re-checked 773ddfd: all four are in.

• A ReplicaPart with a non-iterable logical_origin or shape now raises InvalidReplicaPolicyError, since the tuple(...) conversion is guarded. The pair budget is decided from len(answer) before _resolve, and the budget tests pass.
• plan_transfers is documented in affine-transfer.rst, the docstrings say answers must agree on every rank, and the test cites "Any consistent plan is exact" by name.

I merged the branch with current master: the merge is clean, and the planner, replica-policy, shard-map, ZeRO-affine, AutoEP-affine and executor suites (including the external-group test) pass together on CPU, with 383 passed. LGTM from the policy side. The cost-aware policy will build directly on this interface.