[PR #762] Fix hybrid combine top-k weight gradients with single reduction #764

New issue

Open

opened 2026-10-27 21:06:03 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:06:03 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/762
Author: @Gstdioh
Created: 9/23/2026
Status: 🔄 Open

Base: main ← Head: codex/fix-hybrid-combine-score-grad

📝 Commits (6)

2de61ed Fix hybrid combine top-k weight source selection
299bcbf Add regression for receiver-local top-k weight gradients
612bf08 Apply repository formatting to combine reduce epilogue
6d7a1e4 Complete top-k weight gradient regression
463065c Keep score-gradient fix minimal without reformatting
655d77e Report failing ranks in score-gradient regression

📊 Changes

2 files changed (+115 additions, -1 deletions)

View changed files

📝 deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh (+9 -1)
➕ tests/elastic/test_combine_topk_weights.py (+106 -0)

📄 Description

With hybrid combine, do_expand=False and allow_multiple_reduction=False, the sender forwards a separate top-k weight row for each destination GPU. The reduce epilogue currently selects the master slot by scale-out rank, so routes on different GPUs in the same scale-out domain can read the wrong row.

For example, receiver-local score-gradient rows [11, 0] and [0, 22] should produce [11, 22], but selecting one domain-wide row produces [0, 22]. This affects combine(..., topk_weights=...), including its use for dispatch backward. A plain dispatch/combine weight roundtrip can hide the bug because the dispatched weights are duplicated.

This change selects the score source master slot by global GPU rank for this mode, matching the forwarded row layout. The existing selection is preserved for other modes, with a compile-time assertion for slot layout. The original header formatting is preserved; the production change is limited to this one block (+9/-1 lines).

A standalone distributed regression uses synthetic receiver-local gradients and covers same-GPU, same-domain/different-GPU, and cross-domain routes, including invalid top-k slots. It supports --allow-multiple-reduction=1 as a control.

Validation: Python syntax, Ruff 0.6.5, YAPF 0.40.2, and git diff --check pass. The header already differs from clang-format on the base branch; whole-file formatting is intentionally excluded to keep this fix focused. Real multi-node GPU validation completed for commit 655d77e6bfcdcf46b065a4596104f6f2d4338503: two nodes with 8 NVIDIA H200 GPUs each (16 ranks), with fresh processes and separate JIT caches for --allow-multiple-reduction=0 and =1. All three cases passed for all ranks in both modes, and both node launchers exited with code 0 for each run. No test or production-code changes were needed. This run used PyTorch 2.9.0 (CUDA 12.4) and nvcc 12.8.61; it does not validate the README's PyTorch >=2.10 baseline or add distributed CI coverage. See the GPU validation results.

Run on each node (replace the rendezvous placeholders):

torchrun --nnodes=2 --nproc-per-node=2 --node-rank=NODE_RANK \
    --master-addr=MASTER_ADDR --master-port=29500 \
    tests/elastic/test_combine_topk_weights.py

The test asserts at least two logical scale-out domains and at least two GPUs per domain. Repeat with --allow-multiple-reduction=1 to exercise the control path.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/762 **Author:** [@Gstdioh](https://github.com/Gstdioh) **Created:** 9/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `codex/fix-hybrid-combine-score-grad` --- ### 📝 Commits (6) - [`2de61ed`](https://github.com/deepseek-ai/DeepEP/commit/2de61ede578ef0498f82401f6ebced529995943c) Fix hybrid combine top-k weight source selection - [`299bcbf`](https://github.com/deepseek-ai/DeepEP/commit/299bcbfadb62d8d8e7a80e4cf0a3a6552dfb3f26) Add regression for receiver-local top-k weight gradients - [`612bf08`](https://github.com/deepseek-ai/DeepEP/commit/612bf08abd96cbb8cfab38a4d2510242956e60ed) Apply repository formatting to combine reduce epilogue - [`6d7a1e4`](https://github.com/deepseek-ai/DeepEP/commit/6d7a1e4cb676d537563a585e44d907ed2f430cb6) Complete top-k weight gradient regression - [`463065c`](https://github.com/deepseek-ai/DeepEP/commit/463065c7c3cb27a0a5f8874510201be2f899ad7b) Keep score-gradient fix minimal without reformatting - [`655d77e`](https://github.com/deepseek-ai/DeepEP/commit/655d77e6bfcdcf46b065a4596104f6f2d4338503) Report failing ranks in score-gradient regression ### 📊 Changes **2 files changed** (+115 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/include/deep_ep/impls/combine_reduce_epilogue.cuh` (+9 -1) ➕ `tests/elastic/test_combine_topk_weights.py` (+106 -0) </details> ### 📄 Description With hybrid combine, `do_expand=False` and `allow_multiple_reduction=False`, the sender forwards a separate top-k weight row for each destination GPU. The reduce epilogue currently selects the master slot by scale-out rank, so routes on different GPUs in the same scale-out domain can read the wrong row. For example, receiver-local score-gradient rows `[11, 0]` and `[0, 22]` should produce `[11, 22]`, but selecting one domain-wide row produces `[0, 22]`. This affects `combine(..., topk_weights=...)`, including its use for dispatch backward. A plain dispatch/combine weight roundtrip can hide the bug because the dispatched weights are duplicated. This change selects the score source master slot by global GPU rank for this mode, matching the forwarded row layout. The existing selection is preserved for other modes, with a compile-time assertion for slot layout. The original header formatting is preserved; the production change is limited to this one block (+9/-1 lines). A standalone distributed regression uses synthetic receiver-local gradients and covers same-GPU, same-domain/different-GPU, and cross-domain routes, including invalid top-k slots. It supports `--allow-multiple-reduction=1` as a control. Validation: Python syntax, Ruff 0.6.5, YAPF 0.40.2, and `git diff --check` pass. The header already differs from clang-format on the base branch; whole-file formatting is intentionally excluded to keep this fix focused. Real multi-node GPU validation completed for commit `655d77e6bfcdcf46b065a4596104f6f2d4338503`: two nodes with 8 NVIDIA H200 GPUs each (16 ranks), with fresh processes and separate JIT caches for `--allow-multiple-reduction=0` and `=1`. All three cases passed for all ranks in both modes, and both node launchers exited with code 0 for each run. No test or production-code changes were needed. This run used PyTorch 2.9.0 (CUDA 12.4) and nvcc 12.8.61; it does not validate the README's PyTorch >=2.10 baseline or add distributed CI coverage. See the [GPU validation results](https://github.com/deepseek-ai/DeepEP/pull/762#discussion_r4079669535). Run on each node (replace the rendezvous placeholders): ```bash torchrun --nnodes=2 --nproc-per-node=2 --node-rank=NODE_RANK \ --master-addr=MASTER_ADDR --master-port=29500 \ tests/elastic/test_combine_topk_weights.py ``` The test asserts at least two logical scale-out domains and at least two GPUs per domain. Repeat with `--allow-multiple-reduction=1` to exercise the control path. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:06:03 +00:00

bsakaguchi referenced this issue

2026-10-27 21:06:10 +00:00

[PR #765] [CLOSED] [Bugfix][NVSHMEM] Support v2 state in legacy IBGDA kernels #766