From e548f65f96697830035a28f9070b40829408ccdb Mon Sep 17 00:00:00 2001
From: Roc <30228238+sljlp@users.noreply.github.com>
Date: Tue, 8 Mar 2022 22:26:02 +0800
Subject: [PATCH] support ema optimizer in sharding optimizers (#39860)

---
 .../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 d04a3a53db..b42f21989a 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:
-- 
GitLab