Fix DP-attention SHM feature finalization race (#29543)
Co-authored-by: Yinghai Lu <yinghai@meta.com>
This commit is contained in:
co-authored by
Yinghai Lu
parent
ae09b8302f
commit
5747ed3b19
@@ -235,26 +235,17 @@ class SchedulerRequestReceiver:
|
|||||||
# so that ShmPointerMMData metadata (not full tensor data) is what
|
# so that ShmPointerMMData metadata (not full tensor data) is what
|
||||||
# gets serialized during broadcast_pyobj.
|
# gets serialized during broadcast_pyobj.
|
||||||
if recv_reqs:
|
if recv_reqs:
|
||||||
# Barrier for the non-DP-attention path only: there is a single
|
if self.model_config.is_multimodal and has_shm_features(recv_reqs):
|
||||||
# broadcast_pyobj on tp_cpu_group where the source rank returns
|
# The broadcast source returns with its original objects while
|
||||||
# the original objects immediately while other ranks are still in
|
# peer ranks may still be unpickling ShmPointerMMData
|
||||||
# pickle.loads (-> __setstate__ -> shm_open). Without a barrier
|
# (-> shm_open). Synchronize the same CPU groups that carried
|
||||||
# the source can call materialize() / shm_unlink before others
|
# SHM-backed work requests before materialize() unlinks them.
|
||||||
# open the segment. recv_reqs is consistent across all ranks
|
if self.server_args.enable_dp_attention:
|
||||||
# here (same broadcast), so the guard is deadlock-free.
|
if self.ps.attn_tp_size > 1:
|
||||||
#
|
barrier(group=self.attn_tp_cpu_group)
|
||||||
# Under DP-attention no barrier is needed: the control_reqs
|
if self.ps.attn_cp_size > 1:
|
||||||
# broadcast on tp_cpu_group (step 3) is a collective that forces
|
barrier(group=self.attn_cp_cpu_group)
|
||||||
# every rank to complete the earlier attn_tp / attn_cp work_reqs
|
elif self.ps.tp_size > 1:
|
||||||
# deserializations (steps 1-2, which call shm_open) before any
|
|
||||||
# rank returns from step 3. POSIX guarantees shm_unlink only
|
|
||||||
# removes the name; already-open handles stay valid.
|
|
||||||
if (
|
|
||||||
not self.server_args.enable_dp_attention
|
|
||||||
and self.ps.tp_size > 1
|
|
||||||
and self.model_config.is_multimodal
|
|
||||||
and has_shm_features(recv_reqs)
|
|
||||||
):
|
|
||||||
barrier(group=self.tp_cpu_group)
|
barrier(group=self.tp_cpu_group)
|
||||||
for req in recv_reqs:
|
for req in recv_reqs:
|
||||||
unwrap_shm_features(req)
|
unwrap_shm_features(req)
|
||||||
|
|||||||
Reference in New Issue
Block a user