def get_ops(self):
try:
from deepspeed.ops.sparse_attention import MatMul, Softmax
except:
raise Exception(
'Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`'
)
if self.shape not in SparseAttention.ops:
sparsity_layout = self.sparsity_config.make_layout()
sparse_dot_sdd_nt = MatMul(sparsity_layout,
self.sparsity_config.block,
'sdd',
trans_a=False,
trans_b=True)
sparse_dot_dsd_nn = MatMul(sparsity_layout,
self.sparsity_config.block,
'dsd',
trans_a=False,
trans_b=False)
sparse_softmax = Softmax(sparsity_layout,
self.sparsity_config.block)
SparseAttention.ops[self.shape] = (sparse_dot_sdd_nt,
sparse_dot_dsd_nn,
sparse_softmax)
return SparseAttention.ops[self.shape]
def run_softmax_sparse(x, scale, dx, kp_mask, attn_mask, layout, block):
from deepspeed.ops.sparse_attention import Softmax
sparse_softmax = Softmax(layout, block, bench=False)
dx = dense_to_sparse(dx, layout, block)
x = dense_to_sparse(x, layout, block)
x.retain_grad()
y = sparse_softmax(x,
scale=scale,
key_padding_mask=kp_mask,
key_padding_mask_mode='add',
attn_mask=attn_mask,
attn_mask_mode='mul')
y.backward(dx)
dx = x.grad.clone()
x.grad.zero_()
return x, dx
def get_ops(self):
from deepspeed.ops.sparse_attention import MatMul, Softmax
if self.shape not in SparseAttention.ops:
sparsity_layout = self.sparsity_config.make_layout()
sparse_dot_sdd_nt = MatMul(sparsity_layout,
self.sparsity_config.block,
'sdd',
trans_a=False,
trans_b=True)
sparse_dot_dsd_nn = MatMul(sparsity_layout,
self.sparsity_config.block,
'dsd',
trans_a=False,
trans_b=False)
sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block)
SparseAttention.ops[self.shape] = (sparse_dot_sdd_nt,
sparse_dot_dsd_nn,
sparse_softmax)
return SparseAttention.ops[self.shape]
def get_ops(self):
if self.seq_len not in SparseSelfAttention.ops:
sparsity_layout = self.sparsity_config.make_layout()
sparse_dot_sdd_nt = MatMul(sparsity_layout,
self.sparsity_config.block,
'sdd',
trans_a=False,
trans_b=True)
sparse_dot_dsd_nn = MatMul(sparsity_layout,
self.sparsity_config.block,
'dsd',
trans_a=False,
trans_b=False)
sparse_softmax = Softmax(sparsity_layout,
self.sparsity_config.block)
SparseSelfAttention.ops[self.seq_len] = (sparse_dot_sdd_nt,
sparse_dot_dsd_nn,
sparse_softmax)
return SparseSelfAttention.ops[self.seq_len]
def get_ops(self, H, L):
import sys
if L not in SparseSelfAttention.ops:
sparsity_layout = self.get_layout(L)
sparse_dot_sdd_nt = MatMul(sparsity_layout,
self.sparsity_config.block,
'sdd',
trans_a=False,
trans_b=True)
sparse_dot_dsd_nn = MatMul(sparsity_layout,
self.sparsity_config.block,
'dsd',
trans_a=False,
trans_b=False)
sparse_softmax = Softmax(sparsity_layout,
self.sparsity_config.block)
SparseSelfAttention.ops[L] = (sparse_dot_sdd_nt, sparse_dot_dsd_nn,
sparse_softmax)
return SparseSelfAttention.ops[L]