Skip to content

Commit

Permalink
fix shared param grad_add op_device is null (#33875)
Browse files Browse the repository at this point in the history
  • Loading branch information
wangxicoding authored Jul 2, 2021
1 parent 9b48199 commit cf4c6fb
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions python/paddle/fluid/backward.py
Original file line number Diff line number Diff line change
Expand Up @@ -462,6 +462,7 @@ def _addup_repetitive_outputs_(op_descs, block_idx):
var_rename_count = collections.defaultdict(int)
renamed_vars = collections.defaultdict(list)
renamed_var_start_idx = collections.defaultdict(list)
var_device = collections.defaultdict(str)
for idx, op_desc in enumerate(op_descs):
op_device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName(
)
Expand Down Expand Up @@ -528,16 +529,19 @@ def _addup_repetitive_outputs_(op_descs, block_idx):
arg_names[arg_idx] = new_name
op_desc.set_output(param_name, arg_names)
renamed_vars[var_name].append(new_name)
# record the latest device, for shared param
var_device[var_name] = op_device

for var_name, inputs in six.iteritems(renamed_vars):
if len(renamed_vars[var_name]) > 1:
if len(renamed_vars[var_name]) > _MAX_ADD_NUM_:
_accumulate_gradients_by_sum_op_(var_name, renamed_vars,
pending_sum_ops, len(op_descs))
_accumulate_gradients_by_sum_op_(
var_name, renamed_vars, pending_sum_ops,
len(op_descs), var_device[var_name])
else:
_accumulate_gradients_by_add_ops_(var_name, renamed_vars,
pending_sum_ops,
len(op_descs))
_accumulate_gradients_by_add_ops_(
var_name, renamed_vars, pending_sum_ops,
len(op_descs), var_device[var_name])

# sum_op descs are sorted according to their insert position
for key, value in collections.OrderedDict(
Expand Down

0 comments on commit cf4c6fb

Please sign in to comment.