【PRIM】Support custom_vjp for reducing video memory #50885

cxxly · 2023-02-24T11:39:35Z

PR types

New features

PR changes

Others

Describe

Pcard-66975

Add new CustomVJP feature. When forward and backward are both decomposite, run backward decomposite firstly.
Enable dy2st to support CustomVJP . @2742195759
Fix _lower_composite bugs when op is not register in decomposite rules.
Fix cast prim api and vjp rules DataType and VarType mapping error.
Delete to_prim from autograd all list and supoort blacklist/whitelist feature.

paddle-bot · 2023-02-24T11:39:39Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

JiabinYang

some comments

JiabinYang · 2023-03-03T07:24:28Z

python/paddle/jit/dy2static/partial_program.py

@@ -293,7 +316,8 @@ def _create_pure_fp16_program(self, is_infer_mode=False):
    @switch_to_static_graph
    def _create_forward_backward_train_program(self):
        whole_program = self._train_program
-        _, forward_end_op_index = self._infer_info('fp32', self._create_program)
+        # _, forward_end_op_index = self._infer_info('fp32', self._create_program)


no comments code

JiabinYang · 2023-03-03T07:26:02Z

python/paddle/jit/dy2static/program_translator.py

+        class PrimHooker(PartialProgramLayerHook):
+            def __init__(self):
+                self.custom_vjps = set()
+                if core._is_fwd_prim_enabled() and core._is_bwd_prim_enabled():


use _is_all_prim_enabled?

JiabinYang · 2023-03-03T07:28:16Z

python/paddle/jit/dy2static/program_translator.py

+                self, partial_program_layer, whole_program, backward_start_idx
+            ):
+                backward_length = (
+                    len(whole_program.block(0).ops) - backward_start_idx


Add comments to show we need support other block later

Done, check block length and raise exception when the length is over 1.

JiabinYang · 2023-03-03T07:32:59Z

python/paddle/jit/dy2static/program_translator.py

+                    to_prim(infer_program.block(0))
+                return infer_program
+
+        partial_program = partial_program_from(concrete_program)


Is this the only entry for dy2st?

done, enable fp16 support.

JiabinYang · 2023-03-03T07:34:33Z

python/paddle/jit/dy2static/utils.py

@@ -1519,7 +1519,7 @@ def _out_grad_names(program_desc, fwd_end_op_index, out_size):
        min(fwd_end_op_index + out_size, program_desc.block(0).op_size()),
    ):
        op = program_desc.block(0).op(i)
-        if op.type() == 'fill_any_like':
+        if op.type() in ['fill_any_like', "fill_constant"]:


JiabinYang · 2023-03-03T07:35:23Z

python/paddle/jit/dy2static/program_translator.py

+
+
+@switch_to_static_graph
+def to_prim(blocks, exclude=frozenset()):


change name maybe，too many to_prim

JiabinYang · 2023-03-03T07:36:37Z

python/paddle/fluid/tests/unittests/dygraph_to_static/test_cinn_prim_mean.py

@@ -99,7 +100,12 @@ def _train(self, use_prim, data, axis, keep_dim):
    def check_prim(self, net, use_prim):
        if not use_prim:
            return
-        fwd_ops = [op.type for op in net.forward.main_program.block(0).ops]
+        fwd_ops = [


leave a comment or seal this into a function, in case of possible changes happened later

Aurelius84 · 2023-03-10T11:00:41Z

python/paddle/jit/dy2static/partial_program.py

@@ -183,6 +194,7 @@ def __init__(
        # Set default mode to train
        self.training = True
        self._infer_info = ProgramInfo()
+        self._forward_end_index_map = {}


这个字段可以集成到 ProgramInfo里吧？

这个PR先合入吧，合入之后我来把这个字段放入到ProgramInfo里面。

…Paddle#50557) * [CINN]Enhance CacheKey hash logic by considering input dtypes * add unittest * fix typo * fix typo * fix map.at * fix find * fix test * fix cinn cache key structure realize * using ordered map for attributes * add test by review advice --------- Co-authored-by: jiangcheng <thisjiang@qq.com>

* [CINN]Enhance CacheKey hash logic by considering input dtypes (PaddlePaddle#50557) * [CINN]Enhance CacheKey hash logic by considering input dtypes * add unittest * fix typo * fix typo * fix map.at * fix find * fix test * fix cinn cache key structure realize * using ordered map for attributes * add test by review advice --------- Co-authored-by: jiangcheng <thisjiang@qq.com> * [prim] enable dygraph_to_static to support custom_vjp * fix code in a dy2static-friendly way. * [dystatic] add hooker for prim --------- Co-authored-by: Aurelius84 <zhangliujie@baidu.com> Co-authored-by: jiangcheng <thisjiang@qq.com> Co-authored-by: cxxly <chenxx_id@163.com>

* [CINN]Enhance CacheKey hash logic by considering input dtypes (PaddlePaddle#50557) --------- Co-authored-by: jiangcheng <thisjiang@qq.com> * [prim] enable dygraph_to_static to support custom_vjp * Pr 50885 (#7) * [CINN]Enhance CacheKey hash logic by considering input dtypes (PaddlePaddle#50557) * [CINN]Enhance CacheKey hash logic by considering input dtypes --------- Co-authored-by: jiangcheng <thisjiang@qq.com> * [prim] enable dygraph_to_static to support custom_vjp * fix code in a dy2static-friendly way. * [dystatic] add hooker for prim --------- Co-authored-by: Aurelius84 <zhangliujie@baidu.com> Co-authored-by: jiangcheng <thisjiang@qq.com> Co-authored-by: cxxly <chenxx_id@163.com> * [prim] enable dygraph_to_static to support custom_vjp * fix cast prim and vjp dtype mapping error bug * [dy2static-ci] fix dy2static ci errors. --------- Co-authored-by: Aurelius84 <zhangliujie@baidu.com> Co-authored-by: jiangcheng <thisjiang@qq.com> Co-authored-by: cxxly <chenxx_id@163.com>

2742195759

LGTM

JiabinYang

LGTM with some questions can be solved later

JiabinYang · 2023-03-13T11:12:16Z

python/paddle/fluid/framework.py

@@ -3751,7 +3751,6 @@ def __init__(self, program, idx):
        self.vars = collections.OrderedDict()  # var_name --> var
        self.ops = list()  # operator list
        self.program = program
-        self.removed_vars = collections.OrderedDict()


why remove this?

pre-commit auto format

JiabinYang · 2023-03-13T11:22:08Z

python/paddle/jit/dy2static/partial_program.py

@@ -701,6 +736,7 @@ def _prepare_attributes(self):
            'program_id',
            self.program_id,
        ]
+


useless blank?

pre-commit auto format

JiabinYang · 2023-03-13T11:22:26Z

python/paddle/jit/dy2static/partial_program.py

@@ -1119,5 +1155,8 @@ def add_build_strategy_for(
        if hasattr(compiled_program._program, 'lr_sheduler'):
            builded_program.lr_sheduler = compiled_program._program.lr_sheduler
    else:
-        builded_program = program
+        # can't just create a new program, we need copy the vardesc.
+        builded_program = paddle.static.Program()


fix the bug when the program only contain a var.

@paddle.jit.to_static def f(x): return x

jzhang533

LGTM

XiaoguangHu01

LGTM

cxxly force-pushed the prim_custom_vjp branch 11 times, most recently from 18cdd67 to d75ef48 Compare March 2, 2023 13:17

JiabinYang reviewed Mar 3, 2023

View reviewed changes

cxxly force-pushed the prim_custom_vjp branch 5 times, most recently from 74fd37a to 156c830 Compare March 10, 2023 10:28

Aurelius84 reviewed Mar 10, 2023

View reviewed changes

2742195759 mentioned this pull request Mar 13, 2023

[Dy2static] Support Train step In to_static #51543

Closed

cxxly force-pushed the prim_custom_vjp branch 2 times, most recently from d67a0cc to 6191f33 Compare March 13, 2023 05:17

Aurelius84 and others added 7 commits March 13, 2023 09:20

[prim] enable dygraph_to_static to support custom_vjp

c4bb949

[prim] enable dygraph_to_static to support custom_vjp

454c563

fix cast prim and vjp dtype mapping error bug

eb775b0

[Prim] enable whitelist and blacklist for custom_vjp

2b203b9

cxxly force-pushed the prim_custom_vjp branch from 6191f33 to 2b203b9 Compare March 13, 2023 09:21

2742195759 approved these changes Mar 13, 2023

View reviewed changes

JiabinYang approved these changes Mar 13, 2023

View reviewed changes

jzhang533 approved these changes Mar 14, 2023

View reviewed changes

XiaoguangHu01 approved these changes Mar 14, 2023

View reviewed changes

cxxly merged commit 300f36c into PaddlePaddle:develop Mar 14, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

【PRIM】Support custom_vjp for reducing video memory #50885

【PRIM】Support custom_vjp for reducing video memory #50885

cxxly commented Feb 24, 2023 •

edited

Loading

paddle-bot bot commented Feb 24, 2023

JiabinYang left a comment

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

JiabinYang Mar 3, 2023

cxxly Mar 6, 2023

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

JiabinYang Mar 3, 2023

cxxly Mar 5, 2023

Aurelius84 Mar 10, 2023

2742195759 Mar 13, 2023

2742195759 left a comment

JiabinYang left a comment

JiabinYang Mar 13, 2023

cxxly Mar 13, 2023

JiabinYang Mar 13, 2023

cxxly Mar 14, 2023

JiabinYang Mar 13, 2023

cxxly Mar 13, 2023

JiabinYang Mar 14, 2023

jzhang533 left a comment

XiaoguangHu01 left a comment



		@switch_to_static_graph
		def to_prim(blocks, exclude=frozenset()):

【PRIM】Support custom_vjp for reducing video memory #50885

【PRIM】Support custom_vjp for reducing video memory #50885

Conversation

cxxly commented Feb 24, 2023 • edited Loading

PR types

PR changes

Describe

paddle-bot bot commented Feb 24, 2023

JiabinYang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

2742195759 left a comment

Choose a reason for hiding this comment

JiabinYang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jzhang533 left a comment

Choose a reason for hiding this comment

XiaoguangHu01 left a comment

Choose a reason for hiding this comment

cxxly commented Feb 24, 2023 •

edited

Loading