pd: add flag CINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shape
#4826
Conversation
There was a problem hiding this comment.
Pull Request Overview
This pull request adds an option to allow compiling dynamic input shapes to improve runtime compatibility when input shapes vary, while preserving the option for slightly enhanced performance with fixed shapes. Key changes include:
- Introducing a new flag CINN_ALLOW_DYNAMIC_SHAPE in the environment configuration.
- Branching the jit.to_static call in training to conditionally handle dynamic shape support.
- Updating backend error messaging in the main parser to include the paddle backend.
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| deepmd/pd/utils/env.py | Added CINN_ALLOW_DYNAMIC_SHAPE with explanatory comments. |
| deepmd/pd/train/training.py | Updated warm_up_linear to branch jit conversion based on the flag. |
| deepmd/main.py | Expanded allowed backend values in error messaging. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
📝 Walkthrough""" WalkthroughThe updates introduce support for the "paddle" backend in argument parsing and error messages, add a new environment variable Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Trainer
participant Env
participant JIT_Compiler
User->>Trainer: Initialize Trainer
Trainer->>Env: Read CINN_ALLOW_DYNAMIC_SHAPE flag
alt CINN_ALLOW_DYNAMIC_SHAPE is True
Trainer->>JIT_Compiler: Compile forward (dynamic shape, backend=None)
else CINN_ALLOW_DYNAMIC_SHAPE is False
Trainer->>JIT_Compiler: Compile forward (static shape, backend='CINN')
end
Suggested reviewers
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Pylint (3.3.7)deepmd/main.pyNo files to lint: exiting. deepmd/pd/train/training.pyNo files to lint: exiting. deepmd/pd/utils/env.pyNo files to lint: exiting. ✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pd/train/training.py (1)
613-656: Verify the conditional compilation logic is correct.The current implementation appears counterintuitive:
- When
CINN_ALLOW_DYNAMIC_SHAPE=True: Compiles withbackend=None(disabling CINN)- When
CINN_ALLOW_DYNAMIC_SHAPE=False: Compiles withbackend="CINN"(enabling CINN)This seems backwards - shouldn't CINN be enabled when we want to allow dynamic shapes? Please verify this logic is intentional and matches the expected behavior described in the PR objectives.
Additionally, there's a minor style improvement:
- for k in label_dict.keys() + for k in label_dict
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/main.py(2 hunks)deepmd/pd/train/training.py(2 hunks)deepmd/pd/utils/env.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pd/train/training.py
633-633: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda, cuda)
🔇 Additional comments (4)
deepmd/pd/utils/env.py (1)
72-78: LGTM: Well-documented environment variable with sensible defaults.The implementation correctly introduces the dynamic shape control flag with clear documentation explaining the performance vs compatibility tradeoff. The default value of
Truemaintains backward compatibility while allowing users to optimize for performance when input shapes are fixed.deepmd/main.py (2)
115-115: LGTM: Error message correctly includes all supported backends.The error message now properly lists "paddle" alongside "tensorflow" and "pytorch", ensuring users get accurate guidance when an invalid backend is specified.
315-315: LGTM: Help text accurately documents Paddle backend output formats.The help text correctly specifies that Paddle backend outputs files with
.jsonand.pdiparamssuffixes, providing clear guidance to users about expected output formats for each backend.deepmd/pd/train/training.py (1)
57-57: LGTM: Import aligns with the new feature.The import of
CINN_ALLOW_DYNAMIC_SHAPEis correctly added to support the conditional compilation logic below.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4826 +/- ##
==========================================
- Coverage 84.57% 84.56% -0.01%
==========================================
Files 699 699
Lines 68036 68042 +6
Branches 3540 3544 +4
==========================================
+ Hits 57540 57541 +1
- Misses 9361 9366 +5
Partials 1135 1135 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
HydrogenSulfate
changed the title
HydrogenSulfate
changed the title
CINN_ALLOW_DYNAMIC_SHAPE for dynamic shape
HydrogenSulfate
changed the title
CINN_ALLOW_DYNAMIC_SHAPE for dynamic shapeCINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shape
HydrogenSulfate
force-pushed
the
add_option_for_dynamic_shape
branch
from
1161031 to
19ff021
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
deepmd/pd/train/training.py (1)
613-613: This addresses a past review concern correctly.The past review comment noted that
backendis set toNonein the dynamic shape branch. This behavior is now intentional and documented - whenCINN_ALLOW_DYNAMIC_SHAPEis True, the code deliberately disables CINN backend to allow dynamic shapes, while still providing detailed input specifications for optimization.
🧹 Nitpick comments (1)
deepmd/pd/train/training.py (1)
613-656: Improve code maintainability and address static analysis hint.The conditional CINN compilation logic is functionally correct but can be improved for better maintainability and code quality.
Issues to address:
- Static analysis hint: Use
key in dictinstead ofkey in dict.keys()on line 633- Code clarity: Add explanatory comments for the complex conditional logic
- Code organization: The input specification building is verbose and could be extracted
# NOTE: to_static + compiler should be before distributed wrapper if CINN: from paddle import ( jit, static, ) backend = "CINN" if CINN else None + # CINN dynamic shape handling: + # - If CINN_ALLOW_DYNAMIC_SHAPE=True: Use backend=None to allow dynamic shapes (CINN disabled) + # - If CINN_ALLOW_DYNAMIC_SHAPE=False: Use CINN backend for fixed shape optimization if CINN_ALLOW_DYNAMIC_SHAPE: # Build spec only for keys present in sample data # NOTE: This is a trick to decide the right input_spec for wrapper.forward _, label_dict, _ = self.get_data(is_train=True) # Define specification templates spec_templates = { "find_box": np.float32(1.0), "find_coord": np.float32(1.0), "find_numb_copy": np.float32(0.0), "numb_copy": static.InputSpec([1, 1], "int64", name="numb_copy"), "find_energy": np.float32(1.0), "energy": static.InputSpec([1, 1], "float64", name="energy"), "find_force": np.float32(1.0), "force": static.InputSpec([1, -1, 3], "float64", name="force"), "find_virial": np.float32(0.0), "virial": static.InputSpec([1, 9], "float64", name="virial"), "natoms": static.InputSpec([1, -1], "int32", name="natoms"), } label_dict_spec = { k: spec_templates[k] - for k in label_dict.keys() + for k in label_dict if k in spec_templates } self.wrapper.forward = jit.to_static( backend=None, input_spec=[ static.InputSpec([1, -1, 3], "float64", name="coord"), # coord static.InputSpec([1, -1], "int32", name="atype"), # atype None, # spin static.InputSpec([1, 9], "float64", name="box"), # box static.InputSpec([], "float64", name="cur_lr"), # cur_lr label_dict_spec, # label, # None, # task_key # False, # inference_only # False, # do_atomic_virial # None, # fparam # None, # aparam ], full_graph=True, )(self.wrapper.forward) else: + # Use CINN backend for fixed shape optimization self.wrapper.forward = jit.to_static(full_graph=True, backend=backend)( self.wrapper.forward )Additional suggestions for future improvements:
- Consider extracting the input specification building into a separate method
- Add error handling for
self.get_data()call that could potentially fail- Consider making the specification templates configurable rather than hardcoded
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/main.py(2 hunks)deepmd/pd/train/training.py(2 hunks)doc/train/training.md(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- doc/train/training.md
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/main.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pd/train/training.py
633-633: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (26)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
🔇 Additional comments (1)
deepmd/pd/train/training.py (1)
57-57: LGTM: Import of dynamic shape flag.The import of
CINN_ALLOW_DYNAMIC_SHAPEis correctly added to support the new conditional compilation behavior.
HydrogenSulfate
force-pushed
the
add_option_for_dynamic_shape
branch
from
19ff021 to
7bf2ee7
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pd/train/training.py (1)
613-656: Well-implemented conditional CINN compilation with minor optimization opportunity.The conditional logic effectively provides the flexibility mentioned in the PR objectives:
- Dynamic shape branch builds detailed input specifications for CINN compilation
- Static shape branch compiles without specs for potentially better performance
The implementation correctly passes the
backendparameter in both branches, addressing the concern from the previous review.Apply this minor optimization suggested by static analysis:
label_dict_spec = { - k: spec_templates[k] - for k in label_dict.keys() - if k in spec_templates + k: spec_templates[k] + for k in label_dict + if k in spec_templates }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/main.py(2 hunks)deepmd/pd/train/training.py(2 hunks)doc/train/training.md(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/main.py
🚧 Files skipped from review as they are similar to previous changes (1)
- doc/train/training.md
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pd/train/training.py
633-633: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (30)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
🔇 Additional comments (1)
deepmd/pd/train/training.py (1)
57-57: LGTM: Clean import of the new environment variable.The import of
CINN_ALLOW_DYNAMIC_SHAPEis properly added to support the conditional compilation logic.
HydrogenSulfate
force-pushed
the
add_option_for_dynamic_shape
branch
from
7bf2ee7 to
e00c050
Compare
HydrogenSulfate
changed the title
CINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shapeCINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shape
HydrogenSulfate
force-pushed
the
add_option_for_dynamic_shape
branch
from
e00c050 to
d9563cf
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pd/train/training.py (1)
633-633: Apply static analysis optimization.Use
key in dictinstead ofkey in dict.keys()for better performance.- for k in label_dict.keys() + for k in label_dict
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/main.py(2 hunks)deepmd/pd/train/training.py(2 hunks)doc/train/training.md(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- doc/train/training.md
- deepmd/main.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pd/train/training.py
633-633: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda120, cuda)
🔇 Additional comments (2)
deepmd/pd/train/training.py (2)
57-57: LGTM!The import of
CINN_ALLOW_DYNAMIC_SHAPEis correctly placed and aligns with the PR objective to add this new environment flag.
613-672: Well-implemented conditional CINN compilation logic.The implementation correctly handles both dynamic and fixed shape scenarios:
- Dynamic shape mode: Builds comprehensive input specifications from sample data to enable CINN compilation with varying input shapes
- Fixed shape mode: Compiles without specifications for optimal performance with consistent input shapes
- User guidance: Informative logging helps users understand the implications and make appropriate choices
The logic properly addresses the PR objective of providing users control over dynamic shape compilation for performance optimization.
Minor consideration: Ensure that
self.get_data(is_train=True)on line 616 has robust error handling in case training data is not yet available during initialization, though this is likely handled by the broader training pipeline setup.
HydrogenSulfate
force-pushed
the
add_option_for_dynamic_shape
branch
from
d9563cf to
0524f5c
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
deepmd/pd/train/training.py (1)
613-613: Verify the intentional backend configuration change.The code sets
backend="CINN"in the dynamic shape path, which differs from the past review comment's concern about setting backend to None. The current implementation appears correct - using CINN backend when dynamic shapes are enabled and the same backend when disabled.However, please confirm this change aligns with the intended CINN compilation behavior for both dynamic and fixed shape scenarios.
🧹 Nitpick comments (1)
deepmd/pd/train/training.py (1)
633-633: Apply static analysis suggestion for better performance.Use
k in label_dictinstead ofk in label_dict.keys()for better performance and readability.- for k in label_dict.keys() + for k in label_dict
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
deepmd/main.py(2 hunks)deepmd/pd/train/training.py(2 hunks)deepmd/pd/utils/env.py(2 hunks)doc/train/training.md(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- doc/train/training.md
- deepmd/main.py
- deepmd/pd/utils/env.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
deepmd/pd/train/training.py (2)
deepmd/pt/train/training.py (1)
get_data(1095-1137)source/tests/pd/model/test_saveload_dpa1.py (1)
get_data(117-134)
🪛 Ruff (0.11.9)
deepmd/pd/train/training.py
633-633: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (27)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
🔇 Additional comments (2)
deepmd/pd/train/training.py (2)
57-57: LGTM: Import addition is correct.The import of
CINN_ALLOW_DYNAMIC_SHAPEfromdeepmd.pd.utils.envis properly added and aligns with the PR objective to introduce this configuration flag.
659-671: LGTM: Helpful log messages for user guidance.The log messages provide clear guidance to users about CINN behavior and performance considerations. The conditional logging based on the flag setting is appropriate.
njzjz
changed the title
CINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shapeCINN_ALLOW_DYNAMIC_SHAPE for better performance with dynamic shape
|
@njzjz hello, this PR is already for review. |
Add flag
CINN_ALLOW_DYNAMIC_SHAPEfor allow compiling dynamic input shape, if set to0orFalse, performance might be slightly enhanced for fixed shape. We set default value toTruefor compativility.Summary by CodeRabbit
Summary by CodeRabbit
New Features
Improvements