✅ Correctness issue

[JoelMathewC]

We believe that there is an issue of correctness, evident from the disparity of training logs between Seastar and PyG-T. This discussion will attempt to resolve that.

The logs are as follows

Seastar Logs
-------------
Epoch ┃ Time (s) ┃ MSE    ┃ GPU Memory(mb)       
     0 │ 1.7256   │ 0.9397 │ 718.0      │ 712.2857   
     1 │ 0.126    │ 0.8745 │ 726.0      │ 721.0     
     2 │ 0.1277   │ 0.8534 │ 726.0      │ 721.0     
     3 │ 0.1275   │ 0.8116 │ 726.0      │ 721.8095 
     4 │ 0.1243   │ 0.7949 │ 726.0      │ 721.8095  
     5 │ 0.1277   │ 0.7936 │ 726.0      │ 721.8095 
     6 │ 0.1304   │ 0.7991 │ 726.0      │ 721.0   
     7 │ 0.1287   │ 0.7964 │ 726.0      │ 721.8095 
     8 │ 0.134    │ 0.7923 │ 726.0      │ 721.0    
     9 │ 0.1389   │ 0.7893 │ 726.0      │ 721.8095 
    10 │ 0.1298   │ 0.7841 │ 726.0      │ 721.0    

PyG-T Logs
------------
Epoch 000 | Time(s) 1.2041 | MSE 0.90 | Used GPU Memory (Max) 838.000 mb
Epoch 001 | Time(s) 0.3418 | MSE 0.82 | Used GPU Memory (Max) 846.000 mb
Epoch 002 | Time(s) 0.3144 | MSE 0.81 | Used GPU Memory (Max) 846.000 mb
Epoch 003 | Time(s) 0.3172 | MSE 0.79 | Used GPU Memory (Max) 846.000 mb
Epoch 004 | Time(s) 0.3106 | MSE 0.77 | Used GPU Memory (Max) 846.000 mb
Epoch 005 | Time(s) 0.3093 | MSE 0.76 | Used GPU Memory (Max) 846.000 mb
Epoch 006 | Time(s) 0.3207 | MSE 0.76 | Used GPU Memory (Max) 846.000 mb
Epoch 007 | Time(s) 0.3076 | MSE 0.75 | Used GPU Memory (Max) 846.000 mb
Epoch 008 | Time(s) 0.3132 | MSE 0.74 | Used GPU Memory (Max) 846.000 mb
Epoch 009 | Time(s) 0.1969 | MSE 0.74 | Used GPU Memory (Max) 846.000 mb
Epoch 010 | Time(s) 0.1904 | MSE 0.74 | Used GPU Memory (Max) 846.000 mb

[JoelMathewC]

GAT correctness

The logs of training GAT is as follows

now_mem :  58080256
Epoch 00000 | Time(s) 2.3272 | train_acc 0.151157 | Used_Memory 55.389648 mb
now_mem :  75464704
Epoch 00001 | Time(s) 0.0282 | train_acc 0.267356 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00002 | Time(s) 0.0285 | train_acc 0.381585 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00003 | Time(s) 0.0278 | train_acc 0.421960 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00004 | Time(s) 0.0275 | train_acc 0.445101 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00005 | Time(s) 0.0228 | train_acc 0.461841 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00006 | Time(s) 0.0226 | train_acc 0.476120 | Used_Memory 71.968750 mb
now_mem :  75464704

The training logs of Seastar GAT (#7) prior to all of our modifications was

now_mem :  49416704
Epoch 00000 | Time(s) 2.2858 | train_acc 0.171429 | Used_Memory 47.127441 mb
now_mem :  61494784
Epoch 00001 | Time(s) 0.0061 | train_acc 0.257143 | Used_Memory 58.645996 mb
now_mem :  69211648
Epoch 00002 | Time(s) 0.0059 | train_acc 0.514286 | Used_Memory 66.005371 mb
now_mem :  69211648
Epoch 00003 | Time(s) 0.0057 | train_acc 0.657143 | Used_Memory 66.005371 mb
now_mem :  69211648
Epoch 00004 | Time(s) 0.0055 | train_acc 0.764286 | Used_Memory 66.005371 mb
now_mem :  69211648
Epoch 00005 | Time(s) 0.0055 | train_acc 0.864286 | Used_Memory 66.005371 mb
now_mem :  69211648
Epoch 00006 | Time(s) 0.0064 | train_acc 0.807143 | Used_Memory 66.005371 mb
now_mem :  69211648

[JoelMathewC]

🐛 Found the Bug

The CSR module expects input edges in a sorted order for the CSR to be constructed. In StaticGraph sorting edge_list for forward and backward graphs were done before building both the CSRs. Since a deepcopy was not used to perform the sorting, the forward edge list itself was sorted in the reverse order. This caused the forward CSR to be incorrect.

Fix
---
def _prepare_edge_lst_bwd(self, edge_list):    
        edge_list_for_t = copy.deepcopy(edge_list)
        edge_list_for_t.sort()
        self.bwd_edge_list = edge_list_for_t

The updated training log is as follows

Epoch 00000 | Time(s) 0.8430 | train_acc 0.200000 | Used_Memory 55.392578 mb
now_mem :  79999488
Epoch 00001 | Time(s) 0.0056 | train_acc 0.414286 | Used_Memory 76.293457 mb
now_mem :  84233728
Epoch 00002 | Time(s) 0.0055 | train_acc 0.650000 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00003 | Time(s) 0.0053 | train_acc 0.635714 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00004 | Time(s) 0.0053 | train_acc 0.671429 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00005 | Time(s) 0.0052 | train_acc 0.764286 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00006 | Time(s) 0.0052 | train_acc 0.714286 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00007 | Time(s) 0.0053 | train_acc 0.757143 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00008 | Time(s) 0.0052 | train_acc 0.764286 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00009 | Time(s) 0.0052 | train_acc 0.700000 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00010 | Time(s) 0.0052 | train_acc 0.778571 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00011 | Time(s) 0.0053 | train_acc 0.800000 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00012 | Time(s) 0.0051 | train_acc 0.764286 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00013 | Time(s) 0.0056 | train_acc 0.828571 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00014 | Time(s) 0.0051 | train_acc 0.885714 | Used_Memory 80.331543 mb
now_mem :  84233728
Epoch 00015 | Time(s) 0.0052 | train_acc 0.892857 | Used_Memory 80.331543 mb
now_mem :  84233728

Note: That this setup is with Val._t referring to the original tensor.

[JoelMathewC]

⚠️ Issue identified

When making Val._t a cloned detached CPU tensor, it is not able to train to the level it trained previously. The log is as shown here

Epoch 00000 | Time(s) 0.7406 | train_acc 0.178571 | Used_Memory 55.389648 mb
now_mem :  75464704
Epoch 00001 | Time(s) 0.0051 | train_acc 0.264286 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00002 | Time(s) 0.0047 | train_acc 0.314286 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00003 | Time(s) 0.0046 | train_acc 0.478571 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00004 | Time(s) 0.0045 | train_acc 0.557143 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00005 | Time(s) 0.0048 | train_acc 0.700000 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00006 | Time(s) 0.0045 | train_acc 0.707143 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00007 | Time(s) 0.0046 | train_acc 0.778571 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00008 | Time(s) 0.0045 | train_acc 0.757143 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00009 | Time(s) 0.0046 | train_acc 0.792857 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00010 | Time(s) 0.0046 | train_acc 0.814286 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00011 | Time(s) 0.0045 | train_acc 0.800000 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00012 | Time(s) 0.0046 | train_acc 0.857143 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00013 | Time(s) 0.0045 | train_acc 0.850000 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00014 | Time(s) 0.0047 | train_acc 0.871429 | Used_Memory 71.968750 mb
now_mem :  75464704
Epoch 00015 | Time(s) 0.0046 | train_acc 0.821429 | Used_Memory 71.968750 mb
now_mem :  75464704

We verified that the backend is getting called and we looked at what the effect would be from setting the requires gradient to True for the return tensor and the gradient tensor and noticed no benefit.

First Clue

When running in the setup with Val._t as the original tensor, we noticed that the run failed when tensors were selectively pushed onto the state-stack. But it ran smoothly when Val._t is a detached clone CPU tensor. So my theory is there is a portion of backprop code that is not accessed when this approach is taken that is preventing us from running into the error mentioned. Identifying what that area of code is and allowing the CPU tensor setup to run into that error as well would maybe get us a chance to have parallel performance with the working setup.

[JoelMathewC]

Problem Isolated

The backward kernel args are empty when detached clone CPU tensor is used. Need to figure out how to populate the backward kernel args. Its too difficult to look at the output and work our way back so the only step forward is to understand the role of Val in seastar.

Check 1

I isolated the tracing operation with monkey patching and noticed that the same Program object is generated by both versions of the code (cloned CPU tensor and GPU tensor). Additionally, I printed out the tensors generated by this phase and the output is as follows

<class 'torch.Tensor'> torch.Size([100, 1]) cpu 4
<class 'torch.Tensor'> torch.Size([1]) cpu 4
<class 'torch.Tensor'> torch.Size([100, 1]) cpu 4
<class 'torch.Tensor'> torch.Size([1]) cpu 4
<class 'torch.Tensor'> torch.Size([100, 16]) cpu 4
<class 'torch.Tensor'> torch.Size([16]) cpu 4
<class 'torch.Tensor'> torch.Size([100, 16]) cpu 4
<class 'torch.Tensor'> torch.Size([16]) cpu 4
<class 'torch.Tensor'> torch.Size([100, 1]) cpu 4
<class 'torch.Tensor'> torch.Size([100, 16]) cpu 4

Note: The tracing was done in a function. The check of stored tensors was done in the main program after the function was executed. Hence, the tensors did not go out of scope and these are the same tensors we were trying to identify in #59

Check 2

After looking through autodiff.py I was able to isolate the issue to the following

output_grad_map = {k:k._grad for k in need_grad_var}

This grad map turns out to be empty for detached clone of CPU tensors. It contains values for GPU tensor version, so populating it could possible fix the issue.

[JoelMathewC]

Understanding monkey-patching

Monkey-patching is performed on all torch.* operators. But when two tensors are added as a+b that is not monkey-patched. To handle this case we wrap the variables in Val objects which have custom implementations for __add__, __mul__ and so on.

[JoelMathewC]

Wrapping up

This may have gone a little of topic but for all purposes GAT correctness has been established with the fix of the deepcopy bug in the StaticGraph class.

[JoelMathewC]

GCN Correctness

PyG GCN uses a glorot initialization for weights and zero initialization for biases. So I modified Seastar to do the same using this. The results saw Seastar and PyG GCN having similar values. PyG does perform a but better but I think that could be due to some other model optimizations/initializations that we need not consider here.

[JoelMathewC]

TGCN correctness

Realized that we were making a very major blunder by assuming that norm in the GCN formula was relating to degree when it was actually relating to the sum of the edge weights of all inedges to a particular node. The reference for that can be found here.

The updated logs for seastar graph training are as follows

    Epoch ┃ Time (s) ┃ MSE    ┃
    ━━━━━━━╇━━━━━━━━━━╇
     0 │ 3.4157   │ 1.1095 
     1 │ 0.1413   │ 0.9653
     2 │ 0.1344   │ 0.8569 
     3 │ 0.1316   │ 0.7999
     4 │ 0.1314   │ 0.8388 
     5 │ 0.1284   │ 0.8242 
     6 │ 0.1312   │ 0.7837
     7 │ 0.1355   │ 0.7655
     8 │ 0.1313   │ 0.7651
     9 │ 0.1364   │ 0.763
    10 │ 0.1298   │ 0.7561 

While the training is not as smooth as that of PyG-T, it works sufficiently well. I believe the smoothness in PyG-T is again attributed to the implementation of PyG GCN and any optimization/initialization there, which we can ignore for now.

Hence that closes this discussion here.

[JoelMathewC]

Ran a second test

Noticed smoother results today

 Epoch ┃ Time (s) ┃ MSE    ┃ (mb)       ┃ (mb)       ┃ (mb)       ┃ Memory (mb) 
━━━━━━━╇━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━
     0 │ 5.204    │ 1.0186 │ 722.0      │ 713.4286   │ 1632.3945  │ 1632.3535   
     1 │ 0.1907   │ 0.8638 │ 730.0      │ 724.9524   │ 1628.9219  │ 1628.731    
     2 │ 0.2012   │ 0.8076 │ 730.0      │ 728.0476   │ 1628.7109  │ 1628.7109   
     3 │ 0.192    │ 0.8128 │ 730.0      │ 728.0476   │ 1628.7109  │ 1628.7109   
     4 │ 0.1916   │ 0.7992 │ 730.0      │ 724.9524   │ 1628.9297  │ 1628.7266   
     5 │ 0.1833   │ 0.7786 │ 730.0      │ 728.0476   │ 1629.1758  │ 1629.1758   
     6 │ 0.1381   │ 0.7693 │ 730.0      │ 728.0476   │ 1629.1758  │ 1629.1756   
     7 │ 0.1322   │ 0.7665 │ 730.0      │ 728.0476   │ 1629.1641  │ 1629.1641   
     8 │ 0.1295   │ 0.7611 │ 730.0      │ 728.0476   │ 1629.1641  │ 1629.1641   
     9 │ 0.1417   │ 0.7541 │ 730.0      │ 728.0476   │ 1629.1641  │ 1629.1641   
    10 │ 0.1291   │ 0.7502 │ 730.0      │ 726.3333   │ 1629.1641  │ 1629.1641