Implement one synchronous data-parallel SGD step over replicated parameters. Let R be the rank count, r index ranks, j index parameters, gr,j be a local gradient, θj be the shared original parameter, and η be learning_rate. A missing gradient contributes zero. The mean gradient gˉj is:
gˉj=R1r=0∑R−1gr,j.The updated parameter θj′ is:
θj′=θj−ηgˉj.Mean gradients are stored in each parameter's dtype before use in the float32 update. Return a Python dictionary with mean_gradients, a list of one PyTorch tensor per parameter, and updated_replicas, a rank-ordered list of parameter-tensor lists. Returned tensors preserve the corresponding parameter's shape, dtype, and device; updated replicas are independent copies.
Input: parameter_replicas = [[[1,2]],[[1,2]]], gradient_replicas = [[[2,4]],[[4,6]]], learning_rate = 0.1
Output: {"mean_gradients":[[3,5]],"updated_replicas":[[[0.7,1.5]],[[0.7,1.5]]]}
Explanation: The local gradients average to [3,5], and the same out-of-place SGD update is copied to both ranks.
Input: parameter_replicas = [[[1, -1]], [[1, -1]]], gradient_replicas = [[[2, 4]], [[0, 0]]], learning_rate = 0.5
Output: {"mean_gradients": [[1, 2]], "updated_replicas": [[[0.5, -2]], [[0.5, -2]]]}
Input: parameter_replicas and gradient_replicas each contain 3 ranks with tensor shapes (2,) and (2, 2), learning_rate = 0.25
Output: {"mean_gradients": [[1, 1], [[2, 2], [2, 2]]], "updated_replicas": 3 parameter lists, each containing tensors of shapes (2,) and (2, 2)}
torch.zeros_like can create an accumulator with a chosen dtype.
Tensor.to converts a tensor to the required storage dtype.
Tensor.clone creates independent tensor storage.
Accepts: array
Accepts: array
Accepts: number
Implement one synchronous data-parallel SGD step over replicated parameters. Let R be the rank count, r index ranks, j index parameters, gr,j be a local gradient, θj be the shared original parameter, and η be learning_rate. A missing gradient contributes zero. The mean gradient gˉj is:
gˉj=R1r=0∑R−1gr,j.The updated parameter θj′ is:
θj′=θj−ηgˉj.Mean gradients are stored in each parameter's dtype before use in the float32 update. Return a Python dictionary with mean_gradients, a list of one PyTorch tensor per parameter, and updated_replicas, a rank-ordered list of parameter-tensor lists. Returned tensors preserve the corresponding parameter's shape, dtype, and device; updated replicas are independent copies.
Input: parameter_replicas = [[[1,2]],[[1,2]]], gradient_replicas = [[[2,4]],[[4,6]]], learning_rate = 0.1
Output: {"mean_gradients":[[3,5]],"updated_replicas":[[[0.7,1.5]],[[0.7,1.5]]]}
Explanation: The local gradients average to [3,5], and the same out-of-place SGD update is copied to both ranks.
Input: parameter_replicas = [[[1, -1]], [[1, -1]]], gradient_replicas = [[[2, 4]], [[0, 0]]], learning_rate = 0.5
Output: {"mean_gradients": [[1, 2]], "updated_replicas": [[[0.5, -2]], [[0.5, -2]]]}
Input: parameter_replicas and gradient_replicas each contain 3 ranks with tensor shapes (2,) and (2, 2), learning_rate = 0.25
Output: {"mean_gradients": [[1, 1], [[2, 2], [2, 2]]], "updated_replicas": 3 parameter lists, each containing tensors of shapes (2,) and (2, 2)}
torch.zeros_like can create an accumulator with a chosen dtype.
Tensor.to converts a tensor to the required storage dtype.
Tensor.clone creates independent tensor storage.
Accepts: array
Accepts: array
Accepts: number
