MSCCL++ Integration (#22734)
Co-authored-by: Caio Rocha <caiorocha@microsof.com> Co-authored-by: empyreus <rjsouza1995@gmail.com>
This commit is contained in:
Caio Rocha
co-authored by
Caio Rocha
empyreus
Caio Rocha
empyreus
parent
9c53031d2b
commit
c2eae96c56
@@ -0,0 +1,55 @@
|
||||
## MSCCL++ All-Reduce Benchmark
|
||||
|
||||
[MSCCL++](https://github.com/microsoft/mscclpp) is a GPU-driven communication library that can replace NCCL for all-reduce operations. It supports CUDA graph capture and is optimized for small-to-medium message sizes commonly seen in tensor-parallel inference.
|
||||
|
||||
Currently supported configurations: **TP=8** (single-node) and **TP=16** (two-node).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
1. If you use the default SGLang Docker image build from `docker/Dockerfile`, [MSCCL++](https://github.com/microsoft/mscclpp) is already installed by default.
|
||||
2. If you are not using that Docker image (or want to install manually), install [MSCCL++](https://github.com/microsoft/mscclpp) from source (requires CMake and a CUDA toolkit):
|
||||
```bash
|
||||
git clone https://github.com/microsoft/mscclpp.git
|
||||
cd mscclpp && mkdir build && cd build
|
||||
cmake .. && make -j && pip install ..
|
||||
```
|
||||
3. Ensure `mscclpp` is importable in your Python environment before running the benchmark or using MSCCL++ for inference.
|
||||
|
||||
### Running the Benchmark
|
||||
|
||||
The benchmark compares all-reduce latency across torch/NCCL (eager), MSCCL++ (eager and graph), and PyNccl (graph) for power-of-two message sizes.
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node 8 \
|
||||
--nnodes 1 \
|
||||
--node_rank 0 \
|
||||
benchmark/kernels/all_reduce/benchmark_mscclpp.py
|
||||
```
|
||||
|
||||
For multi-node (TP=16):
|
||||
```bash
|
||||
export WORLD_SIZE=2
|
||||
export MASTER_ADDR=<master-ip>
|
||||
export MASTER_PORT=12345
|
||||
|
||||
# Run on each node with the appropriate RANK (0 or 1):
|
||||
torchrun --nproc_per_node 8 \
|
||||
--nnodes $WORLD_SIZE \
|
||||
--node_rank $RANK \
|
||||
--master_addr $MASTER_ADDR \
|
||||
--master_port $MASTER_PORT \
|
||||
benchmark/kernels/all_reduce/benchmark_mscclpp.py
|
||||
```
|
||||
|
||||
### Inference with MSCCL++
|
||||
|
||||
Use the `--enable-mscclpp` flag to select MSCCL++ as the all-reduce backend during CUDA-graph-captured inference:
|
||||
|
||||
```bash
|
||||
python -m sglang.launch_server \
|
||||
--model-path Qwen/Qwen3-8B \
|
||||
--tp-size 8 \
|
||||
--enable-mscclpp
|
||||
```
|
||||
|
||||
> **Note:** MSCCL++ performs auto-tuning on first initialization, which may add a few seconds to startup time. The tuned configurations are cached for the lifetime of the process.
|
||||
@@ -24,6 +24,7 @@ from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.device_communicators.pymscclpp import PyMscclppCommunicator
|
||||
from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
cleanup_dist_env_and_memory,
|
||||
get_tensor_model_parallel_group,
|
||||
graph_capture,
|
||||
initialize_model_parallel,
|
||||
@@ -51,10 +52,12 @@ def pynccl_allreduce(
|
||||
|
||||
def _bench_graph_time(func, inp_randn, warmup_loop=2, graph_loop=10, test_loop=10):
|
||||
graph_input = inp_randn.clone()
|
||||
graph_input_snapshot = inp_randn.clone()
|
||||
with graph_capture() as graph_capture_context:
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=graph_capture_context.stream):
|
||||
for _ in range(graph_loop):
|
||||
graph_input.copy_(graph_input_snapshot)
|
||||
graph_out = func(graph_input)
|
||||
|
||||
graph.replay()
|
||||
@@ -222,3 +225,7 @@ if __name__ == "__main__":
|
||||
prof_dir = f"prof/msccl"
|
||||
os.makedirs(prof_dir, exist_ok=True)
|
||||
ctx.export_chrome_trace(f"{prof_dir}/trace_rank{dist.get_rank()}.json.gz")
|
||||
|
||||
pymscclpp_comm.destroy()
|
||||
dist.barrier()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
@@ -21,6 +21,7 @@ ARG GITHUB_ARTIFACTORY=github.com
|
||||
ARG INSTALL_FLASHINFER_JIT_CACHE=0
|
||||
ARG FLASHINFER_VERSION=0.6.12
|
||||
ARG MOONCAKE_VERSION=0.3.11.post1
|
||||
ARG MSCCLPP_VERSION=sglang-v0.9.1
|
||||
#if need other arg please add in MOONCAKE_COMPILE_ARG
|
||||
ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=ON"
|
||||
|
||||
@@ -443,6 +444,7 @@ ARG USE_LATEST_SGLANG
|
||||
ARG GITHUB_ARTIFACTORY
|
||||
ARG MOONCAKE_VERSION
|
||||
ARG MOONCAKE_COMPILE_ARG
|
||||
ARG MSCCLPP_VERSION
|
||||
|
||||
WORKDIR /sgl-workspace
|
||||
|
||||
@@ -523,6 +525,24 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \
|
||||
fi
|
||||
|
||||
# Install MSCCL++ Python dependencies and package (builds extension via CMake through pip)
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
git clone --depth=1 --branch ${MSCCLPP_VERSION} https://${GITHUB_ARTIFACTORY}/microsoft/mscclpp.git /tmp/mscclpp \
|
||||
&& case "${CUDA_VERSION}" in \
|
||||
12.*) \
|
||||
CMAKE_ARGS="-DMSCCLPP_BYPASS_GPU_CHECK=ON -DMSCCLPP_USE_CUDA=ON -DMSCCLPP_GPU_ARCHS=80,90,100,100a,103,103a" \
|
||||
python3 -m pip install "/tmp/mscclpp[cuda12]"; \
|
||||
;; \
|
||||
13.*) \
|
||||
CMAKE_ARGS="-DMSCCLPP_BYPASS_GPU_CHECK=ON -DMSCCLPP_USE_CUDA=ON -DMSCCLPP_GPU_ARCHS=80,90,100,100a,103,103a" \
|
||||
python3 -m pip install "/tmp/mscclpp[cuda13]"; \
|
||||
;; \
|
||||
*) \
|
||||
echo "Unsupported CUDA version for MSCCL++: ${CUDA_VERSION}" && exit 1; \
|
||||
;; \
|
||||
esac \
|
||||
&& rm -rf /tmp/mscclpp
|
||||
|
||||
# Install essential Python packages (use constraints to prevent conflicts)
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
python3 -m pip install -c /sgl-workspace/constraints.txt \
|
||||
|
||||
@@ -65,7 +65,10 @@ import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed.parallel_state import destroy_distributed_environment
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
destroy_distributed_environment,
|
||||
destroy_model_parallel,
|
||||
)
|
||||
from sglang.srt.entrypoints.engine import _set_envs_and_config
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
||||
from sglang.srt.layers.moe import initialize_moe_config
|
||||
@@ -936,6 +939,7 @@ def latency_test(
|
||||
fout.write(json.dumps(result) + "\n")
|
||||
|
||||
if server_args.tp_size > 1:
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
|
||||
|
||||
@@ -957,12 +961,17 @@ def main(server_args, bench_args):
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
# Calculate local ranks for multi-node setup
|
||||
nranks_per_node = server_args.tp_size // server_args.nnodes
|
||||
local_rank_start = server_args.node_rank * nranks_per_node
|
||||
local_rank_end = local_rank_start + nranks_per_node
|
||||
|
||||
if server_args.tp_size == 1:
|
||||
work_func(server_args, port_args, bench_args, 0, 0)
|
||||
else:
|
||||
workers = []
|
||||
for tp_rank in range(server_args.tp_size):
|
||||
with maybe_reindex_device_id(tp_rank) as gpu_id:
|
||||
for tp_rank in range(local_rank_start, local_rank_end):
|
||||
with maybe_reindex_device_id(tp_rank - local_rank_start) as gpu_id:
|
||||
proc = multiprocessing.Process(
|
||||
target=work_func,
|
||||
args=(
|
||||
|
||||
@@ -328,7 +328,8 @@ class CustomAllreduce:
|
||||
|
||||
def close(self):
|
||||
if not self.disabled and self._ptr:
|
||||
ops.dispose(self._ptr)
|
||||
if ops is not None:
|
||||
ops.dispose(self._ptr)
|
||||
if _is_cuda:
|
||||
self.free_shared_buffer(self.meta_ptrs)
|
||||
self.free_shared_buffer(self.buffer_ptrs)
|
||||
|
||||
@@ -16,8 +16,6 @@ _is_musa = is_musa()
|
||||
|
||||
IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip or _is_musa
|
||||
IS_QUICK_AR_AVAILABLE = _is_hip
|
||||
# TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
|
||||
IS_MSCCLPP_AR_AVAILABLE = _is_cuda
|
||||
|
||||
try:
|
||||
import sgl_kernel.allreduce as _custom_ar
|
||||
@@ -26,7 +24,6 @@ except ImportError as e:
|
||||
logger.warning("Failed to import from custom_ar with %r", e)
|
||||
IS_CUSTOM_AR_AVAILABLE = False
|
||||
IS_QUICK_AR_AVAILABLE = False
|
||||
IS_MSCCLPP_AR_AVAILABLE = False
|
||||
|
||||
# region IS_CUSTOM_AR_AVAILABLE
|
||||
|
||||
@@ -167,44 +164,3 @@ elif _is_hip:
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region IS_MSCCLPP_AR_AVAILABLE
|
||||
|
||||
if not IS_MSCCLPP_AR_AVAILABLE:
|
||||
pass
|
||||
|
||||
elif _is_cuda:
|
||||
|
||||
def mscclpp_generate_unique_id() -> bytes:
|
||||
return _custom_ar.mscclpp_generate_unique_id()
|
||||
|
||||
def mscclpp_init_context(
|
||||
unique_id: bytes,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
scratch: torch.Tensor,
|
||||
put_buffer: torch.Tensor,
|
||||
nranks_per_node: int,
|
||||
rank_to_node: List[int],
|
||||
rank_to_ib: List[int],
|
||||
context_selection: int,
|
||||
) -> int:
|
||||
return _custom_ar.mscclpp_init_context(
|
||||
unique_id,
|
||||
rank,
|
||||
world_size,
|
||||
scratch,
|
||||
put_buffer,
|
||||
nranks_per_node,
|
||||
rank_to_node,
|
||||
rank_to_ib,
|
||||
context_selection,
|
||||
)
|
||||
|
||||
def mscclpp_allreduce(
|
||||
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
|
||||
) -> None:
|
||||
return _custom_ar.mscclpp_allreduce(context, inp, out, nthreads, nblocks)
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
@@ -1,108 +1,247 @@
|
||||
import bisect
|
||||
import importlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from enum import IntEnum
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup, ReduceOp
|
||||
|
||||
import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.srt.compilation.piecewise_context_manager import (
|
||||
get_pcg_capture_stream,
|
||||
is_in_pcg_torch_compile,
|
||||
is_in_piecewise_cuda_graph,
|
||||
)
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_is_hip = is_hip()
|
||||
|
||||
|
||||
class MscclContextSelection(IntEnum):
|
||||
MSCCL1SHOT1NODELL = 1
|
||||
MSCCL1SHOT2NODELL = 2
|
||||
|
||||
|
||||
def mscclpp_is_weak_contiguous(inp: torch.Tensor):
|
||||
return inp.is_contiguous() or (
|
||||
inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
|
||||
== inp.numel() * inp.element_size()
|
||||
)
|
||||
|
||||
|
||||
def mscclpp_convert_to_bytes(size_str):
|
||||
"""
|
||||
Converts a human-readable size string (e.g., "1MB", "2.5kb", "3 GB")
|
||||
into the equivalent number of bytes using binary units.
|
||||
|
||||
Args:
|
||||
size_str (str): A string representing size with unit (KB, MB, GB).
|
||||
|
||||
Returns:
|
||||
int: Number of bytes.
|
||||
"""
|
||||
size_str = size_str.strip().lower()
|
||||
|
||||
if not size_str:
|
||||
raise ValueError("Empty input string")
|
||||
|
||||
# Extract numeric part and unit
|
||||
for i in range(len(size_str)):
|
||||
if not size_str[i].isdigit() and size_str[i] != ".":
|
||||
break
|
||||
num_str = size_str[:i]
|
||||
unit = size_str[i:].strip()
|
||||
|
||||
try:
|
||||
num = float(num_str)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid numeric value in '{size_str}'")
|
||||
|
||||
# Conversion factors
|
||||
if unit == "b":
|
||||
return int(num)
|
||||
elif unit == "kb":
|
||||
return int(num * 1024)
|
||||
elif unit == "mb":
|
||||
return int(num * 1024 * 1024)
|
||||
elif unit == "gb":
|
||||
return int(num * 1024 * 1024 * 1024)
|
||||
else:
|
||||
raise ValueError(f"Unsupported unit: {unit}, support B, KB, MB, GB only")
|
||||
|
||||
|
||||
def mscclpp_bench_time(func, test_niter: int = 10, warmup_niter: int = 2):
|
||||
# warmup
|
||||
for _ in range(warmup_niter):
|
||||
func()
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier()
|
||||
start_event.record()
|
||||
for _ in range(test_niter):
|
||||
func()
|
||||
end_event.record()
|
||||
end_event.synchronize()
|
||||
func_cost_us = start_event.elapsed_time(end_event) / test_niter * 1000
|
||||
return func_cost_us
|
||||
|
||||
|
||||
class PyMscclppCommunicator:
|
||||
_SUPPORTED_WORLD_SIZES = [8, 16]
|
||||
_MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB"))
|
||||
_SUPPORTED_WORLD_SIZES = [8, 16, 32]
|
||||
_SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16]
|
||||
|
||||
# max_bytes: max supported mscclpp allreduce size
|
||||
# in A100 mscclpp is faster than nccl only under condition of msg size smaller than1MB
|
||||
def _is_symm_mem_enabled(self) -> bool:
|
||||
try:
|
||||
return get_global_server_args().enable_symm_mem
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def _is_weak_contiguous(self, inp: torch.Tensor):
|
||||
return inp.is_contiguous() or (
|
||||
inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
|
||||
== inp.numel() * inp.element_size()
|
||||
)
|
||||
|
||||
def _get_tuned_config(self, size):
|
||||
if size <= 512:
|
||||
target_size = 512
|
||||
elif size > 256 * 1024 * 1024:
|
||||
target_size = 256 * 1024 * 1024
|
||||
else:
|
||||
target_size = 1 << (size - 1).bit_length()
|
||||
return self.best_configs.get(target_size)
|
||||
|
||||
def _create_dsl_algorithms(self):
|
||||
dsl_algos_config = []
|
||||
n_nodes = self.world_size // self.nranks_per_node
|
||||
if n_nodes == 2 or n_nodes == 4:
|
||||
for tbg in [1, 2, 4, 8]:
|
||||
for num_threads_per_block in [256, 512, 768, 1024]:
|
||||
spec = self.mscclpp.language.AlgoSpec(
|
||||
name=f"allreduce_{n_nodes}node_{tbg}TBG_{num_threads_per_block}TPB",
|
||||
collective=self.mscclpp.language.collectives.AllReduce(
|
||||
self.world_size, 1, True
|
||||
),
|
||||
nranks_per_node=self.nranks_per_node,
|
||||
world_size=self.world_size,
|
||||
in_place=True,
|
||||
instances=1,
|
||||
protocol="LL",
|
||||
auto_sync=False,
|
||||
num_threads_per_block=num_threads_per_block,
|
||||
reuse_resources=True,
|
||||
use_double_scratch_buffer=True,
|
||||
min_message_size=tbg * (1 << 10),
|
||||
max_message_size=8 << 20,
|
||||
tags={"default": 1},
|
||||
)
|
||||
algo = self.mscclpp.compile(
|
||||
self.def_algo.allreduce_multi_nodes,
|
||||
spec,
|
||||
self.rank,
|
||||
thread_block_group_size=tbg,
|
||||
)
|
||||
dsl_algos_config.append((algo, [0], [0]))
|
||||
return dsl_algos_config
|
||||
|
||||
def _create_native_algorithms(self):
|
||||
navitve_algorithms_config = []
|
||||
dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
|
||||
data_type=str(torch.float16)
|
||||
)
|
||||
self.scratch_buffer = torch.utils.dlpack.from_dlpack(dlpack)
|
||||
self.flag_buffer = torch.ones(128, dtype=torch.uint32, device="cuda")
|
||||
algos = self.mscclpp_ext.AlgorithmCollectionBuilder().build_default_algorithms(
|
||||
scratch_buffer=self.scratch_buffer.data_ptr(),
|
||||
scratch_buffer_size=self.scratch_buffer.nbytes,
|
||||
rank=self.rank,
|
||||
)
|
||||
|
||||
for algo in algos:
|
||||
if algo.name == "default_allreduce_nvls_packet":
|
||||
algo.set_message_size_range(0, 512 << 10)
|
||||
navitve_algorithms_config.append(
|
||||
(algo, [4, 8, 12, 16], [256, 512, 768, 1024])
|
||||
)
|
||||
if algo.name == "default_allreduce_packet":
|
||||
algo.set_message_size_range(0, 2 << 20)
|
||||
navitve_algorithms_config.append(
|
||||
(algo, [14, 21, 28, 42, 56], [256, 512, 768, 1024])
|
||||
)
|
||||
if algo.name == "default_allreduce_rsag_zero_copy":
|
||||
algo.set_message_size_range(512 << 10, 4 << 30)
|
||||
navitve_algorithms_config.append(
|
||||
(algo, [32, 48, 64, 128], [256, 512, 768, 1024])
|
||||
)
|
||||
if (
|
||||
self.symm_mem_enabled
|
||||
and algo.name == "default_allreduce_nvls_zero_copy"
|
||||
):
|
||||
algo.set_message_size_range(512 << 10, 4 << 30)
|
||||
navitve_algorithms_config.append(
|
||||
(algo, [4, 8, 12, 16, 32], [256, 512, 768, 1024])
|
||||
)
|
||||
|
||||
return navitve_algorithms_config
|
||||
|
||||
def _create_algorithms(self):
|
||||
if self.world_size == 8:
|
||||
self.algos_config = self._create_native_algorithms()
|
||||
self._tune(5, 10, 20, self.algos_config)
|
||||
elif self.world_size == 16 or self.world_size == 32:
|
||||
self.dsl_algos_config = self._create_dsl_algorithms()
|
||||
self._tune(5, 10, 20, self.dsl_algos_config)
|
||||
|
||||
def _get_time(
|
||||
self,
|
||||
algo,
|
||||
tune_tensor,
|
||||
size,
|
||||
nb,
|
||||
nt,
|
||||
n_warmup,
|
||||
n_graph_launches,
|
||||
n_ops_per_graph,
|
||||
):
|
||||
# Check if the algorithm can run with the given configuration
|
||||
if self._run_algo(algo, tune_tensor, size, nb, nt, True) != 0:
|
||||
return float("inf")
|
||||
|
||||
# Warmup iterations to stabilize performance
|
||||
for _ in range(n_warmup):
|
||||
self._run_algo(algo, tune_tensor, size, nb, nt, True)
|
||||
|
||||
# Warmup on capture stream
|
||||
capture_stream = torch.cuda.Stream()
|
||||
capture_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(capture_stream):
|
||||
self._run_algo(algo, tune_tensor, size, nb, nt, True)
|
||||
capture_stream.synchronize()
|
||||
|
||||
# Capture the algorithm execution in a CUDA graph
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g, stream=capture_stream):
|
||||
for _ in range(n_ops_per_graph):
|
||||
self._run_algo(algo, tune_tensor, size, nb, nt, True)
|
||||
|
||||
# Measure the execution time of the captured graph
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
start_event.record(capture_stream)
|
||||
with torch.cuda.stream(capture_stream):
|
||||
for _ in range(n_graph_launches):
|
||||
g.replay()
|
||||
end_event.record(capture_stream)
|
||||
end_event.synchronize()
|
||||
elapsed = start_event.elapsed_time(end_event)
|
||||
|
||||
# Synchronize timing results across all ranks to ensure consistent algorithm selection
|
||||
# replicate n times such due to algo limitations
|
||||
time_tensor = torch.full(
|
||||
(self.world_size,), elapsed, dtype=torch.float64, device="cuda"
|
||||
).to(dtype=torch.float32)
|
||||
torch.cuda.current_stream().wait_stream(capture_stream)
|
||||
if self.rank == 0:
|
||||
avg_time = time_tensor[self.rank].item() / self.world_size
|
||||
tensor = torch.tensor([avg_time])
|
||||
else:
|
||||
tensor = torch.empty(1)
|
||||
dist.broadcast(tensor, src=0, group=self.group)
|
||||
avg_time = tensor.item()
|
||||
|
||||
return avg_time
|
||||
|
||||
def _tune(self, n_warmup, n_graph_launches, n_ops_per_graph, algos_config):
|
||||
sizes = [1 << i for i in range(9, 24)]
|
||||
dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
|
||||
data_type=str(torch.float16)
|
||||
)
|
||||
tune_tensor = torch.utils.dlpack.from_dlpack(dlpack)
|
||||
|
||||
for size in sizes:
|
||||
best_time = float("inf")
|
||||
best_config = None
|
||||
for i in range(len(algos_config)):
|
||||
algo, candidates_nblocks, candidates_nthreads = algos_config[i]
|
||||
if (
|
||||
size >= algo.message_size_range[0]
|
||||
and size <= algo.message_size_range[1]
|
||||
):
|
||||
for nb in candidates_nblocks:
|
||||
for nt in candidates_nthreads:
|
||||
avg_time = self._get_time(
|
||||
algo,
|
||||
tune_tensor,
|
||||
size,
|
||||
nb,
|
||||
nt,
|
||||
n_warmup,
|
||||
n_graph_launches,
|
||||
n_ops_per_graph,
|
||||
)
|
||||
if avg_time < best_time:
|
||||
best_time = avg_time
|
||||
best_config = (algo, nb, nt)
|
||||
if best_config:
|
||||
self.best_configs[size] = best_config
|
||||
|
||||
torch.cuda.synchronize()
|
||||
for algo, _, _ in algos_config:
|
||||
algo.reset()
|
||||
|
||||
def _run_algo(self, algo, tensor, size, nblocks, nthreads, sym_mem_enabled=False):
|
||||
return algo.execute(
|
||||
comm=self.comm.communicator,
|
||||
executor=self.executor,
|
||||
input_buffer=tensor.data_ptr(),
|
||||
output_buffer=tensor.data_ptr(),
|
||||
input_size=size,
|
||||
output_size=size,
|
||||
dtype=self.dtype_to_mscclpp_dtype(tensor.dtype),
|
||||
op=self.mscclpp.ReduceOp.SUM,
|
||||
stream=torch.cuda.current_stream().cuda_stream,
|
||||
nblocks=nblocks,
|
||||
nthreads_per_block=nthreads,
|
||||
symmetric_memory=sym_mem_enabled,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
group: ProcessGroup,
|
||||
device: Union[int, str, torch.device],
|
||||
max_bytes=_MAX_BYTES,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
"""Args:
|
||||
group: the process group to work on. If None, it will use the
|
||||
default process group.
|
||||
device: the device to bind the CustomAllreduce to. If None,
|
||||
@@ -114,11 +253,16 @@ class PyMscclppCommunicator:
|
||||
self._IS_CAPTURING = False
|
||||
self.disabled = True
|
||||
|
||||
if not ops.IS_MSCCLPP_AR_AVAILABLE:
|
||||
# disable because of missing mscclpp library
|
||||
# e.g. in a non-cuda environment
|
||||
try:
|
||||
self.mscclpp = importlib.import_module("mscclpp")
|
||||
self.mscclpp_ext = importlib.import_module("mscclpp.ext")
|
||||
self.def_algo = importlib.import_module("mscclpp.default_algos")
|
||||
except ImportError:
|
||||
self.available = False
|
||||
self.mscclpp = None
|
||||
return
|
||||
|
||||
self.available = True
|
||||
self.group = group
|
||||
|
||||
assert (
|
||||
@@ -161,137 +305,83 @@ class PyMscclppCommunicator:
|
||||
assert isinstance(device, torch.device)
|
||||
self.device = device
|
||||
|
||||
self.max_bytes = max_bytes
|
||||
self.rank = rank
|
||||
self.world_size = world_size
|
||||
|
||||
if dist.get_rank(group) == 0:
|
||||
unique_id = [ops.mscclpp_generate_unique_id()]
|
||||
else:
|
||||
unique_id = [None]
|
||||
dist.broadcast_object_list(unique_id, src=self.ranks[0], group=self.group)
|
||||
self.unique_id = unique_id[0]
|
||||
self.rank_to_node, self.rank_to_ib = list(range(world_size)), list(
|
||||
range(world_size)
|
||||
self.comm = self.mscclpp.CommGroup(
|
||||
torch_group=self.group, rank=rank, size=world_size
|
||||
)
|
||||
for r in range(world_size):
|
||||
self.rank_to_node[r] = r // 8
|
||||
self.rank_to_ib[r] = self.rank % 8
|
||||
self.executor = self.mscclpp.Executor(self.comm.communicator)
|
||||
self.symm_mem_enabled = self._is_symm_mem_enabled()
|
||||
self.best_configs = {}
|
||||
self._create_algorithms()
|
||||
|
||||
self._context = None
|
||||
self.context_selection = None
|
||||
self.msg_size_for_finetune = [
|
||||
2**i for i in range(10, math.floor(math.log2(self.max_bytes)) + 1)
|
||||
]
|
||||
self.msg_size2best_config = {}
|
||||
if world_size == 8:
|
||||
self.context_selection = MscclContextSelection.MSCCL1SHOT1NODELL
|
||||
elif world_size == 16:
|
||||
self.context_selection = MscclContextSelection.MSCCL1SHOT2NODELL
|
||||
if not _is_hip:
|
||||
self.scratch = torch.empty(
|
||||
self.max_bytes * 8,
|
||||
dtype=torch.uint8,
|
||||
device=self.device,
|
||||
)
|
||||
self.put_buffer = torch.empty(
|
||||
self.max_bytes * 8 // self.nranks_per_node,
|
||||
dtype=torch.uint8,
|
||||
device=self.device,
|
||||
)
|
||||
self._context = ops.mscclpp_init_context(
|
||||
self.unique_id,
|
||||
self.rank,
|
||||
self.world_size,
|
||||
self.scratch,
|
||||
self.put_buffer,
|
||||
self.nranks_per_node,
|
||||
self.rank_to_node,
|
||||
self.rank_to_ib,
|
||||
int(self.context_selection),
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError("HIP Mscclpp is not supported yet.")
|
||||
|
||||
self.msg_size2best_config = {}
|
||||
self.pre_tune_config()
|
||||
if dist.get_rank(group) == 0:
|
||||
msg_size2best_config = [self.msg_size2best_config]
|
||||
else:
|
||||
msg_size2best_config = [None]
|
||||
dist.broadcast_object_list(
|
||||
msg_size2best_config, src=self.ranks[0], group=self.group
|
||||
)
|
||||
self.msg_size2best_config = msg_size2best_config[0]
|
||||
|
||||
# PyMscclpp is enabled only in cuda graph
|
||||
self.disabled = True
|
||||
|
||||
def pre_tune_config(self, dtype=torch.bfloat16) -> bool:
|
||||
logger.debug(f"start to pre-tune configs for rank {self.rank}")
|
||||
nthreads_to_try = [256, 512, 1024]
|
||||
nblocks_to_try = [21, 42, 84]
|
||||
inp_randn = torch.ones(
|
||||
self.msg_size_for_finetune[-1] // dtype.itemsize, dtype=dtype, device="cuda"
|
||||
)
|
||||
oup_randn = torch.empty_like(inp_randn)
|
||||
for msg_size in self.msg_size_for_finetune:
|
||||
mock_inp, mock_outp = (
|
||||
inp_randn[: msg_size // dtype.itemsize],
|
||||
oup_randn[: msg_size // dtype.itemsize],
|
||||
)
|
||||
best_config, best_time = None, None
|
||||
for nthreads in nthreads_to_try:
|
||||
for nblocks in nblocks_to_try:
|
||||
cur_cost = mscclpp_bench_time(
|
||||
lambda: ops.mscclpp_allreduce(
|
||||
self._context, mock_inp, mock_outp, nthreads, nblocks
|
||||
)
|
||||
)
|
||||
if best_time is None or cur_cost < best_time:
|
||||
best_config = (nthreads, nblocks)
|
||||
best_time = cur_cost
|
||||
self.msg_size2best_config[msg_size] = best_config
|
||||
if self.rank == 0:
|
||||
logger.debug(
|
||||
f"for msg_size {msg_size}, best_config: {best_config}, best_time: {best_time}us"
|
||||
)
|
||||
def destroy(self):
|
||||
self.algos_config = None
|
||||
self.best_configs = None
|
||||
self.executor = None
|
||||
self.scratch_buffer = None
|
||||
self.flag_buffer = None
|
||||
self.comm = None
|
||||
|
||||
def should_mscclpp_allreduce(
|
||||
self, inp: torch.Tensor, op: ReduceOp = ReduceOp.SUM
|
||||
) -> bool:
|
||||
if self.disabled or self._context is None:
|
||||
if (
|
||||
self.disabled
|
||||
or self.world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES
|
||||
):
|
||||
return False
|
||||
if inp.dtype not in PyMscclppCommunicator._SUPPORTED_DTYPE:
|
||||
return False
|
||||
if not mscclpp_is_weak_contiguous(inp):
|
||||
if not self._is_weak_contiguous(inp):
|
||||
return False
|
||||
# only support sum op
|
||||
if op != ReduceOp.SUM:
|
||||
if op is not ReduceOp.SUM:
|
||||
return False
|
||||
if inp.numel() * inp.element_size() > self.max_bytes:
|
||||
if self._get_tuned_config(inp.numel() * inp.element_size()) is None:
|
||||
return False
|
||||
# mscclpp must not be used during any piecewise CUDA graph phase
|
||||
# (compile, capture, or replay) as it changes the allreduce dispatch
|
||||
# path and triggers recompilation.
|
||||
if (
|
||||
is_in_piecewise_cuda_graph()
|
||||
or is_in_pcg_torch_compile()
|
||||
or get_pcg_capture_stream() is not None
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def all_reduce(self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM):
|
||||
if self._IS_CAPTURING:
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
self.graph_input_set.add((tensor.dtype, tensor.numel()))
|
||||
msg_size = tensor.numel() * tensor.itemsize
|
||||
index = bisect.bisect_left(self.msg_size_for_finetune, msg_size)
|
||||
msg_size_finetune = self.msg_size_for_finetune[index]
|
||||
nthreads, nblocks = self.msg_size2best_config[msg_size_finetune]
|
||||
result = torch.empty_like(tensor)
|
||||
ops.mscclpp_allreduce(self._context, tensor, result, nthreads, nblocks)
|
||||
return result
|
||||
def dtype_to_mscclpp_dtype(self, dtype: torch.dtype):
|
||||
if dtype == torch.float16:
|
||||
return self.mscclpp.DataType.float16
|
||||
elif dtype == torch.float32:
|
||||
return self.mscclpp.DataType.float32
|
||||
elif dtype == torch.int32:
|
||||
return self.mscclpp.DataType.int32
|
||||
elif dtype == torch.bfloat16:
|
||||
return self.mscclpp.DataType.bfloat16
|
||||
else:
|
||||
raise ValueError(f"Unknown data type: {dtype}")
|
||||
|
||||
def all_reduce(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
stream: torch.cuda.Stream = None,
|
||||
):
|
||||
assert op == torch.distributed.ReduceOp.SUM
|
||||
nbytes = tensor.numel() * tensor.element_size()
|
||||
algo, nblocks, nthreads = self._get_tuned_config(nbytes)
|
||||
self._run_algo(algo, tensor, nbytes, nblocks, nthreads, self.symm_mem_enabled)
|
||||
return tensor
|
||||
|
||||
@contextmanager
|
||||
def change_state(
|
||||
self,
|
||||
enable: Optional[bool] = None,
|
||||
):
|
||||
if enable is None:
|
||||
if enable is None or self.available is False:
|
||||
# guess a default value when not specified
|
||||
# DO: Decided if raise an exception here or not
|
||||
enable = self.available
|
||||
|
||||
old_disable = self.disabled
|
||||
|
||||
@@ -605,7 +605,15 @@ class GroupCoordinator:
|
||||
if self.npu_communicator is not None and not self.npu_communicator.disabled:
|
||||
return self.npu_communicator.all_reduce(input_)
|
||||
|
||||
if self.pynccl_comm is not None and self.is_symmetric_memory_enabled():
|
||||
should_use_pymscclpp_allreduce = (
|
||||
self.pymscclpp_comm is not None
|
||||
and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
|
||||
)
|
||||
if (
|
||||
self.pynccl_comm is not None
|
||||
and self.is_symmetric_memory_enabled()
|
||||
and not should_use_pymscclpp_allreduce
|
||||
):
|
||||
self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
|
||||
with self.pynccl_comm.change_state(enable=True):
|
||||
self.pynccl_comm.all_reduce(input_)
|
||||
@@ -615,6 +623,7 @@ class GroupCoordinator:
|
||||
if (
|
||||
self.ca_comm is not None
|
||||
and not self.ca_comm.disabled
|
||||
and not should_use_pymscclpp_allreduce
|
||||
and self.ca_comm.should_custom_ar(input_)
|
||||
):
|
||||
outplace_all_reduce_method = "ca"
|
||||
@@ -624,11 +633,7 @@ class GroupCoordinator:
|
||||
and self.qr_comm.should_quick_allreduce(input_)
|
||||
):
|
||||
outplace_all_reduce_method = "qr"
|
||||
elif (
|
||||
self.pymscclpp_comm is not None
|
||||
and not self.pymscclpp_comm.disabled
|
||||
and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
|
||||
):
|
||||
elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce:
|
||||
outplace_all_reduce_method = "pymscclpp"
|
||||
elif (
|
||||
self.torch_symm_mem_comm is not None
|
||||
@@ -1473,6 +1478,8 @@ class GroupCoordinator:
|
||||
self.cpu_group = None
|
||||
if self.pynccl_comm is not None:
|
||||
self.pynccl_comm = None
|
||||
if self.pymscclpp_comm is not None:
|
||||
self.pymscclpp_comm.destroy()
|
||||
if self.ca_comm is not None:
|
||||
self.ca_comm = None
|
||||
if self.mq_broadcaster is not None:
|
||||
|
||||
@@ -86,24 +86,6 @@ FetchContent_Declare(
|
||||
)
|
||||
FetchContent_Populate(repo-flash-attention)
|
||||
|
||||
# mscclpp
|
||||
FetchContent_Declare(
|
||||
repo-mscclpp
|
||||
URL https://${GITHUB_ARTIFACTORY}/microsoft/mscclpp/archive/51eca89d20f0cfb3764ccd764338d7b22cd486a6.tar.gz
|
||||
URL_HASH SHA256=b064de701da5253e32f4031d16b01a61d78c369327d585ff7bfb521bbe742677
|
||||
)
|
||||
FetchContent_Populate(repo-mscclpp)
|
||||
|
||||
# mscclpp's own CMakeLists.txt hardcodes a github.com FetchContent for nlohmann/json.
|
||||
# Patch it to route through GITHUB_ARTIFACTORY so it follows the same mirror as our deps.
|
||||
file(READ "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" _mscclpp_cmakelists)
|
||||
string(REPLACE
|
||||
"https://github.com/nlohmann/json"
|
||||
"https://${GITHUB_ARTIFACTORY}/nlohmann/json"
|
||||
_mscclpp_cmakelists "${_mscclpp_cmakelists}")
|
||||
file(WRITE "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" "${_mscclpp_cmakelists}")
|
||||
unset(_mscclpp_cmakelists)
|
||||
|
||||
# ccache option
|
||||
option(ENABLE_CCACHE "Whether to use ccache" ON)
|
||||
find_program(CCACHE_FOUND ccache)
|
||||
@@ -257,7 +239,6 @@ endif()
|
||||
# NOTE: Please sort the filenames alphabetically
|
||||
set(SOURCES
|
||||
"csrc/allreduce/custom_all_reduce.cu"
|
||||
"csrc/allreduce/mscclpp_allreduce.cu"
|
||||
"csrc/attention/cutlass_mla_kernel.cu"
|
||||
"csrc/attention/merge_attn_states.cu"
|
||||
"csrc/attention/vertical_slash_index.cu"
|
||||
@@ -329,7 +310,6 @@ set(INCLUDES
|
||||
${repo-cutlass_SOURCE_DIR}/tools/util/include
|
||||
${repo-flashinfer_SOURCE_DIR}/include
|
||||
${repo-flashinfer_SOURCE_DIR}/csrc
|
||||
${repo-mscclpp_SOURCE_DIR}/include
|
||||
${repo-cutlass_SOURCE_DIR}/examples/77_blackwell_fmha
|
||||
${repo-cutlass_SOURCE_DIR}/examples/common
|
||||
${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src
|
||||
@@ -379,19 +359,8 @@ else()
|
||||
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1")
|
||||
endif()
|
||||
|
||||
# mscclpp option
|
||||
set(MSCCLPP_USE_CUDA ON)
|
||||
set(MSCCLPP_BYPASS_GPU_CHECK ON)
|
||||
set(MSCCLPP_BUILD_TESTS OFF)
|
||||
set(MSCCLPP_BUILD_PYTHON_BINDINGS OFF)
|
||||
set(MSCCLPP_BUILD_APPS_NCCL OFF)
|
||||
add_subdirectory(
|
||||
${repo-mscclpp_SOURCE_DIR}
|
||||
${CMAKE_CURRENT_BINARY_DIR}/mscclpp-build
|
||||
)
|
||||
|
||||
target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt mscclpp_static)
|
||||
target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt mscclpp_static)
|
||||
target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
|
||||
target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
|
||||
|
||||
# sparse flash attention
|
||||
target_compile_definitions(common_ops_sm90_build PRIVATE
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/library.h>
|
||||
|
||||
#include "mscclpp_allreduce.cuh"
|
||||
|
||||
enum MscclContextSelection {
|
||||
MSCCL1NODELL = 1,
|
||||
MSCCL2NODELL = 2,
|
||||
};
|
||||
|
||||
class MscclContext {
|
||||
public:
|
||||
MscclContextSelection selection_;
|
||||
std::shared_ptr<sglang::Msccl1NodeLLcontext> msccl_1nodeLL_context;
|
||||
std::shared_ptr<sglang::Msccl2NodeLLcontext> msccl_2nodeLL_context;
|
||||
MscclContext(MscclContextSelection selection) : selection_(selection) {}
|
||||
template <typename T>
|
||||
void allreduce(
|
||||
cudaStream_t stream, T* input, T* output, const size_t input_numel, int threads = 512, int block_limit = 21) {
|
||||
if (selection_ == MSCCL1NODELL) {
|
||||
msccl_1nodeLL_context->allreduce<T>(stream, input, output, input_numel, threads, block_limit);
|
||||
} else if (selection_ == MSCCL2NODELL) {
|
||||
msccl_2nodeLL_context->allreduce<T>(stream, input, output, input_numel, threads, block_limit);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
using fptr_t = int64_t;
|
||||
static_assert(sizeof(void*) == sizeof(fptr_t));
|
||||
|
||||
torch::Tensor _unique_id2tensor(const mscclpp::UniqueId& unique_id) {
|
||||
auto options = torch::TensorOptions().dtype(torch::kByte).device(torch::kCPU);
|
||||
auto tensor = torch::empty({static_cast<int64_t>(unique_id.size())}, options);
|
||||
std::memcpy(tensor.data_ptr<uint8_t>(), unique_id.data(), unique_id.size());
|
||||
return tensor;
|
||||
}
|
||||
|
||||
// Function to convert vector of int32_t back to array of uint8_t
|
||||
mscclpp::UniqueId _tensor2unique_id(const torch::Tensor& tensor) {
|
||||
mscclpp::UniqueId unique_id;
|
||||
std::memcpy(unique_id.data(), tensor.data_ptr<uint8_t>(), unique_id.size());
|
||||
return unique_id;
|
||||
}
|
||||
|
||||
torch::Tensor mscclpp_generate_unique_id() {
|
||||
mscclpp::UniqueId unique_id = mscclpp::TcpBootstrap::createUniqueId();
|
||||
return _unique_id2tensor(unique_id);
|
||||
}
|
||||
|
||||
fptr_t mscclpp_init_context(
|
||||
const torch::Tensor& unique_id,
|
||||
const int64_t rank,
|
||||
const int64_t world_size,
|
||||
torch::Tensor& scratch,
|
||||
torch::Tensor& put_buffer,
|
||||
const int64_t nranks_per_node,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib,
|
||||
const int64_t context_selection) {
|
||||
MscclContext* context_ptr = new MscclContext(static_cast<MscclContextSelection>(context_selection));
|
||||
mscclpp::UniqueId uid = _tensor2unique_id(unique_id);
|
||||
if (context_selection == MSCCL1NODELL) {
|
||||
void* scratch_ptr = reinterpret_cast<void*>(scratch.data_ptr());
|
||||
const size_t scratch_bytes = scratch.numel() * scratch.element_size();
|
||||
context_ptr->msccl_1nodeLL_context = std::make_shared<sglang::Msccl1NodeLLcontext>(
|
||||
uid, rank, world_size, scratch_ptr, scratch_bytes, nranks_per_node, rank_to_node, rank_to_ib);
|
||||
} else if (context_selection == MSCCL2NODELL) {
|
||||
void* scratch_ptr = reinterpret_cast<void*>(scratch.data_ptr());
|
||||
const size_t scratch_bytes = scratch.numel() * scratch.element_size();
|
||||
void* put_buffer_ptr = reinterpret_cast<void*>(put_buffer.data_ptr());
|
||||
const size_t put_buffer_bytes = put_buffer.numel() * put_buffer.element_size();
|
||||
context_ptr->msccl_2nodeLL_context = std::make_shared<sglang::Msccl2NodeLLcontext>(
|
||||
uid,
|
||||
rank,
|
||||
world_size,
|
||||
scratch_ptr,
|
||||
scratch_bytes,
|
||||
put_buffer_ptr,
|
||||
put_buffer_bytes,
|
||||
nranks_per_node,
|
||||
rank_to_node,
|
||||
rank_to_ib);
|
||||
} else {
|
||||
throw std::runtime_error("invalid context selection");
|
||||
}
|
||||
return (fptr_t)context_ptr;
|
||||
}
|
||||
|
||||
bool _mscclpp_is_weak_contiguous(torch::Tensor& t) {
|
||||
return t.is_contiguous() ||
|
||||
(t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
|
||||
}
|
||||
void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out, int64_t nthreads, int64_t nblocks) {
|
||||
MscclContext* context = reinterpret_cast<MscclContext*>(_context);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
|
||||
auto stream = c10::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
|
||||
TORCH_CHECK_EQ(inp.numel(), out.numel());
|
||||
TORCH_CHECK(_mscclpp_is_weak_contiguous(out));
|
||||
TORCH_CHECK(_mscclpp_is_weak_contiguous(inp));
|
||||
switch (out.scalar_type()) {
|
||||
case at::ScalarType::Float: {
|
||||
context->allreduce<float>(
|
||||
stream,
|
||||
reinterpret_cast<float*>(inp.data_ptr()),
|
||||
reinterpret_cast<float*>(out.data_ptr()),
|
||||
inp.numel(),
|
||||
nthreads,
|
||||
nblocks);
|
||||
break;
|
||||
}
|
||||
case at::ScalarType::Half: {
|
||||
context->allreduce<half>(
|
||||
stream,
|
||||
reinterpret_cast<half*>(inp.data_ptr()),
|
||||
reinterpret_cast<half*>(out.data_ptr()),
|
||||
inp.numel(),
|
||||
nthreads,
|
||||
nblocks);
|
||||
break;
|
||||
}
|
||||
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
|
||||
case at::ScalarType::BFloat16: {
|
||||
context->allreduce<__nv_bfloat16>(
|
||||
stream,
|
||||
reinterpret_cast<__nv_bfloat16*>(inp.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
|
||||
inp.numel(),
|
||||
nthreads,
|
||||
nblocks);
|
||||
break;
|
||||
}
|
||||
#endif
|
||||
default:
|
||||
throw std::runtime_error("custom allreduce only supports float32, float16 and bfloat16");
|
||||
}
|
||||
}
|
||||
@@ -1,779 +0,0 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// Licensed under the MIT license.
|
||||
#pragma once
|
||||
#ifdef USE_ROCM
|
||||
#include <hip/hip_fp16.h>
|
||||
#else
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#endif
|
||||
|
||||
#include <mscclpp/concurrency_device.hpp>
|
||||
#include <mscclpp/core.hpp>
|
||||
#include <mscclpp/memory_channel.hpp>
|
||||
#include <mscclpp/memory_channel_device.hpp>
|
||||
#include <mscclpp/nvls_device.hpp>
|
||||
#include <mscclpp/port_channel.hpp>
|
||||
#include <mscclpp/port_channel_device.hpp>
|
||||
|
||||
// comment this for test_mscclpp_allreduce.cu
|
||||
#include "utils.h"
|
||||
|
||||
namespace sglang {
|
||||
|
||||
__device__ mscclpp::DeviceSyncer deviceSyncer;
|
||||
__device__ mscclpp::DeviceSyncer allGatherDeviceSyncer;
|
||||
__device__ mscclpp::DeviceSyncer reduceScatterDeviceSyncer;
|
||||
__device__ mscclpp::DeviceSyncer ibDeviceSyncer;
|
||||
|
||||
template <typename To, typename From>
|
||||
__forceinline__ __device__ To bit_cast(const From& src) {
|
||||
static_assert(sizeof(To) == sizeof(From), "Size mismatch for bit_cast");
|
||||
|
||||
union {
|
||||
From f;
|
||||
To t;
|
||||
} u;
|
||||
u.f = src;
|
||||
return u.t;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ T add_elements(T a, T b) {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ __half2 add_elements(__half2 a, __half2 b) {
|
||||
return __hadd2(a, b);
|
||||
}
|
||||
|
||||
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
|
||||
template <>
|
||||
__forceinline__ __device__ __nv_bfloat162 add_elements(__nv_bfloat162 a, __nv_bfloat162 b) {
|
||||
return __hadd2(a, b);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ int4 add_vectors_helper(int4 a, int4 b) {
|
||||
int4 ret;
|
||||
ret.w = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.w), bit_cast<T, int>(b.w)));
|
||||
ret.x = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.x), bit_cast<T, int>(b.x)));
|
||||
ret.y = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.y), bit_cast<T, int>(b.y)));
|
||||
ret.z = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.z), bit_cast<T, int>(b.z)));
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ int4 add_vectors(int4 a, int4 b) {
|
||||
return add_vectors_helper<T>(a, b);
|
||||
}
|
||||
|
||||
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
|
||||
template <>
|
||||
__forceinline__ __device__ int4 add_vectors<__nv_bfloat16>(int4 a, int4 b) {
|
||||
return add_vectors_helper<__nv_bfloat162>(a, b);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ int4 add_vectors<__half>(int4 a, int4 b) {
|
||||
return add_vectors_helper<__half2>(a, b);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ uint2 add_vectors_helper(uint2 a, uint2 b) {
|
||||
uint2 ret;
|
||||
ret.x = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.x), bit_cast<T, int>(b.x)));
|
||||
ret.y = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.y), bit_cast<T, int>(b.y)));
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ uint2 add_vectors(uint2 a, uint2 b) {
|
||||
return add_vectors_helper<T>(a, b);
|
||||
}
|
||||
|
||||
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
|
||||
template <>
|
||||
__forceinline__ __device__ uint2 add_vectors<__nv_bfloat16>(uint2 a, uint2 b) {
|
||||
return add_vectors_helper<__nv_bfloat162>(a, b);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ uint2 add_vectors<__half>(uint2 a, uint2 b) {
|
||||
return add_vectors_helper<__half2>(a, b);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ int add_vectors_helper(int a, int b) {
|
||||
return bit_cast<int, T>(add_elements(bit_cast<T, int>(a), bit_cast<T, int>(b)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ int add_vectors(int a, int b) {
|
||||
return add_vectors_helper<T>(a, b);
|
||||
}
|
||||
|
||||
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
|
||||
template <>
|
||||
__forceinline__ __device__ int add_vectors<__nv_bfloat16>(int a, int b) {
|
||||
return add_vectors_helper<__nv_bfloat162>(a, b);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ int add_vectors<__half>(int a, int b) {
|
||||
return add_vectors_helper<__half2>(a, b);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------
|
||||
// allreduce_LL_1node using LLPacket, origin allreduce2
|
||||
// -------------------------------------------------------
|
||||
|
||||
__device__ uint64_t globalFlag = 1;
|
||||
|
||||
template <typename TYPE>
|
||||
__global__ void __launch_bounds__(1024, 1) allreduce_LL_1node(
|
||||
mscclpp::MemoryChannelDeviceHandle* memChans,
|
||||
TYPE* buff,
|
||||
TYPE* scratch,
|
||||
void* resultBuff,
|
||||
int rank,
|
||||
int worldSize,
|
||||
size_t nelems) {
|
||||
nelems = nelems / (sizeof(int) / sizeof(TYPE));
|
||||
// This version of allreduce only works for single nodes
|
||||
const int nPeers = worldSize - 1;
|
||||
const size_t nPkts = nelems / 2;
|
||||
const int nelemsPerRank = nelems / worldSize;
|
||||
const int nPktsPerRank = nelemsPerRank / 2;
|
||||
// flag for packets. Initially 1
|
||||
const uint32_t flag = (uint32_t)globalFlag;
|
||||
// thread block & channel info
|
||||
const int nBlocksPerPeer = gridDim.x / nPeers;
|
||||
const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
|
||||
const int peerIdx = blockIdx.x / nBlocksPerPeer;
|
||||
const int remoteRank = peerIdx < rank ? peerIdx : peerIdx + 1;
|
||||
mscclpp::MemoryChannelDeviceHandle memChan = memChans[peerIdx];
|
||||
const int tid = threadIdx.x + localBlockIdx * blockDim.x;
|
||||
// double buffering
|
||||
size_t scratchBaseOffset = (flag & 1) ? 0 : nPkts * sizeof(mscclpp::LLPacket);
|
||||
void* scratchBuff = (void*)((char*)scratch + scratchBaseOffset);
|
||||
size_t scratchOffset = scratchBaseOffset + rank * nPktsPerRank * sizeof(mscclpp::LLPacket);
|
||||
size_t scratchResultOffset =
|
||||
(flag & 1) ? 2 * nPkts * sizeof(mscclpp::LLPacket) : 3 * nPkts * sizeof(mscclpp::LLPacket);
|
||||
size_t srcOffset = remoteRank * nelemsPerRank * sizeof(int);
|
||||
uint2* src = (uint2*)((char*)buff + rank * nelemsPerRank * sizeof(int));
|
||||
uint2* dst = (uint2*)((char*)resultBuff + rank * nelemsPerRank * sizeof(int));
|
||||
|
||||
// step 1: write to scratch buffer
|
||||
memChan.putPackets(scratchOffset, srcOffset, nelemsPerRank * sizeof(int), tid, blockDim.x * nBlocksPerPeer, flag);
|
||||
// step 2: get data from scratch buffer, reduce data and write result to remote scratch buffer
|
||||
for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerRank; idx += blockDim.x * gridDim.x) {
|
||||
uint2 data = make_uint2(0, 0);
|
||||
for (int index = 0; index < nPeers; index++) {
|
||||
const int remoteRank = index < rank ? index : index + 1;
|
||||
mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + remoteRank * nPktsPerRank;
|
||||
uint2 val = dstPkt[idx].read(flag);
|
||||
data = add_vectors<TYPE>(val, data);
|
||||
}
|
||||
data = add_vectors<TYPE>(data, src[idx]);
|
||||
dst[idx] = data;
|
||||
|
||||
mscclpp::LLPacket packet;
|
||||
packet.data1 = data.x;
|
||||
packet.flag1 = flag;
|
||||
packet.data2 = data.y;
|
||||
packet.flag2 = flag;
|
||||
size_t offset = scratchResultOffset / sizeof(mscclpp::LLPacket) + (idx + rank * nPktsPerRank);
|
||||
for (int index = 0; index < nPeers; index++) {
|
||||
memChans[index].write(offset, packet);
|
||||
}
|
||||
}
|
||||
// step 3: get data result from scratch buffer
|
||||
mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)((char*)scratch + scratchResultOffset);
|
||||
const int dstOffset = remoteRank * nPktsPerRank;
|
||||
uint2* result = (uint2*)((char*)resultBuff + remoteRank * nelemsPerRank * sizeof(int));
|
||||
for (int idx = threadIdx.x + localBlockIdx * blockDim.x; idx < nPktsPerRank; idx += blockDim.x * nBlocksPerPeer) {
|
||||
uint2 data = dstPkt[idx + dstOffset].read(flag);
|
||||
result[idx].x = data.x;
|
||||
result[idx].y = data.y;
|
||||
}
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
globalFlag += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------
|
||||
// allreduce_LL_2node using LLPacket, origin allreduce5
|
||||
// -------------------------------------------------------
|
||||
|
||||
template <typename TYPE>
|
||||
__global__ void __launch_bounds__(1024, 1) allreduce_LL_2node(
|
||||
mscclpp::MemoryChannelDeviceHandle* memChans,
|
||||
mscclpp::PortChannelDeviceHandle* portChans,
|
||||
TYPE* buff,
|
||||
TYPE* scratch,
|
||||
TYPE* putBuff,
|
||||
TYPE* resultBuff,
|
||||
int rank,
|
||||
int nRanksPerNode,
|
||||
int worldSize,
|
||||
size_t nelems) {
|
||||
nelems = nelems / (sizeof(int) / sizeof(TYPE));
|
||||
// This version of allreduce only works for single nodes
|
||||
const int nPeersInNode = nRanksPerNode - 1;
|
||||
const int nPkts = nelems / 2;
|
||||
const int nelemsPerLocalRank = nelems / nRanksPerNode;
|
||||
const int nPktsPerLocalRank = nelemsPerLocalRank / 2;
|
||||
const int localRankId = rank % nRanksPerNode;
|
||||
// flag for packets. Initially 1
|
||||
const uint32_t flag = (uint32_t)globalFlag;
|
||||
// thread block & channel info
|
||||
const int nBlocksPerPeer = gridDim.x / nPeersInNode;
|
||||
const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
|
||||
const int peerIdx = blockIdx.x / nBlocksPerPeer;
|
||||
const int remoteRankIdx = peerIdx < localRankId ? peerIdx : peerIdx + 1;
|
||||
mscclpp::MemoryChannelDeviceHandle memChan = memChans[peerIdx];
|
||||
mscclpp::PortChannelDeviceHandle portChan = portChans[localRankId];
|
||||
const int tid = threadIdx.x + localBlockIdx * blockDim.x;
|
||||
// double buffering
|
||||
size_t scratchBaseOffset = (flag & 1) ? 0 : nPkts * sizeof(mscclpp::LLPacket);
|
||||
size_t putBaseOffset = (flag & 1) ? 0 : nPktsPerLocalRank * sizeof(mscclpp::LLPacket);
|
||||
void* scratchBuff = (void*)((char*)scratch + scratchBaseOffset);
|
||||
size_t scratchOffset = scratchBaseOffset + localRankId * nPktsPerLocalRank * sizeof(mscclpp::LLPacket);
|
||||
size_t scratchResultOffset =
|
||||
(flag & 1) ? 2 * nPkts * sizeof(mscclpp::LLPacket) : 3 * nPkts * sizeof(mscclpp::LLPacket);
|
||||
size_t srcOffset = remoteRankIdx * nelemsPerLocalRank * sizeof(int);
|
||||
uint2* src = (uint2*)((char*)buff + localRankId * nelemsPerLocalRank * sizeof(int));
|
||||
uint2* dst = (uint2*)((char*)resultBuff + localRankId * nelemsPerLocalRank * sizeof(int));
|
||||
|
||||
// step 1: write to scratch buffer
|
||||
if (nRanksPerNode > 1) {
|
||||
memChan.putPackets(
|
||||
scratchOffset, srcOffset, nelemsPerLocalRank * sizeof(int), tid, blockDim.x * nBlocksPerPeer, flag);
|
||||
}
|
||||
// step 2: get data from scratch buffer, do local reduce-scatter in each node.
|
||||
mscclpp::LLPacket* putPkt = (mscclpp::LLPacket*)((char*)putBuff + putBaseOffset);
|
||||
for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerLocalRank; idx += blockDim.x * gridDim.x) {
|
||||
uint2 data = make_uint2(0, 0);
|
||||
for (int index = 0; index < nPeersInNode; index++) {
|
||||
const int remoteRank = index < localRankId ? index : index + 1;
|
||||
mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + remoteRank * nPktsPerLocalRank;
|
||||
uint2 val = dstPkt[idx].read(flag);
|
||||
data = add_vectors<TYPE>(val, data);
|
||||
}
|
||||
data = add_vectors<TYPE>(data, src[idx]);
|
||||
putPkt[idx].write(data.x, data.y, flag);
|
||||
dst[idx] = data;
|
||||
}
|
||||
deviceSyncer.sync(gridDim.x);
|
||||
// step 3. send local reduced data to remote node.
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
portChan.put(scratchOffset, putBaseOffset, nPktsPerLocalRank * sizeof(mscclpp::LLPacket));
|
||||
if ((flag & 63) == 0) {
|
||||
portChan.flush();
|
||||
}
|
||||
}
|
||||
// step 4. try to read the data from scratch buffer and write to local peers
|
||||
mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + localRankId * nPktsPerLocalRank;
|
||||
for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerLocalRank; idx += blockDim.x * gridDim.x) {
|
||||
uint2 res = dst[idx];
|
||||
uint2 val = dstPkt[idx].read(flag);
|
||||
res = add_vectors<TYPE>(res, val);
|
||||
|
||||
mscclpp::LLPacket packet;
|
||||
packet.data1 = res.x;
|
||||
packet.flag1 = flag;
|
||||
packet.data2 = res.y;
|
||||
packet.flag2 = flag;
|
||||
size_t offset = scratchResultOffset / sizeof(mscclpp::LLPacket) + (idx + localRankId * nPktsPerLocalRank);
|
||||
for (int index = 0; index < nPeersInNode; index++) {
|
||||
memChans[index].write(offset, packet);
|
||||
}
|
||||
dst[idx] = res;
|
||||
}
|
||||
|
||||
// step 5: get data result from scratch buffer
|
||||
dstPkt = (mscclpp::LLPacket*)((char*)scratch + scratchResultOffset);
|
||||
const int dstOffset = remoteRankIdx * nPktsPerLocalRank;
|
||||
uint2* result = (uint2*)((char*)resultBuff + remoteRankIdx * nelemsPerLocalRank * sizeof(int));
|
||||
if (nRanksPerNode > 1) {
|
||||
for (int idx = threadIdx.x + localBlockIdx * blockDim.x; idx < nPktsPerLocalRank;
|
||||
idx += blockDim.x * nBlocksPerPeer) {
|
||||
uint2 data = dstPkt[idx + dstOffset].read(flag);
|
||||
result[idx] = data;
|
||||
}
|
||||
}
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
globalFlag += 1;
|
||||
}
|
||||
}
|
||||
|
||||
static const mscclpp::Transport IBs[] = {
|
||||
mscclpp::Transport::IB0,
|
||||
mscclpp::Transport::IB1,
|
||||
mscclpp::Transport::IB2,
|
||||
mscclpp::Transport::IB3,
|
||||
mscclpp::Transport::IB4,
|
||||
mscclpp::Transport::IB5,
|
||||
mscclpp::Transport::IB6,
|
||||
mscclpp::Transport::IB7};
|
||||
|
||||
class MscclCommGroup {
|
||||
public:
|
||||
std::shared_ptr<mscclpp::Communicator> comm_;
|
||||
const size_t rank_;
|
||||
const size_t world_size_;
|
||||
const std::vector<int64_t> rank_to_node_;
|
||||
const std::vector<int64_t> rank_to_ib_;
|
||||
MscclCommGroup(
|
||||
mscclpp::UniqueId unique_id,
|
||||
const size_t rank,
|
||||
const size_t world_size,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib)
|
||||
: rank_(rank), world_size_(world_size), rank_to_node_(rank_to_node), rank_to_ib_(rank_to_ib) {
|
||||
auto bootstrap = std::make_shared<mscclpp::TcpBootstrap>(rank, world_size);
|
||||
bootstrap->initialize(unique_id);
|
||||
comm_ = std::make_shared<mscclpp::Communicator>(bootstrap);
|
||||
}
|
||||
template <typename T>
|
||||
void allreduce(cudaStream_t stream, T* output, size_t input_numel, int threads = 512, int block_limit = 21) {
|
||||
throw std::runtime_error("you should not call allreduce of a base context");
|
||||
}
|
||||
bool is_same_node(int r1, int r2) {
|
||||
return rank_to_node_[r1] == rank_to_node_[r2];
|
||||
}
|
||||
|
||||
void make_connection(
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& same_node_connections,
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& cross_node_connections) {
|
||||
same_node_connections.clear();
|
||||
cross_node_connections.clear();
|
||||
std::unordered_map<int, mscclpp::NonblockingFuture<std::shared_ptr<mscclpp::Connection>>> conn_futures;
|
||||
for (int r = 0; r < world_size_; ++r) {
|
||||
if (r == rank_) continue;
|
||||
mscclpp::Transport transport = is_same_node(r, rank_) ? mscclpp::Transport::CudaIpc : IBs[rank_to_ib_[r]];
|
||||
conn_futures.emplace(r, comm_->connectOnSetup(r, 0, transport));
|
||||
}
|
||||
comm_->setup();
|
||||
for (int r = 0; r < world_size_; ++r) {
|
||||
if (r == rank_) continue;
|
||||
if (is_same_node(r, rank_)) {
|
||||
same_node_connections.emplace(r, conn_futures[r].get());
|
||||
} else {
|
||||
cross_node_connections.emplace(r, conn_futures[r].get());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void make_memory_channels_with_scratch(
|
||||
void* tensor_ptr,
|
||||
const size_t tensor_bytes,
|
||||
void* scratch_ptr,
|
||||
const size_t scratch_bytes,
|
||||
const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& semaphores,
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories,
|
||||
std::unordered_map<int, mscclpp::MemoryChannel>& channels) {
|
||||
channels.clear();
|
||||
make_semaphores<mscclpp::MemoryDevice2DeviceSemaphore>(connections, semaphores);
|
||||
register_tensor_with_connections(scratch_ptr, scratch_bytes, connections, registered_memories);
|
||||
for (const auto& [peer, _] : connections) {
|
||||
channels.emplace(
|
||||
peer, mscclpp::MemoryChannel(semaphores[peer], registered_memories[peer], tensor_ptr, scratch_ptr));
|
||||
}
|
||||
}
|
||||
void make_port_channels_with_scratch(
|
||||
std::shared_ptr<mscclpp::ProxyService> proxyService,
|
||||
void* tensor_ptr,
|
||||
const size_t tensor_bytes,
|
||||
void* scratch_ptr,
|
||||
const size_t scratch_bytes,
|
||||
const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Host2DeviceSemaphore>>& semaphores,
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories,
|
||||
std::unordered_map<int, mscclpp::PortChannel>& channels) {
|
||||
channels.clear();
|
||||
make_semaphores<mscclpp::Host2DeviceSemaphore>(connections, semaphores);
|
||||
|
||||
mscclpp::TransportFlags flags;
|
||||
for (const auto& [_, conn] : connections) {
|
||||
flags |= conn->transport();
|
||||
}
|
||||
auto local_reg_memory = comm_->registerMemory(tensor_ptr, tensor_bytes, flags);
|
||||
|
||||
register_tensor_with_connections(scratch_ptr, scratch_bytes, connections, registered_memories);
|
||||
std::unordered_map<int, mscclpp::SemaphoreId> semaphore_ids;
|
||||
std::unordered_map<int, size_t> memory_ids;
|
||||
memory_ids[rank_] = proxyService->addMemory(local_reg_memory);
|
||||
for (const auto& [peer, memory] : registered_memories) {
|
||||
if (peer == rank_) continue;
|
||||
memory_ids[peer] = proxyService->addMemory(memory);
|
||||
}
|
||||
for (const auto& [peer, semaphore] : semaphores) {
|
||||
semaphore_ids[peer] = proxyService->addSemaphore(semaphore);
|
||||
}
|
||||
|
||||
for (const auto& [peer, _] : connections) {
|
||||
channels.emplace(peer, proxyService->portChannel(semaphore_ids[peer], memory_ids[peer], memory_ids[rank_]));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename SemaphoreType>
|
||||
void make_semaphores(
|
||||
const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
|
||||
std::unordered_map<int, std::shared_ptr<SemaphoreType>>& semaphores) {
|
||||
semaphores.clear();
|
||||
for (const auto& [peer, conn] : connections) {
|
||||
semaphores[peer] = std::make_shared<SemaphoreType>(*comm_, conn);
|
||||
}
|
||||
comm_->setup();
|
||||
}
|
||||
|
||||
void register_tensor_with_connections(
|
||||
void* tensor_ptr,
|
||||
size_t tensor_bytes,
|
||||
const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories) {
|
||||
registered_memories.clear();
|
||||
mscclpp::TransportFlags all_transports;
|
||||
for (const auto& [_, connection] : connections) {
|
||||
all_transports |= connection->transport();
|
||||
}
|
||||
mscclpp::RegisteredMemory buf_reg_mem = comm_->registerMemory(tensor_ptr, tensor_bytes, all_transports);
|
||||
registered_memories[rank_] = buf_reg_mem;
|
||||
|
||||
std::unordered_map<int, mscclpp::NonblockingFuture<mscclpp::RegisteredMemory>> remote_mem_futures;
|
||||
for (const auto& [r, connection] : connections) {
|
||||
comm_->sendMemoryOnSetup(buf_reg_mem, r, 0);
|
||||
auto remoteMemory = comm_->recvMemoryOnSetup(r, 0);
|
||||
remote_mem_futures.emplace(r, remoteMemory);
|
||||
}
|
||||
comm_->setup();
|
||||
for (auto& [r, mem_feature] : remote_mem_futures) {
|
||||
registered_memories.emplace(r, mem_feature.get());
|
||||
}
|
||||
}
|
||||
|
||||
void make_device_memory_handle_base_on_new_ptr(
|
||||
const std::unordered_map<int, mscclpp::MemoryChannel>& old_memory_channels,
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory>& registered_sm_memories,
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memory_semaphores,
|
||||
std::unordered_map<int, mscclpp::MemoryChannel>& memory_channels,
|
||||
mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>& device_memory_handle,
|
||||
void* input,
|
||||
void* scratch,
|
||||
const cudaStream_t stream) {
|
||||
memory_channels.clear();
|
||||
for (const auto& [peer, channel] : old_memory_channels) {
|
||||
memory_channels.emplace(
|
||||
peer, mscclpp::MemoryChannel(memory_semaphores[peer], registered_sm_memories[peer], input, scratch));
|
||||
}
|
||||
std::vector<mscclpp::MemoryChannel> memory_channels_list;
|
||||
for (int r = 0; r < world_size_; r++) {
|
||||
if (r == rank_) continue;
|
||||
if (is_same_node(r, rank_)) {
|
||||
memory_channels_list.push_back(memory_channels[r]);
|
||||
}
|
||||
}
|
||||
std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
|
||||
std::transform(
|
||||
memory_channels_list.begin(),
|
||||
memory_channels_list.end(),
|
||||
memory_channel_handlers.begin(),
|
||||
[](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
|
||||
mscclpp::gpuMemcpyAsync<mscclpp::MemoryChannelDeviceHandle>(
|
||||
device_memory_handle.data(),
|
||||
memory_channel_handlers.data(),
|
||||
memory_channel_handlers.size(),
|
||||
stream,
|
||||
cudaMemcpyHostToDevice);
|
||||
}
|
||||
};
|
||||
|
||||
class Msccl1NodeLLcontext {
|
||||
private:
|
||||
std::shared_ptr<MscclCommGroup> comm_group_ = nullptr;
|
||||
void* scratch_;
|
||||
const size_t scratch_bytes_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> same_node_connections_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> cross_node_connections_;
|
||||
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory> registered_sm_memories_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memory_semaphores_;
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels_;
|
||||
mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> d_memHandles_;
|
||||
std::unordered_map<void*, std::unordered_map<int, mscclpp::MemoryChannel>> input_ptr2memory_channels_;
|
||||
std::unordered_map<void*, mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>> input_ptr2d_memHandles_;
|
||||
cudaStream_t h2d_stream;
|
||||
const size_t nranks_per_node_;
|
||||
|
||||
public:
|
||||
Msccl1NodeLLcontext(
|
||||
mscclpp::UniqueId unique_id,
|
||||
const size_t rank,
|
||||
const size_t world_size,
|
||||
void* scratch,
|
||||
const size_t scratch_bytes,
|
||||
const size_t nranks_per_node,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib)
|
||||
: scratch_(scratch),
|
||||
scratch_bytes_(scratch_bytes),
|
||||
nranks_per_node_(nranks_per_node),
|
||||
d_memHandles_(nranks_per_node - 1) {
|
||||
CHECK_CUDA_SUCCESS(cudaStreamCreateWithFlags(&h2d_stream, cudaStreamNonBlocking));
|
||||
comm_group_ = std::make_shared<MscclCommGroup>(unique_id, rank, world_size, rank_to_node, rank_to_ib);
|
||||
comm_group_->make_connection(same_node_connections_, cross_node_connections_);
|
||||
comm_group_->make_memory_channels_with_scratch(
|
||||
scratch_,
|
||||
scratch_bytes_,
|
||||
scratch_,
|
||||
scratch_bytes_,
|
||||
same_node_connections_,
|
||||
memory_semaphores_,
|
||||
registered_sm_memories_,
|
||||
memory_channels_);
|
||||
std::vector<mscclpp::MemoryChannel> memory_channels_list;
|
||||
for (int r = 0; r < comm_group_->world_size_; r++) {
|
||||
if (r == comm_group_->rank_) continue;
|
||||
memory_channels_list.push_back(memory_channels_[r]);
|
||||
}
|
||||
std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
|
||||
std::transform(
|
||||
memory_channels_list.begin(),
|
||||
memory_channels_list.end(),
|
||||
memory_channel_handlers.begin(),
|
||||
[](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
|
||||
mscclpp::gpuMemcpy<mscclpp::MemoryChannelDeviceHandle>(
|
||||
d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice);
|
||||
}
|
||||
|
||||
~Msccl1NodeLLcontext() {
|
||||
CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void allreduce(cudaStream_t stream, T* input, T* output, size_t input_numel, int nthreads = 512, int nblocks = 21) {
|
||||
dim3 nthrs(nthreads);
|
||||
dim3 nblks(nblocks);
|
||||
cudaStreamCaptureStatus capturing_status;
|
||||
CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &capturing_status));
|
||||
mscclpp::MemoryChannelDeviceHandle* memChans;
|
||||
if (capturing_status != cudaStreamCaptureStatusActive) {
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
|
||||
comm_group_->make_device_memory_handle_base_on_new_ptr(
|
||||
memory_channels_,
|
||||
registered_sm_memories_,
|
||||
memory_semaphores_,
|
||||
memory_channels,
|
||||
d_memHandles_,
|
||||
input,
|
||||
scratch_,
|
||||
h2d_stream);
|
||||
CHECK_CUDA_SUCCESS(cudaStreamSynchronize(h2d_stream));
|
||||
memChans = d_memHandles_.data();
|
||||
} else {
|
||||
void* input_void_ptr = reinterpret_cast<void*>(input);
|
||||
if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) {
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
|
||||
mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> device_memory_handle(comm_group_->world_size_ - 1);
|
||||
comm_group_->make_device_memory_handle_base_on_new_ptr(
|
||||
memory_channels_,
|
||||
registered_sm_memories_,
|
||||
memory_semaphores_,
|
||||
memory_channels,
|
||||
device_memory_handle,
|
||||
input,
|
||||
scratch_,
|
||||
h2d_stream);
|
||||
input_ptr2memory_channels_.emplace(input_void_ptr, memory_channels);
|
||||
input_ptr2d_memHandles_.emplace(input_void_ptr, device_memory_handle);
|
||||
}
|
||||
auto it = input_ptr2d_memHandles_.find(input_void_ptr);
|
||||
memChans = it->second.data();
|
||||
}
|
||||
allreduce_LL_1node<T><<<nblks, nthrs, 0, stream>>>(
|
||||
memChans, (T*)input, (T*)scratch_, output, comm_group_->rank_, comm_group_->world_size_, input_numel);
|
||||
|
||||
cudaError_t status = cudaGetLastError();
|
||||
if (status != cudaSuccess) {
|
||||
printf("rank: %lu failed to launch allreduce_LL_1node: %s\n", comm_group_->rank_, cudaGetErrorString(status));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
class Msccl2NodeLLcontext {
|
||||
private:
|
||||
std::shared_ptr<MscclCommGroup> comm_group_ = nullptr;
|
||||
void* scratch_;
|
||||
const size_t scratch_bytes_;
|
||||
void* put_buffer_;
|
||||
const size_t put_buffer_bytes_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> same_node_connections_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> cross_node_connections_;
|
||||
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory> registered_sm_memories_;
|
||||
std::unordered_map<int, mscclpp::RegisteredMemory> registered_port_memories_;
|
||||
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memory_semaphores_;
|
||||
std::unordered_map<int, std::shared_ptr<mscclpp::Host2DeviceSemaphore>> port_semaphores_;
|
||||
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels_;
|
||||
std::unordered_map<int, mscclpp::PortChannel> port_channels_;
|
||||
|
||||
mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> d_memHandles_;
|
||||
mscclpp::GpuBuffer<mscclpp::PortChannelDeviceHandle> d_portHandles_;
|
||||
|
||||
std::shared_ptr<mscclpp::ProxyService> proxyService;
|
||||
cudaStream_t h2d_stream;
|
||||
const size_t nranks_per_node_;
|
||||
|
||||
std::unordered_map<void*, std::unordered_map<int, mscclpp::MemoryChannel>> input_ptr2memory_channels_;
|
||||
std::unordered_map<void*, mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>> input_ptr2d_memHandles_;
|
||||
|
||||
public:
|
||||
Msccl2NodeLLcontext(
|
||||
mscclpp::UniqueId unique_id,
|
||||
const size_t rank,
|
||||
const size_t world_size,
|
||||
void* scratch,
|
||||
const size_t scratch_bytes,
|
||||
void* put_buffer,
|
||||
const size_t put_buffer_bytes,
|
||||
const size_t nranks_per_node,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib)
|
||||
: scratch_(scratch),
|
||||
scratch_bytes_(scratch_bytes),
|
||||
put_buffer_(put_buffer),
|
||||
put_buffer_bytes_(put_buffer_bytes),
|
||||
nranks_per_node_(nranks_per_node),
|
||||
d_memHandles_(nranks_per_node - 1),
|
||||
d_portHandles_(world_size - nranks_per_node) {
|
||||
CHECK_CUDA_SUCCESS(cudaStreamCreateWithFlags(&h2d_stream, cudaStreamNonBlocking));
|
||||
comm_group_ = std::make_shared<MscclCommGroup>(unique_id, rank, world_size, rank_to_node, rank_to_ib);
|
||||
proxyService = std::make_shared<mscclpp::ProxyService>();
|
||||
proxyService->startProxy();
|
||||
comm_group_->make_connection(same_node_connections_, cross_node_connections_);
|
||||
comm_group_->make_memory_channels_with_scratch(
|
||||
scratch_,
|
||||
scratch_bytes_,
|
||||
scratch_,
|
||||
scratch_bytes_,
|
||||
same_node_connections_,
|
||||
memory_semaphores_,
|
||||
registered_sm_memories_,
|
||||
memory_channels_);
|
||||
comm_group_->make_port_channels_with_scratch(
|
||||
proxyService,
|
||||
put_buffer_,
|
||||
put_buffer_bytes_,
|
||||
scratch_,
|
||||
scratch_bytes_,
|
||||
cross_node_connections_,
|
||||
port_semaphores_,
|
||||
registered_port_memories_,
|
||||
port_channels_);
|
||||
std::vector<mscclpp::MemoryChannel> memory_channels_list;
|
||||
std::vector<mscclpp::PortChannel> port_channels_list;
|
||||
for (int r = 0; r < comm_group_->world_size_; r++) {
|
||||
if (r == comm_group_->rank_) continue;
|
||||
if (comm_group_->is_same_node(r, comm_group_->rank_)) {
|
||||
memory_channels_list.push_back(memory_channels_[r]);
|
||||
} else {
|
||||
port_channels_list.push_back(port_channels_[r]);
|
||||
}
|
||||
}
|
||||
std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
|
||||
std::transform(
|
||||
memory_channels_list.begin(),
|
||||
memory_channels_list.end(),
|
||||
memory_channel_handlers.begin(),
|
||||
[](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
|
||||
mscclpp::gpuMemcpy<mscclpp::MemoryChannelDeviceHandle>(
|
||||
d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice);
|
||||
|
||||
std::vector<mscclpp::PortChannelDeviceHandle> port_channel_handlers(port_channels_list.size());
|
||||
std::transform(
|
||||
port_channels_list.begin(),
|
||||
port_channels_list.end(),
|
||||
port_channel_handlers.begin(),
|
||||
[](const mscclpp::PortChannel& channel) { return channel.deviceHandle(); });
|
||||
mscclpp::gpuMemcpy<mscclpp::PortChannelDeviceHandle>(
|
||||
d_portHandles_.data(), port_channel_handlers.data(), port_channel_handlers.size(), cudaMemcpyHostToDevice);
|
||||
}
|
||||
|
||||
~Msccl2NodeLLcontext() {
|
||||
CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream));
|
||||
if (proxyService) {
|
||||
proxyService->stopProxy();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
allreduce(cudaStream_t stream, T* input, T* output, const size_t input_numel, int nthreads = 512, int nblocks = 21) {
|
||||
dim3 nthrs(nthreads);
|
||||
dim3 nblks(nblocks);
|
||||
cudaStreamCaptureStatus capturing_status;
|
||||
CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &capturing_status));
|
||||
mscclpp::MemoryChannelDeviceHandle* memChans;
|
||||
if (capturing_status != cudaStreamCaptureStatusActive) {
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
|
||||
comm_group_->make_device_memory_handle_base_on_new_ptr(
|
||||
memory_channels_,
|
||||
registered_sm_memories_,
|
||||
memory_semaphores_,
|
||||
memory_channels,
|
||||
d_memHandles_,
|
||||
input,
|
||||
scratch_,
|
||||
h2d_stream);
|
||||
CHECK_CUDA_SUCCESS(cudaStreamSynchronize(h2d_stream));
|
||||
memChans = d_memHandles_.data();
|
||||
} else {
|
||||
void* input_void_ptr = reinterpret_cast<void*>(input);
|
||||
if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) {
|
||||
std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
|
||||
mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> device_memory_handle(7);
|
||||
comm_group_->make_device_memory_handle_base_on_new_ptr(
|
||||
memory_channels_,
|
||||
registered_sm_memories_,
|
||||
memory_semaphores_,
|
||||
memory_channels,
|
||||
device_memory_handle,
|
||||
input,
|
||||
scratch_,
|
||||
h2d_stream);
|
||||
input_ptr2memory_channels_.emplace(input_void_ptr, memory_channels);
|
||||
input_ptr2d_memHandles_.emplace(input_void_ptr, device_memory_handle);
|
||||
}
|
||||
auto it = input_ptr2d_memHandles_.find(input_void_ptr);
|
||||
memChans = it->second.data();
|
||||
}
|
||||
allreduce_LL_2node<T><<<nblks, nthrs, 0, stream>>>(
|
||||
memChans,
|
||||
d_portHandles_.data(),
|
||||
(T*)input,
|
||||
(T*)scratch_,
|
||||
(T*)put_buffer_,
|
||||
output,
|
||||
comm_group_->rank_,
|
||||
nranks_per_node_,
|
||||
comm_group_->world_size_,
|
||||
input_numel);
|
||||
|
||||
cudaError_t status = cudaGetLastError();
|
||||
if (status != cudaSuccess) {
|
||||
printf("rank: %lu failed to launch allreduce_LL_2node: %s\n", comm_group_->rank_, cudaGetErrorString(status));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -1,153 +0,0 @@
|
||||
/*
|
||||
* this file is used to test mscclpp_allreduce.cu using mpirun
|
||||
* this file is adapted from https://github.com/flashinfer-ai/flashinfer/blob/v0.2.5/src/test_sum_all_reduce.cu
|
||||
usage:
|
||||
cd PATH-TO-THIS-FILE
|
||||
export MPI_HOME=/usr/local/mpi
|
||||
# export MPI_HOME=/opt/hpcx/ompi/
|
||||
export MSCCLPP_HOME=/workspace/test/mscclpp
|
||||
nvcc -O2 -arch=native -std=c++17 test_mscclpp_allreduce.cu \
|
||||
-o test_mscclpp_allreduce -D_GLIBCXX_USE_CXX11_ABI=0 \
|
||||
-I${MSCCLPP_HOME}/include -L${MSCCLPP_HOME}/build -lmscclpp \
|
||||
-lnccl -I${MPI_HOME}/include -L${MPI_HOME}/lib -lmpi
|
||||
|
||||
/opt/hpcx/ompi/bin/
|
||||
mpirun --allow-run-as-root -H 127.0.0.1:8 -np 8 \
|
||||
--map-by ppr:8:node \
|
||||
--mca btl_openib_warn_no_device_params_found 0 \
|
||||
--mca btl_tcp_if_include bond0 \
|
||||
--allow-run-as-root -np 8 \
|
||||
-x NCCL_RUNTIME_CONNECT=0 -x NCCL_IB_GID_INDEX=3 -x NCCL_DEBUG=WARN \
|
||||
-x LD_PRELOAD=${MSCCLPP_HOME}/build/libmscclpp.so ./test_mscclpp_allreduce
|
||||
*/
|
||||
#include <mpi.h>
|
||||
#include <thrust/detail/raw_pointer_cast.h>
|
||||
#include <thrust/device_vector.h>
|
||||
#include <thrust/host_vector.h>
|
||||
|
||||
#ifndef CHECK_CUDA_SUCCESS
|
||||
#define CHECK_CUDA_SUCCESS(cmd) \
|
||||
do { \
|
||||
cudaError_t e = cmd; \
|
||||
if (e != cudaSuccess) { \
|
||||
printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \
|
||||
exit(EXIT_FAILURE); \
|
||||
} \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include "mscclpp_allreduce.cuh"
|
||||
|
||||
template <typename T>
|
||||
bool isclose(T a, T b, float rtol = 1e-5, float atol = 1e-8) {
|
||||
return fabs(a - b) <= (atol + rtol * fabs(b));
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
// init mpi
|
||||
MPI_Init(&argc, &argv);
|
||||
printf("MPI Initialized.\n");
|
||||
int nranks, rank;
|
||||
|
||||
// get work size and rank id
|
||||
MPI_Comm_size(MPI_COMM_WORLD, &nranks);
|
||||
MPI_Comm_rank(MPI_COMM_WORLD, &rank);
|
||||
cudaSetDevice(rank);
|
||||
printf("nranks: %d, rank: %d\n", nranks, rank);
|
||||
|
||||
// init host and device buffers
|
||||
using T = float;
|
||||
using ReduceT = float;
|
||||
const size_t num_elems = 2 * 1024 * 1024;
|
||||
std::vector<T> host_buf(num_elems);
|
||||
for (uint32_t i = 0; i < num_elems; ++i) {
|
||||
host_buf[i] = T(i + rank);
|
||||
}
|
||||
thrust::device_vector<T> device_buf(host_buf);
|
||||
const size_t buf_size_in_bytes = num_elems * sizeof(T);
|
||||
std::vector<T> host_result_buf(num_elems);
|
||||
thrust::device_vector<T> device_result_buf(host_result_buf);
|
||||
|
||||
std::vector<T> host_scratch_buf(num_elems * 8);
|
||||
for (uint32_t i = 0; i < num_elems; ++i) {
|
||||
host_scratch_buf[i] = 1;
|
||||
}
|
||||
thrust::device_vector<T> device_scratch_buf(host_scratch_buf);
|
||||
std::vector<T> host_put_buf(num_elems);
|
||||
thrust::device_vector<T> device_put_buf(host_put_buf);
|
||||
|
||||
mscclpp::UniqueId unique_id;
|
||||
if (rank == 0) unique_id = mscclpp::TcpBootstrap::createUniqueId();
|
||||
MPI_Bcast(&unique_id, sizeof(unique_id), MPI_BYTE, 0, MPI_COMM_WORLD);
|
||||
|
||||
std::vector<int64_t> rank_to_node(nranks);
|
||||
std::vector<int64_t> rank_to_ib(nranks);
|
||||
for (int i = 0; i < nranks; i++) {
|
||||
rank_to_node[i] = i / 8;
|
||||
rank_to_ib[i] = i % 8;
|
||||
}
|
||||
|
||||
cudaStream_t s;
|
||||
CHECK_CUDA_SUCCESS(cudaStreamCreate(&s));
|
||||
CHECK_CUDA_SUCCESS(cudaStreamSynchronize(s));
|
||||
if (nranks == 8) {
|
||||
auto context = std::make_shared<sglang::Msccl1NodeLLcontext>(
|
||||
unique_id,
|
||||
rank,
|
||||
nranks,
|
||||
thrust::raw_pointer_cast(device_scratch_buf.data()),
|
||||
buf_size_in_bytes * 8,
|
||||
rank_to_node,
|
||||
rank_to_ib);
|
||||
printf("rank: %d, Msccl1NodeLLcontext setup.\n", rank);
|
||||
MPI_Barrier(MPI_COMM_WORLD);
|
||||
context->allreduce<T>(
|
||||
s,
|
||||
thrust::raw_pointer_cast(device_buf.data()),
|
||||
thrust::raw_pointer_cast(device_result_buf.data()),
|
||||
device_buf.size());
|
||||
} else if (nranks == 16) {
|
||||
// TODO: this branch is untested since there is something wrong with mpirun in my test machince
|
||||
auto context = std::make_shared<sglang::Msccl2NodeLLcontext>(
|
||||
unique_id,
|
||||
rank,
|
||||
nranks,
|
||||
thrust::raw_pointer_cast(device_scratch_buf.data()),
|
||||
buf_size_in_bytes * 8,
|
||||
thrust::raw_pointer_cast(device_put_buf.data()),
|
||||
buf_size_in_bytes,
|
||||
rank_to_node,
|
||||
rank_to_ib);
|
||||
printf("rank: %d, Msccl2NodeLLcontext setup.\n", rank);
|
||||
MPI_Barrier(MPI_COMM_WORLD);
|
||||
context->allreduce<T>(
|
||||
s,
|
||||
thrust::raw_pointer_cast(device_buf.data()),
|
||||
thrust::raw_pointer_cast(device_result_buf.data()),
|
||||
device_buf.size());
|
||||
}
|
||||
|
||||
// check result correctness
|
||||
thrust::host_vector<T> host_buf_result = device_result_buf;
|
||||
size_t num_results_error_atol_1e_3_rtol_1e_3 = 0;
|
||||
bool nan_detected = false;
|
||||
|
||||
for (uint32_t i = 0; i < num_elems; ++i) {
|
||||
T expected = T(i * nranks + (nranks - 1) * nranks / 2);
|
||||
if (std::isnan(float(host_buf_result[i]))) {
|
||||
nan_detected = true;
|
||||
}
|
||||
if (!isclose(float(host_buf_result[i]), float(expected), 1e-3, 1e-3)) {
|
||||
num_results_error_atol_1e_3_rtol_1e_3++;
|
||||
}
|
||||
}
|
||||
float result_accuracy = 1. - float(num_results_error_atol_1e_3_rtol_1e_3) / float(num_elems);
|
||||
|
||||
printf("rank: %d, nan_detected: %d accuracy: %f\n", rank, nan_detected, result_accuracy);
|
||||
|
||||
CHECK_CUDA_SUCCESS(cudaStreamDestroy(s));
|
||||
MPI_Finalize();
|
||||
return 0;
|
||||
}
|
||||
@@ -38,15 +38,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
"int reg_buffer_sz_bytes) -> ()");
|
||||
m.impl("all_reduce", torch::kCUDA, &all_reduce);
|
||||
|
||||
m.def("mscclpp_generate_unique_id", &mscclpp_generate_unique_id);
|
||||
m.def(
|
||||
"mscclpp_init_context(Tensor unique_id, int rank, int world_size, Tensor scratch, Tensor put_buffer, "
|
||||
"int nranks_per_node, int[] rank_to_node, int[] rank_to_ib, int context_selection) -> int");
|
||||
m.impl("mscclpp_init_context", torch::kCUDA, &mscclpp_init_context);
|
||||
|
||||
m.def("mscclpp_allreduce(int context, Tensor inp, Tensor! out, int nthreads, int nblocks) -> ()");
|
||||
m.impl("mscclpp_allreduce", torch::kCUDA, &mscclpp_allreduce);
|
||||
|
||||
/*
|
||||
* From csrc/attention
|
||||
*/
|
||||
|
||||
@@ -89,20 +89,6 @@ std::tuple<std::vector<int64_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta
|
||||
void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs);
|
||||
void register_graph_buffers(
|
||||
fptr_t _fa, const std::vector<std::vector<int64_t>>& handles, const std::vector<std::vector<int64_t>>& offsets);
|
||||
|
||||
// mscclpp
|
||||
torch::Tensor mscclpp_generate_unique_id();
|
||||
fptr_t mscclpp_init_context(
|
||||
const torch::Tensor& unique_id,
|
||||
const int64_t rank,
|
||||
const int64_t world_size,
|
||||
torch::Tensor& scratch,
|
||||
torch::Tensor& put_buffer,
|
||||
const int64_t nranks_per_node,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib,
|
||||
const int64_t context_selection);
|
||||
void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out, int64_t nthreads, int64_t nblocks);
|
||||
#endif
|
||||
|
||||
/*
|
||||
|
||||
@@ -92,28 +92,6 @@ if torch.version.hip is not None:
|
||||
def qr_max_size() -> int:
|
||||
return torch.ops.sgl_kernel.qr_max_size.default()
|
||||
|
||||
# mscclpp
|
||||
def mscclpp_generate_unique_id() -> bytes:
|
||||
raise NotImplementedError()
|
||||
|
||||
def mscclpp_init_context(
|
||||
unique_id: bytes,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
scratch: torch.Tensor,
|
||||
put_buffer: torch.Tensor,
|
||||
nranks_per_node: int,
|
||||
rank_to_node: List[int],
|
||||
rank_to_ib: List[int],
|
||||
context_selection: int,
|
||||
) -> int:
|
||||
raise NotImplementedError()
|
||||
|
||||
def mscclpp_allreduce(
|
||||
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
|
||||
) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
else:
|
||||
|
||||
def init_custom_ar(
|
||||
@@ -150,36 +128,3 @@ else:
|
||||
|
||||
def meta_size() -> int:
|
||||
return torch.ops.sgl_kernel.meta_size.default()
|
||||
|
||||
def mscclpp_generate_unique_id() -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.mscclpp_generate_unique_id.default()
|
||||
|
||||
def mscclpp_init_context(
|
||||
unique_id: torch.Tensor,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
scratch: torch.Tensor,
|
||||
put_buffer: torch.Tensor,
|
||||
nranks_per_node: int,
|
||||
rank_to_node: List[int],
|
||||
rank_to_ib: List[int],
|
||||
context_selection: int,
|
||||
) -> int:
|
||||
return torch.ops.sgl_kernel.mscclpp_init_context.default(
|
||||
unique_id,
|
||||
rank,
|
||||
world_size,
|
||||
scratch,
|
||||
put_buffer,
|
||||
nranks_per_node,
|
||||
rank_to_node,
|
||||
rank_to_ib,
|
||||
context_selection,
|
||||
)
|
||||
|
||||
def mscclpp_allreduce(
|
||||
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
|
||||
) -> None:
|
||||
torch.ops.sgl_kernel.mscclpp_allreduce.default(
|
||||
context, inp, out, nthreads, nblocks
|
||||
)
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import socket
|
||||
import unittest
|
||||
from enum import IntEnum
|
||||
from typing import Any
|
||||
|
||||
import sgl_kernel.allreduce as custom_ops
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
class MscclContextSelection(IntEnum):
|
||||
MSCCL1SHOT1NODELL = 1
|
||||
MSCCL1SHOT2NODELL = 2
|
||||
|
||||
|
||||
def _run_correctness_worker(world_size, rank, distributed_init_port, test_sizes):
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method=distributed_init_method,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
group = dist.group.WORLD
|
||||
cpu_group = torch.distributed.new_group(list(range(world_size)), backend="gloo")
|
||||
if rank == 0:
|
||||
unique_id = [custom_ops.mscclpp_generate_unique_id()]
|
||||
else:
|
||||
unique_id = [None]
|
||||
dist.broadcast_object_list(
|
||||
unique_id, src=0, device=torch.device("cpu"), group=cpu_group
|
||||
)
|
||||
unique_id = unique_id[0]
|
||||
rank_to_node, rank_to_ib = list(range(world_size)), list(range(world_size))
|
||||
for r in range(world_size):
|
||||
rank_to_node[r] = r // 8
|
||||
rank_to_ib[r] = rank % 8
|
||||
MAX_BYTES = 2**20
|
||||
scratch = torch.empty(
|
||||
MAX_BYTES * 8, dtype=torch.bfloat16, device=torch.cuda.current_device()
|
||||
)
|
||||
put_buffer = torch.empty(
|
||||
MAX_BYTES, dtype=torch.bfloat16, device=torch.cuda.current_device()
|
||||
)
|
||||
print(f"[{rank}] start mscclpp_context init")
|
||||
nranks_per_node = torch.cuda.device_count()
|
||||
selection = int(MscclContextSelection.MSCCL1SHOT1NODELL)
|
||||
mscclpp_context = custom_ops.mscclpp_init_context(
|
||||
unique_id,
|
||||
rank,
|
||||
world_size,
|
||||
scratch,
|
||||
put_buffer,
|
||||
nranks_per_node,
|
||||
rank_to_node,
|
||||
rank_to_ib,
|
||||
selection,
|
||||
)
|
||||
try:
|
||||
test_loop = 10
|
||||
for sz in test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
if sz * dtype.itemsize > MAX_BYTES:
|
||||
continue
|
||||
if rank == 0:
|
||||
print(f"mscclpp allreduce test sz {sz}, dtype {dtype}")
|
||||
for _ in range(test_loop):
|
||||
inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device)
|
||||
inp1_ref = inp1.clone()
|
||||
out1 = torch.empty_like(inp1)
|
||||
custom_ops.mscclpp_allreduce(
|
||||
mscclpp_context, inp1, out1, nthreads=512, nblocks=21
|
||||
)
|
||||
dist.all_reduce(inp1_ref, group=group)
|
||||
torch.testing.assert_close(out1, inp1_ref)
|
||||
finally:
|
||||
dist.barrier(group=group)
|
||||
dist.destroy_process_group(group=group)
|
||||
|
||||
|
||||
def get_open_port() -> int:
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
except OSError:
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
|
||||
s.bind(("::1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def multi_process_parallel(
|
||||
world_size: int, test_target: Any, target_args: tuple = ()
|
||||
) -> None:
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
procs = []
|
||||
distributed_init_port = get_open_port()
|
||||
for i in range(world_size):
|
||||
proc_args = (world_size, i, distributed_init_port) + target_args
|
||||
proc = mp.Process(target=test_target, args=proc_args, name=f"Worker-{i}")
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
|
||||
for i in range(world_size):
|
||||
procs[i].join()
|
||||
assert (
|
||||
procs[i].exitcode == 0
|
||||
), f"Process {i} failed with exit code {procs[i].exitcode}"
|
||||
|
||||
|
||||
class TestMSCCLAllReduce(unittest.TestCase):
|
||||
test_sizes = [
|
||||
512,
|
||||
2560,
|
||||
4096,
|
||||
5120,
|
||||
7680,
|
||||
32768,
|
||||
262144,
|
||||
524288,
|
||||
]
|
||||
world_sizes = [8]
|
||||
|
||||
def test_correctness(self):
|
||||
for world_size in self.world_sizes:
|
||||
available_gpus = torch.cuda.device_count()
|
||||
if world_size > available_gpus:
|
||||
print(
|
||||
f"Skipping world_size={world_size}, found {available_gpus} and now ray is not supported here"
|
||||
)
|
||||
continue
|
||||
|
||||
print(f"Running test for world_size={world_size}")
|
||||
multi_process_parallel(
|
||||
world_size, _run_correctness_worker, target_args=(self.test_sizes,)
|
||||
)
|
||||
print(f"custom allreduce tp = {world_size}: OK")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,196 +0,0 @@
|
||||
"""For Now, MSCCL is only supported on TP16 and TP8 case
|
||||
|
||||
if [[ $RANK -eq 0 ]]; then
|
||||
ray start --block --head --port=6379 &
|
||||
python3 test_mscclpp.py;
|
||||
else
|
||||
ray start --block --address=${MASTER_ADDR}:6379;
|
||||
fi
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import unittest
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.communication_op import ( # noqa
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_group,
|
||||
graph_capture,
|
||||
initialize_model_parallel,
|
||||
set_custom_all_reduce,
|
||||
set_mscclpp_all_reduce,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
def get_open_port() -> int:
|
||||
# try ipv4
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
except OSError:
|
||||
# try ipv6
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def multi_process_parallel(
|
||||
world_size: int,
|
||||
master_addr: str,
|
||||
cls: Any,
|
||||
test_target: Any,
|
||||
) -> None:
|
||||
|
||||
# Using ray helps debugging the error when it failed
|
||||
# as compared to multiprocessing.
|
||||
# NOTE: We need to set working_dir for distributed tests,
|
||||
# otherwise we may get import errors on ray workers
|
||||
|
||||
ray.init(log_to_driver=True)
|
||||
|
||||
distributed_init_port = get_open_port()
|
||||
refs = []
|
||||
for rank in range(world_size):
|
||||
refs.append(
|
||||
test_target.remote(
|
||||
cls, world_size, master_addr, rank, distributed_init_port
|
||||
)
|
||||
)
|
||||
ray.get(refs)
|
||||
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
class TestMSCCLAllReduce(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
random.seed(42)
|
||||
# 1KB to 1MB
|
||||
cls.test_sizes = [512, 4096, 32768, 262144, 524288]
|
||||
cls.world_sizes = [8]
|
||||
TEST_TP16 = int(os.getenv("SGL_MSCCLPP_TEST_TP16", "0"))
|
||||
if TEST_TP16:
|
||||
cls.world_sizes = [16]
|
||||
cls.test_loop = 10
|
||||
|
||||
def test_graph_allreduce(self):
|
||||
TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
|
||||
for world_size in self.world_sizes:
|
||||
if world_size not in [8, 16]:
|
||||
continue
|
||||
multi_process_parallel(
|
||||
world_size, TEST_MASTER_ADDR, self, self.graph_allreduce
|
||||
)
|
||||
|
||||
def test_eager_allreduce(self):
|
||||
TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
|
||||
for world_size in self.world_sizes:
|
||||
if world_size not in [8, 16]:
|
||||
continue
|
||||
multi_process_parallel(
|
||||
world_size, TEST_MASTER_ADDR, self, self.eager_allreduce
|
||||
)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(self, world_size, master_addr, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
|
||||
set_mscclpp_all_reduce(True)
|
||||
set_custom_all_reduce(False)
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank % torch.cuda.device_count(),
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.cuda.synchronize()
|
||||
del data
|
||||
|
||||
for sz in self.test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.test_loop):
|
||||
with graph_capture() as graph_capture_context:
|
||||
# use integers so result matches NCCL exactly
|
||||
inp1 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
inp2 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(
|
||||
graph, stream=graph_capture_context.stream
|
||||
):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
# the input buffer is immediately modified to test
|
||||
# synchronization
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
torch.testing.assert_close(out2, inp2)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(self, world_size, master_addr, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
|
||||
set_mscclpp_all_reduce(True)
|
||||
set_custom_all_reduce(False)
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
for sz in self.test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.test_loop):
|
||||
inp1 = torch.randint(
|
||||
1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
|
||||
)
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
dist.all_reduce(inp1, group=group)
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user