| 7 | |
| 8 | |
| 9 | def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1): |
| 10 | if bucket_size_mb > 0: |
| 11 | bucket_size_bytes = bucket_size_mb * 1024 * 1024 |
| 12 | buckets = _take_tensors(tensors, bucket_size_bytes) |
| 13 | else: |
| 14 | buckets = OrderedDict() |
| 15 | for tensor in tensors: |
| 16 | tp = tensor.type() |
| 17 | if tp not in buckets: |
| 18 | buckets[tp] = [] |
| 19 | buckets[tp].append(tensor) |
| 20 | buckets = buckets.values() |
| 21 | |
| 22 | for bucket in buckets: |
| 23 | flat_tensors = _flatten_dense_tensors(bucket) |
| 24 | dist.all_reduce(flat_tensors) |
| 25 | flat_tensors.div_(world_size) |
| 26 | for tensor, synced in zip( |
| 27 | bucket, _unflatten_dense_tensors(flat_tensors, bucket)): |
| 28 | tensor.copy_(synced) |
| 29 | |
| 30 | |
| 31 | def allreduce_grads(params, coalesce=True, bucket_size_mb=-1): |