File size: 4,119 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | from __future__ import annotations
from contextlib import contextmanager, nullcontext
from typing import Iterator
import torch
import torch.distributed as dist
from diffulex.config import Config
from diffulex.distributed.parallel_state import fetch_parallel_state
from diffulex.logger import get_logger
logger = get_logger(__name__)
_VLLM_TP_GROUP = None
_VLLM_TP_GROUP_FAILED = False
def _tp_group_ranks_for_vllm() -> list[list[int]]:
state = fetch_parallel_state()
tp_size = int(state.tp_size)
if tp_size <= 1:
return [[int(state.global_rank)]]
dp_size = int(state.dp_size)
return [
list(range(dp_rank * tp_size, (dp_rank + 1) * tp_size))
for dp_rank in range(dp_size)
]
def get_vllm_tp_group():
"""Return a vLLM GroupCoordinator matching Diffulex TP ranks.
dInfer routes custom all-reduce and CUDA graph capture through the vLLM /
sglang group coordinator. Diffulex owns process-group initialization, so we
build the coordinator on top of the already initialized torch distributed
group and fall back silently when vLLM is unavailable.
"""
global _VLLM_TP_GROUP, _VLLM_TP_GROUP_FAILED
if _VLLM_TP_GROUP is not None:
return _VLLM_TP_GROUP
if _VLLM_TP_GROUP_FAILED:
return None
try:
from vllm.distributed.parallel_state import GroupCoordinator, set_custom_all_reduce
set_custom_all_reduce(True)
_VLLM_TP_GROUP = GroupCoordinator(
group_ranks=_tp_group_ranks_for_vllm(),
local_rank=int(torch.cuda.current_device()),
torch_distributed_backend=dist.get_backend(),
use_device_communicator=True,
group_name="tp",
)
return _VLLM_TP_GROUP
except Exception:
_VLLM_TP_GROUP_FAILED = True
logger.debug("Failed to initialize vLLM TP coordinator.", exc_info=True)
return None
@contextmanager
def vllm_current_config(config: Config) -> Iterator[None]:
"""Temporarily install a minimal vLLM config during module construction."""
try:
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
parallel_config = ParallelConfig(
tensor_parallel_size=int(config.tensor_parallel_size),
pipeline_parallel_size=1,
data_parallel_size=int(config.data_parallel_size),
enable_expert_parallel=bool(int(config.expert_parallel_size) > 1),
disable_custom_all_reduce=False,
distributed_timeout_seconds=int(config.distributed_timeout_seconds),
)
with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)):
yield
except Exception:
logger.debug("Using Diffulex model init without vLLM current config.", exc_info=True)
yield
@contextmanager
def vllm_graph_capture(stream: torch.cuda.Stream, pool) -> Iterator[torch.cuda.Stream]:
"""Enter vLLM graph-capture side contexts when available."""
try:
from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
set_graph_pool_id(pool)
except Exception:
logger.debug("Failed to set vLLM graph pool id.", exc_info=True)
group = get_vllm_tp_group()
if group is None or not hasattr(group, "graph_capture"):
with torch.cuda.stream(stream):
yield stream
return
try:
context = getattr(group, "graph_capture")()
with context as graph_context:
yield getattr(graph_context, "stream", stream)
except Exception:
logger.debug("vLLM graph_capture context failed; using raw CUDA stream.", exc_info=True)
with torch.cuda.stream(stream):
yield stream
def reset_vllm_compat_state() -> None:
global _VLLM_TP_GROUP, _VLLM_TP_GROUP_FAILED
group = _VLLM_TP_GROUP
_VLLM_TP_GROUP = None
_VLLM_TP_GROUP_FAILED = False
if group is not None and hasattr(group, "destroy"):
try:
group.destroy()
except Exception:
logger.debug("Failed to destroy vLLM TP coordinator.", exc_info=True)
|