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,27 +235,18 @@ class SchedulerRequestReceiver:
|
||||
# so that ShmPointerMMData metadata (not full tensor data) is what
|
||||
# gets serialized during broadcast_pyobj.
|
||||
if recv_reqs:
|
||||
# Barrier for the non-DP-attention path only: there is a single
|
||||
# broadcast_pyobj on tp_cpu_group where the source rank returns
|
||||
# the original objects immediately while other ranks are still in
|
||||
# pickle.loads (-> __setstate__ -> shm_open). Without a barrier
|
||||
# the source can call materialize() / shm_unlink before others
|
||||
# open the segment. recv_reqs is consistent across all ranks
|
||||
# here (same broadcast), so the guard is deadlock-free.
|
||||
#
|
||||
# Under DP-attention no barrier is needed: the control_reqs
|
||||
# broadcast on tp_cpu_group (step 3) is a collective that forces
|
||||
# every rank to complete the earlier attn_tp / attn_cp work_reqs
|
||||
# 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)
|
||||
if self.model_config.is_multimodal and has_shm_features(recv_reqs):
|
||||
# The broadcast source returns with its original objects while
|
||||
# peer ranks may still be unpickling ShmPointerMMData
|
||||
# (-> shm_open). Synchronize the same CPU groups that carried
|
||||
# SHM-backed work requests before materialize() unlinks them.
|
||||
if self.server_args.enable_dp_attention:
|
||||
if self.ps.attn_tp_size > 1:
|
||||
barrier(group=self.attn_tp_cpu_group)
|
||||
if self.ps.attn_cp_size > 1:
|
||||
barrier(group=self.attn_cp_cpu_group)
|
||||
elif self.ps.tp_size > 1:
|
||||
barrier(group=self.tp_cpu_group)
|
||||
for req in recv_reqs:
|
||||
unwrap_shm_features(req)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user