zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[EP] Fix the tensorSize of the barrier op (#1222) 

* Fix the tensorSize of the barrier op

* Fix
This commit is contained in:
Xun Sun 2025-12-17 15:50:12 +08:00 committed by GitHub
parent dbb211d1d4
commit 44dde2dfa2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mooncake-ep/src/mooncake_backend.cpp
View File

@ -14,6 +14,7 @@ constexpr const char* SYNC_OP_ERROR_MSG = "Expecting async op but got sync op.";
constexpr const char* REDUCE_OP_ERROR_MSG = "Only support SUM.";
constexpr const char* SPARSE_ERROR_MSG = "Sparse op not supported.";
constexpr const char* REDUCE_DTYPE_ERROR_MSG = "Unsupported reduce dtype: ";
constexpr int kBarrierDummyTensorSize = 1;
std::string MooncakeBackend::hostIp_ = "127.0.0.1";
TransferEngine MooncakeBackend::engine_ = TransferEngine(true);
@ -469,8 +470,10 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::barrier(
const c10d::BarrierOptions& opts) {
TORCH_CHECK(isCpu_, "Barrier is available only for CPU.")
return worker_.putTaskCpu(
c10d::OpType::BARRIER, 0, 0, &meta_, [=](void*, size_t, size_t) {},
[=](void*, size_t, size_t) {});
// a non-zero tensorSize is required to ensure the worker task for the
// barrier is created
c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, &meta_,
[=](void*, size_t, size_t) {}, [=](void*, size_t, size_t) {});
}
void MooncakeBackend::shutdown() {

7
mooncake-wheel/tests/test_mooncake_backend_cpu.py
View File

@ -26,6 +26,10 @@ def worker(rank, world_size, results, collective):
dist.all_gather(gathered, tensor)
results[rank] = [t.item() for t in gathered]
elif collective == "barrier":
dist.barrier()
results[rank] = "ok"
else:
raise ValueError(f"Unsupported collective: {collective}")
@ -66,6 +70,9 @@ class TestMooncakeBackend(unittest.TestCase):
# Expected gather = [0, 1, 2, 3]
self._spawn_and_check("all_gather", lambda size: list(range(size)))
def test_barrier(self):
self._spawn_and_check("barrier", lambda size: "ok")
if __name__ == "__main__":
unittest.main()