[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;
  • 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.
• 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.

What it does not change
plan_transfer behaves exactly as before; it now 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 Theorem 3 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;
  • the wrong number of answers;
  • splits under min_split_bytes, and more parts than the budget allows.
• The maps are validated before the policy runs.

With DS_ACCELERATOR=cpu, the new file passes together with test_affine_transfer_plan.py, test_autoep_affine.py and test_affine_shard_map.py (272 tests), 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 and the license check are clean.

How it fits with other PRs

• [AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885's affine_transfer.py changes apply cleanly on top of this branch (git apply --check).
• @jinyouzhi: this restructures plan_transfer into _prepare / _emit ahead of your affine_geometry.py move. I'm happy to rebase in whichever order suits you.
• Docs: the planner/executor page is affine-transfer.rst in [checkpoint] Execute affine transfer plans with bounded peer staging #8809. I'll add a plan_transfers section there once one of the two lands.

cc @Achyuthan-S @delock