This commit is contained in:
Binyang Li
2026-06-22 21:50:39 +00:00
parent 5e0c1de254
commit bd0f15b4ef
5 changed files with 92 additions and 32 deletions

View File

@@ -85,8 +85,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int):
def main():
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
NUM_MAX_NVL_PEERS = _detect_local_world_size()
assert (
num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS
@@ -162,7 +165,7 @@ def main():
print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True)
buf = ep.ExpertParallelRuntime(
group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False
ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False
)
print(
f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()} "

View File

@@ -65,8 +65,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int):
def main():
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
# Small settings for functional check
num_tokens = 128
hidden = 1024
@@ -122,7 +125,7 @@ def main():
)
print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True)
buf = ep.ExpertParallelRuntime(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False)
buf = ep.ExpertParallelRuntime(ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False)
print(f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()}", flush=True)
assert buf.is_available()

View File

@@ -78,8 +78,11 @@ def init_dist():
def main():
args = parse_args()
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
# Shrink the "bf16 precision" anchor to keep values small.
rank_offset = 128
assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor"
@@ -107,7 +110,7 @@ def main():
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = -1
moe_comm = ep.MoECommunicator(
comm=group,
comm=ep_group,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden,