Skip to content

GitLab

Commit 143a8372 authored by Kai Chen's avatar Kai Chen
Browse files

bug fix for all_reduce when coalesce is False

parent df14830a
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
mmdet/core/utils/dist_utils.py
View file @ 143a8372
...@@ -47,11 +47,12 @@ def all_reduce_coalesced(tensors): ...@@ -47,11 +47,12 @@ def all_reduce_coalesced(tensors):
buckets[tp] = [] buckets[tp] = []
buckets[tp].append(tensor) buckets[tp].append(tensor)
world_size = dist.get_world_size()
for tp in buckets: for tp in buckets:
bucket = buckets[tp] bucket = buckets[tp]
coalesced = _flatten_dense_tensors(bucket) coalesced = _flatten_dense_tensors(bucket)
dist.all_reduce(coalesced) dist.all_reduce(coalesced)
coalesced /= dist.get_world_size() coalesced.div_(world_size)
for buf, synced in zip(bucket, for buf, synced in zip(bucket,
_unflatten_dense_tensors(coalesced, bucket)): _unflatten_dense_tensors(coalesced, bucket)):
...@@ -66,8 +67,9 @@ def reduce_grads(model, coalesce=True): ...@@ -66,8 +67,9 @@ def reduce_grads(model, coalesce=True):
if coalesce: if coalesce:
all_reduce_coalesced(grads) all_reduce_coalesced(grads)
else: else:
world_size = dist.get_world_size()
for tensor in grads: for tensor in grads:
dist.all_reduce(tensor) dist.all_reduce(tensor.div_(world_size))
class DistOptimizerHook(OptimizerHook): class DistOptimizerHook(OptimizerHook):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment