tf.distribute.NcclAllReduce
Reduction using NCCL all-reduce.
tf.distribute.NcclAllReduce(
num_packs=1
)
| Args |
|---|
num_packs | values will be packed in this many splits. num_packs should be greater than or equals 0. When it is zero, no packing will be done. |
| Raises |
|---|
ValueError if num_packs is negative. |
Methods
batch_reduce
View source
batch_reduce(
reduce_op, value_destination_pairs, experimental_hints=None
)
Reduce PerReplica objects in a batch.
Reduce each first element in value_destination_pairs to each second element which indicates the destinations.
This can be faster than multiple individual reduces because we can fuse several tensors into one or multiple packs before reduction.
| Args |
|---|
reduce_op | An instance of tf.distribute.ReduceOp that indicates how the per_replica_value will be reduced. |
value_destination_pairs | A list or a tuple of PerReplica objects (or tensors with device set if there is one device) and destinations. |
experimental_hints | A tf.distrbute.experimental.CollectiveHints. Hints to perform collective operations. |
| Returns |
|---|
| a list of Mirrored objects. |
| Raises |
|---|
ValueError | if value_destination_pairs is not an iterable of tuples of PerReplica objects and destinations. |
broadcast
View source
broadcast(
tensor, destinations
)
Broadcast the tensor to destinations.
| Args |
|---|
tensor | the tensor to broadcast. |
destinations | the broadcast destinations. |
| Returns |
|---|
| a Mirrored object. |
reduce
View source
reduce(
reduce_op, per_replica_value, destinations, experimental_hints=None
)
Reduce per_replica_value to destinations.
It runs the reduction operation defined by reduce_op and put the result on destinations.
| Args |
|---|
reduce_op | An instance of tf.distribute.ReduceOp that indicates how per_replica_value will be reduced. |
per_replica_value | A tf.distribute.DistributedValues object or a tensor with device set. |
destinations | the reduction destinations. |
experimental_hints | A tf.distrbute.experimental.CollectiveHints. Hints to perform collective operations. |
| Returns |
|---|
| a Mirrored object. |
| Raises |
|---|
ValueError | if per_replica_value can't be converted to a PerReplica object or if destinations aren't strings, Variables or DistributedValues |