[Serde][Dy2St] Use cached PartialProgramLayer in jit.load

[SigureMo]

PR Category

Execute Infrastructure

PR Types

Bug fixes

Description

jit.load 运行阶段，当 PartialProgramLayer 已经被构造后，不应该重新构造，应该直接复用，否则这会导致：

• PartialProgramLayer 构造阶段需要每次都重新跑，这包含了前反向拆分等诸多编译期逻辑，不应该在热启动链路执行
• scope 是在 PartialProgramLayer 上 cache 的，如果 PartialProgramLayer 每次重新构造，那么 scope 必然重新构造，这会导致执行器 cache 无法命中，每次都重新构造执行器，造成非常大的启动开销

同样代码（ResNet50），动转静直接跑 timeline 如下：

jit.save 再 jit.load 本 PR 修复前 timeline 如下（慢了十几倍）：

本 PR 修复后 timeline 如下（和动转静直接跑基本对齐了）：

另外，因为有一个 TranslatedLayer load 多个 program 的情况，因此需要存多个 PartialProgramLayer，根据 method_name 存在 _partial_program_layers 中

[paddle-bot]

你的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.

[Copilot]

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (1)

python/paddle/jit/pir_translated_layer.py:324

• [nitpick] The parameter name 'input' shadows Python's built-in function. Consider renaming it to 'input_data' to improve clarity.

def _run_dygraph(instance, input, program_holder, method_name):