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