zero3: SDMA allgather via mori (sdma_allgather) - #7999
Merged
Merged
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 66673546b5
ℹ️ 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".
|
|
||
|
|
||
| def _dist_allgather_fn(input_tensor: Tensor, output_tensor: Tensor, group=None): | ||
| work = mori.allgather_into_tensor(input_tensor, output_tensor) |
There was a problem hiding this comment.
When ZeRO is initialized with a non-WORLD data-parallel group, or with a secondary zero-param group, _all_gather passes that group down as ds_process_group (partition_parameters.py:1463-1471), but this new SDMA call ignores the group argument and uses mori's WORLD-backed default process group. In those model/tensor-parallel configurations mori gathers from more ranks than the caller allocated output_tensor for, which can corrupt fetched parameters or write past the expected buffer; fall back unless group is WORLD or make mori initialize/use the matching group.
Useful? React with 👍 / 👎.
Comment on lines
+1282
to
+1283
| handle = instrument_w_nvtx(dist.allgather_fn)( | ||
| flat_tensor, partitions[rank_in_group], group=ds_process_group, async_op=True) |
There was a problem hiding this comment.
With the default stage3_allgather_sequential=false, any ZeRO-3 fetch containing more than one parameter takes _all_gather_dtype, but this path now calls dist.allgather_fn directly instead of _dist_allgather_fn. As a result, enabling sdma_allgather has no effect for the common coalesced prefetch path (including the added sample config, which does not enable sequential allgather), so the advertised optimization is skipped for most multi-parameter buckets.
Useful? React with 👍 / 👎.
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Signed-off-by: wuyl1 <yangwu@amd.com>
Move all mori-specific code (handle, dtype map, transit buffer sizing,
PG-registration helper, Work wrapper) out of partition_parameters.py
into a dedicated runtime/comm backend module:
deepspeed/runtime/comm/mori.py
mori.init(max_numel) # one-shot, idempotent, exception-safe
mori.is_enabled() # cheap predicate
mori.allgather_into_tensor(in, out)
-> Work-compatible object on success, None on fallback
The new backend uses mori_cpp.AllGatherIntoTensor (NCCL/RCCL-style
flat->flat C++ dispatcher) instead of the old mori.ccl.AllgatherSdma
templated Python wrapper, so DeepSpeed no longer has to pre-convert
numel into uint32 lane counts or template the C++ class on dtype.
partition_parameters.py is now agnostic to the SDMA path:
def _dist_allgather_fn(input_tensor, output_tensor, group=None):
work = mori.allgather_into_tensor(input_tensor, output_tensor)
if work is not None:
return work
return instrument_w_nvtx(dist.allgather_fn)(...)
Init failure (mori missing, non-AMD/ROCm runtime, shmem init error)
leaves the handle unset and logs a single rank-0 warning, so the SDMA
path silently no-ops and dist.allgather_fn (RCCL/NCCL) takes over —
no hard fail.
Net change: partition_parameters.py shrinks by 79 lines; one new
self-contained module under runtime/comm/.
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
_SdmaWork.wait() previously blocked the CPU on _event.synchronize()
before issuing the stream-level dependency. RCCL's Work.wait() only
records a stream-level wait (cudaStreamWaitEvent / hipStreamWaitEvent)
and does NOT block the CPU, which is what the ZeRO-3 prefetch pipeline
relies on: while bucket N is in flight on the GPU, the CPU is free to
queue bucket N+1 so it can overlap with the trailing compute of N.
The CPU-blocking variant turned out to be a per-step critical-path tax
that wiped out SDMA's headroom on workloads that issue many small
allgathers per step. Concretely, on Qwen3-32B + ZeRO-3 + seq_len=128,
8x MI300X, ~6400 prefetch buckets per step:
before: SDMA 1014 ms / step (1009 tok/s)
RCCL 932 ms / step (1099 tok/s) -> SDMA -8.0%
after: SDMA 927 ms / step (1104 tok/s)
RCCL 929 ms / step (1100 tok/s) -> within noise
Loss curve is bit-identical with and without the CPU sync, so this is
purely a CPU-pipelining fix. is_completed() is unchanged (it polls
via _event.query() without blocking, same as before).
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Move the zero3_overlap demo dir into examples/sdma_allgather/ (the name
that matches the feature being demoed) and add a Qwen3-32B + ZeRO-3
trainer that reproduces the +9.93% end-to-end speedup of this PR on
8x MI300X with the default DeepSpeed bucket sizes.
Layout:
ds_config_zero3_{sdma,nosdma}.json ZeRO-3 + bf16 + DS-default buckets
run_gpt_sdma_{on,off}.sh GPT-7B-ish demo (existing trainer)
run_qwen3_sdma_{on,off}.sh Qwen3-32B demo (new trainer)
train_qwen3_zero3.py self-contained Qwen3 trainer
README.md feature overview + repro steps
train_zero3.py unchanged (renamed only)
test_sdma_allgather_zero3.py unchanged (renamed only)
train_qwen3_zero3.py inlines a minimal wikitext-103 dataloader so the
benchmark has no dependency on external benchmark repos. Loading via
AutoConfig + from_config keeps the example weight-free; only the model
config and tokenizer are pulled from HuggingFace.
The configs use DeepSpeed's default ZeRO-3 bucket sizes
(stage3_prefetch_bucket_size = 5e7, etc.) so the published numbers
in README.md are reproducible without any tuning.
Verified on 8x MI300X, two fresh rounds:
Qwen3-32B + ZeRO-3 + DP=8, seq_len=1024, micro_bs=1, 100 steps
SDMA off : 1402.5 ms / step (5841 tok/s)
SDMA on : 1263.2 ms / step (6486 tok/s) -> +9.93% e2e
GPT-7B + ZeRO-3 + DP=8, 100 steps -> +5.9% e2e
Loss curves match across the two backends, peak memory is identical
(96.45 GB), per-step jitter is 1.4-2.7%, so the ~140 ms gap is well
above noise.
Drops:
examples/zero3_overlap/run.sh superseded by run_gpt_*
examples/zero3_overlap/ds_config_zero3.json superseded by *_sdma.json
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Collaborator
|
Hi @inkcherry @wuyl1, thanks for your contributed code. I have read this PR and your discussion with @tjruwase in RFC 7884. And I have some high level comments before going into detailed review.
The following are current diagram and the ideal diagram |
Per @delock's review (PR deepspeedai#7999) and @tjruwase's RFC 7884 feedback, the SDMA allgather should be an internal optimisation inside deepspeed.comm, not a user-facing ZeRO-3 config knob. This commit reshapes the integration accordingly: - TorchBackend.all_gather_into_tensor() now attempts a mori SDMA fast-path before delegating to torch.distributed.all_gather_into_tensor. Detection is automatic at backend init (AMD/ROCm + mori importable + shmem init succeeds); failure leaves the standard RCCL/NCCL path untouched. - ZeRO-3 partition_parameters.py becomes mori-agnostic: _dist_allgather_fn is just a thin wrapper around dist.allgather_fn. The dedicated mori init in Init.__init__ is gone. - DeepSpeedZeroConfig.sdma_allgather and sdma_allgather_max_numel are removed (no ds_config field — transparent). An opt-out env var DS_DISABLE_SDMA_ALLGATHER=1 is provided for A/B baselines and debugging, plus DS_SDMA_ALLGATHER_MAX_NUMEL=N to override the transit buffer size (default 64M elements). - mori.supports() now refuses the SDMA path when the call is on a non-WORLD process group, when the shard is larger than the pre-allocated transit buffer, or when the dtype isn't in the mori dtype map. This fixes the Codex P1 review note: a ZeRO instance initialised with a sub-PG no longer dispatches to a WORLD-bound mori handle. - Because ZeRO-3's coalesced allgather path (_all_gather_dtype, the default with stage3_allgather_sequential=False) calls dist.allgather_fn directly, it now ALSO benefits from SDMA. This fixes the Codex P2 review note that the original hook only intercepted the sequential path. Examples: - ds_config_zero3.json is now a single file (no _sdma / _nosdma split). - run_*_sdma_off.sh export DS_DISABLE_SDMA_ALLGATHER=1. - run_*_sdma_on.sh export MORI_ENABLE_SDMA=1 (still required for uncached transit buffers on the mori side). - test_sdma_allgather_zero3.py exercises the transparent path via dist.allgather_fn instead of the old private hook. Quick test on 8x MI300X (GPT-7B-ish, ZeRO-3 default coalesced path, 100 steps, wikitext-2): off (DS_DISABLE_SDMA_ALLGATHER=1) : 705.0 ms / step (11.5 samples/s) on (MORI_ENABLE_SDMA=1) : 649.8 ms / step (12.5 samples/s) -> +8.5% e2e, mori log "AllgatherSdma initialized: PE 0 of 8", loss bit-identical at every checkpoint. The +8.5% gain on the coalesced path is direct evidence Codex P2 is fixed: before this commit the same config would have shown ~0% gain because the SDMA hook only sat on the sequential path. Signed-off-by: inkcherry <mingzhi.liu@amd.com>
…ched Follow-up to the TorchBackend refactor (606f309) in response to @delock's review. Two changes: 1. User-facing control surface is a single opt-in env var. DS_SDMA_ALLGATHER=1 is now the only switch. Default is OFF even when mori is installed, so users explicitly choose the hardware-specific fast-path (per the discussion under PR deepspeedai#7999: users want to keep control over a HW-specific optimisation even when the runtime can auto-detect it). When set, mori.init() auto-exports MORI_ENABLE_SDMA=1 on the user's behalf (via os.environ.setdefault), so the SDMA kernel gets its uncached transit buffers without the user having to know mori's internal env var. DS_SDMA_ALLGATHER_MAX_NUMEL=N stays as the transit-buffer-size override. The previous opt-out DS_DISABLE_SDMA_ALLGATHER is gone (default off already covers the A/B baseline scenario). mori.init() is now a silent no-op when the user did not opt in, so backend init on machines that happen to have mori installed no longer prints anything. 2. ZeRO-3 hot path is bit-identical to upstream. partition_parameters.py is now 0 diff against the upstream pristine version (commit 57b10d5). All earlier cosmetic churn from the original SDMA branch (AllGatherHandle multi-line signature, gathered=... temp var, total_numel=... temp var, _dist_allgather_fn inlining, and notably collapsing the `if original_dtype == allgather_dtype:` fast-path that upstream explicitly keeps separate "for safety" re: tied-parameter post_accumulate_grad_hook re-entry) is reverted. The SDMA acceleration entirely lives inside TorchBackend.all_gather_into_tensor; ZeRO-3 just calls dist.allgather_fn and lands on the fast-path transparently when it is enabled and the call is on the WORLD group with a fitting shard size and supported dtype. Net core-file footprint of this PR is now just deepspeed/comm/torch.py (TorchBackend) + the new deepspeed/runtime/comm/mori.py module. Examples and README updated accordingly: run_*_sdma_on.sh export DS_SDMA_ALLGATHER=1 (no MORI_ENABLE_SDMA needed); run_*_sdma_off.sh just don't export anything (default). Smoke test on 8x MI300X (GPT-7B-ish, ZeRO-3 default coalesced path, 100 steps, wikitext-2): off (default, no env var) : 686.2 ms / step (11.8 samples/s) on (DS_SDMA_ALLGATHER=1) : 645.9 ms / step (12.5 samples/s) -> +6.2% step time, +5.9% throughput mori reports "AllgatherSdma initialized: PE X of 8" on all ranks of the on run and nothing on the off run; loss curves overlap (off step 50 loss=7.28 vs on step 50 loss=7.25, within data-loader jitter). Co-authored-by: wuyl1 <yangwu@amd.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Forgot to put a concrete number in the GPT column. Both off and on runs report Max_MA = 12.12 GB on rank 0 (DeepSpeed's memory_status report), so the table now matches the Qwen3 column's format. Signed-off-by: inkcherry <mingzhi.liu@amd.com>
inkcherry
added a commit
to inkcherry/DeepSpeed
that referenced
this pull request
May 13, 2026
…ched Follow-up to the TorchBackend refactor (606f309) in response to @delock's review. Two changes: 1. User-facing control surface is a single opt-in env var. DS_SDMA_ALLGATHER=1 is now the only switch. Default is OFF even when mori is installed, so users explicitly choose the hardware-specific fast-path (per the discussion under PR deepspeedai#7999: users want to keep control over a HW-specific optimisation even when the runtime can auto-detect it). When set, mori.init() auto-exports MORI_ENABLE_SDMA=1 on the user's behalf (via os.environ.setdefault), so the SDMA kernel gets its uncached transit buffers without the user having to know mori's internal env var. DS_SDMA_ALLGATHER_MAX_NUMEL=N stays as the transit-buffer-size override. The previous opt-out DS_DISABLE_SDMA_ALLGATHER is gone (default off already covers the A/B baseline scenario). mori.init() is now a silent no-op when the user did not opt in, so backend init on machines that happen to have mori installed no longer prints anything. 2. ZeRO-3 hot path is bit-identical to upstream. partition_parameters.py is now 0 diff against the upstream pristine version (commit 57b10d5). All earlier cosmetic churn from the original SDMA branch (AllGatherHandle multi-line signature, gathered=... temp var, total_numel=... temp var, _dist_allgather_fn inlining, and notably collapsing the `if original_dtype == allgather_dtype:` fast-path that upstream explicitly keeps separate "for safety" re: tied-parameter post_accumulate_grad_hook re-entry) is reverted. The SDMA acceleration entirely lives inside TorchBackend.all_gather_into_tensor; ZeRO-3 just calls dist.allgather_fn and lands on the fast-path transparently when it is enabled and the call is on the WORLD group with a fitting shard size and supported dtype. Net core-file footprint of this PR is now just deepspeed/comm/torch.py (TorchBackend) + the new deepspeed/runtime/comm/mori.py module. Examples and README updated accordingly: run_*_sdma_on.sh export DS_SDMA_ALLGATHER=1 (no MORI_ENABLE_SDMA needed); run_*_sdma_off.sh just don't export anything (default). Smoke test on 8x MI300X (GPT-7B-ish, ZeRO-3 default coalesced path, 100 steps, wikitext-2): off (default, no env var) : 686.2 ms / step (11.8 samples/s) on (DS_SDMA_ALLGATHER=1) : 645.9 ms / step (12.5 samples/s) -> +6.2% step time, +5.9% throughput mori reports "AllgatherSdma initialized: PE X of 8" on all ranks of the on run and nothing on the off run; loss curves overlap (off step 50 loss=7.28 vs on step 50 loss=7.25, within data-loader jitter). Co-authored-by: wuyl1 <yangwu@amd.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
baseline_pp.py was a stray 257 KB UTF-16 + CRLF copy of deepspeed/runtime/zero/partition_parameters.py left at the repo root by an earlier development snapshot. It is not referenced from anywhere in the codebase and shows up as a binary blob in the PR diff against upstream master. Drop it. Signed-off-by: inkcherry <mingzhi.liu@amd.com>
delock
reviewed
May 13, 2026
| import torch | ||
| from deepspeed.accelerator import get_accelerator | ||
| from deepspeed.ops.op_builder import NotImplementedBuilder | ||
| try: |
Collaborator
There was a problem hiding this comment.
What is the condition of this fallback? Is this fallback related to Mori library?
Contributor
Author
There was a problem hiding this comment.
removed try/except here.
delock
reviewed
May 13, 2026
| @@ -0,0 +1,229 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
Collaborator
There was a problem hiding this comment.
Should replace Microsoft Corporation in licence header with DeepSpeed team.
delock
reviewed
May 13, 2026
|
|
||
| User-visible controls (env vars, no ``ds_config`` field): | ||
|
|
||
| * ``DS_SDMA_ALLGATHER=1`` opt in to the SDMA path. Required: |
Collaborator
There was a problem hiding this comment.
Can we use DS_DMA_ALLGATHER and DS_DMA_ALLGATHER_MAX_NUMEL as behavior control environment variable? I'm thinking about in the future other accelerators wants to use similiar techniques, and can share the same environment variable name that is not brand binded.
Contributor
Author
There was a problem hiding this comment.
I initially wanted to use a config file, but considering that currently only the AMD path is available, I set it as an environment variable.
I suggest adding it to the config file in the future if more accelerators support this approach， because environment variables tend to be more specific.
Collaborator
There was a problem hiding this comment.
@inkcherry I feel that SDMA implementation should be a drop in replacement of RCCL implementation. So maybe it should be turn on by default. In comments, I saw descriptions that Mori allgather may not be bit by bit identical to RCCL implementation. However allgather has no computation, how could it be not bit by bit identical?
I agree for *MAX_NUMEL option, if in the future more accelerator use it, we need to move it as a config controled variable and give it a generic name.
Contributor
Author
There was a problem hiding this comment.
For now, set it as an environment variable. we don't want DeepSpeed's collective backend to silently change behaviour based on which extra packages happen to be installed — mori may end up on the system as a transitive dep of some other framework, not because the user explicitly chose to enable SDMA.
Allgather itself will not have changes. Due to timing variations (data dependencies are correct), end-to-end results might have bit-level differences at the tail due to the order of additions (which can be considered noise, and this is normal in training). However, the loss remains unchanged. Thanks, I have revised the wording to prevent ambiguity.
delock
requested changes
May 13, 2026
delock
reviewed
May 13, 2026
| # SDMA unsafe (non-WORLD group, oversized shard, unsupported dtype, | ||
| # mori unavailable) yields None and we fall through to RCCL/NCCL. | ||
| from deepspeed.runtime.comm import mori as _mori | ||
| sdma_work = _mori.allgather_into_tensor(input_tensor, output_tensor, group=group) |
Collaborator
There was a problem hiding this comment.
What is the overhead here if the hardware does not support mori? I'm a little bit concerned in potential overhead into this common path. How long does it take on a CUDA system without Mori library, for example?
Contributor
Author
There was a problem hiding this comment.
Refactored to a lazy check at TorchBackend.init
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
delock
approved these changes
May 13, 2026
Three small fixes for the format pipeline that started failing after
the recent comment cleanup:
- mori.py: restore the standard DeepSpeed license header. The
earlier "update" commit collapsed
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
#
# DeepSpeed Team
into a single "# Copyright (c) DeepSpeed Team." line, which makes
scripts/check-license.py fail (it greps for the literal lines
"^# SPDX-License-Identifier: Apache-2.0$" and "^# DeepSpeed Team$").
Header is back in the canonical form used everywhere else in the
repo.
- mori.py: yapf-reformat three call sites the editor had wrapped
too aggressively for the project's 119-column limit
(allgather_into_tensor signature, the _handle(...) call, the
rank-0 warning guard).
- .pre-commit-config.yaml: add deepspeed/runtime/comm/mori.py to
the check-torchdist exclude list. Like its sibling
deepspeed/runtime/comm/coalesced_collectives.py (already on the
list), mori is an intentionally-low-level comm backend that
needs torch.distributed primitives (group.WORLD identity,
is_initialized(), get_rank()) directly; the higher-level
deepspeed.comm wrappers aren't a substitute here.
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Collaborator
|
Hi @inkcherry can you resolve merge conflicts? Thanks! |
mori is a backend helper for TorchBackend (selected at backend init, exposed via TorchBackend.all_gather_into_tensor), so it belongs under deepspeed/comm/ alongside the other backend implementations rather than under deepspeed/runtime/comm/ which holds higher-level comm utilities. This also drops the explicit deepspeed/runtime/comm/mori.py entry from the check-torchdist exclude list, since deepspeed/comm/ is already covered by the broader parent-directory exclude pattern. Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Brings the SDMA allgather branch up to date with upstream master (57 commits since the last sync, including the v0.19.0 release). Only conflict was the Latest News section of the top-level README, which now lists [2026/05] SDMA above the [2026/03] entries added upstream. Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Contributor
Author
@delock resolved. |
After merging upstream/master, pre-commit run --all-files surfaced
three latent violations in the SDMA example scripts that the prior
DeepSpeedAI CI runs hadn't been exercising:
- check-license : missing the standard DeepSpeed Apache-2.0 header
on all three example .py files.
- check-torchdist : direct use of torch.distributed.get_world_size,
`import torch.distributed as dist`, `import torch.distributed as
torch_dist`. Routed through deepspeed.comm (`from deepspeed
import comm as dist`), matching the rest of the examples in the
repo.
- check-torchcuda : torch.cuda.set_device / manual_seed_all /
synchronize / reset_peak_memory_stats / max_memory_allocated /
device_count / Event used directly. Routed through
`get_accelerator()` so the demos stay portable across CUDA, ROCm
and other DeepSpeed accelerators (which matters here -- the
SDMA fast-path is a ROCm-only feature in the first place).
yapf also re-wrapped a few argparse / torch.empty call sites to
the project's 119-column limit.
No behaviour change: the examples produce the same output and the
benchmark numbers in examples/sdma_allgather/README.md still apply.
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
nathon-lee
pushed a commit
to nathon-lee/DeepSpeed_woo
that referenced
this pull request
Jul 1, 2026
## Summary RFC: deepspeedai#7884 Wire `sdma_allgather` into ZeRO-3's parameter prefetch path (`_dist_allgather_fn`). When enabled, ZeRO-3 allgather routes through `mori_cpp.AllGatherIntoTensor` (intra-node SDMA copy on AMD MI300), with a transparent fallback to `dist.allgather_fn` (RCCL/NCCL) on init failure. End-to-end demo + repro steps + verified numbers live in [`examples/sdma_allgather/README.md`](examples/sdma_allgather/README.md). Headline (8x MI300X, DeepSpeed default ZeRO-3 buckets, 100 steps): | | GPT-7B-ish | Qwen3-32B | |---|---|---| | SDMA off | 697.7 ms / step | 1402.5 ms / step | | SDMA on | 622.0 ms / step | 1263.2 ms / step | | **gain** | **+10.85 %** | **+9.93 %** | Loss curves match off ↔ on, peak memory unchanged. Speedup is workload-dependent — gains shrink (or invert) when allgather can't be overlapped with compute Co-authored-by: wuyl1 <yangwu@amd.com> --------- Signed-off-by: wuyl1 <yangwu@amd.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com> Co-authored-by: wuyl1 <yangwu@amd.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>
nathon-lee
pushed a commit
to nathon-lee/DeepSpeed_woo
that referenced
this pull request
Jul 1, 2026
## Summary RFC: deepspeedai#7884 Wire `sdma_allgather` into ZeRO-3's parameter prefetch path (`_dist_allgather_fn`). When enabled, ZeRO-3 allgather routes through `mori_cpp.AllGatherIntoTensor` (intra-node SDMA copy on AMD MI300), with a transparent fallback to `dist.allgather_fn` (RCCL/NCCL) on init failure. End-to-end demo + repro steps + verified numbers live in [`examples/sdma_allgather/README.md`](examples/sdma_allgather/README.md). Headline (8x MI300X, DeepSpeed default ZeRO-3 buckets, 100 steps): | | GPT-7B-ish | Qwen3-32B | |---|---|---| | SDMA off | 697.7 ms / step | 1402.5 ms / step | | SDMA on | 622.0 ms / step | 1263.2 ms / step | | **gain** | **+10.85 %** | **+9.93 %** | Loss curves match off ↔ on, peak memory unchanged. Speedup is workload-dependent — gains shrink (or invert) when allgather can't be overlapped with compute Co-authored-by: wuyl1 <yangwu@amd.com> --------- Signed-off-by: wuyl1 <yangwu@amd.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com> Co-authored-by: wuyl1 <yangwu@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
RFC: #7884
Wire
sdma_allgatherinto ZeRO-3's parameter prefetch path(
_dist_allgather_fn). When enabled, ZeRO-3 allgather routes throughmori_cpp.AllGatherIntoTensor(intra-node SDMA copy on AMD MI300), with atransparent fallback to
dist.allgather_fn(RCCL/NCCL) on init failure.End-to-end demo + repro steps + verified numbers live in
examples/sdma_allgather/README.md.Headline (8x MI300X, DeepSpeed default ZeRO-3 buckets, 100 steps):
Loss curves match off ↔ on, peak memory unchanged.
Speedup is workload-dependent — gains shrink (or invert) when allgather can't be overlapped with compute
Co-authored-by: wuyl1 yangwu@amd.com