Skip to content

GitLab

Unverified Commit 4a1aa97e authored by Thor Johnsen's avatar Thor Johnsen Committed by GitHub
Browse files

Merge pull request #845 from NVIDIA/distopt_bug_fix 

Bug fix
parents 8abb6908 3ccdfaa3
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 3 deletions
apex/contrib/optimizers/distributed_fused_adam.py
View file @ 4a1aa97e
...@@ -264,8 +264,10 @@ class DistributedFusedAdam(torch.optim.Optimizer): ...@@ -264,8 +264,10 @@ class DistributedFusedAdam(torch.optim.Optimizer):
grp = torch.distributed.new_group(ranks=ranks) grp = torch.distributed.new_group(ranks=ranks)
if torch.distributed.get_rank() in ranks: if torch.distributed.get_rank() in ranks:
self._rs_pg.append(grp) self._rs_pg.append(grp)
if self._compute_L2_grad_norm and torch.distributed.get_rank() in ranks: if self._compute_L2_grad_norm:
self._l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks) l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
if torch.distributed.get_rank() in ranks:
self._l2_grad_norm_pg = l2_grad_norm_pg
torch.distributed.all_reduce(self._overflow_buf,group=self._l2_grad_norm_pg) torch.distributed.all_reduce(self._overflow_buf,group=self._l2_grad_norm_pg)
self._rs_st = [torch.cuda.Stream() for _ in range(self._num_rs_pg)] self._rs_st = [torch.cuda.Stream() for _ in range(self._num_rs_pg)]
for rs_pg in self._rs_pg: for rs_pg in self._rs_pg:
......
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