Repository: dmlc/dgl

Issue:
NCCL based SparseEmbedding fails for pytorch < 1.7, due to missing `delete_key` method in TCPStore
## 🐛 Bug

NCCL based SparseEmbedding fails for pytorch < 1.7, due to missing `delete_key` method in TCPStore:
```
AttributeError: ‘torch.distributed.TCPStore’ object has not attribute ‘delete_key’
```

## To Reproduce

Steps to reproduce the behavior:

1. Install DGL 0.7 and pytorch 1.6.
1. Run the `/examples/pytorch/rgcn/entity_classify_mp.py` example using a GPU and the `--dgl-sparse` flag.
1. This will result in:
```
AttributeError: ‘torch.distributed.TCPStore’ object has not attribute ‘delete_key’
```

## Environment

 - DGL Version (e.g., 1.0): 0.7-cu111 from pip
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): pytorch 1.6
 - OS (e.g., Linux): ubuntu
 - How you installed DGL (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source):
 - Python version: 3.6
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:


NCCL based SparseEmbedding fails for pytorch < 1.7, due to missing `delete_key` method in TCPStore
## 🐛 Bug

NCCL based SparseEmbedding fails for pytorch < 1.7, due to missing `delete_key` method in TCPStore:
```
AttributeError: ‘torch.distributed.TCPStore’ object has not attribute ‘delete_key’
```

## To Reproduce

Steps to reproduce the behavior:

1. Install DGL 0.7 and pytorch 1.6.
1. Run the `/examples/pytorch/rgcn/entity_classify_mp.py` example using a GPU and the `--dgl-sparse` flag.
1. This will result in:
```
AttributeError: ‘torch.distributed.TCPStore’ object has not attribute ‘delete_key’
```

## Environment

 - DGL Version (e.g., 1.0): 0.7-cu111 from pip
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): pytorch 1.6
 - OS (e.g., Linux): ubuntu
 - How you installed DGL (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source):
 - Python version: 3.6
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:




Discussion:

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
@@ -35,6 +35,7 @@ def __init__(self, params, lr):
         # otherwise it will crash the training
         self.shmem_buffer_holder = []
 
+        assert len(params) > 0, 'Empty parameters'
         # if we are using shared memory for communication
         for emb in params:
             assert isinstance(emb, NodeEmbedding), \
@@ -50,6 +51,7 @@ def __init__(self, params, lr):
                     'MultiGPU world_size for each embedding should be same.'
         assert not self._rank is None
         assert not self._world_size is None
+        self._nccl_root_id = 'SparseGradOptimizer.nccl_root_id'
 
     def step(self):
         ''' The step function.
@@ -111,17 +113,14 @@ def _comm_setup(self):
                     # root process broadcasts nccl id
                     nccl_id = nccl.UniqueId()
                     uid = str(nccl_id)
-                    store.set('nccl_root_id', uid)
+                    store.set(self._nccl_root_id, uid)
                 else:
-                    uid = store.get('nccl_root_id')
+                    uid = store.get(self._nccl_root_id)
                     nccl_id = nccl.UniqueId(uid)
                 # needs to be set for nccl to work
                 self._comm = nccl.Communicator(self._world_size,
                                                self._rank,
                                                nccl_id)
-                if self._rank == 0:
-                    # clear the store entry for future communicators
-                    store.delete_key('nccl_root_id')
                 th.distributed.barrier()
 
     def _shared_setup(self):


Test patch: