Repository: dmlc/dgl

Issue:
OOM issue with dgl.unbatch on GPU
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

The GPU memory keeps increasing while conducting `dgl.unbatch` on batched graphs on GPU and copying splitted graphs to CPU. 

## To Reproduce

Run the following script, and we will notice that the allocated GPU memory keeps increasing. If a much larger graph dataset is used, this will evently cause OOM crash.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```python
import dgl
import torch
from dgl.dataloading import GraphDataLoader


dataset = dgl.data.QM9EdgeDataset()
dataloader = GraphDataLoader(dataset, batch_size=64)

graph_list = []

for batch_graph, _ in dataloader:
    batch_graph = batch_graph.to('cuda:0')
    split_graphs = dgl.unbatch(batch_graph)

    graph_list.extend([graph.cpu() for graph in split_graphs])

    print(f'memory allocation: {torch.cuda.memory_allocated() / 1024**2} Mb')

    torch.cuda.empty_cache()
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

How could the GPU memory increase be prevented?

## Environment

 - DGL Version: 1.1.2+cu117
 - Backend Library & Version: Pytorch 1.13.1
 - OS (e.g., Linux): Linux
 - How you installed DGL (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source):
 - Python version: 3.9
 - CUDA/cuDNN version (if applicable): 11.7
 - GPU models and configuration (e.g. V100): V100
 - Any other relevant information: reproduce the same phenomenon on a Windows machine (dgl 1.1.2+cu116 and python3.10)

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi @wondey-sh, have you estimated the expected memory usage?
> Hi @wondey-sh, have you estimated the expected memory usage?

Hi @czkkkkkk , the example code is used in part of my model inference code, and I expect that it can be run in a g4dn.xlarge machine. However, OOM happens even when `batch_size` is 1. The graphs in my dataset is much larger than the example code, but are perfectly fine for training with `batch_size=16`. Thanks.
I see. I reproduced the issue. We will dive deep into it.

Resolving patch:
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
--- a/python/dgl/heterograph.py
+++ b/python/dgl/heterograph.py
@@ -5661,6 +5661,9 @@ def to(self, device, **kwargs):  # pylint: disable=invalid-name
         If the graph is already on the specified device, the function directly returns it.
         Otherwise, it returns a cloned graph on the specified device.
 
+        Note that data of node and edge features are not moved to the specified
+        device before being accessed or `materialize_data()` is called.
+
         Parameters
         ----------
         device : Framework-specific device context object
@@ -5752,6 +5755,21 @@ def cpu(self):
         """
         return self.to(F.cpu())
 
+    def materialize_data(self):
+        """Materialize the graph data on the current device.
+
+        This method is a no-op if the graph data is already materialized.
+
+        Returns
+        -------
+        DGLGraph
+            The graph on the current device.
+        """
+        for frame in itertools.chain(self._node_frames, self._edge_frames):
+            for col in frame._columns.values():
+                col.data  # pylint: disable=pointless-statement
+        return self
+
     def pin_memory_(self):
         """Pin the graph structure and node/edge data to the page-locked memory for
         GPU zero-copy access.


Test patch: