🫙 Improving Storage

[JoelMathewC]

Naive

The state stack used to push all resulting tensors from the current timestamp onto the state stack. That gave the following benchmark

Epoch	Time (s)	MSE	Max GPU Memory(mb)
0	1.8506	1.0943	2716.875
1	0.165	1.0462	2772.875
2	0.1377	1.0386	2772.875
3	0.1557	1.0366	2772.875

Modifying the code to only store those tensors that will be used by the backend resulted in an improvement.

Epoch	Time (s)	MSE	Max GPU Memory(mb)
0	1.8506	1.0943	2500.875
1	0.165	1.0462	2516.875
2	0.1377	1.0386	2516.875
3	0.1557	1.0366	2516.875

Further modifying the code to not account for the CUDA context object of NaiveGraph

Epoch	Time (s)	MSE	Max GPU Memory(mb)
0	1.8506	1.0943	2394.0
1	0.165	1.0462	2410.0
2	0.1377	1.0386	2410.0
3	0.1557	1.0366	2410.0

The modification was made here

In ExeState.reset
-----------------
for bu in bunits:
  for arg in bu.kernel_args():
     if arg.id in input_map.keys():
         self.bwd_common_tensor_list.append(arg.id)

In Executor.forward_cb
-----------------------
self.ts.tensor_map_stack.push({key: self.ts.current_tensor_map[key] for key in self.ts.bwd_common_tensor_list})

[JoelMathewC]

Trying to log the stored tensors

I printed the storage log of three timestamps. I noticed that most of the tensors reported repeatedly have the shape of the node_feat matrix passed as input to the Seastar GCN kernel.

The following tensor is repeated multiple times.

<class 'torch.Tensor'> torch.Size([55000, 48]) cuda:0

[JoelMathewC]

Looking at the evidence again

After taking a look at the logs again I noticed that tensors of size (55000,48) are duplicated in the first timestamp, but for every timestamp following that the tensors logged are only those associated with the actual temporal layers which is unavoidable.

I made a val_registry storing all the val objects and noticed that this did not increment the memory at all. This implies one of the following

1. The val references for some reason were never deleted
2. The tensors associated with the val reference was never deleted (maybe something else is still referencing it)

[JoelMathewC]

A fix

By modifying the Val._t to store a CPU based detached clone of the tensor we were able to bring down the memory by a considerable margin.

Epoch	Time (s)	MSE	Max GPU Memory(mb)
0	1.8506	1.0943	2102.0
1	0.165	1.0462	2122.0
2	0.1377	1.0386	2122.0
3	0.1557	1.0366	2122.0

Analysis

So it seems that when we do not use a cloned detached version pytorch maintains some intermediate tensors in the backend that are causing the memory overhead. We came to this conclusion by using a GPU based detached clone of the tensor and noticed that there was only small increase in memory to 2196.0 the rest of it must have been from the intermediate results stored by pytorch. A possible reference to a case where something like this happens is there here.

Additionally, it may be noted that we will continue to proceed with CPU based detached clone since that will provide for extra memory on GPU for GPU based computations.

Note: This can be mentioned as a modification in the paper

UPDATE: As per #62 this approach will affect the correctness of the program.

[JoelMathewC]

Understanding the storage logs

After building an isolated version of the tracing system as mentioned in #62, the following observations were made for a GCN.

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

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

It is my belief that the CentralNode attributes and the val objects are in some sort of reference cycle preventing their deallocation even after they run out of scope. Using a _destroy_central_node as follows gives the output

def _destroy_central_node(self, cen, nfeats, efeats):
        if nfeats:
            for k, _ in nfeats.items():
                delattr(cen, k)
                for n in cen.innbs:
                    delattr(n, k)
        if efeats:
            for k, _ in efeats.items():
                for e in cen.inedges:
                    delattr(e, k)

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

Using _destroy_central_node allows us to get rid of a few intermediate outputs that were stored unnecessarily.

[JoelMathewC]

Got cloned CPU tensors to work

Made the following changes

class Val(abc.ABC):
    @abc.abstractmethod
    def __init__(self, tensor, id, fprog):
        self._t = tensor.clone().detach().cpu().requires_grad_(tensor.requires_grad)
    ...

class TorchVal(Val):
    def __init__(self, backend, tensor, val_type, id, fprog, reduce_dim):
        ...
        ...
        if reduce_dim:
            self._v = self._t.clone().detach().requires_grad_(False).mean(dim=0)
        else:
            self._v = self._t.clone().detach().requires_grad_(False)

This allows tensors who have requires_grad=True to have a cloned CPU tensor with those same properties. This generated the correct output. Need to verify the storage benefits. However, it must be noted that those replicated tensors from from the gc.get_objects() output is still there.

[JoelMathewC]

An unexpected output

With our modified code, when I run the code there seems to be no benefit from using CPU tensors, which is quite odd since we did see the benefit a while back.

[JoelMathewC]

Storage in GCN

The logs of the GCN storage is as shown below.

PyG GCN
----------
Epoch 00000 | Time(s) 0.5773 | train_acc 0.128508 | Used_Memory 24.181152 mb
Epoch 00001 | Time(s) 0.0024 | train_acc 0.323978 | Used_Memory 32.647461 mb
Epoch 00002 | Time(s) 0.0020 | train_acc 0.347612 | Used_Memory 32.647461 mb

Seastar GCN
-------------
Epoch 00000 | Time(s) 0.8960 | train_acc 0.086164 | Used_Memory 23.895996 mb
Epoch 00001 | Time(s) 0.0030 | train_acc 0.308715 | Used_Memory 32.527832 mb
Epoch 00002 | Time(s) 0.0024 | train_acc 0.307730 | Used_Memory 32.527832 mb

It is my belief that if we can remove the intermediate outputs generated by Seastar we should be able to have a memory improvement of close to 1 MB. Which will scale to larger values for larger graphs.

[JoelMathewC]

Understanding the repetitions

There are a series of repeated tensors for a single GCN layer when using Seastar. They are repeated 4 times, these can be interpreted as

1. Original Tensor
2. Output Tensor
3. Intermediate tensor generated during FuncWrapper.apply()
4. Intermediate Tensor generated after output is passed to rest of the Seastar code

[JoelMathewC]

pyNVML acting oddly

It may be noted that pyNVML does act quite oddly and seems to account for some unnecessary memory when making measurements. We need to look into this more clearly when doing benchmarking.

[JoelMathewC]

An example of odd behaviour is as shown. This is probably because pytorch takes control over a certain portion of the memory that isnt being used but is probably cached by the allocator to increase speed. This begs to ask the questions how do we efficiently measure storage??

# UNCOMMENT THIS 
# if torch.cuda.is_available():
#   print("🎉 CUDA is available")

# (OR) THIS AND RUN
# torch.cuda.empty_cache()

nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0)
initial = nvmlDeviceGetMemoryInfo(handle).used
print("Initial Mem: ", initial/(1024**2))

# THEN TRY UNCOMMENTING THIS
# if torch.cuda.is_available():
#   print("🎉 CUDA is available")

# (OR) THIS AND RUNNING
# torch.cuda.empty_cache()

a = torch.randn([55000,48]).to(torch.float32).cuda()
print("Expected Size: ", (a.nelement() * a.element_size())/(1024**2), " MB")
print("Measuerd Size (torch): ", torch.cuda.memory_allocated(0)/(1024**2), " MB")
print("Measuerd Size (pynvml): ", (nvmlDeviceGetMemoryInfo(handle).used - initial)/(1024**2), " MB")
cuda.cuMemAlloc(10560000)
print("Expected Size: ", (2*10560000)/(1024**2), " MB")
print("Measuerd Size (torch): ", torch.cuda.memory_allocated(0)/(1024**2), " MB")
print("Measuerd Size (pynvml): ", (nvmlDeviceGetMemoryInfo(handle).used - initial)/(1024**2), " MB")

[JoelMathewC]

Fixed It

Adding the following line before measuring the initial memory should account for the pytorch CUDA context object.

a = torch.tensor([2]).cuda()

[JoelMathewC]

Best way to measure space occupied in seastar

We will need separate scripts for speed and memory analysis since the pynvml functions seem to take time in the order of 0.001 and that could affect our performance.

# To account for the initial CUDA Context object for pynvml
tmp = StaticGraph([(0,0)], [1], 1)

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
g = StaticGraph(cora.get_edges(), edge_weight, features.shape[0])
graph_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem


# In train
now_mem = torch.cuda.max_memory_allocated(0) + graph_mem

[JoelMathewC]

Did Pytorch update increase the memory usage of Seastar

I wonder if updating to Pytorch 2.0 caused an increase in memory. This is because old logs of Seastar seem to show that it takes 16MB to run a GCN but now it takes 32MB.

This reference will come in handy since it deals with the part where Pytorch interfaces with Seastar.

[JoelMathewC]

Some Progress

Didnt seem to find anything solid for now. But was able to get 200 MB lesser memory by setting the requires grad in out_tensors generated in executor to false.

[JoelMathewC]

This was resolved.

[JoelMathewC]

Closing this out.