Eagerly accumulate embedding grads into fp32 buffer (#6958)

Signed-off-by: Tim Moon <[email protected]>

nemo/core/optim/distributed_adam.py

@@ -77,31 +77,37 @@ def __init__(self, params, disable_distributed_parameters=False, **kwargs):
  distopt_param_groups = param_groups
  dtype = kwargs['dtype'] if 'dtype' in kwargs else torch.float32
  grad_sync_dtype = kwargs['grad_sync_dtype'] if 'grad_sync_dtype' in kwargs else dtype
- needs_fp32_optimizer = any(
- getattr(param, '_with_fp32_optimizer', False)
- for param in itertools.chain.from_iterable(param_group['params'] for param_group in param_groups)
- )
- if (dtype != torch.float32 or grad_sync_dtype != torch.float32) and needs_fp32_optimizer:
+ needs_fp32_optimizer = dtype != torch.float32 or grad_sync_dtype != torch.float32
+ if needs_fp32_optimizer:
+ needs_fp32_optimizer = any(
+ any(getattr(param, '_with_fp32_optimizer', False) for param in param_group['params'])
+ for param_group in param_groups
+ )
+ if needs_fp32_optimizer:
 
  # Find params that require explicit FP32 optimizer
  distopt_param_groups = []
  fp32_param_groups = []
  self._fp32_optim_main_params = collections.OrderedDict()
  for param_group in param_groups:
- distopt_param_group = {key: val for key, val in param_group.items() if key != 'params'}
+ distopt_param_group = param_group.copy()
  distopt_param_group['params'] = []
- fp32_param_group = {key: val for key, val in param_group.items() if key != 'params'}
+ fp32_param_group = param_group.copy()
  fp32_param_group['params'] = []
  for model_param in param_group['params']:
  if getattr(model_param, '_with_fp32_optimizer', False):
  main_param = model_param.detach().clone().float()
+ model_param.main_grad = main_param.grad
  fp32_param_group['params'].append(main_param)
  self._fp32_optim_main_params[model_param] = main_param
  else:
  distopt_param_group['params'].append(model_param)
  distopt_param_groups.append(distopt_param_group)
  fp32_param_groups.append(fp32_param_group)
 
+ # Add callback hook so grads accumulate into FP32 buffer
+ self._fp32_register_post_backward_hooks()
+
  # Construct explicit FP32 optimizer
  adamw_kwargs = {}
  for name in ('lr', 'betas', 'eps', 'weight_decay', 'amsgrad'):
@@ -113,6 +119,30 @@ def __init__(self, params, disable_distributed_parameters=False, **kwargs):
  # Construct distributed optimizer
  super().__init__(distopt_param_groups, **kwargs)
 
+ def _fp32_register_post_backward_hooks(self):
+ """Attach hooks for FP32 gradients"""
+
+ # Helper function to avoid issues with late binding closures
+ def make_post_backward_hook(param):
+ def post_backward_hook(*unused):
+ self._fp32_optim_grad_sync_needed = True
+ if hasattr(param, 'main_grad'):
+ with torch.no_grad():
+ if param.grad is not None:
+ param.main_grad += param.grad
+ param.grad = None
+
+ return post_backward_hook
+
+ # Construct hooks and register with params
+ self._fp32_grad_accs = []
+ for param in self._fp32_optim_main_params.keys():
+ param_tmp = param.expand_as(param)
+ grad_acc = param_tmp.grad_fn.next_functions[0][0]
+ hook = make_post_backward_hook(param)
+ grad_acc.register_hook(hook)
+ self._fp32_grad_accs.append(grad_acc)
+
  def _make_post_backward_hook(self, param, param_group_id, param_id):
  def hook(*unused):
  if getattr(param, '_pre_forward_hook_is_enabled', False):