未验证 提交 8bb509d5 编写于 作者: B Baibaifan 提交者: GitHub

fix_stage3_fp16 (#39171)

上级 9059ef69
...@@ -393,6 +393,7 @@ class ShardingStage3(nn.Layer): ...@@ -393,6 +393,7 @@ class ShardingStage3(nn.Layer):
else: else:
param.bw_storage.scale_(scale=self._world_size_scaling) param.bw_storage.scale_(scale=self._world_size_scaling)
param.fw_storage = _VarBaseWrapper(param) param.fw_storage = _VarBaseWrapper(param)
assert param.fw_storage.grad is None
param.fw_storage._copy_gradient_from(param.bw_storage) param.fw_storage._copy_gradient_from(param.bw_storage)
update_list.append(param) update_list.append(param)
return update_list return update_list
...@@ -495,10 +496,9 @@ class ShardingStage3(nn.Layer): ...@@ -495,10 +496,9 @@ class ShardingStage3(nn.Layer):
def _redefine_opt_step(self): def _redefine_opt_step(self):
params_slice_func = self._update_params_slice params_slice_func = self._update_params_slice
opt_step = self._optim.step opt_step = self._optim.step
update_scaler = self._optim.update_scaler
def _opt_step(self): def _opt_step(self):
if not update_scaler: if not self.update_scaler:
params_slice_func() params_slice_func()
if self.offload: if self.offload:
with device_guard(device="cpu"): with device_guard(device="cpu"):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册