Repository navigation
Conversation
Add schedule_transfers, which returns what execute_transfers would run for a batch without buffers or communication: for each process, the rounds per peer and direction with their element and chunk counts and whether each is sent or received in place, the staging arena each peer direction allocates, the local copies, and the round count TransferStats reports. A replica policy or cost model can then price a plan before running it, without importing the executor's private helpers (deepspeedai#8967 asked for this). execute_transfers now builds its rounds and arenas from the same private computation, so the two cannot drift apart. The in-place decision is made from geometry, not from the buffer: _region is an as_strided view of a packed shard, so the view is contiguous exactly when (chunk_shape, strides) is row-major packed, extent-1 axes excepted. plan_digest, the batch fingerprint and TransferStats are unchanged, and build_operations and _packed_rounds keep their names and signatures. Tests: test_schedule_is_what_the_executor_runs (2 and 4 ranks, mixed packed and strided regions, 16-element messages) checks rounds, sent, received and local copy bytes against the executor's TransferStats, that both sides of every message agree, the arena sizes, and that the batch exercises both the in-place and the staged path. test_packed_regions_are_exactly_the_contiguous_views compares the geometric rule with Tensor.is_contiguous over 2000 random shapes and strides. A single-process case covers local-only batches and the argument checks. Forcing every round to staged, or dropping the extent-1 exception, fails these tests. On macOS arm64 CPU (Python 3.12, torch 2.15.0), the executor, external-group, planner and shard-map suites pass (263). All affine checkpoint tests pass except the four TestAffineUniversalCheckpointResume cases, which fail the same way on master here (uninstalled package version). pre-commit (yapf, flake8, codespell, license) is clean. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9ee3cc392d
ℹ️ 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".
| if not isinstance(element_size, int) or isinstance(element_size, bool) or element_size <= 0: | ||
| raise AffineTransferExecutionError('element_size must be a positive integer') |
There was a problem hiding this comment.
schedule_transfers([], endpoints, StagingBudget(0), 4) currently bypasses build_operations (where the only budget-size check lives) and returns schedules, while execute_transfers rejects that budget unconditionally at affine_transfer_executor.py:499-501; similarly, an integral-valued float budget can be accepted for an unsplit plan even though execution rejects non-integers. In these inputs a cost model can receive a supposedly executable schedule that the executor will refuse, so mirror the executor's integer and minimum-element validation before _indexed_operations.
Useful? React with 👍 / 👎.
| torch.rand((), generator=generator) < 0.7) else int(torch.randint(1, 9, (1, ), generator=generator)) | ||
| for stride in row_major) | ||
| view = torch.as_strided(storage, shape, strides, storage_offset=3) | ||
| assert _is_packed(shape, strides) == view.is_contiguous(), (shape, strides) |
There was a problem hiding this comment.
This test calls the newly added private _is_packed helper directly, so a correct executor refactor that removes or replaces the helper while preserving every public ScheduledRound.direct result will fail collection or this assertion. Cover the same geometry cases through schedule_transfers and assert the public round metadata instead, so the test protects the contract rather than an implementation detail.
AGENTS.md reference: AGENTS.md:L76-L79
Useful? React with 👍 / 👎.
|
Reviewed from the cost-model side, since this is the input the model reads. It has what the run term needs. Checked:
Two small things:
A nit: LGTM from the cost-model side. |
…e executor would refuse schedule_transfers now checks the staging budget the same way execute_transfers does, so an empty batch can no longer be scheduled under a budget the executor refuses, and refuses a process the endpoints don't name instead of returning an all-zero schedule. _is_packed treats a zero-extent region as packed, matching Tensor.is_contiguous. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Thanks for checking it from the cost-model side. Both points and the nit are in the new commit: |
|
Thanks @0z5a, I checked e22ab73. Merged with current master, the executor, external-group, planner, shard-map, ZeRO-affine and AutoEP-affine suites pass on CPU with 4 ranks (322 passed). The two updated tests fail against the previous head's executor, as expected. LGTM. With this in, a replica policy can read every count the cost model's run term needs from |
…ysis Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
The analyzer rebuilt the executor's rounds with the private _packed_rounds and build_operations, and re-derived the direct-send rule and arena sizes, so staging_bytes could drift from what execute_transfers allocates. It now reads rounds, direct flags, arenas and local copies from deepspeedai#9027's public schedule_transfers. Per-chunk sizes are not part of that schedule, so the report keeps the chunk count and drops the chunk-size distribution. The summary also records process_numbering: processes are the transport-group ranks TransferEndpoints binds, which join with a calibration only when it measured the same group. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
gpu-test-labeler
Bot
added
the
skip-gpu-test
Why
A replica policy (#8936) or cost model (#8252) needs to know what a batch will cost before running it: messages per peer, rounds, and staging memory. #8967 builds that report today by importing the executor's private
_packed_roundsandbuild_operations. It then re-derives the in-place decision and the arena sizing, so its numbers would drift silently whenever the executor's staging changes. It asked for a public API instead.What
schedule_transfers(plans, endpoints, budget, element_size, process=None)inaffine_transfer_executorreturns, for one process or for every process the endpoints name, theProcessSchedulethatexecute_transferswould run. It uses no buffers and no communication:ScheduledRound: elements, chunks, and whether the round is sent or received in place);TransferStats.roundsreports.execute_transfersnow builds its rounds and arenas from the same private computation, so the report and the run cannot disagree. The in-place decision is made from geometry rather than from the buffer._regionis anas_stridedview of a packed shard, so it is contiguous exactly when(chunk_shape, strides)is row-major packed, ignoring extent-1 axes.Unchanged:
plan_digest, the batch fingerprint andTransferStats, so mixed versions still agree on the schedule;build_operationsand_packed_rounds, which keep their names and signatures until [AffineIR] Analyze affine transfer batches before execution (Plan analysis & calibration Part-1) #8967 switches.Tests
test_schedule_is_what_the_executor_runs(2 and 4 ranks; mixed packed and strided regions with 16-element messages) checks:TransferStats;test_packed_regions_are_exactly_the_contiguous_viewscompares the geometric rule withTensor.is_contiguousover 2000 random shapes and strides, including extent-1 axes.On macOS arm64 CPU (Python 3.12, torch 2.15.0):
TestAffineUniversalCheckpointResumecases, which fail the same way on master in this environment (uninstalled package version).This merges cleanly with #8936, #9010 and #9011. With #8966 and #8967 the only conflict is that both append a section to
affine-transfer.rst.