optimize-offload support adamw op type (#36432)

66c58fa3 · duanboqiang · GitHub · 8256f6fa · 66c58fa3
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py ...tributed/fleet/meta_optimizers/sharding/offload_helper.py +1 -1

未找到文件。
--- a/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py
+++ b/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py
@@ -284,7 +284,7 @@ class OffloadHelper(object):
                break
            vars_name = []
-            if op.type == "adam":
+            if op.type == "adam" or op.type == "adamw":
                # {Moment1Out = [''], Moment2Out = [''], ParamOut = ['']} =
                # adam(inputs={Moment1 = [''], Moment2 = [''], Param = ['']})
                vars_name.append(op.desc.input("Moment1")[0])