RGCN generates nan in PyTorch 1.8 but not in PyTorch 1.7.x

[nv-dlasalle]

🐛 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 ( 09ef2c2)
• 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:

[nv-dlasalle]

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.

[classicsong]

I tested it with cuda 10.1. The training progress looks good.
Will check cuda 11.1 later.

[BarclayII]

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.

[nv-dlasalle]

I think i've tracked this down to the sparse_optim.py copying data from the GPU to the CPU with non_blocking=True, and then running CPU operations before the copy finishes.

[classicsong]

I think i've tracked this down to the sparse_optim.py copying data from the GPU to the CPU with non_blocking=True, and then running CPU operations before the copy finishes.

It seems PyTorch does not check the data access order. Is this a breaking change in PyTorch 1.8?