From edd97f94464a53153ac8ec38b62c9a64a8f25893 Mon Sep 17 00:00:00 2001 From: WangXi Date: Mon, 14 Mar 2022 16:06:20 +0800 Subject: [PATCH] [hybrid fix] fix sharding save offload (#40477) --- .../distributed/fleet/meta_optimizers/sharding/utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py index b42f21989ab..1a3a8a4883d 100755 --- a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py +++ b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py @@ -900,11 +900,12 @@ 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 + # now only Momentum and adam are compatible with sharding, + # support EMA optimizer with '_ema_0', + # support offload with '@offload_0' and '.cast_fp16' checks = [ "_moment1_0", "_moment2_0", "_beta1_pow_acc_0", "_beta2_pow_acc_0", - "_velocity_0", "_ema_0" + "_velocity_0", "_ema_0", "@offload_0", ".cast_fp16" ] for check in checks: if var.name.endswith(check) and var.persistable: -- GitLab