pd: add CINN compiler for dpa2, dpa1 training #4514
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
📝 WalkthroughWalkthroughThe changes modify two modules. In Changes
Sequence Diagram(s)sequenceDiagram
participant Runner as training.run()
participant Env as CINN Flag
participant Paddle as Paddle API
participant NV as nvprof_context
participant Model as Model.forward()
Runner->>Env: Check if CINN is enabled
alt CINN is enabled
Runner->>Paddle: Import jit, static, BuildStrategy
Runner->>Paddle: Wrap Model.forward() using jit.to_static
end
Runner->>NV: Enter profiling context
Runner->>Model: Execute data retrieval and prediction
Model-->>Runner: Return predictions
Runner->>NV: Exit profiling context
sequenceDiagram
participant OS as OS Environment
participant Env as Environment Setup
participant TB as to_bool()
participant Paddle as PaddlePaddle
OS->>TB: Provide "JIT" value
TB-->>Env: Return boolean for JIT
OS->>TB: Provide "CINN" value
TB-->>Env: Return boolean for CINN
alt CINN is true
Env->>Paddle: Check for CINN support
alt Not Supported
Env->>Env: Raise assertion error
end
end
Env->>Env: In enable_prim(), if JIT true, enable primitives in JIT mode
Env->>Env: Otherwise, enable primitives in eager mode with composite op blacklist
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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 (2)
deepmd/pd/train/training.py (2)
402-405: Remove or clarify commented-out code.These lines comment out a previously raised
NotImplementedErrorand a potentialpaddle.jit.to_staticcall. If you no longer need this logic, removing it altogether might reduce confusion and keep the file tidy. Otherwise, add a comment explaining why these lines are kept for future reference.-# if JIT: -# raise NotImplementedError("JIT is not supported yet when training with Paddle") -# self.model = paddle.jit.to_static(self.model)Also applies to: 406-406
925-926: Consider removing the extra commented-out code.This snippet appears to comment out a JIT debugging break. If it's no longer needed, removing it can avoid potential confusion.
-# if JIT: -# break🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 925-926: Commented-out code
This comment appears to contain commented-out code.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pd/train/training.py(8 hunks)deepmd/pd/utils/env.py(2 hunks)
🧰 Additional context used
🪛 GitHub Check: CodeQL
deepmd/pd/utils/env.py
[notice] 61-61: Unused global variable
The global variable 'CINN' is not used.
deepmd/pd/train/training.py
[notice] 925-926: Commented-out code
This comment appears to contain commented-out code.
🔇 Additional comments (7)
deepmd/pd/train/training.py (4)
56-57: Check the new imports for consistency.
The addition of CINN and DEFAULT_PRECISION is consistent with improvements to JIT compilation and precision handling. Nothing problematic is observed here; just ensure that CINN is successfully imported where used and that DEFAULT_PRECISION is consistently applied.
636-648: JIT and CINN integration logic looks good.
You're conditionally enabling JIT using jit.to_static and passing the build_strategy.build_cinn_pass = CINN. This is a clean approach, ensuring that CINN is only used if enabled. Just verify upstream usage to avoid unexpected behaviors if CINN is disabled at runtime.
688-692: Validate data fetching performance within profiling context.
Wrapping the data loading with nvprof_context helps profile overhead. Ensure that exceptions thrown within this block are properly handled so that the profiler is closed gracefully.
706-706: Precision usage for learning rate.
Using paddle.full([], pref_lr, DEFAULT_PRECISION) enforces consistent floating-point precision for the learning rate. This is beneficial for uniformity, especially in multi-task scenarios. Good practice!
Also applies to: 765-765, 815-815
deepmd/pd/utils/env.py (3)
36-50: Robust input validation in to_bool.
Excellent job handling integers, booleans, and string values thoroughly, with clear error messages and suitable lowercasing of strings. This ensures minimal confusion for environment variable parsing.
60-61: CINN usage is valid despite static analysis hint.
While static analysis tools suggest that CINN is “unused,” it is in fact used within training.py (via build_strategy.build_cinn_pass). You can safely disregard that hint.
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 61-61: Unused global variable
The global variable 'CINN' is not used.
167-182: Flexible prim enabling for JIT and eager modes.
Switching between JIT mode and eager mode, while setting blacklists only in eager mode, is a neat approach to leverage performance improvements. This design offers more control and clarity to developers using either mode.
HydrogenSulfate
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
deepmd/pd/train/training.py (3)
709-709: Confirm precision alignment for learning rate.
Replacing a float withpaddle.full([], pref_lr, DEFAULT_PRECISION)ensures consistent floating-point precision. Verify that using a 0-dimensional tensor is intentional and that the rest of the forward pass expects a scalar.
768-768: Avoid code duplication in validation logic.
This pattern repeats the same approach for creating a scalar tensor frompref_lr. Consider encapsulating the logic to reduce repetitive calls and simplify maintenance.
818-818: Maintain consistent approach for multi-task calls.
Reusingpaddle.full([], pref_lr, DEFAULT_PRECISION)in multiple task keys is valid. However, centralizing this snippet could prevent future errors in multi-task scenarios.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pd/train/training.py(6 hunks)deepmd/pd/utils/env.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pd/utils/env.py
🔇 Additional comments (2)
deepmd/pd/train/training.py (2)
56-57: Use environment-driven feature flags carefully.
Importing CINN and DEFAULT_PRECISION helps toggle experimental optimizations and ensure matching data types throughout training. This is acceptable, but confirm that these environment variables or configurations are properly documented and tested to avoid accidental misconfigurations.
691-695: Validate data fetching within profiler context.
The usage of nvprof_context around get_data() is beneficial for profiling. Make sure exceptions (e.g., StopIteration) in data fetching are handled gracefully to preserve consistent profiling measurements.
HydrogenSulfate
force-pushed
the
add_paddle_cinn_dpa2
branch
from
95b201d to
5cdd421
Compare
HydrogenSulfate
force-pushed
the
add_paddle_cinn_dpa2
branch
from
5cdd421 to
7ca2a9e
Compare
HydrogenSulfate
changed the title
HydrogenSulfate
changed the title
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4514 +/- ##
==========================================
- Coverage 84.60% 84.57% -0.03%
==========================================
Files 680 680
Lines 64473 64500 +27
Branches 3540 3540
==========================================
+ Hits 54546 54554 +8
- Misses 8786 8804 +18
- Partials 1141 1142 +1 ☔ View full report in Codecov by Sentry. |
|
@njzjz I have a question about the code at deepmd-kit/deepmd/pt/model/network/layernorm.py Lines 97 to 99 in 8d4c27b numel()==0? I haven't noticed this happening during Python training and testing. Could you please help me with this? Thank you!
|
When using LAMMPS with MPI. |
|
See also #2668. Some users may encounter the situation where a processor has no atom. |
Thank you for your response. Although 0-size tensors are not very common, we have indeed encountered similar issues with some object detection models. We are going to support for training and inference with 0-size tensors in Paddle in the near future. |
HydrogenSulfate
changed the title
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/pd/utils/env.py (2)
40-60: Consider using pattern matching for a more concise implementation.The function is well-implemented with proper type checking and error handling. However, it could be more concise using Python 3.10+ pattern matching.
Here's a more concise version using pattern matching:
-def to_bool(flag: int | bool | str) -> bool: - if isinstance(flag, int): - if flag not in [0, 1]: - raise ValueError(f"flag must be either 0 or 1, but received {flag}") - return bool(flag) - - elif isinstance(flag, str): - flag = flag.lower() - if flag not in ["1", "0", "true", "false"]: - raise ValueError( - f"flag must be either '0', '1', 'true', 'false', but received '{flag}'" - ) - return flag in ["1", "true"] - - elif isinstance(flag, bool): - return flag - - else: - raise ValueError( - f"flag must be either int, bool, or str, but received {type(flag).__name__}" - ) +def to_bool(flag: int | bool | str) -> bool: + match flag: + case bool(): + return flag + case int() if flag in [0, 1]: + return bool(flag) + case str() as s if s.lower() in ["0", "1", "true", "false"]: + return s.lower() in ["1", "true"] + case int(): + raise ValueError(f"flag must be either 0 or 1, but received {flag}") + case str(): + raise ValueError( + f"flag must be either '0', '1', 'true', 'false', but received '{flag}'" + ) + case _: + raise ValueError( + f"flag must be either int, bool, or str, but received {type(flag).__name__}" + )
177-195: Enhance logging to indicate the specific mode.The function correctly handles both JIT/CINN and eager modes, but the log message could be more informative.
Consider updating the log message to indicate the specific mode:
- log.info(f"{'Enable' if enable else 'Disable'} prim in eager and static mode.") + mode = "JIT/CINN" if JIT or CINN else "eager" + log.info(f"{'Enable' if enable else 'Disable'} prim in {mode} mode.")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pd/utils/env.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (19)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- 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 library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (2)
deepmd/pd/utils/env.py (2)
2-4: LGTM! Good use of type hints.The
annotationsimport enables modern type hint syntax, improving code readability and type safety.
63-70: LGTM! Good defensive programming with CINN support check.The code properly validates CINN support at startup with a clear error message, preventing runtime issues.
Note: The past review comment about unused
CINNvariable is outdated as the variable is now used intraining.py.
We verified paddle CINN compiler in DPA-2 example(single A100-SXM (40G), cada11.8, Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz x 160).
To enable CINN compiler in training, add one flag:
CINN=1before training command, e.g.CINN=1 dp --pd train input_torch_medium.json.Curves:
dpa2
Performance
We tested with torch==2.6.0.dev20241219+cu118
se_atten
Performance
We tested with torch==2.6.0.dev20241219+cu118
Accuracy details:
dpa2
Pytorch:
Paddle(eager mode):
Paddle(CINN compiler)
se_atten
Pytorch
Paddle(eager_mode)
Paddle(CINN compliler)
TODO:
Summary by CodeRabbit
New Features
Refactor