[diffusion] fix: fix NCCL deadlock in ulysses sp when sequence length has remainder (#24694)
Signed-off-by: storyicon <storyicon@foxmail.com>
This commit is contained in:
@@ -478,6 +478,10 @@ class USPAttention(nn.Module):
|
||||
k = _usp_input_all_to_all(k, head_dim=2)
|
||||
v = _usp_input_all_to_all(v, head_dim=2)
|
||||
|
||||
# If NCCL timeout/deadlock occurs here, check whether
|
||||
# attn_mask is inconsistent across SP ranks (None on some, Tensor on
|
||||
# others), which causes all_gather participant mismatch. Upstream
|
||||
# mask builders must ensure all ranks produce the same mask type.
|
||||
gathered_mask = sequence_model_parallel_all_gather(
|
||||
attn_mask.contiguous(), dim=1
|
||||
)
|
||||
|
||||
@@ -1050,12 +1050,9 @@ class LTX2DenoisingStage(DenoisingStage):
|
||||
if valid is None:
|
||||
return None
|
||||
valid = int(valid)
|
||||
if valid <= 0 or valid >= int(seq_len):
|
||||
return None
|
||||
mask = torch.ones(
|
||||
(batch_size, int(seq_len)), device=device, dtype=torch.float32
|
||||
)
|
||||
mask[:, valid:] = 0.0
|
||||
mask = torch.ones((batch_size, int(seq_len)), device=device, dtype=torch.bool)
|
||||
if valid < int(seq_len):
|
||||
mask[:, max(0, valid) :] = False
|
||||
return mask
|
||||
|
||||
@staticmethod
|
||||
|
||||
Reference in New Issue
Block a user