mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-17 01:07:19 +00:00
WIP
This commit is contained in:
@@ -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()} "
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user