207 lines
6.6 KiB
Python
207 lines
6.6 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import math
|
||
import random
|
||
from typing import Tuple, Optional, List
|
||
from einops import rearrange
|
||
|
||
from block_sparse_attn import block_sparse_attn_func
|
||
|
||
out_path = '/cpfs01/shared/xpixel/zhuangjunhao/swin_base/DiffSynth-Studio/diffsynth/models/out_mask/'
|
||
class WindowPartition3D:
|
||
"""Partition / reverse‑partition helpers for 5‑D tensors (B,F,H,W,C)."""
|
||
|
||
@staticmethod
|
||
def partition(x: torch.Tensor, win: Tuple[int, int, int]):
|
||
B, F, H, W, C = x.shape
|
||
wf, wh, ww = win
|
||
assert F % wf == 0 and H % wh == 0 and W % ww == 0, "Dims must divide by window size."
|
||
|
||
x = x.view(B,
|
||
F // wf, wf,
|
||
H // wh, wh,
|
||
W // ww, ww,
|
||
C)
|
||
# (B, nf, nh, nw, wf, wh, ww, C) -> (B*nf*nh*nw, wf*wh*ww, C)
|
||
x = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous()
|
||
return x.view(-1, wf * wh * ww, C)
|
||
|
||
@staticmethod
|
||
def reverse(windows: torch.Tensor,
|
||
win: Tuple[int, int, int],
|
||
orig: Tuple[int, int, int]):
|
||
F, H, W = orig
|
||
wf, wh, ww = win
|
||
nf, nh, nw = F // wf, H // wh, W // ww
|
||
B = windows.size(0) // (nf * nh * nw)
|
||
x = windows.view(B, nf, nh, nw, wf, wh, ww, -1)
|
||
x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous()
|
||
return x.view(B, F, H, W, -1)
|
||
|
||
|
||
def generate_causal_draft_block_mask(batch_size, nheads, seqlen, local_num, window_size, q_w, k_w, topk=10, device='cuda'):
|
||
assert batch_size == 1, "batch_size must be 1"
|
||
# 生成 causal_mask
|
||
causal_mask = generate_causal_block_mask(batch_size, nheads, seqlen, local_num, window_size, device).squeeze(0)
|
||
|
||
# q_w 和 k_w 的形状为 (nf*nh*nw, wf*wh*ww, D)
|
||
# avgpool_q 和 avgpool_k 的形状为 (nf*nh*nw, D) avgpool wf*wh*ww
|
||
avgpool_q = torch.mean(q_w, dim=1)
|
||
avgpool_k = torch.mean(k_w, dim=1)
|
||
# avgpool_q = torch.randn(avgpool_q.shape, device=avgpool_q.device, dtype=avgpool_q.dtype)
|
||
# avgpool_k = torch.randn(avgpool_k.shape, device=avgpool_k.device, dtype=avgpool_k.dtype)
|
||
# print(avgpool_q.shape, avgpool_k.shape)
|
||
|
||
avgpool_q = rearrange(avgpool_q, 's (h d) -> s h d', h=nheads) # (nf*nh*nw, h, D)
|
||
avgpool_k = rearrange(avgpool_k, 's (h d) -> s h d', h=nheads) # (nf*nh*nw, h, D)
|
||
|
||
|
||
q_heads = avgpool_q.permute(1, 0, 2) # (h, nf*nh*nw, D)
|
||
k_heads = avgpool_k.permute(1, 0, 2) # (h, nf*nh*nw, D)
|
||
|
||
D = avgpool_q.shape[-1]
|
||
|
||
scores = torch.einsum("hld,hmd->hlm", q_heads, k_heads) / math.sqrt(D) # (h, nf*nh*nw, nf*nh*nw)
|
||
|
||
# causal_mask 中为 False 的区域设置为 -inf True 的区域设置为 0
|
||
causal_mask = causal_mask.to(torch.float32)
|
||
causal_mask = causal_mask.masked_fill(causal_mask == False, -float('inf'))
|
||
causal_mask = causal_mask.masked_fill(causal_mask == True, 0)
|
||
|
||
scores = scores + causal_mask
|
||
|
||
attn_map = torch.softmax(scores, dim=-1)
|
||
# print(attn_map)
|
||
assert attn_map.shape == causal_mask.shape, "attn_map.shape must be equal to causal_mask.shape"
|
||
|
||
|
||
attn_map = rearrange(attn_map, 'h (it s1) s2 -> (h it) s1 s2', it=seqlen)
|
||
|
||
loop_num = attn_map.shape[0]
|
||
|
||
mask = torch.zeros_like(attn_map, dtype=torch.bool)
|
||
|
||
# process each head independently
|
||
for h in range(loop_num):
|
||
head_scores = attn_map[h] # [S1, S2]
|
||
# print(head_scores)
|
||
flat = head_scores.flatten() # [S1*S2]
|
||
n = flat.numel()
|
||
k = topk # number of smallest to exclude
|
||
|
||
if k == 0:
|
||
mask[h] = True
|
||
continue
|
||
if k >= n:
|
||
# nothing to keep
|
||
continue
|
||
|
||
# threshold = max of the k smallest scores
|
||
# print(flat, k)
|
||
# 计算非0元素个数
|
||
# non_zero_num = (flat > 0).sum()
|
||
# print('non_zero_num', non_zero_num, k)
|
||
# threshold = torch.topk(flat, k, largest=False).values.max()
|
||
# print(torch.sort(flat, descending=True)[0])
|
||
threshold = torch.sort(flat, descending=True)[0][k]
|
||
|
||
# build head mask
|
||
# print(threshold)
|
||
mask[h] = head_scores > threshold
|
||
# print(mask[h])
|
||
|
||
mask = rearrange(mask, '(h it) s1 s2 -> h (it s1) s2', it=seqlen)
|
||
|
||
# save mask 为图像
|
||
import matplotlib.pyplot as plt
|
||
for i in range(mask.shape[0]):
|
||
plt.imshow(mask[i].float().int().cpu().numpy())
|
||
plt.savefig(out_path + f'mask_{i}.png')
|
||
|
||
mask = mask.unsqueeze(0).repeat(batch_size, 1, 1, 1)
|
||
|
||
return mask
|
||
|
||
|
||
def generate_causal_block_mask(batch_size, nheads, seqlen, local_num, window_size, device='cuda'):
|
||
# i: 当前 query 的位置，j: key 的位置
|
||
i = torch.arange(seqlen, device=device).view(-1, 1) # [seqlen, 1]
|
||
j = torch.arange(seqlen, device=device).view(1, -1) # [1, seqlen]
|
||
|
||
# 条件：j <= i 且 j >= i - local_num + 1
|
||
causal_mask = (j <= i) & (j >= i - local_num + 1) # [seqlen, seqlen]
|
||
|
||
# print(causal_mask.float().int())
|
||
|
||
causal_mask = causal_mask.unsqueeze(1).unsqueeze(-1).repeat(1, window_size, 1, window_size)
|
||
causal_mask = rearrange(causal_mask, 'a n1 b n2 -> (a n1) (b n2)')
|
||
# save causal_mask 为图像
|
||
# import matplotlib.pyplot as plt
|
||
# plt.imshow(causal_mask.float().int())
|
||
# plt.savefig('causal_mask.png')
|
||
# 扩展 batch 维度
|
||
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0).repeat(batch_size, nheads, 1, 1)
|
||
# causal_mask[:] = True
|
||
return causal_mask
|
||
|
||
|
||
win = (2, 8, 8)
|
||
B = 1
|
||
nheads = 12
|
||
f = 16
|
||
h = 32
|
||
w = 56
|
||
dim = 128
|
||
device = 'cuda'
|
||
dtype = torch.bfloat16
|
||
|
||
seqlen = f//win[0]
|
||
|
||
token_num = f*h*w
|
||
|
||
local_num = random.randint(2, seqlen)
|
||
# local_num = 4
|
||
print(local_num)
|
||
# print(seqlen)
|
||
|
||
window_size = win[0]*h*w//128
|
||
|
||
q = torch.randn(B, token_num, nheads*dim, device=device, dtype=dtype)
|
||
k = torch.randn(B, token_num, nheads*dim, device=device, dtype=dtype)
|
||
k = q+torch.randn(B, token_num, nheads*dim, device=device, dtype=dtype)
|
||
v = torch.randn(B, token_num, nheads*dim, device=device, dtype=dtype)
|
||
|
||
q = q.view(B, f, h, w, nheads*dim)
|
||
k = k.view(B, f, h, w, nheads*dim)
|
||
v = v.view(B, f, h, w, nheads*dim)
|
||
|
||
q_w = WindowPartition3D.partition(q, win) # (B*nf*nh*nw, wf*wh*ww, D)
|
||
k_w = WindowPartition3D.partition(k, win) # (B*nf*nh*nw, wf*wh*ww, D)
|
||
v_w = WindowPartition3D.partition(v, win) # (B*nf*nh*nw, wf*wh*ww, D)
|
||
|
||
square_num = window_size*window_size
|
||
topk_ratio = (random.uniform(0., 1.0)**2)*(local_num-1.2)+1.2
|
||
topk = int(square_num*topk_ratio)
|
||
print(topk/square_num)
|
||
# topk = int(window_size*window_size*1.5)
|
||
attention_mask = generate_causal_draft_block_mask(B, nheads, seqlen, local_num, window_size, q_w, k_w, topk=topk, device=q.device)
|
||
|
||
print(attention_mask.shape)
|
||
print(attention_mask.float().mean())
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|