Open zhao1157 opened 4 years ago
grad = tf.add_n(grads)
is a sum
without an all-reduce. An all-reduce is simply an add_n
except you get the output tensor on every device instead of just one device.
Note the reason parameter server mode takes a mean of the gradients instead of a sum is that use_mean=True
is passed here. Passing use_mean=True
causes the gradients to be divided by the number of replicas after summing them.
@reedwm Thanks for getting back to me about tf.add_n()
. The reason I asked this question is that I heard nccl all-reduce
consumes far less memory than tf.add_n()
does, so I wondered if there is a function I could use to achieve the functionality of tf.add_n()
with less memory.
On a single device, add_n
uses the least possible amount memory, as it only allocates its output tensor. The reason add_n
may use more memory than nccl all-reduce is that if add_n
is used to implement an all-reduce, all tensors must be copied to a single device, so that device must allocate all the tensors. I haven't checked, but it's likely nccl all-reduce does not require a single device to allocate enough space for all the tensors.
However, with parameter server mode, add_n
isn't used to implement an all-reduce, so the issue doesn't occur. There is no way to replace add_n
with some other reduction to make it take less memory. All the per-GPU tensors are transferred to the parameter server device(s), so the parameter server device(s) must have enough memory to hold all the per-GPU tensors.
@reedwm Thanks for your very helpful clarification.
I haven't checked, but it's likely nccl all-reduce does not require a single device to allocate enough space for all the tensors.
I know ring all-reduce
does the sum sequentially on each device instead of having all tensors transferred to one single device, thus requiring less memory. Do you know the difference between nccl all-reduce
and ring all-reduce
?
https://github.com/tensorflow/benchmarks/blob/5d03cf8e356d2ae17df440cdb612c378cbacf5ef/scripts/tf_cnn_benchmarks/variable_mgr_util.py#L575
@reedwm In parameter server mode, I managed to replace
grad = tf.add_n(grads)
bynccl all reduce
:I tried to figure out a way to accomplish the
sum
without usingall-reduce
since I only need one copy of the sum, notnumber-of-gpus
copies of sum. Intf
, is there areduce
API I can use? Thanks.