pd: suppport CINN for se_e2_a inference #4770
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR aims to add support for CINN in se_e2_a inference and adjust the order between the to_static (CINN) wrapper and the distributed wrapper.
- Added CINN configuration checks in the C++ source file.
- Reordered and updated the CINN-related compilation wrappers in the Python training file.
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| source/api_cc/src/DeepPotPD.cc | Added CINN configuration checks with logging of compilation info. |
| deepmd/pd/train/training.py | Revised the order of CINN-related to_static wrapping before the distributed wrapper and fixed code block placement. |
📝 Walkthrough""" WalkthroughThe CINN JIT compilation logic for Paddle models is moved from the Changes
Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (29)
✨ 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 (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
deepmd/pd/train/training.py (5)
609-609: Remove redundant backend assignment.The line
backend = "CINN" if CINN else Noneis redundant since this code block only executes whenCINNisTrue. The backend will always be"CINN".Apply this diff to simplify the code:
- backend = "CINN" if CINN else None + backend = "CINN"
631-635: Clean up commented code.The commented lines should either be removed if they're no longer needed, or properly documented if they represent future functionality.
Consider removing these commented lines or adding a TODO comment explaining their purpose:
}, # label, - # None, # task_key - # False, # inference_only - # False, # do_atomic_virial - # None, # fparam - # None, # aparam ],
641-643: Fix typo in log message.There's a typo: "traning" should be "training".
Apply this diff to fix the typo:
log.info( "Enable CINN during training, there may be some additional " - "compilation time in the first traning step." + "compilation time in the first training step." )
610-638: Consider validating input specification compatibility.The
input_specdefinition is quite detailed and hard-coded. Consider adding validation to ensure that the actual runtime inputs match this specification, as mismatches could cause compilation errors.The current implementation assumes specific tensor shapes and types. If the actual model inputs differ from this specification, CINN compilation might fail. Consider adding runtime validation or making the input specification more flexible based on the actual model requirements.
604-607: Consider moving imports to module level.While the conditional imports work correctly, consider moving the paddle imports to the module level for better performance if CINN checking happens frequently.
The current approach:
if CINN: from paddle import ( jit, static, )Could be optimized to:
# At module level try: from paddle import jit, static PADDLE_JIT_AVAILABLE = True except ImportError: PADDLE_JIT_AVAILABLE = False # In the conditional block if CINN and PADDLE_JIT_AVAILABLE: # Use jit and static
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pd/train/training.py(1 hunks)source/api_cc/src/DeepPotPD.cc(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
🔇 Additional comments (1)
deepmd/pd/train/training.py (1)
602-603: LGTM on the wrapper ordering comment.The comment correctly explains the important sequencing requirement for CINN compilation before distributed wrappers.
67a024f to
9dd2daf
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4770 +/- ##
==========================================
- Coverage 84.79% 84.78% -0.01%
==========================================
Files 698 698
Lines 67746 67756 +10
Branches 3540 3544 +4
==========================================
+ Hits 57444 57447 +3
- Misses 9171 9175 +4
- Partials 1131 1134 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
2851a82 to
6479157
Compare
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Style