gistfile1.txt

addmm(bias, input, weight)

bias: replicated

output = input * weight

input shard(1), weight, shard(0) -> partial tensor
output -> partial -> replicated?

output + bias -> partial? only do on one rank
bias.grad = output -> partial
output.grad = bias -> replicated

forward of dist op:
disttensor + dist tensor = out_dist_tensor
out_dist_tensor.grad_fn = addbackward -> 

dist_tensor0.mm(dist_tensor_1) = out_dist_tensor
out_dist_tensor.grad_fn = mmbackward -> mm
dist_tensor0.grad -> mm(output_grad, dist_tensor_1)

addbackward -> 

bias = 1