Repository navigation
Conversation
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. |
|
This looks like a solid, focused change. |
![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: 35f904af97
ℹ️ 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".
| local_rank = int(os.environ.get('LOCAL_RANK', dist.get_rank() if dist.is_initialized() else 0)) | ||
| device = local_rank % max(get_accelerator().device_count(), 1) | ||
| get_accelerator().set_device(device) | ||
| return torch.device('cuda', device) |
There was a problem hiding this comment.
When cpu-torch-latest runs this newly added file (the workflow invokes all of unit/ in .github/workflows/cpu-torch-latest.yml), DistributedTest initializes CPU/Gloo because CPU_Accelerator.is_available() is true, so these classes are not skipped; the first _full() then allocates on this unconditional CUDA device and fails with the CPU-only PyTorch build. Mark the distributed classes CUDA-only before launch while leaving the *CPU classes runnable.
Useful? React with 👍 / 👎.
| if get_accelerator().device_count() < self.world_size: | ||
| pytest.skip(f'needs {self.world_size} visible accelerators, found {get_accelerator().device_count()}') |
There was a problem hiding this comment.
On CUDA hosts exposing fewer than eight devices, this check runs too late to skip: DistributedTest first spawns world_size = 8, and DistributedTest._dist_run calls set_device(local_rank) before invoking the test method. Ranks beyond the available device count therefore fail during initialization, so ordinary 2/4-GPU full-suite runs fail or hang instead of reaching this pytest.skip; move the resource check to collection or launch time.
Useful? React with 👍 / 👎.
| return cdb.new_group(ranks) | ||
| if timeout is None: | ||
| return cdb.new_group(ranks) | ||
| return cdb.new_group(ranks, timeout=timeout) |
There was a problem hiding this comment.
When oneCCL is available, init_distributed can install CCLBackend; calling the newly advertised dist.new_group(..., timeout=...) reaches this keyword call, but deepspeed/comm/ccl.py still defines new_group(self, ranks) only, so every timeout-bearing subgroup creation raises TypeError instead of creating a group. Extend that sibling override (and the base interface) or explicitly reject unsupported timeouts.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
| if not segment.shape or len(segment.shape) != ndim: | ||
| raise AffineTransferExecutionError('segment dimensionality does not match its logical shape') |
There was a problem hiding this comment.
When source and target maps describe a 0-D scalar (for example replicated_map((), 1)), plan_transfer legitimately emits a one-element TransferSegment whose shape is (), but this truthiness check rejects it before either a local or P2P copy. Consequently the generic executor cannot execute valid planner output for scalar parameters; validate dimensionality with len(segment.shape) != ndim without treating a zero-dimensional shape as an empty segment.
Useful? React with 👍 / 👎.
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
35f904a to
4a4bee9
Compare
| with a non-writable stock ``loss_function`` raises an actionable error rather | ||
| than being silently left unmodified. | ||
|
|
||
| Executing affine transfers |
There was a problem hiding this comment.
The doc (including the part above) is in autoEP but should be able to work for autoTP as well. Maybe we need a better place. Please propose document placement in slack channel and modify accordingly.
There was a problem hiding this comment.
Agreed, none of this is AutoEP-specific. AutoTP layers publish the same affine maps, and both the planner and the executor only see maps and buffers.
My proposal for the placement is a new Affine Shard Transfers page (docs/code-docs/source/affine-transfer.rst) in the Checkpointing API section, next to Model Checkpointing. This push moves the section there, with a new subsection on transport and staging. As a result this PR no longer touches autoep.rst. I'll also raise the placement in the Slack channel as you suggested, and if the discussion settles on a different place, I'll move it again.
The Chunked causal-LM loss text above this line was added by #8653, not by this PR. It isn't AutoEP-specific either, but I'd rather not mix it into this PR. I can move it in a separate docs PR if its owners agree.
| @pytest.mark.skipif(get_accelerator().device_name() != 'cuda' or get_accelerator().device_count() < 8, | ||
| reason='requires eight visible CUDA devices before distributed worker launch') | ||
| class TestB2DisjointEndpoints(DistributedTest): | ||
| world_size = 8 |
There was a problem hiding this comment.
we don't have runner with 8 ranks. There are either 2 GPU cards on modal runner or 4 ranks with cpu-torch-latest-multirank nightly. Suggest not confine the test to cuda and allow run this test in CPU multirank nightly instead.
There was a problem hiding this comment.
Done. The multi-rank classes no longer check for CUDA. When there are too few devices, the harness's own device-count gate skips them before any worker launches.
The disjoint-endpoint case is now parametrized over 2 and 4 ranks, with readers on the lower half of the ranks and writers on the upper half. With four ranks, each writer receives from two readers. The [2] case fits the two-GPU runners, and both cases run in the CPU multi-rank nightly. The test also asserts that every rank's sent and received bytes equal the source and target bytes it holds, with no local copies, so every element crosses the transport exactly once.
Run locally on CPU/gloo (macOS arm64, Python 3.12.13, torch 2.14.1), with this file on top of current master:
| Environment | Result |
|---|---|
LOCAL_SIZE=4 DS_MULTIRANK_ONLY=1 DS_DISABLE_REUSE_DIST_ENV=1, --forked (nightly shape) |
11 multi-rank cases selected, 11 passed |
LOCAL_SIZE=2, --forked |
48 passed, the 4-rank case skipped before launch |
no LOCAL_SIZE, --forked (cpu-torch-latest shape) |
38 passed, 11 multi-rank cases skipped before launch |
| if operations: | ||
| for handle in dist.batch_isend_irecv(operations): | ||
| block_stream = getattr(handle, 'block_current_stream', None) | ||
| if device.type == 'cuda' and callable(block_stream): |
There was a problem hiding this comment.
Check is_synchronized_device instead of check if device type is cuda. We should check device capability instead of device type.
There was a problem hiding this comment.
Done. A single predicate, stream_ordered = initialized and not get_accelerator().is_synchronized_device(), now gates both the per-round handle path and the final current_stream().synchronize(). The final synchronize was a device-type check as well. The initialized term is there because without a process group the buffers are host tensors. On CUDA this selects the same branches as before, so the H100 results in the description still describe this code. The CPU multi-rank runs now cover the synchronized branch: forcing the other branch there makes them fail.
One caveat for other accelerators. Work.block_current_stream() orders the device stream only on backends that override it. WorkNCCL and WorkXCCL both implement it as synchronize(); I checked this at the torch 2.12.1 and 2.14.1 tags and the torch-xpu-ops commits they pin. The base Work implementation instead uses a CUDA stream baton. Forcing this branch on CPU/gloo shows what happens on a non-synchronized accelerator whose backend does not override the method: the two-rank cases fail at the first round with RuntimeError: Failed to create StreamBlock, and the four-rank case stalled instead of failing. I did not find such an override for HCCL or CNCL. This PR has not been validated on those devices, and before this change they took the wait() path.
|
Transport assumption worth documenting: staging is not inherent to affine transfer The executor hands the transport one contiguous tensor per peer per round, so strided regions are packed through an arena. A future GPUDirect-RDMA path would instead want the per-round |
|
Hi @0z5a , thanks for your contribution. I have left my comments. Thanks! |
Address review on the affine transfer executor: - Decide stream-ordered completion with is_synchronized_device() instead of the device type, for both the per-round handle and the final stream synchronization. - Run the multi-rank executor tests on any accelerator. The disjoint-endpoint case now uses two or four ranks, so it fits the two-device PR runners and the four-rank CPU nightly. - Move the executor documentation from the AutoEP page to a new Checkpointing API page, and record that the staging arena and StagingBudget belong to the current torch point-to-point transport. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Thanks @delock. Pushed ba9e8b6, which addresses all four points. Details are in the inline threads.
No new GPU run for this push. The CUDA branches are unchanged, and the H100 rows in the description are from the previous head, which is marked there. |
|
@0z5a the modified code looks good to me. Let me know when you get the document placement discussion settle on slack. |
autotp_rollout_e2e.py trains Qwen3-style dense models with AutoTP + ZeRO and moves the weights into rollout engines with a different TP size through a Universal checkpoint, through gathered parameters, and through ZeRO storage affine maps (plan_transfer + execute_transfers), checking every path bit for bit and recording time and peak memory per sync. A --replica_policy chooses which ZeRO-0/1/2 replica each target block is read from; the default reads locally when possible and otherwise from the least-loaded holder. Needs DeepSpeed with deepspeedai/DeepSpeed#8809 and deepspeedai/DeepSpeed#8885. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…pspeedai#8800) ## Summary deepspeedai#8623 plans shard-to-shard copies between two `ParamAffineMap`s. Today, though, the only AutoEP maps come from checkpoint metadata, so nothing at runtime can feed the planner. This PR closes that gap: it builds the source map from the placement the live `AutoEPMoELayer` actually holds, and builds the target map from a target placement descriptor. Both maps then go straight to `plan_transfer`. ``` AutoEP runtime placement -> source / target ParamAffineMap -> plan_transfer (deepspeedai#8623) -> copy descriptors ``` Maps of what each process actually stores under expert data parallelism (EDP) and ZeRO are left to a follow-up, which builds them by nesting these layout maps with ZeRO's storage maps. Related: deepspeedai#8252 (§8.4 AutoEP expert placement), deepspeedai#8230, deepspeedai#8544. ## Changes - `AutoEPMoELayer.expert_affine_maps()` returns the affine maps of the live `w1`/`w2`/`w3`. - Map ranks are EP-local ranks. The docstring explains how to translate them with `dist.get_global_rank(self.ep_group, rank)`. - Each shard is the unpartitioned packed `[local_experts, ...]` tensor. ZeRO-3/EDP fragments stay outside the map, and the shape comes from `ds_shape`. - No collectives. - `AutoEPMoELayer.plan_expert_transfer(target_placement, limits=None, *, replica_selector=None)` returns one `TransferPlan` per expert parameter. - It only plans: it moves no data, swaps no parameters and migrates no optimizer state. - `autoep_affine.plan_autoep_placement_transfer(source_placement, target_placement, logical_shape, ...)` is the layer-independent form. Both placements go through the existing `autoep_placement_to_affine_map`, so no second placement→map conversion is introduced. - Placements are compared by explicit rank IDs rather than list order (`autoep_placements_equal`), and each parameter's published `DS_AUTOEP_UC_META` placement and affine map must equal the live lowering. - The live placement must match how `repack_expert_weights` actually packed the weights; today that is contiguous and uniform. A descriptor that disagrees raises `NotImplementedError`, rather than producing a map that addresses shards the weights do not have. - A placement the descriptor can express is not necessarily one the runtime can execute. The target placement is only used for planning. ## Not in scope - Executing plans or transporting data (deepspeedai#8809), parameter swap, optimizer-state migration. - Storage maps of EDP replicas and ZeRO-3 fragments (follow-up PR), and replica selection policy (deepspeedai#8252). - Placement generation/versioning: the runtime does not change placement after construction yet. - Runtime support for non-uniform placements. ## Tests `tests/unit/v1/moe/test_autoep_unit.py`: - `test_live_expert_maps_address_the_packed_runtime_weights`: for each rank, the map's shard shape and contents match the real packed `w1/w2/w3`. The reference is the full tensor from `repack_expert_weights(ep_size=1)`, which does not depend on any affine map. - `test_plan_expert_transfer_moves_live_shards_into_target_packed_order`: from EP=2 into 3 ranks with reordered, uneven, empty and replicated placements; executing the plan reproduces the target shards, every element is written once, and planning never calls `rebuild`/`extract`. - `test_live_expert_maps_reject_a_placement_the_weights_do_not_have`, `test_live_expert_maps_identify_placement_ranks_by_id`, `test_live_expert_maps_reject_a_published_map_the_placement_does_not_produce`. `tests/unit/checkpoint/test_autoep_affine.py`: `test_placements_equal_identifies_ranks_by_id`. Results (Intel Arc Pro B70): - `test_autoep_unit.py`, `test_autoep_affine.py`, `test_autoep_affine_universal.py`, `test_affine_transfer_plan.py`, `test_affine_shard_map.py`, `test_autoep_zero3_metadata.py`: 432 passed, 6 skipped. - `test_autoep_checkpoint.py` (8× B70): 32 passed, 2 skipped. ## Contribution checklist **Related work.** - This PR builds on deepspeedai#8623 (`plan_transfer`). - Plan execution is deepspeedai#8809; replica policy is deepspeedai#8252. - deepspeedai#8844 and deepspeedai#8848 change Universal Checkpointing for AutoEP. They do not produce runtime maps. A search for AutoEP placement and transfer and for affine maps found no overlapping PR. **Commands and results.** Hardware and software: Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu. ``` pytest -p no:randomly tests/unit/v1/moe/test_autoep_unit.py tests/unit/checkpoint/test_autoep_affine.py \ tests/unit/checkpoint/test_autoep_affine_universal.py tests/unit/checkpoint/test_affine_transfer_plan.py \ tests/unit/checkpoint/test_affine_shard_map.py tests/unit/checkpoint/test_autoep_zero3_metadata.py pytest -p no:randomly tests/unit/v1/moe/test_autoep_checkpoint.py ``` Results are listed under Tests above. **Numerics, performance, memory.** Training is not affected. The new methods only run when called and use no collectives. **Compatibility matrix.** | | Validated | Not validated | |---|---|---| | ZeRO | Independent: maps describe the unpartitioned packed tensors | Storage under EDP/ZeRO (follow-up) | | Parallelism | EP=2; target placements on 3 ranks | | | Precision | Dtype-independent: maps describe shapes and offsets | | | Accelerator | XPU | CUDA, HPU | **AI assistance.** This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
RL frameworks build their trainer <-> engine weight-update group
outside the default world: each process keeps its own default group
(world size 1 for a TP=1 trainer or engine) and joins a separate group
created with _new_process_group_helper on its own store, as slime,
SGLang and vLLM do. torch's rank bookkeeping for such a group maps
every member's default-world rank 0 to group rank 0, so
dist.get_rank(group) returned 0 on both ends and the preflight refused
the whole batch ("target location 0 is not owned by this process").
When a group is passed, read the rank and size from the group itself.
Every other path is unchanged.
Add a CPU reproduction: two processes, each with a world-size-1
default group, joined by an externally created gloo group. It fails
before this change and passes after it. The new test file builds that
group with torch.distributed directly, so it joins the existing
check-torchdist exclusion next to tests/unit/comm/test_dist.py.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Pushed 99aef48 with one more fix, found while wiring the executor into an RL trainer↔engine weight-update group. The bug. RL frameworks build that group outside the default world. Each trainer or engine process keeps its own default group (world size 1 at TP=1) and joins a separate two-process group created with The fix. When a group is passed, The test. With the fix, the executor also syncs a full Qwen3-8B (399 tensors, 16.4 GB) between two GPU hosts over NCCL into a running vLLM engine's parameters, through exactly such a group, bitwise-checked after every sync. |
|
|
||
|
|
||
| def _worker(rank, port, tmp): | ||
| import torch.distributed as tdist |
There was a problem hiding this comment.
@0z5a can you added one line here note torch.distributed usage here is intended?
There was a problem hiding this comment.
Done in e38477d: the comment above the import now says the direct torch.distributed use is intended, since the test builds the weight-update group outside deepspeed.comm the way RL frameworks do. Thanks!
|
Hi @0z5a , I'm fine with check-dist change. I leave a comment for a minor change though, can you take a look? Thanks! |
The test builds the trainer <-> engine weight-update group the way RL frameworks do, outside deepspeed.comm, so calling torch.distributed directly is intended. Say so next to the import, as requested in review. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
… states (deepspeedai#8885) ## Summary Describes, as affine maps in global ranks, where a live ZeRO engine stores each element of a parameter and its optimizer states, so `plan_transfer` can read from or write into training storage directly (e.g. training → rollout weight sync) without gathering full parameters first. - `zero_affine.py` - `zero_param_storage_affine_map`: low-precision replicas under ZeRO-0/1/2 (DP or expert-DP group), `ds_tensor` partitions under ZeRO-3, padding excluded. - `zero_optimizer_storage_affine_maps`: fp32 master + optimizer states (ZeRO-1/2, bf16 optimizer, ZeRO-3). - `zero_param_storage_shard` / `zero_optimizer_storage_shard`: flat views of ZeRO's own buffers matching the maps. - `autotp_storage_affine_maps`: AutoTP layout × ZeRO storage, one all-gather over the TP group. - `affine.py`: `nest_affine_maps` composes a layout map (TP/EP shard placement) with per-shard storage maps. A process's stored run of an N-D shard is described as at most `2·ndim−1` N-D boxes in the parameter's own shape, so plans stay small (row-parallel shards are not cut per row). - `affine_transfer.py`: `plan_transfer` accepts an N-D piece in a 1-D (flat) shard only when the piece is dense there; every other rank mismatch is still refused. - `affine_ir_spec.md`: phase-2 section on layout vs storage maps. - Fix `test_convert_checkpoint` `stat().size` → `st_size` (from deepspeedai#8515). No change to ZeRO, AutoTP or the executor. Offloaded optimizer states raise `NotImplementedError`; MiCS / hpZ are not covered. ## Tests - `tests/unit/checkpoint/test_zero_affine.py`, `test_affine_transfer_plan.py`: ZeRO-0/1/2/3 param + optimizer maps checked against ZeRO's own buffers (bitwise rebuild), AutoTP nest, N-D box count, flat-shard acceptance/refusal. - ≤2 ranks: Intel Arc B70 (XPU) ×2, pass. >2 ranks: CPU (gloo), pass. - End-to-end (deepspeedai/DeepSpeedExamples#1021): Qwen3-8B, 8× NVIDIA L20, AutoTP-2 training engine → AutoTP-4 rollout engine via `autotp_storage_affine_maps` + `plan_transfer` + executor (deepspeedai#8809), ZeRO-2 and ZeRO-3, weights bitwise equal after every sync. Plans are built once and reused. | Qwen3-8B, 8×L20 | ZeRO-2 gather | ZeRO-2 affine `lowest` | ZeRO-2 affine `balanced` | ZeRO-3 gather | ZeRO-3 affine | |---|---|---|---|---|---| | maps + plan (once) | - | 0.63 + 0.57 s | 0.70 + 0.70 s | - | 0.82 + 0.60 s | | first sync | 1.01 s | 1.16 s | 0.48 s | 1.71 s | 1.18 s | | repeated syncs | 0.78-0.82 s | 1.05-1.35 s | 0.49-0.55 s | 1.22-1.28 s | 0.45-0.74 s | | device memory added | 3.48 GB | 0.30 GB | 0.16 GB | 3.48 GB | 0.66 GB | | sent / local copy | - | 35.2 / 5.0 GB | 15.1 / 25.1 GB | - | 34.0 / 6.3 GB | | executor rounds | - | 96 | 41 | - | 32 | Each map has 3,192 pieces on both sides; before N-D boxes (one piece per stored row) the same ZeRO-3 maps had 592k source / 2.36M target pieces and planning took 207 s. Under ZeRO-2 every DP rank holds a replica: `lowest` is the planner's default holder choice, `balanced` is a `replica_selector` in the example (local copy first, else the least-loaded holder). --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Hi @0z5a can you take a look at CI failure in merge queue? Thanks! |
The merged #8623 planner describes affine shard copies. This PR executes those plans directly into existing frozen target shards with bounded peer staging, without reconstructing the full logical tensor.
Refs #8230, #8252 and 0z5a#1. This is the executor portion of that work; #8800 supplies live AutoEP metadata separately. The diff is limited to 4 files, 1014 added lines: executor, contract tests, a new API documentation page and its toctree entry. Historical benchmark drivers and evidence remain in the fork and are linked below.
execute_transfer/execute_transfersaccept explicit source and target rank bindings, preserve strided regions and storage offsets, and collectively reject invalid buffers, aliases or request-order disagreement before writes. Shared plans reuse their schedule within a call. On accelerators that are not synchronized with the host (is_synchronized_device()is false), work handles order the caller stream and drain every 32 rounds and before return; on synchronized devices such as CPU each handle is waited immediately.Review fixes, first round: valid scalar plans now execute; the unrelated
new_group(timeout=...)changes are removed, preserving the existing CCL interface.Review fixes, second round: stream ordering is gated on
is_synchronized_device()rather than the device type; the multi-rank tests are no longer confined to CUDA and rely on the harness's device-count gate before launch; the disjoint-endpoint case uses 2 or 4 ranks so it fits the two-GPU PR runners and the four-rank CPU nightly; the documentation moved from the AutoEP page to an Affine Shard Transfers page under Checkpointing API, and records that the staging arena andStagingBudgetbelong to the current torch point-to-point transport.Validation on this exact executor and test source:
LOCAL_SIZE=4 DS_MULTIRANK_ONLY=1 DS_DISABLE_REUSE_DIST_ENV=1 --forked, on current masterLOCAL_SIZE=2 DS_DISABLE_REUSE_DIST_ENV=1 --forkedLOCAL_SIZE,--forkedThe second review round was not rerun on GPU: on CUDA the capability predicate selects the same branches as the device-type check it replaces, so the H100 row still describes the executor code, but the reshaped disjoint-endpoint test has run only on CPU.
On the same two H100s, restoring the original scalar predicate causes all four scalar regressions (single/batched, local/cross-process) to fail with the dimension-rejection defect. The fixed full file passes, including FP32/FP16/BF16 transfer, mixed storage windows, collective refusal, request ordering and a shared-plan batch spanning 48 rounds. The final pytest run exits zero and the assigned GPU contexts are gone. Startup/resource failures are retained separately.
Historical eight-H100 non-scalar correctness evidence records 304 BF16/FP32 checks plus mixed-layout and shared-plan checks. Historical paired performance reports 1.289x geometric speedup, 95% CI [1.2703, 1.3080], for a 1-GiB synthetic frozen-weight handoff. These are prior non-scalar runs, not a new eight-GPU qualification of this rebased and reduced branch or a full-model performance claim.
Scope: scale-1 plans and frozen packed weights. No optimizer migration, live AutoEP swap, CPU/NVMe offload integration, training convergence/resumption, serving or non-CUDA accelerator qualification. CPU/gloo runs the contract tests only;
Work.block_current_stream()orders the device stream on NCCL and XCCL, while other non-synchronized backends that do not override it are untested. The staging budget applies per peer direction; a payload-time transport failure may partially modify targets. Maintenance contact: @0z5a.AI assistance was used to prepare this change, its tests and this description. I reviewed every changed line and ran the tests listed above.