Repository: dmlc/dgl

Issue:
RGCN generates nan in PyTorch 1.8 but not in PyTorch 1.7.x
## 🐛 Bug

Running `/examples/pytorch/rgcn/entity_classify_mp.py` results in the training accuracy staying near zero and the loss becomes `nan`.

## To Reproduce

Steps to reproduce the behavior:

1. Build from source, and run the example according to the readme:
```
$ OMP_NUM_THREADS=1 python3 entity_classify_mp.py -d ogbn-mag --testing --fanout='30,30' --batch-size 1024 --n-hidden 128 --lr 0.01 --num-worker 1 --eval-batch-size 8 --low-mem --gpu 0 --dropout 0.7 --use-self-loop --n-bases 2 --n-epochs 3 --dgl-sparse --sparse-lr 0.08
Using backend: pytorch
Namespace(batch_size=1024, dataset='ogbn-mag', dgl_sparse=True, dropout=0.7, eval_batch_size=8, fanout='30,30', gpu='0', l2norm=0, layer_norm=False, low_mem=True, lr=0.01, n_bases=2, n_epochs=3, n_hidden=128, n_layers=2, node_feats=False, num_workers=1, sparse_lr=0.08, use_self_loop=True, validation=False)
num_tasks = 1
task_type = multiclass classification
eval_metric = acc
num_classes = 349
is_hetero = True
Opening 'dataset/ogbn_mag_dgl/processed/dgl_data_processed'.
Number of relations: 8
Number of class: 349
Number of train: 629571
Number of valid: 64879
Number of test: 41939
0:author
1:field_of_study
2:institution
3:paper
start training...
Train Accuracy: 0.0117 | Train Loss: 5.9985
Train Accuracy: 0.0107 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0010 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0010 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0078 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0049 | Train Loss: nan
Train Accuracy: 0.0000 | Train Loss: nan
Train Accuracy: 0.0029 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0020 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
Train Accuracy: 0.0049 | Train Loss: nan
Train Accuracy: 0.0039 | Train Loss: nan
```

## Expected behavior
However, when pytorch's implementation is used:
```
$ OMP_NUM_THREADS=1 python3 entity_classify_mp.py -d ogbn-mag --testing --fanout='30,30' --batch-size 1024 --n-hidden 128 --lr 0.01 --num-worker 1 --eval-batch-size 8 --low-mem --gpu 0 --dropout 0.7 --use-self-loop --n-bases 2 --n-epochs 3 --sparse-lr 0.08
Using backend: pytorch
Namespace(batch_size=1024, dataset='ogbn-mag', dgl_sparse=False, dropout=0.7, eval_batch_size=8, fanout='30,30', gpu='0', l2norm=0, layer_norm=False, low_mem=True, lr=0.01, n_bases=2, n_epochs=3, n_hidden=128, n_layers=2, node_feats=False, num_workers=1, sparse_lr=0.08, use_self_loop=True, validation=False)
num_tasks = 1
task_type = multiclass classification
eval_metric = acc
num_classes = 349
is_hetero = True
Opening 'dataset/ogbn_mag_dgl/processed/dgl_data_processed'.
Number of relations: 8
Number of class: 349
Number of train: 629571
Number of valid: 64879
Number of test: 41939
0:author
1:field_of_study
2:institution
3:paper
start training...
Train Accuracy: 0.0107 | Train Loss: 5.9672
Train Accuracy: 0.0166 | Train Loss: 5.7614
Train Accuracy: 0.0215 | Train Loss: 5.5110
Train Accuracy: 0.0332 | Train Loss: 5.4334
Train Accuracy: 0.0361 | Train Loss: 5.2588
Train Accuracy: 0.0469 | Train Loss: 5.1994
Train Accuracy: 0.0557 | Train Loss: 5.1170
Train Accuracy: 0.0518 | Train Loss: 5.0814
Train Accuracy: 0.0674 | Train Loss: 5.0409
Train Accuracy: 0.0791 | Train Loss: 4.9510
Train Accuracy: 0.0615 | Train Loss: 5.0149
Train Accuracy: 0.0693 | Train Loss: 5.0005
Train Accuracy: 0.0801 | Train Loss: 4.7365
Train Accuracy: 0.0928 | Train Loss: 4.8522
Train Accuracy: 0.0840 | Train Loss: 4.8094
Train Accuracy: 0.1055 | Train Loss: 4.7392
Train Accuracy: 0.0967 | Train Loss: 4.8593
Train Accuracy: 0.0830 | Train Loss: 4.8212
Train Accuracy: 0.0957 | Train Loss: 4.6933
Train Accuracy: 0.1094 | Train Loss: 4.6382
Train Accuracy: 0.1055 | Train Loss: 4.6573
Train Accuracy: 0.1016 | Train Loss: 4.6128
Train Accuracy: 0.1211 | Train Loss: 4.6512
```

## Environment

 - DGL Version (e.g., 1.0): master ( 09ef2c2c3d73b06941acb2e4ac73000c1d524aef)
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): pytorch 1.8
 - OS (e.g., Linux): ubuntu 18.04
 - How you installed DGL (`conda`, `pip`, source): source
 - Build command you used (if compiling from source): cmake .. -DUSE_CUDA=ON -DBUILD_TORCH=ON -DBUILD_CPP_TEST=ON -DCMAKE_BUILD_TYPE=Debug -DCMAKE_VERBOSE_MAKEFILE=1 && make -j13
 - Python version: 3.6
 - CUDA/cuDNN version (if applicable): 11.1
 - GPU models and configuration (e.g. V100): TitanV
 - Any other relevant information:





Discussion:
More information, is that it does not seem to occur when running on the CPU, and I've only been able to reproduce it on the ogbn-mag dataset.
I tested it with cuda 10.1. The training progress looks good.
Will check cuda 11.1 later.
Me and @zheng-da observed the same phenomenon.  Using PyTorch 1.8 will yield nan while PyTorch 1.7+ does not have this problem.

In my case this is not related to SparseEmbedding and OGBN-MAG.  NaNs occur even when node features are available.

Resolving patch:
diff --git a/python/dgl/optim/pytorch/sparse_optim.py b/python/dgl/optim/pytorch/sparse_optim.py
--- a/python/dgl/optim/pytorch/sparse_optim.py
+++ b/python/dgl/optim/pytorch/sparse_optim.py
@@ -598,6 +598,12 @@ def update(self, idx, grad, emb):
             exec_dev = grad.device
             state_dev = state_step.device
 
+            # only perform async copies cpu -> gpu, or gpu-> gpu, but block
+            # when copying to the cpu, so as to ensure the copy is finished
+            # before operating on the data on the cpu
+            state_nonblock = state_dev != th.device('cpu')
+            exec_nonblock = exec_dev != th.device('cpu')
+
             # There can be duplicated indices due to sampling.
             # Thus unique them here and average the gradient here.
             grad_indices, inverse, cnt = th.unique(idx,
@@ -605,9 +611,9 @@ def update(self, idx, grad, emb):
                                                    return_counts=True)
             state_idx = grad_indices.to(state_dev)
             state_step[state_idx] += 1
-            state_step = state_step[state_idx].to(exec_dev, non_blocking=True)
-            orig_mem = state_mem[state_idx].to(exec_dev, non_blocking=True)
-            orig_power = state_power[state_idx].to(exec_dev, non_blocking=True)
+            state_step = state_step[state_idx].to(exec_dev, non_blocking=exec_nonblock)
+            orig_mem = state_mem[state_idx].to(exec_dev, non_blocking=exec_nonblock)
+            orig_power = state_power[state_idx].to(exec_dev, non_blocking=exec_nonblock)
 
             grad_values = th.zeros((grad_indices.shape[0], grad.shape[1]), device=exec_dev)
             grad_values.index_add_(0, inverse, grad)
@@ -617,8 +623,10 @@ def update(self, idx, grad, emb):
             grad_power = grad_values * grad_values
             update_mem = beta1 * orig_mem + (1.-beta1) * grad_mem
             update_power = beta2 * orig_power + (1.-beta2) * grad_power
-            state_mem[state_idx] = update_mem.to(state_dev, non_blocking=True)
-            state_power[state_idx] = update_power.to(state_dev, non_blocking=True)
+            state_mem[state_idx] = update_mem.to(state_dev,
+                                                 non_blocking=state_nonblock)
+            state_power[state_idx] = update_power.to(state_dev,
+                                                     non_blocking=state_nonblock)
 
             update_mem_corr = update_mem / (1. - th.pow(th.tensor(beta1, device=exec_dev),
                                                         state_step)).unsqueeze(1)


Test patch: