Repository: dmlc/dgl

Issue:
[Example][Bug] Random crash of a multi-gpu run in graphsage example folder
## 🐛 Bug

I currently encounter a random crash in one of graphsage multi-gpu runs. By "random", I mean, it crashed at random batch id (within the first epoch) or it can run through but with `Loss = nan` reported. 

This random crash(issue) is disappeared when `CUDA_LAUNCH_BLOCKING=1` is set, which may indicate there missing a device-synchronization somewhere in the script or code.

## To Reproduce

With merging the fix in https://github.com/dmlc/dgl/commit/c56e27a80748e3264cbc75a39095262ae6cb4989, here are two examples running with 8GPUs (V100). 

- Example 1 (in example/graphsage/advanced/)
`yes |  python3 train_sampling_unsupervised.py --dataset ogbn-products --num-epochs 1 --gpu 0,1,2,3,4,5,6,7`

A sample error msg and stack traces:
```
[0]Epoch 00000 | Step 00000 | Loss 11.6581 | Speed (samples/sec) nan|nan | Load nan| train nan | GPU 1985.0 MB
[0]Epoch 00000 | Step 00020 | Loss 3.4909 | Speed (samples/sec) 5815.4997|5815.4997 | Load 1.7095| train 0.0131 | GPU 2180.7 MB
[0]Epoch 00000 | Step 00040 | Loss 1.7599 | Speed (samples/sec) 5724.2924|5724.2924 | Load 1.7376| train 0.0129 | GPU 2180.8 MB
[0]Epoch 00000 | Step 00060 | Loss 1.5126 | Speed (samples/sec) 5773.5686|5773.5686 | Load 1.7229| train 0.0128 | GPU 2181.5 MB
[0]Epoch 00000 | Step 00080 | Loss nan | Speed (samples/sec) 5726.9199|5726.9199 | Load 1.7380| train 0.0128 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00100 | Loss nan | Speed (samples/sec) 5692.8335|5692.8335 | Load 1.7493| train 0.0128 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00120 | Loss nan | Speed (samples/sec) 5680.3568|5680.3568 | Load 1.7527| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00140 | Loss nan | Speed (samples/sec) 5669.1963|5669.1963 | Load 1.7557| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00160 | Loss nan | Speed (samples/sec) 5655.7341|5655.7341 | Load 1.7599| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00180 | Loss nan | Speed (samples/sec) 5631.4966|5631.4966 | Load 1.7676| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00200 | Loss nan | Speed (samples/sec) 5618.9977|5618.9977 | Load 1.7718| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00220 | Loss nan | Speed (samples/sec) 5607.6820|5607.6820 | Load 1.7757| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00240 | Loss nan | Speed (samples/sec) 5598.7065|5598.7065 | Load 1.7784| train 0.0129 | GPU 2183.2 MB
[0]Epoch 00000 | Step 00260 | Loss nan | Speed (samples/sec) 5588.3978|5588.3978 | Load 1.7817| train 0.0128 | GPU 2183.2 MB
[W CUDAGuardImpl.h:113] Warning: CUDA warning: an illegal memory access was encountered (function destroyEvent)
terminate called after throwing an instance of 'c10::CUDAError'
  what():  CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.

-- Process 3 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/workspace/examples/graphsage/advanced/train_sampling_unsupervised.py", line 135, in run
    for step, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(dataloader):
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 498, in __next__
    self._next_non_threaded() if not self.use_thread else self._next_threaded()
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 480, in _next_non_threaded
    batch, feats, stream_event = _prefetch(batch, self.dataloader, self.stream)
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 324, in _prefetch
    batch = recursive_apply(batch, lambda x: x.to(dataloader.device, non_blocking=True))
  File "/opt/conda/lib/python3.8/site-packages/dgl/utils/internal.py", line 996, in recursive_apply
    return [recursive_apply(v, fn, *args, **kwargs) for v in data]
  File "/opt/conda/lib/python3.8/site-packages/dgl/utils/internal.py", line 996, in <listcomp>
    return [recursive_apply(v, fn, *args, **kwargs) for v in data]
  File "/opt/conda/lib/python3.8/site-packages/dgl/utils/internal.py", line 998, in recursive_apply
    return fn(data, *args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 324, in <lambda>
    batch = recursive_apply(batch, lambda x: x.to(dataloader.device, non_blocking=True))
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
```
- Example 2 (in example/graphsage/advanced/)
`yes |   CUDA_LAUNCH_BLOCKING=1  python3 train_sampling_unsupervised.py --dataset ogbn-products --num-epochs 1 --gpu 0,1,2,3,4,5,6,7`

ran through with the output:
```
[0]Epoch 00000 | Step 00000 | Loss 16.5902 | Speed (samples/sec) nan|nan | Load nan| train nan | GPU 1984.8 MB
[0]Epoch 00000 | Step 00020 | Loss 3.8702 | Speed (samples/sec) 5421.0691|5421.0691 | Load 1.7165| train 0.1309 | GPU 2182.5 MB
[0]Epoch 00000 | Step 00040 | Loss 2.1864 | Speed (samples/sec) 5392.3610|5392.3610 | Load 1.6664| train 0.1910 | GPU 2182.5 MB
[0]Epoch 00000 | Step 00060 | Loss 1.4191 | Speed (samples/sec) 5399.0877|5399.0877 | Load 1.6672| train 0.1876 | GPU 2182.5 MB
[0]Epoch 00000 | Step 00080 | Loss 1.1865 | Speed (samples/sec) 5423.6634|5423.6634 | Load 1.6573| train 0.1892 | GPU 2182.5 MB
[0]Epoch 00000 | Step 00100 | Loss 0.9519 | Speed (samples/sec) 5487.2282|5487.2282 | Load 1.6413| train 0.1845 | GPU 2182.5 MB
...
[0]Epoch 00000 | Step 01500 | Loss 0.4924 | Speed (samples/sec) 5776.5354|5776.5354 | Load 1.6084| train 0.1252 | GPU 2184.5 MB
[0]Epoch 00000 | Step 01520 | Loss 0.4939 | Speed (samples/sec) 5777.3222|5777.3222 | Load 1.6084| train 0.1249 | GPU 2184.5 MB
[0]Epoch 00000 | Step 01540 | Loss 0.4853 | Speed (samples/sec) 5777.6396|5777.6396 | Load 1.6085| train 0.1247 | GPU 2184.5 MB
Epoch Time(s): 2691.0062
Avg epoch time: -0.0
```

## Environment

 - DGL Version (e.g., 1.0): 0.9 (build from source HEAD@ 9d9295b7fe337393b789c07ef2cd8eab2ed902d3)
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): torch 1.13
 - CUDA/cuDNN version (if applicable): 11.7
 - GPU models and configuration (e.g. V100):  8xV100




Discussion:
This should be caused by the incorrect usage of CUDA streams in the dataloader. I had realized the potential risk here but didn't observe bugs so I left it unfixed. https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html
https://github.com/dmlc/dgl/blob/c56e27a80748e3264cbc75a39095262ae6cb4989/python/dgl/dataloading/dataloader.py#L319-L326
Line 324 is added in #4177 but I believe Line 321-322 have similar issues too.

A temporary solution is setting `use_alternative_streams=False` for the dataloader.

Resolving patch:
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
--- a/python/dgl/dataloading/dataloader.py
+++ b/python/dgl/dataloading/dataloader.py
@@ -28,7 +28,6 @@
 from .. import backend as F
 from ..distributed import DistGraph
 from ..multiprocessing import call_once_and_share
-from ..cuda import stream as dgl_stream
 
 PYTORCH_VER = LooseVersion(torch.__version__)
 PYTHON_EXIT_STATUS = False
@@ -305,6 +304,18 @@ def _await_or_return(x):
     else:
         return x
 
+def _record_stream(x, stream):
+    if stream is None:
+        return x
+    if isinstance(x, torch.Tensor):
+        x.record_stream(stream)
+        return x
+    elif isinstance(x, _PrefetchedGraphFeatures):
+        node_feats = recursive_apply(x.node_feats, _record_stream, stream)
+        edge_feats = recursive_apply(x.edge_feats, _record_stream, stream)
+        return _PrefetchedGraphFeatures(node_feats, edge_feats)
+    else:
+        return x
 
 def _prefetch(batch, dataloader, stream):
     # feats has the same nested structure of batch, except that
@@ -316,12 +327,21 @@ def _prefetch(batch, dataloader, stream):
     #
     # Once the futures are fetched, this function waits for them to complete by
     # calling its wait() method.
-    with torch.cuda.stream(stream), dgl_stream(stream):
+    if stream is not None:
+        current_stream = torch.cuda.current_stream()
+        current_stream.wait_stream(stream)
+    else:
+        current_stream = None
+    with torch.cuda.stream(stream):
         # fetch node/edge features
         feats = recursive_apply(batch, _prefetch_for, dataloader)
         feats = recursive_apply(feats, _await_or_return)
-        # transfer input nodes/seed nodes/sampled subgraph
+        feats = recursive_apply(feats, _record_stream, current_stream)
+        # transfer input nodes/seed nodes
+        # TODO(Xin): sampled subgraph is transferred in the default stream
+        # because heterograph doesn't support .record_stream() for now
         batch = recursive_apply(batch, lambda x: x.to(dataloader.device, non_blocking=True))
+        batch = recursive_apply(batch, _record_stream, current_stream)
     stream_event = stream.record_event() if stream is not None else None
     return batch, feats, stream_event
 


Test patch: