Skip to content

adasum integration #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
klipto wants to merge 1 commit into
base: master
Choose a base branch
from
Open

adasum integration #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
359 changes: 358 additions & 1 deletion PyTorch/LanguageModeling/BERT/optimization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import math
import torch
import horovod.torch as hvd
from torch.optim import Optimizer
from torch.optim.optimizer import required
from torch.nn.utils import clip_grad_norm_
Expand Down Expand Up @@ -255,7 +256,7 @@ def step(self, closure=None):
# if self.max_steps != -1:
schedule_fct = SCHEDULES[self.schedule]
lr_scheduled = self.learning_rate * schedule_fct(self.step_count / self.max_steps, self.warmup)
if torch.distributed.get_rank() == 0:
if hvd.rank() == 0 :#torch.distributed.get_rank() == 0:
print("Step {} LR {}".format(self.step_count, lr_scheduled))
# else:
# lr_scheduled = self.learning_rate
Expand Down Expand Up @@ -393,3 +394,359 @@ def step(self, closure=None):

return loss

from contextlib import contextmanager

def find_duplicates(lst):
seen = set()
dups = set()
for el in lst:
if el in seen:
dups.add(el)
seen.add(el)
return dups

class DistributedOptimizer(torch.optim.Optimizer):
def __init__(self, optimizer, compression,
backward_passes_per_step=1):
params = [p
for param_group in optimizer.param_groups
for p in param_group['params']]
super(self.__class__, self).__init__(params, dict(DistributedOptimizer.__dict__))
self.optimizer = optimizer
self._compression = compression

named_parameters = [('allreduce.noname.%s' % i, v)
for param_group in self.param_groups
for i, v in enumerate(param_group['params'])]

# make sure that named_parameters are tuples
if any([not isinstance(p, tuple) for p in named_parameters]):
raise ValueError('named_parameters should be a sequence of '
'tuples (name, parameter), usually produced by '
'model.named_parameters().')

dups = find_duplicates([k for k, _ in named_parameters])
if len(dups) > 0:
raise ValueError('Parameter names in named_parameters must be unique. '
'Found duplicates: %s' % ', '.join(dups))

all_param_ids = {id(v)
for param_group in self.param_groups
for v in param_group['params']}
named_param_ids = {id(v) for k, v in named_parameters}
unnamed_param_ids = all_param_ids - named_param_ids
if len(unnamed_param_ids):
raise ValueError('named_parameters was specified, but one or more model '
'parameters were not named. Python object ids: '
'%s' % ', '.join(str(id) for id in unnamed_param_ids))

self._parameter_names = {v: k for k, v in sorted(named_parameters)}
self.backward_passes_per_step = backward_passes_per_step
self._allreduce_delay = {v: self.backward_passes_per_step
for _, v in sorted(named_parameters)}
self._handles = {}
self._grad_accs = []
self._requires_update = set()
self._synchronized = False
self._should_synchronize = True
if hvd.size() > 1:
self._register_hooks()

def set_backward_passes_per_step(self, passes):
self.backward_passes_per_step = passes
for p in self._allreduce_delay:
self._allreduce_delay[p] = self.backward_passes_per_step

def _register_hooks(self):
for param_group in self.param_groups:
for p in param_group['params']:
if p.requires_grad:
p.grad = p.data.new(p.size()).zero_()
self._requires_update.add(p)
p_tmp = p.expand_as(p)
grad_acc = p_tmp.grad_fn.next_functions[0][0]
grad_acc.register_hook(self._make_hook(p))
self._grad_accs.append(grad_acc)

def _allreduce_grad_async(self, p):
name = self._parameter_names.get(p)
tensor = p.grad
tensor_compressed, ctx = self._compression.compress(tensor)
handle = hvd.allreduce_async_(tensor_compressed, op=hvd.Adasum, name=name)
return handle, ctx

def _make_hook(self, p):
def hook(*ignore):
if p in self._handles and self._handles[p][0] is not None:
if self._allreduce_delay[p] <= 0:
raise AssertionError(
"Gradients were computed more than "
"backward_passes_per_step times before call "
"to step(). Increase backward_passes_per_step to "
"accumulate gradients locally.")
assert not p.grad.requires_grad
assert self._allreduce_delay[p] > 0
handle, ctx = None, None
self._allreduce_delay[p] -= 1
if self._allreduce_delay[p] == 0:
handle, ctx = self._allreduce_grad_async(p)
self._handles[p] = (handle, ctx)
return hook

def synchronize(self):
missing_p = self._requires_update - set(self._handles.keys())
for p in missing_p:
handle, ctx = self._allreduce_grad_async(p)
self._handles[p] = (handle, ctx)

for p, value in self._handles.items():
handle, ctx = value
if handle is None:
handle, ctx = self._allreduce_grad_async(p)
self._handles[p] = (handle, ctx)
for p, (handle, _) in self._handles.items():
output = hvd.synchronize(handle)
self._allreduce_delay[p] = self.backward_passes_per_step
p.grad.copy_(self._compression.decompress(output, ctx).data)
self._handles.clear()

self._synchronized = True

@contextmanager
def skip_synchronize(self):
"""
A context manager used to specify that optimizer.step() should
not perform synchronization.

It's typically used in a following pattern:

.. code-block:: python

optimizer.synchronize()
with optimizer.skip_synchronize():
optimizer.step()
"""
self._should_synchronize = False
try:
yield
finally:
self._should_synchronize = True

def step(self, closure=None):
if self._should_synchronize:
if self._synchronized:
warnings.warn("optimizer.step() called without "
"optimizer.skip_synchronize() context after "
"optimizer.synchronize(). This can cause training "
"slowdown. You may want to consider using "
"optimizer.skip_synchronize() context if you use "
"optimizer.synchronize() in your code.")
self.synchronize()
self._synchronized = False
return self.optimizer.step(closure)
#return super(self.__class__, self).step(closure)

def zero_grad(self):
if self._handles:
raise AssertionError("optimizer.zero_grad() was called after loss.backward() "
"but before optimizer.step() or optimizer.synchronize(). "
"This is prohibited as it can cause a race condition.")
return self.optimizer.zero_grad()
#return super(self.__class__, self).zero_grad()

from apex.fp16_utils.loss_scaler import DynamicLossScaler
class DistributedAdasumOptimizer(torch.optim.Optimizer):
def __init__(self, optimizer, compression,
backward_passes_per_step=1):
params = [p
for param_group in optimizer.param_groups
for p in param_group['params']]
super(self.__class__, self).__init__(params, dict(DistributedAdasumOptimizer.__dict__))

self._compression = compression
self.optimizer = optimizer

named_parameters = [('allreduce.noname.%s' % i, v)
for param_group in self.param_groups
for i, v in enumerate(param_group['params'])]

# make sure that named_parameters are tuples
if any([not isinstance(p, tuple) for p in named_parameters]):
raise ValueError('named_parameters should be a sequence of '
'tuples (name, parameter), usually produced by '
'model.named_parameters().')

dups = find_duplicates([k for k, _ in named_parameters])
if len(dups) > 0:
raise ValueError('Parameter names in named_parameters must be unique. '
'Found duplicates: %s' % ', '.join(dups))

all_param_ids = {id(v)
for param_group in self.param_groups
for v in param_group['params']}
named_param_ids = {id(v) for k, v in named_parameters}
unnamed_param_ids = all_param_ids - named_param_ids
if len(unnamed_param_ids):
raise ValueError('named_parameters was specified, but one or more model '
'parameters were not named. Python object ids: '
'%s' % ', '.join(str(id) for id in unnamed_param_ids))

self._parameter_names = {v: k for k, v in sorted(named_parameters)}
self.backward_passes_per_step = backward_passes_per_step
self._allreduce_delay = {v: self.backward_passes_per_step
for _, v in sorted(named_parameters)}
self._handles = {}
self._grad_accs = []
self._requires_update = set()
self._synchronized = False
self._should_synchronize = True

self._starting_models = {
p : torch.zeros_like(p, requires_grad=False)
for _, p in named_parameters
}

self._scalers = {
p : DynamicLossScaler()
for _, p in named_parameters
}

self._is_first = True

#self._register_hooks()

def set_backward_passes_per_step(self, passes):
self.backward_passes_per_step = passes
for p in self._allreduce_delay:
self._allreduce_delay[p] = self.backward_passes_per_step

def _register_hooks(self):
for param_group in self.param_groups:
for p in param_group['params']:
if p.requires_grad:
p.grad = p.data.new(p.size()).zero_()
self._requires_update.add(p)
p_tmp = p.expand_as(p)
grad_acc = p_tmp.grad_fn.next_functions[0][0]
grad_acc.register_hook(self._make_hook(p))
self._grad_accs.append(grad_acc)

def _allreduce_grad_async(self, p):
# Delta optimizer implements this logic:
# start = current.copy()
# step() -> computes 'current - \alpha.f(g)' where f is
# optimizer logic and g is the gradient
# delta = current-start
# allreduce_(delta)
# start += delta
# current = start
# In order to suppport this logic using function hook to improve performance,
# we do:
# delta = (start - \alpha.f(g)) - start
# = -\alpha.f(g)
# set start to zero and step computes -\alpha.f(g)
# where f is the underlying optimizer logic

name = self._parameter_names.get(p)
start = self._starting_models[p]

stashed_params = []
for group in self.param_groups:
stashed_params.append(group['params'])
# only want to step on p
if any([p is v for v in group['params']]):
group['params'] = [p]
else:
group['params'] = []

start.data.copy_(p)

self.optimizer.step()

# compute delta = curr - start
p.data.sub_(start)

# allreduce as before
scaler = self._scalers[p]
p.data.mul_(scaler.loss_scale)
tensor_compressed, ctx = self._compression.compress(p)
handle = hvd.allreduce_async_(tensor_compressed.data, name=name, op=hvd.Adasum)

# reset stashed parameters
for stashed, group in zip(stashed_params, self.param_groups):
group['params'] = stashed

return handle, ctx

def _make_hook(self, p):
def hook(*ignore):
if p in self._handles and self._handles[p][0] is not None:
if self._allreduce_delay[p] <= 0:
raise AssertionError(
"Gradients were computed more than "
"backward_passes_per_step times before call "
"to step(). Increase backward_passes_per_step to "
"accumulate gradients locally.")
assert not p.grad.requires_grad
assert self._allreduce_delay[p] > 0
handle, ctx = None, None
self._allreduce_delay[p] -= 1
if self._allreduce_delay[p] == 0:
handle, ctx = self._allreduce_grad_async(p)
self._handles[p] = (handle, ctx)
return hook

def synchronize(self):
pass

@contextmanager
def skip_synchronize(self):
raise AssertionError("Skipping synchronization is not supported when using Adasum optimizer.")

def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()

if self._is_first:
self._is_first = False
for group in self.param_groups:
for p in group['params']:
self._starting_models[p].data.copy_(p.data)

self.optimizer.step()

handles = []
for group in self.param_groups:
for p in group['params']:
name = self._parameter_names.get(p)
#scaler = self._scalers[p]
start = self._starting_models[p]
p.data.sub_(start)
#p.data.mul_(scaler.loss_scale)
#tensor_compressed, ctx = self._compression.compress(p)
tensor_compressed, ctx = p, None
handle = hvd.allreduce_async_(tensor_compressed.data, name=name, op=hvd.Adasum)
handles.append((handle, p, ctx))

for handle, p, ctx in handles:
start = self._starting_models[p]
#scaler = self._scalers[p]
delta = hvd.synchronize(handle)
#has_overflow = not (torch.isfinite(delta.data).all().item())
#if not has_overflow:
# delta = self._compression.decompress(delta, ctx)
# delta.data.div_(scaler.loss_scale)
# start.data.add_(delta.data)
start.data.add_(delta.data)
p.data.copy_(start)
#scaler.update_scale(has_overflow)

return loss

def zero_grad(self):
if self._handles:
raise AssertionError("optimizer.zero_grad() was called after loss.backward() "
"but before optimizer.step() or optimizer.synchronize(). "
"This is prohibited as it can cause a race condition.")
return self.optimizer.zero_grad()
Loading