Solved by mrshenli in post #2 
              
              
                Hey @Kunlin_Yang, say, param_rref is the RRef of the corresponding parameter, you can do sth like the following: 
import torch.distributed.autograd as dist_autograd
...

def update_grad(param_rref, cid):
    grad = dist_autograd.get_gradients(cid)[param_rref.local_value()]
    # do inplace update on…