Commit 561a3c43 authored by Ayush Dubey's avatar Ayush Dubey Committed by TensorFlower Gardener
Browse files

Dynamic subdivisions in collective ring reduce.

Before this change, a CollectiveOp user was required to specify subdiv_offsets
for the RingReduce algorithm.  During ring reduction, we created chunks of the
tensor to exchange between devices.  If the chunks were too large, or if the
hardware supported multiple data exchanges in parallel, the user could further
subdivide the chunk by specifying more than 1 subdiv offset.  Each subdiv
offset corresponded to another subdivision of the chunk, so effectively the
total number of tensor chunks is number of devices * number of subdivs.

After this change, we can dynamically infer the number of subdivisions based on
a target chunk size.  In ring_reducer.cc, we start with 1 subdiv, and keep
increasing until chunk size is less than MAX_CHUNK_SIZE.  Currently,
MAX_CHUNK_SIZE is set at 4 MB, although it may make sense to change this based
on specific hardware.

As a part of this change, a user can now provide an empty subdiv_offset list.
If empty, we dynamically add subdivisions based on the above algorithm.  If
non-empty, we take the user-specified subdivions.

PiperOrigin-RevId: 214815959
parent 9a0a768d
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment