I come up with a simple solution:
def weights_sharing(weights, weights_grad, group_info):
    #share weights
    group_len = group_info.size(0)
    weights_size = weights.size(1)
    weights_collection = weights.gather(0,group_info.expand(weights_size,group_len).transpose(0,1))
    averge_weight = weights_collection.mean(dim=0)
    for i in group_info.numpy():
        weights[i] = averge_weight

    #share gradient
    grad_collection = weights_grad.gather(0, group_info.expand(weights_size, group_len).transpose(0,1))
    averge_grad = grad_collection.mean(dim=0)
    for i in group_info.numpy():
        weights_grad[i] = averge_grad

sample usage:
linear = nn.Linear(10,5)
x=torch.randn(5,10)
y=torch.randn(5,5)
y_c = linear(x)
loss = (y_c-y).pow(2).mean()
loss.backward()
weight_sharing(linear.weight.data, linear.weight.grad.data, group_info)
#then update the parameter as usual

I still have a question, since the weight sharing happens after backward pass, I think it won’t affect computation graph. I think this is true, but not sure about it.
Any discussion is welcome.