[Distributed] turn off recording on embeddings in the inference. (#1861)

* turn on/off recording in sparse embedding.

* add test.

python/dgl/backend/backend.py

@@ -1474,6 +1474,11 @@ def is_no_grad(x):
     """
     pass
 
+def is_recording():
+    """ Test if the execution is recording gradients.
+    """
+    pass
+
 class record_grad(object):
     """Context manager that records the gradients"""
     def __init__(self):

python/dgl/backend/mxnet/tensor.py

@@ -605,6 +605,9 @@ def grad(x):
 def is_no_grad(x):
     return (x != 0).sum() == 0
 
+def is_recording():
+    return mx.autograd.is_recording()
+
 record_grad = mx.autograd.record
 
 class no_grad(object):

python/dgl/backend/pytorch/tensor.py

@@ -517,6 +517,9 @@ def grad(x):
 def is_no_grad(x):
     return x.grad is None or (x.grad == 0).all()
 
+def is_recording():
+    return th.is_grad_enabled()
+
 class record_grad(object):
     def __init__(self):
         pass

python/dgl/backend/tensorflow/tensor.py

@@ -685,6 +685,9 @@ def grad(x):
 def is_no_grad(x):
     return cgrad.is_no_grad(x)
 
+def is_recording():
+    raise NotImplementedError("Tensorflow doesn't support is_recording")
+
 no_grad = None
 
 initialize_context()

python/dgl/distributed/sparse_emb.py

@@ -47,8 +47,10 @@ class DistEmbedding:
 
     def __call__(self, idx):
         idx = utils.toindex(idx).tousertensor()
-        emb = F.attach_grad(self._tensor[idx])
-        self._trace.append((idx, emb))
+        emb = self._tensor[idx]
+        if F.is_recording():
+            emb = F.attach_grad(emb)
+            self._trace.append((idx, emb))
         return emb
 
 class SparseAdagradUDF:

tests/distributed/test_dist_graph_store.py

@@ -142,6 +142,10 @@ def check_dist_graph(g, num_nodes, num_edges):
         assert np.all(F.asnumpy(grad_sum[rest]) == np.zeros((len(rest), 1)))
 
         emb = DistEmbedding(g, g.number_of_nodes(), 1, 'emb2', emb_init)
+        with F.no_grad():
+            feats1 = emb(nids)
+        assert np.all(F.asnumpy(feats1) == 0)
+
         optimizer = SparseAdagrad([emb], lr=lr)
         with F.record_grad():
             feats1 = emb(nids)
@@ -151,7 +155,8 @@ def check_dist_graph(g, num_nodes, num_edges):
             loss = F.sum(feats + 1, 0)
         loss.backward()
         optimizer.step()
-        feats = emb(nids)
+        with F.no_grad():
+            feats = emb(nids)
         assert_almost_equal(F.asnumpy(feats), np.ones((len(nids), 1)) * math.sqrt(2) * -lr)
         rest = np.setdiff1d(np.arange(g.number_of_nodes()), F.asnumpy(nids))
         feats1 = emb(rest)