From 4847a40deb7f2140b3a2dedac209966d91e2a68d Mon Sep 17 00:00:00 2001 From: sljlp Date: Wed, 23 Feb 2022 19:23:00 +0800 Subject: [PATCH] support ema optimizer in sharding optimizers --- .../paddle/distributed/fleet/meta_optimizers/sharding/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py index d04a3a53db3e2..b42f21989abd7 100755 --- a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py +++ b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py @@ -901,9 +901,10 @@ def save_persistables(exe, dirname, main_program, filename=None): def is_opt_vars(var): # NOTE(JZ-LIANG): The checks should be updated when add new compatible optimizer # now only Momentum and adam are compatible with sharding + # support EMA optimizer checks = [ "_moment1_0", "_moment2_0", "_beta1_pow_acc_0", "_beta2_pow_acc_0", - "_velocity_0" + "_velocity_0", "_ema_0" ] for check in checks: if var.name.endswith(check) and var.persistable: