<jupyter_start><jupyter_text>Table of Contents1&nbsp;&nbsp;Dense vs Sparse Self-Attention1.1&nbsp;&nbsp;Check equality of sparse implementation with reference2&nbsp;&nbsp;Benchmarking Sparse vs Dense Self-Attention3&nbsp;&nbsp;scatter_add vs index_add4&nbsp;&nbsp;embedding lookup vs index_select4.1&nbsp;&nbsp;Scattered GRU input<jupyter_code>%load_ext autoreload
%autoreload 2
%matplotlib inline

from pathlib import Path
import sys, os

# make this file executable from anywhere
#if __name__ == '__main__':

# in editor
#full_path = os.path.realpath(__file__)

# in jupyter
full_path = !pwd
full_path = full_path[0]


print(full_path)
REPO_ROOT = full_path.rsplit('ProGraML', maxsplit=1)[0] + 'ProGraML'
print(REPO_ROOT)
#insert at 1, 0 is the script path (or '' in REPL)
sys.path.insert(1, REPO_ROOT)
REPO_ROOT = Path(REPO_ROOT)
from matplotlib import pyplot as plt
import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric import utils
from torch_geometric.utils import softmax as scatter_softmax
from torch_geometric.data import DataLoader
from deeplearning.ml4pl.models.ggnn.modeling import SelfAttentionMessageLayer
from deeplearning.ml4pl.models.ggnn.configs import ProGraMLBaseConfig, GraphTransformerConfig
from deeplearning.ml4pl.models.ggnn.modeling import (
    SelfAttentionMessageLayer, TransformerUpdateLayer, NodeEmbeddings
)
config = GraphTransformerConfig()
ne = NodeEmbeddings(config)
x = torch.randint(8568,(10000,))
ns = ne(x)
ns
torch.std_mean(ns.view(-1))
plt.hist(ns.view(-1).detach().numpy(), bins=100)
plt.show()

tu = TransformerUpdateLayer(config)
new = tu(ns, ns)
torch.std_mean(new)
# att = SelfAttentionMessageLayer()
# real data
from deeplearning.ml4pl.poj104.dataset import POJ104Dataset
ds_dir = REPO_ROOT / 'deeplearning/ml4pl/poj104/classifyapp_data'
valid_data = POJ104Dataset(ds_dir, 'val')


tu()<jupyter_output><empty_output><jupyter_text>Dense vs Sparse Self-Attention<jupyter_code>def print_state_dict(mod):
    for n, t in mod.state_dict().items():
        print(n, t.size())

def num_parameters(mod) -> int:
    """Compute the number of trainable parameters in a nn.Module and its children."""
    num_params = sum(param.numel() for param in mod.parameters(recurse=True) if param.requires_grad)
    return num_params, f"{num_params * 4 / 1e6:.3f}MB"
config = BaseConfig.from_dict({
    'edge_type_count': 3,
    'backward_edges': True,
    'hidden_size': 200,
    'transformer_attn_bias': True,
    'transformer_num_heads': 8,
    'transformer_attn_dropout': 0.0,
})<jupyter_output><empty_output><jupyter_text>Check equality of sparse implementation with reference<jupyter_code>sparse_attn = SelfAttentionMessageLayer(config).to(device='cuda')

#print_state_dict(sparse_attn)
#print("")
#print(sparse_attn)
#print("")
#num_parameters(sparse_attn)
# get dummy data

num_n = 5
heads = 8
hidd = 200

random_emb = nn.Parameter(torch.randn(9000, hidd, device='cuda'), requires_grad=True)

ones = torch.ones(num_n, num_n, device='cuda')

adj = torch.tril(ones, diagonal=0)
#adj = ones
print("adj matrix: edges go from (row -> column) if [row, column] == 1")
print(adj)

attn_mask = torch.tril(ones * float('-inf'), diagonal=-1)

print("")
print("attn_mask.t(): We print the transposed attention mask,")
print("bc. the reference implementation indexes the attn_mask <target, src>")
print(attn_mask.t())

edge_index = utils.dense_to_sparse(adj)[0]
print(edge_index.size())
#print(edge_index.t()[:10])

x = torch.randint(9000, (num_n,))
print(x.size())

node_states = random_emb[x]
# get sparse attn on dummy data
sa_node_states, sa_weights = sparse_attn(edges=edge_index, node_states=node_states)
sa_sum = torch.sum(sa_node_states)

print(sa_sum)


in_proj = sparse_attn.qkv_in_proj.weight
in_proj_bias = sparse_attn.qkv_in_proj.bias
out_proj = sparse_attn.out_proj.weight
out_proj_bias = sparse_attn.out_proj.bias


ns = node_states.unsqueeze(1)
da_node_states, da_weights = F.multi_head_attention_forward(
                ns, ns, ns, 200, 8,
                in_proj, in_proj_bias,
                None, None, False,
                0.0, out_proj, out_proj_bias,
                training=False,
                key_padding_mask=None, need_weights=True, 
                attn_mask=attn_mask)
da_sum = torch.sum(da_node_states)


# reformat
da_node_states = da_node_states.squeeze()
da_weights = da_weights.squeeze().t()

print(da_sum)

print(da_weights)


sa_weights_matrix = torch.zeros(num_n, num_n, device='cuda')
for i, (s, t) in enumerate(edge_index.t()):
    sa_weights_matrix[s, t] = sa_weights[i]
print(sa_weights_matrix)
#print(sa_weights)
# assert similarity of sums of outputs
check1 = torch.abs(sa_sum - da_sum)
print((check1 < 1e-7).item())
print(torch.max(check1).item())
# assert similarity of attention weights for complete graphs
check2 = torch.abs(da_weights - sa_weights_matrix )
print(torch.all(check2 < 1e-7).item())
print(torch.max(check2).item())

# assert similarity of self-attention output (new node states)
check3 = torch.abs(da_node_states - sa_node_states)
print(torch.all(check3 < 1e-7).item())
print(torch.max(check3).item())<jupyter_output><empty_output><jupyter_text>Benchmarking Sparse vs Dense Self-Attention<jupyter_code>def satt(x, emb_table, edge_index):
    node_states = emb_table[x]
    s = sparse_attn(edges=edge_index, node_states=node_states)
    loss = torch.sum(s[0])
    loss.backward()
    torch.cuda.synchronize()

def datt(x, emb_table, attn_mask):
    query = emb_table[x.unsqueeze(1)]
    s = dense_attn(query, query, query, attn_mask=attn_mask, need_weights=False)
    loss = torch.sum(s[0])
    loss.backward()
    torch.cuda.synchronize()

def benchmark_on_causal(s,n):
    for num_n in [2**i for i in range(s,n+1)]:
        print(num_n)
        with torch.no_grad():
            causal_mask = (torch.triu(torch.ones(num_n, num_n, device='cuda')) - 1) * 1e4
            edge_index = utils.dense_to_sparse(causal_mask)[0]
            print(edge_index.size())
            x = torch.randint(9000, (num_n,))
        torch.cuda.synchronize()
        #print("sparse")
        %timeit satt(x, random_emb, edge_index)
        
        #print("dense")
        try:
            %timeit datt(x, random_emb, causal_mask)
        except RuntimeError:
            print("Dense OOM")
        del causal_mask
        del edge_index
        del x

benchmark_on_causal(8, 14)

def benchmark_on_real(s,n):
    for bsz in [2**i for i in range(s,n+1)]:
        dl = DataLoader(valid_data, batch_size=bsz)
        print(bsz)
        for graph in dl:
            graph.to(device='cuda')
            print(graph)
            break

        with torch.no_grad():
            x = graph.x[:,0]
            edge_index = graph.edge_index

        torch.cuda.synchronize()
        %timeit satt(x, random_emb, edge_index)
        
        try:
            with torch.no_grad():
                mask = (1 - utils.to_dense_adj(graph.edge_index).squeeze() ) * -10000
                print(f"Size of adjacency matrix: {mask.size()[0]**2 * 4 / 1e6:,.2f} MB")
            torch.cuda.synchronize()    
            %timeit datt(x, random_emb, mask)
        except RuntimeError:
            print("Dense OOM")
            
        print("----"*10)
benchmark_on_real(0, 10)
dl = DataLoader(valid_data, batch_size=64)
for graph in dl:
    graph.to(device='cuda')
    break
print(graph)

with torch.no_grad():
    adj = (1 - to_dense_adj(graph.edge_index).squeeze() ) * -10000
adj.size(), adj.device
print(f"Size of adjacency matrix: {adj.size()[0]**2 * 4 / 1e6:,.2f} MB")<jupyter_output><empty_output><jupyter_text>scatter_add vs index_add<jupyter_code>torch.utils.
target = torch.randn(2000,200, requires_grad=True, device='cuda')
mbs = torch.randn(10000, 200, device='cuda', requires_grad=True)
index = torch.randint(0, 2000, [10000], device='cuda')
sa = torch.zeros_like(target)
sa.scatter_add_(0, index.unsqueeze(1).repeat(1, 200), mbs)

s = torch.zeros_like(target)
s.index_add_(0, index, mbs)

torch.max(torch.abs(sa - s)), torch.argmax(torch.abs(sa - s))
def ia():
    s = torch.zeros_like(target)
    s.index_add_(0, index, mbs)
    loss = torch.sum(s)
    loss.backward()
    torch.cuda.synchronize()

def sa():
    s = torch.zeros_like(target)
    s.scatter_add_(0, index.unsqueeze(1).repeat(1, 200), mbs)
    loss = torch.sum(s)
    loss.backward()
    torch.cuda.synchronize()
    
    
torch.cuda.synchronize()    
%timeit ia()
%timeit sa()<jupyter_output><empty_output><jupyter_text>embedding lookup vs index_select<jupyter_code>bs = 2000
mask = torch.randn(bs, device='cuda')
mask2 = mask>0
index = torch.randint(0, 2000, [10000], dtype=torch.long, device='cuda')

target = torch.randn(bs,200, requires_grad=True, device='cuda')
pred = torch.randn(bs,200, requires_grad=True, device='cuda')

h = torch.nn.GRUCell(200, 200).to(device='cuda')
g = torch.nn.GRUCell(200, 200).to(device='cuda')
torch.equal(F.embedding(index, target), torch.index_select(target, dim=0, index=index))
torch.index_select(target, dim=0, index=index).size()
def emb():
    selected = F.embedding(index, target)
    loss = torch.sum(selected)
    loss.backward()
    torch.cuda.synchronize()

def index_select():
    selected = torch.index_select(target, dim=0, index=index)
    loss = torch.sum(selected)
    loss.backward()
    torch.cuda.synchronize()

torch.cuda.synchronize()

%timeit emb()
%timeit index_select()<jupyter_output><empty_output><jupyter_text>Scattered GRU input<jupyter_code>def fn1():
    a = torch.zeros_like(target,device='cuda')
    a[mask2] += g(target[mask2])
    a[~mask2] += h(target[~mask2])
    
    loss = torch.sum(a)
    loss.backward()
    torch.cuda.synchronize()
    return loss

l = fn1()

def fn2():
    a = torch.zeros_like(target,device='cuda')
    a += mask2.view(-1, 1) * g(target)
    a += ~mask2.view(-1, 1) * h(target)
    loss = torch.sum(a)
    loss.backward()
    torch.cuda.synchronize()
    return loss

ll = fn2()
    
torch.cuda.synchronize()

%timeit fn1()
%timeit fn2()
print(ll)
print(l)
def fn1():
    loss = torch.sum(mask.view(-1, 1) * (target - pred) ** 2)
    loss.backward()
    torch.cuda.synchronize()
def fn2():
    loss = torch.sum((target[mask2,:] - pred[mask2,:]) ** 2)
    loss.backward()
    torch.cuda.synchronize()

torch.cuda.synchronize()

%timeit fn1()
%timeit fn2()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13735732249951635, 'fraction_numerical': 0.024182627200619075, 'mean_word_length': 4.143781094527363, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '269770', 'license': 'Apache-2.0', 'n_tokens_mistral': 4016, 'n_tokens_neox': 3858, 'n_words': 747, 'path': 'deeplearning/ml4pl/poj104/Runtime_Tests_Scattered_Ops.ipynb', 'repo_name': 'Zacharias030/ProGraML', 'size': '10338'}