From d5cc7bff1bc777f0da77dd969b6c4fc22bfbcafa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=AD=A3?= <2042519524@qq.com> Date: Wed, 2 Jun 2021 17:40:22 +0800 Subject: [PATCH] update mp (#33194) * update mp --- python/paddle/distributed/collective.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/python/paddle/distributed/collective.py b/python/paddle/distributed/collective.py index 4f3a6f47689..5775a734c87 100644 --- a/python/paddle/distributed/collective.py +++ b/python/paddle/distributed/collective.py @@ -1009,16 +1009,18 @@ def _parallel_linear(x, name=name) linear_out = linear(x) - startup_block = paddle.static.default_startup_program().global_block() - main_block = paddle.static.default_main_program().global_block() - startup_block.vars[linear.weight.name].is_distributed = True - main_block.vars[linear.weight.name].is_distributed = True + startup_block = paddle.static.default_startup_program().current_block() + main_block = paddle.static.default_main_program().current_block() + startup_block._find_var_recursive(linear.weight.name).is_distributed = True + main_block._find_var_recursive(linear.weight.name).is_distributed = True + # set is_distributed for splited bias # if a linear layer is splited by row, each rank would hold a complete bias and they should be the same in each rank. # if a linear layer is splited by col, the bias would also be split into each rank as its weight if axis == 1 and linear._bias_attr != False: - startup_block.vars[linear.bias.name].is_distributed = True - main_block.vars[linear.bias.name].is_distributed = True + startup_block._find_var_recursive( + linear.bias.name).is_distributed = True + main_block._find_var_recursive(linear.bias.name).is_distributed = True if not gather_out: return linear_out -- GitLab