nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-10 12:49:38 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Merge remote-tracking branch 'upstream/main' into hotfix/fix-sp

Browse source
This commit is contained in:
FrankLeeeee 2024-06-24 09:08:21 +00:00
commit 6bb2c599b6
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
opensora/models/layers/blocks.py
View file

@ -499,7 +499,7 @@ class SeqParallelMultiHeadCrossAttention(MultiHeadCrossAttention):
# shape: # shape:
# q, k, v: [B, SUB_N, NUM_HEADS, HEAD_DIM] # q, k, v: [B, SUB_N, NUM_HEADS, HEAD_DIM]
q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim) q = self.q_linear(x).view(B, -1, self.num_heads, self.head_dim)
kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim) kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
kv = split_forward_gather_backward(kv, get_sequence_parallel_group(), dim=3, grad_scale="down") kv = split_forward_gather_backward(kv, get_sequence_parallel_group(), dim=3, grad_scale="down")
k, v = kv.unbind(2) k, v = kv.unbind(2)