[jinyouzhi]

Analyze affine transfer batches before execution

Stacked on #8809 (executor) and #8966 (calibration path labels). Only the last commit is new.

Why

Plans that move the same bytes can differ a lot in time (up to 1.75× on 8× B70 between replica choices), and executor measurements show that the round count under a StagingBudget dominates sync time. A policy (#8935/#8936) or cost model (#8252) needs these facts before it runs a plan.

What

deepspeed/checkpoint/affine_transfer_analysis.py:

• analyze_transfers(plans, endpoints, element_size, budget, path_of) is a pure function with no communication. It replays the executor's build_operations and round packing and reports:
  • volume;
  • segment, chunk and message size distributions;
  • rounds;
  • per process: sent/received bytes, peers, messages, pack/unpack bytes, staging arena;
  • per-pair and per-path bytes.
• expected_stats(report, process) predicts that process's TransferStats.
• report_to_dict returns a JSON-ready summary.
• Messages use the same path_of labels as calibrate_links, so per-path totals can be priced with fit_links.

It does not choose replicas, rank plans, or model time.

Note: it imports the executor's private _packed_rounds / build_operations. @0z5a, would a public dry-run schedule API on the executor be preferable?

Tests

• Hand-counted rounds and bytes on an uneven handoff.
• Staging of strided sources.
• A batch with a shared plan.
• Path and pair labels in the JSON summary.
• 4 ranks: expected_stats == the executor's TransferStats, for a strided and a replicated parameter under budgets of 3 and 64 elements.

Results: 6 passed on 4× L20 (nccl) and on 4× Intel Arc Pro B70 (xccl). Mutation checks (dropping the receiver's round count, or the local bytes in the path summary) are caught.

[0z5a]

Read the last commit (72081c1b). The report is what a policy or cost model needs: per-process sent and received bytes, rounds and messages give the bandwidth and per-message terms, and expected_stats matching the executor's TransferStats on 4 ranks is a good anchor. It passes on CPU (gloo, 4 ranks) here.

I checked the schedule reconstruction against the executor at e38477d7. The direct-send test (_is_packed versus region.is_contiguous()), the per-peer, per-direction arena sized by the largest staged round, and rounds as the maximum over a process's peer queues all match today. Two structural suggestions:

• Share the schedule instead of re-deriving it. The module imports the executor's private _packed_rounds, and build_operations isn't in its __all__ either. It then re-implements the direct-versus-staged decision and the arena sizing. Those will drift when the executor's staging changes. staging_bytes in particular isn't covered by the expected_stats comparison, since TransferStats has no staging field. Would a public schedule helper in the executor work for you, returning each peer's rounds and whether each round is sent directly? Both sides could then call it.
• Ordering. This stacks on #8966 only for LOCAL and default_path_of. Defining those here, or in a small shared module that calibration imports, would let this land first.

Minor: the report numbers processes by transport-group rank, while calibration uses default-group ranks, so the path labels join only when the two groups coincide. The docstring says so, but recording that assumption in the report itself would stop a consumer from mixing the two.

[0z5a]

@jinyouzhi on the schedule question: yes, I'd rather the executor own it, and I can add it. #8809 is in master now, and the executor is unchanged from the e38477d7 copy this module mirrors, so the report still matches today.

What I have in mind is a public, tensor-free helper in affine_transfer_executor, say schedule_transfers(plans, endpoints, budget, element_size, process=None). For one process, or for all of them, it returns, per peer and direction, the packed rounds with their element counts, whether each round goes direct, and the arena size, plus the local copies. The direct decision needs only geometry: _region is an as_strided view of a contiguous flat shard, so region.is_contiguous() depends only on (chunk_shape, strides), the same rule as your _is_packed. execute_transfers would build its rounds and arenas from that result, and this module could drop _packed_rounds, build_operations and its own direct-send and arena logic, so staging_bytes can't drift from what the executor allocates. That's now #9027, and this one can switch to it once it lands.

Separately, this head is stacked on 17aa897e, before #8966's d00ba303. Once it's rebased, it needs the external-group test fix from my #8966 comment.

[jinyouzhi]

Rebased onto #8966's 57a47ff, which carries the test fix. I'll switch to schedule_transfers and drop the mirrored scheduling once #9027 lands.

[jinyouzhi]

@0z5a Switched to schedule_transfers (#9027 is merged into this branch until it lands): the module no longer imports _packed_rounds / build_operations or re-derives direct sends and arenas, so staging_bytes is the executor's arena. Reports are identical to the previous version on 300 random batches, except the per-chunk size distribution, which the schedule doesn't carry (the chunk count stays). The summary now records process_numbering, and the docs say to calibrate on the transfer's group. Since #8966 is close, I kept LOCAL / default_path_of there rather than moving them. Tests: 72 pass on 4×H20.