Compatible upgrade of sparse_momentum for master param (#107)

PaddlePaddle · Mar 1, 2022 · 731eb3f · 731eb3f
1 parent ef0a32b
commit 731eb3f
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/dynamic/utils/hybrid_optimizer.py b/dynamic/utils/hybrid_optimizer.py
@@ -85,12 +85,20 @@ def _append_optimize_op(self, block, param_and_grad):
  if getattr(param_and_grad[0], 'is_sparse_grad', None):
  index = getattr(param_and_grad[0], 'index', None)
  axis = getattr(param_and_grad[0], 'axis', None)
- _, _ = paddle._C_ops.sparse_momentum(
- param_and_grad[0], param_and_grad[1], velocity_acc, index, lr,
- param_and_grad[0], velocity_acc, 'mu', self._momentum,
- 'use_nesterov', self._use_nesterov, 'regularization_method',
- self._regularization_method, 'regularization_coeff',
- self._regularization_coeff, 'axis', axis)
+ try:
+ _, _ = paddle._C_ops.sparse_momentum(
+ param_and_grad[0], param_and_grad[1], velocity_acc, index, lr,
+ param_and_grad[0], velocity_acc, 'mu', self._momentum,
+ 'use_nesterov', self._use_nesterov, 'regularization_method',
+ self._regularization_method, 'regularization_coeff',
+ self._regularization_coeff, 'axis', axis)
+ except:
+ _, _, _ = paddle._C_ops.sparse_momentum(
+ param_and_grad[0], param_and_grad[1], velocity_acc, index, lr, master_weight,
+ param_and_grad[0], velocity_acc, master_weight, 'mu', self._momentum,
+ 'use_nesterov', self._use_nesterov, 'regularization_method',
+ self._regularization_method, 'regularization_coeff',
+ self._regularization_coeff, 'axis', axis, 'multi_precision', find_master)
  else:
  _, _, _ = paddle._C_ops.momentum(
  param_and_grad[0], param_and_grad[1], velocity_acc, lr,