[NPU] fix pp 2 hang on npu (#38249)
This commit is contained in:
liupeng374
@@ -1811,6 +1811,161 @@ class GroupCoordinator:
|
||||
tensor_dict[key] = value
|
||||
return tensor_dict
|
||||
|
||||
def send_recv_tensor_dict(
|
||||
self,
|
||||
send_tensor_dict: Dict[str, Union[torch.Tensor, Any]],
|
||||
send_dst: Optional[int] = None,
|
||||
recv_src: Optional[int] = None,
|
||||
send_all_gather_group: Optional["GroupCoordinator"] = None,
|
||||
recv_all_gather_group: Optional["GroupCoordinator"] = None,
|
||||
) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
|
||||
"""Send tensor dict to *send_dst* and simultaneously recv from *recv_src*.
|
||||
|
||||
Uses ``batch_isend_irecv`` to submit all send/recv operations
|
||||
atomically, avoiding deadlock on backends (e.g. NPU/HCCL) where
|
||||
``isend`` may block until a matching ``recv`` is posted.
|
||||
|
||||
NOTE: ``send_dst`` / ``recv_src`` are local ranks within this group.
|
||||
"""
|
||||
if not torch.distributed.is_initialized() or self.world_size == 1:
|
||||
return None
|
||||
|
||||
if send_dst is None:
|
||||
send_dst = (self.rank_in_group + 1) % self.world_size
|
||||
if recv_src is None:
|
||||
recv_src = (self.rank_in_group - 1) % self.world_size
|
||||
|
||||
assert send_dst < self.world_size, f"Invalid send_dst rank ({send_dst})"
|
||||
assert recv_src < self.world_size, f"Invalid recv_src rank ({recv_src})"
|
||||
|
||||
group = self.device_group
|
||||
metadata_group = self.cpu_group
|
||||
|
||||
# ---- 1. Exchange metadata via batch_isend_irecv on CPU group ----
|
||||
send_metadata_list, send_tensor_list = _split_tensor_dict(send_tensor_dict)
|
||||
|
||||
send_meta_bytes = pickle.dumps(send_metadata_list)
|
||||
send_meta_size = torch.tensor([len(send_meta_bytes)], dtype=torch.long)
|
||||
send_meta_data = torch.frombuffer(send_meta_bytes, dtype=torch.uint8)
|
||||
|
||||
recv_meta_size = torch.empty(1, dtype=torch.long)
|
||||
|
||||
meta_ops = [
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.isend,
|
||||
send_meta_size,
|
||||
self.ranks[send_dst],
|
||||
group=metadata_group,
|
||||
),
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.irecv,
|
||||
recv_meta_size,
|
||||
self.ranks[recv_src],
|
||||
group=metadata_group,
|
||||
),
|
||||
]
|
||||
reqs = torch.distributed.batch_isend_irecv(meta_ops)
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
|
||||
recv_meta_len = recv_meta_size.item()
|
||||
recv_meta_data = torch.empty(recv_meta_len, dtype=torch.uint8)
|
||||
|
||||
meta_ops = [
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.isend,
|
||||
send_meta_data,
|
||||
self.ranks[send_dst],
|
||||
group=metadata_group,
|
||||
),
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.irecv,
|
||||
recv_meta_data,
|
||||
self.ranks[recv_src],
|
||||
group=metadata_group,
|
||||
),
|
||||
]
|
||||
reqs = torch.distributed.batch_isend_irecv(meta_ops)
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
|
||||
recv_metadata_list = pickle.loads(recv_meta_data.numpy())
|
||||
|
||||
# ---- 2. Prepare recv buffers and collect all tensor ops ----
|
||||
recv_tensor_dict: Dict[str, Any] = {}
|
||||
tensor_ops: List[torch.distributed.P2POp] = []
|
||||
recv_tensor_info: List[
|
||||
Tuple[str, torch.Tensor, bool, Optional[torch.Size]]
|
||||
] = []
|
||||
|
||||
for key, value in recv_metadata_list:
|
||||
if isinstance(value, TensorMetadata):
|
||||
tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
|
||||
if tensor.numel() == 0:
|
||||
recv_tensor_dict[key] = tensor
|
||||
continue
|
||||
|
||||
use_all_gather = (
|
||||
recv_all_gather_group is not None
|
||||
and tensor.numel() % recv_all_gather_group.world_size == 0
|
||||
)
|
||||
orig_shape = None
|
||||
if use_all_gather:
|
||||
orig_shape = tensor.shape
|
||||
tensor = tensor.reshape(recv_all_gather_group.world_size, -1)[
|
||||
recv_all_gather_group.rank_in_group
|
||||
]
|
||||
|
||||
comm_group = metadata_group if tensor.is_cpu else group
|
||||
tensor_ops.append(
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.irecv,
|
||||
tensor,
|
||||
self.ranks[recv_src],
|
||||
group=comm_group,
|
||||
)
|
||||
)
|
||||
recv_tensor_info.append((key, tensor, use_all_gather, orig_shape))
|
||||
else:
|
||||
recv_tensor_dict[key] = value
|
||||
|
||||
# Add send ops
|
||||
for tensor in send_tensor_list:
|
||||
if tensor.numel() == 0:
|
||||
continue
|
||||
send_t = tensor
|
||||
if (
|
||||
send_all_gather_group is not None
|
||||
and send_t.numel() % send_all_gather_group.world_size == 0
|
||||
):
|
||||
send_t = send_t.reshape(send_all_gather_group.world_size, -1)[
|
||||
send_all_gather_group.rank_in_group
|
||||
]
|
||||
comm_group = metadata_group if send_t.is_cpu else group
|
||||
tensor_ops.append(
|
||||
torch.distributed.P2POp(
|
||||
torch.distributed.isend,
|
||||
send_t,
|
||||
self.ranks[send_dst],
|
||||
group=comm_group,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- 3. Batch exchange all tensors ----
|
||||
if tensor_ops:
|
||||
reqs = torch.distributed.batch_isend_irecv(tensor_ops)
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
|
||||
# ---- 4. Post-process received tensors (all_gather if needed) ----
|
||||
for key, tensor, use_all_gather, orig_shape in recv_tensor_info:
|
||||
if use_all_gather:
|
||||
tensor = recv_all_gather_group.all_gather(tensor, dim=0)
|
||||
tensor = tensor.reshape(orig_shape)
|
||||
recv_tensor_dict[key] = tensor
|
||||
|
||||
return recv_tensor_dict
|
||||
|
||||
def barrier(self):
|
||||
"""Barrier synchronization among the group.
|
||||
NOTE: don't use `device_group` here! `barrier` in NCCL is
|
||||
|
||||
@@ -32,8 +32,9 @@ from sglang.srt.sampling.sampling_observer_pp import (
|
||||
pop_auxiliary_output_from_pp_tensors,
|
||||
)
|
||||
from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj
|
||||
from sglang.srt.utils.common import is_xpu
|
||||
from sglang.srt.utils.common import is_npu, is_xpu
|
||||
|
||||
_is_npu = is_npu()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -1025,6 +1026,20 @@ class SchedulerPPMixin:
|
||||
batch_result = None
|
||||
send_output_work = []
|
||||
|
||||
# On NPU (HCCL), isend/irecv may block until a matching peer op is
|
||||
# posted, so the parity-based send-first/recv-first ordering used
|
||||
# for NPU is replaced by batch_isend_irecv which submits all
|
||||
# send/recv operations atomically.
|
||||
if _is_npu and self.ps.pp_size == 2:
|
||||
return self._pp2_only_send_recv_output_tensors_npu(
|
||||
next_first_rank_mb_id,
|
||||
next_mb_id,
|
||||
mbs,
|
||||
mb_metadata,
|
||||
last_rank_comm_queue,
|
||||
pp_outputs,
|
||||
)
|
||||
|
||||
# On CUDA, isend is async: it enqueues to the stream and returns,
|
||||
# so every rank can send first safely. On some backends isend is
|
||||
# effectively blocking and does not return until the peer posts a
|
||||
@@ -1075,6 +1090,118 @@ class SchedulerPPMixin:
|
||||
|
||||
return next_pp_outputs, batch_result, d2h_event, send_output_work
|
||||
|
||||
def _pp2_only_send_recv_output_tensors_npu(
|
||||
self: Scheduler,
|
||||
next_first_rank_mb_id: int,
|
||||
next_mb_id: int,
|
||||
mbs: List[ScheduleBatch],
|
||||
mb_metadata: List[PPBatchMetadata],
|
||||
last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]],
|
||||
pp_outputs: PPProxyTensors | None,
|
||||
) -> Tuple[
|
||||
Optional[PPProxyTensors],
|
||||
Optional[GenerationBatchResult],
|
||||
Optional[torch.Event],
|
||||
List[P2PWork],
|
||||
]:
|
||||
"""NPU-specific output tensor send/recv using batch_isend_irecv.
|
||||
|
||||
Pairs the send of output tensors to the next stage with the recv
|
||||
of output tensors from the previous stage in a single
|
||||
``batch_isend_irecv`` call, avoiding the deadlock that can occur
|
||||
on HCCL when separate isend/irecv calls block waiting for each
|
||||
other.
|
||||
"""
|
||||
next_pp_outputs = None
|
||||
d2h_event = None
|
||||
batch_result = None
|
||||
send_output_work = []
|
||||
|
||||
all_gather_group = (
|
||||
self.attn_tp_group if self.require_attn_tp_allgather else None
|
||||
)
|
||||
|
||||
# ---- Prepare send dict ----
|
||||
# On NPU, always send something (full output or a lightweight skip
|
||||
# marker) so the peer's recv in batch_isend_irecv always has a
|
||||
# matching send. Without this, SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM
|
||||
# would cause asymmetric skip decisions between adjacent ranks
|
||||
# (send target and recv target are different micro-batches), leading
|
||||
# to deadlock or forcing the user to disable the optimisation
|
||||
# entirely (≈10 % throughput loss).
|
||||
send_dict: Optional[Dict[str, torch.Tensor]] = None
|
||||
if self.pp_group.is_last_rank:
|
||||
target_send = mbs[next_first_rank_mb_id]
|
||||
if target_send is not None:
|
||||
q_event, pp_outputs_to_send = last_rank_comm_queue.popleft()
|
||||
if not target_send.forward_mode.is_prebuilt():
|
||||
if _pp_can_skip_output_comm(target_send):
|
||||
send_dict = {"__msg_type__": "output", "__skip__": True}
|
||||
else:
|
||||
self.device_module.current_stream().wait_event(q_event)
|
||||
send_dict = dict(pp_outputs_to_send.tensors)
|
||||
send_dict["__msg_type__"] = "output"
|
||||
elif pp_outputs:
|
||||
if pp_outputs.tensors.get("__skip__"):
|
||||
send_dict = {"__msg_type__": "output", "__skip__": True}
|
||||
else:
|
||||
send_dict = dict(pp_outputs.tensors)
|
||||
send_dict["__msg_type__"] = "output"
|
||||
|
||||
# ---- Determine recv conditions ----
|
||||
# NOTE: skip_recv is NOT computed here. The skip decision is made
|
||||
# by the sender (last rank) and communicated via the __skip__
|
||||
# marker. The receiver always participates in the
|
||||
# batch_isend_irecv and inspects the marker after recv.
|
||||
target_recv = mbs[next_mb_id]
|
||||
should_recv = (
|
||||
target_recv is not None and not target_recv.forward_mode.is_prebuilt()
|
||||
)
|
||||
|
||||
def _handle_recv_dict(recv_dict):
|
||||
nonlocal next_pp_outputs, batch_result, d2h_event
|
||||
if recv_dict.get("__skip__"):
|
||||
# _pp_make_skip_output_result returns next_pp_outputs=None
|
||||
# (correct for the non-NPU path where the skip propagates
|
||||
# via pp_outputs=None). On NPU we must propagate the skip
|
||||
# marker through the ring, so override it here.
|
||||
_, batch_result, d2h_event = self._pp_make_skip_output_result(
|
||||
target_recv, mb_metadata[next_mb_id]
|
||||
)
|
||||
next_pp_outputs = PPProxyTensors(recv_dict)
|
||||
else:
|
||||
next_pp_outputs = PPProxyTensors(recv_dict)
|
||||
with self.copy_stream_ctx:
|
||||
self.copy_stream.wait_stream(self.schedule_stream)
|
||||
batch_result = self._pp_prep_batch_result(
|
||||
target_recv, mb_metadata[next_mb_id], next_pp_outputs
|
||||
)
|
||||
d2h_event = self.device_module.Event()
|
||||
d2h_event.record(self.device_module.current_stream())
|
||||
|
||||
# ---- Execute communication ----
|
||||
if send_dict is not None and should_recv:
|
||||
# Paired send + recv via batch_isend_irecv
|
||||
with torch.profiler.record_function("send_recv_res_dict"):
|
||||
recv_dict = self.pp_group.send_recv_tensor_dict(
|
||||
send_tensor_dict=send_dict,
|
||||
send_all_gather_group=all_gather_group,
|
||||
recv_all_gather_group=all_gather_group,
|
||||
)
|
||||
_handle_recv_dict(recv_dict)
|
||||
elif send_dict is not None:
|
||||
# Send only (recv not needed — target is None or prebuilt)
|
||||
send_output_work = self._pp_send_dict_to_next_stage(
|
||||
send_dict, async_send=True, msg_type="output"
|
||||
)
|
||||
elif should_recv:
|
||||
# Recv only (no send needed)
|
||||
with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
|
||||
recv_dict = self._pp_recv_dict_from_prev_stage()
|
||||
_handle_recv_dict(recv_dict)
|
||||
|
||||
return next_pp_outputs, batch_result, d2h_event, send_output_work
|
||||
|
||||
def _pp_launch_batch(
|
||||
self: Scheduler,
|
||||
mb_id: int,
|
||||
|
||||
Reference in New Issue
Block a user