Repository navigation
build(pt): reduce the CUDA compile time of the PyTorch operators - #6044
Conversation
Compiling libdeepmd_op_pt dominates the installation of a CUDA build.
The DPA4 training operators alone generated 114 translation units: 90
rotation shards over degree x radial rank x {float, double, bfloat16}
and 24 value-path shards over degree x {float, double, half,
bfloat16}. Every shard parsed <ATen/ATen.h>, and every operator host
parsed the complete <torch/torch.h> C++ frontend. This commit removes
that cost without changing the kernels that training runs.
DPA4 training dtypes
- The rotation and value-path grids instantiate float32 and bfloat16,
the working types of float32 training and of bfloat16 autocast.
float64 is instantiated only under the new CMake option
DEEPMD_ENABLE_DPA4_FP64 (OFF by default): the model binds the fused
path for float32 parameters only, so the float64 kernels serve the
numerical consistency tests alone. float16 is no longer instantiated:
autocast runs in bfloat16, and the rotation kernels that every
backward composes never had a float16 form, so a float16 value path
could not complete a training step.
- One dispatch macro (DPA4_SEZM_DISPATCH_TYPES in sezm_train_ops.cuh)
serves every training entry. The value-path entry rejects any other
dtype with a message naming the option, and the new operator
sezm_so2_value_supports_float64() (float64_available() in Python)
reports whether the build carries the float64 kernels.
- test_float64_agrees_with_eager_to_reduction_order skips without the
float64 kernels, the CUDA CI job builds with the option through
CMAKE_ARGS, and doc/install/install-from-source.md documents it.
Double accumulation in the float64 kernels
- Every DPA4 training kernel computes in acc_type<scalar_t>: float for
float32 and bfloat16, double for float64. This now includes the
rotation and segment-sum kernels, the pointwise kernels of the mixing
traversal, and the scalar head parameters (inverse temperature, label
smoothing, norm epsilon) of the resident forward and the competition
head, all of which previously computed in float even in the float64
instantiations. A float64 evaluation therefore agrees with the eager
reference to reduction order, 1e-15 to 3e-11 relative across the
parity shapes, where a float32 kernel stage leaves disagreements on
the order of 1e-7. The float64 test tolerance tightens from 5e-6 to
1e-12: its shapes measure 1e-14, and a float32-rounded label
smoothing alone leaves 5e-10.
- The 384-thread top-degree variant of the first-order rotation
backward does not stage its edge rows in shared memory when it runs
in double: its batched-reduction scratch alone occupies 46,560 bytes,
and the staged rows would exceed the 48 KiB static limit. The
condition is the static footprint itself, so every float32 and
bfloat16 variant keeps the staging.
- The float32 and bfloat16 instantiations compile to unchanged SASS,
compared function by function on representative rotation and
value-path shards and on the three training hosts. The exception is
mixing_entry_bwd_kernel, which divides by the competition weight it
already holds instead of loading it a second time, with identical
results.
ATen-free kernel shards
- rotate_mix_train/kernels.cuh and so2_conv_train/kernels.cuh carry
device code and launchers only. Input validation and the runtime
(lmax, rank) dispatch move into the host files, launch_so2_value_fwd
returns a cudaError_t for the host to raise, and acc_type moves into
the new header sezm_types.cuh, so the shards include only
<c10/util/BFloat16.h>. <ATen/ATen.h> alone adds about 33 s to every
translation unit that includes it.
Lean framework headers
- The operator hosts and shared headers of DPA4, DPA4C, DPA1, the graph
fitting and the edge force/virial operators include <torch/types.h>
and <torch/library.h> instead of <torch/torch.h>, which keeps the C++
frontend (nn, optim, data, serialization) out of every CUDA unit. The
header set compiles in 34-38 s instead of 50-55 s, with a peak
compiler footprint of 1.5 GB instead of 2.6 GB.
- dpa1_graph_descriptor.cu, the longest unit of the library, compiles
with nvcc --threads=2 like the fused DPA4 convolution units, so its
native and portable-PTX device passes run concurrently.
Measured effect
Single translation units compiled alone on a 192-core host (CUDA 13.3,
sm_120 SASS plus compute_75 PTX):
rotation shard, degree 1, rank 0, float32 38.2 s -> 3.4 s
rotation shard, degree 6, rank 4, bfloat16 73.1 s -> 39.2 s
value-path shard, degree 1, float32 50.4 s -> 16.2 s
value-path shard, degree 6, bfloat16 61.4 s -> 28.2 s
dpa1_graph_descriptor.cu 196.6 s -> 129.3 s
A full default rebuild of the operator library with 64 jobs now takes
3,999 CPU-s over 134 units; the rotation grid accounts for 512 CPU-s
(60 units) and the value-path grid for 246 CPU-s (12 units), against
5,132 CPU-s (90 units) and 1,500 CPU-s (24 units) in the last full
build before this change, which ran up to 94 concurrent jobs. The
critical path of a parallel build now runs through the fused DPA4
inference convolution units (up to 201 s).
Tests
Default build: test_so2_value_train.py (the two float64 cases skip),
the grid-pair and segment-softmax kernel tests, the DPA4 and SeZM
descriptor training-path, CUDA, Triton and vacuum-reference tests, and
the DPA4C, DPA1 and edge force/virial suites pass. With
DEEPMD_ENABLE_DPA4_FP64=ON, test_so2_value_train.py passes in full,
including the float64 cases at the 1e-12 tolerance.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (3)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughDPA4 training kernels now use shared accumulator types and dtype dispatch. A build option controls float64 kernel generation. The SO(2) operator reports whether float64 kernels are available, and CUDA workflows and tests check the enabled and default build configurations. ChangesDPA4 training kernels
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Change: Other Suggested reviewers: Merge Risk: ⚪ Minimal · up to The default training path remains float32/bfloat16, with float64 validation kernels available through an opt-in build setting. The supplied evidence establishes no current behavior or deployment failure, so the change appears mergeable. Security Architecture ReviewSecurity architecture risk: 🔵 Low · up to The changes are bounded to CUDA training operators and their build-time capabilities. The inspected public operators reject unsupported working dtypes before computation, and their supported callers preserve private-buffer ownership. No introduced security attack path was established, but incomplete tensor-contract and deployment coverage prevents a minimal-risk assessment. Retained concerns Security review detailsSecurity Blast Radius
Trust Boundaries and Controls
Resilience and Maintainability Implications
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 41.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 36 functions across 18 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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. Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6044 +/- ##
==========================================
- Coverage 77.85% 77.60% -0.26%
==========================================
Files 1170 1170
Lines 140912 140914 +2
Branches 5056 5056
==========================================
- Hits 109708 109356 -352
- Misses 29318 29675 +357
+ Partials 1886 1883 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed all 28 changed files: CMake-generated dtype grids, declaration/dispatch consistency, moved host validation, rotation/mixing accumulation through second order, resident-forward scalar precision, shared-memory staging, narrowed framework headers, capability query, tests, documentation and CUDA workflow configuration. The caller's float32 eligibility gate and bfloat16 autocast registrations were also checked.
No new high-confidence correctness blocker found in static review. The float64 CMake option consistently controls shard generation, extern declarations, dispatch and the capability query; double matrix multiplications retain their ATen fallback. The top-degree static shared-memory condition accounts for widened reduction scratch and avoids staging when it would exceed 48 KiB. The removed half dispatch is consistent with the supported fused training entry points.
Validation remains the main limitation. Although all seven exact-head workflow conclusions succeed, both real GPU jobs were skipped; only the Test CUDA aggregate job ran (https://github.com/deepmodeling/deepmd-kit/actions/runs/37214771171). The PR description also says source-branch numerical/build measurements were not rerun after cherry-picking onto this base. No CUDA compilation or numerical suite was run locally, so this review does not independently validate those measurements or the tightened 1e-12 bound. Please run the actual GPU jobs on this head before treating the change as validated: the Python CUDA job enables FP64, while the C++ CUDA installation exercises the default-off build.
Agent: dot
Reviewed head: c10b237
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Static review only: no CUDA toolchain was available to me, and no CI job on this head compiled the changed CUDA sources (see the inline note). What I checked by inspection:
- Instantiation coverage is consistent before and after. Rotation goes from 90 shards (f32, f64, bf16) to 60 by default and back to 90 with
DEEPMD_ENABLE_DPA4_FP64=ON; the value path goes from 24 (f32, f64, f16, bf16) to 12 by default and 18 with the option, f16 dropped for good. The hostexterndeclarations,DPA4_SEZM_DISPATCH_TYPESand the generated shards are gated by the same macro, which is a PRIVATE definition on the one target that owns both, so host and shards cannot disagree;check_value_inputsrejects the other dtypes before dispatch. f16 never worked end to end (the rotate kernels never dispatched Half, autocast is bf16 only,_is_supportedrequires float32 weights), so dropping it changes nothing reachable. - float32 and bfloat16 arithmetic is unchanged (
acc_typestays float,exp_a<float>keeps__expf);__launch_bounds__unchanged; everyCHECK_LAUNCHsite survives the move of validation into the host files (5 rotate, 6 mixing, 3 so2). - The float64 top-degree rotation backward no longer stages edge rows in shared memory; at L=6 rank 4 the staged layout would exceed 48 KiB, so that is the right call, and float32/bf16 still stage.
- The
--threads=2and--use_fast_mathCOMPILE_OPTIONSgo to disjoint file sets;torch/types.hplustorch/library.hcovers everytorch::symbol used in the narrowed hosts; the doc entry matches the option;CMAKE_ARGSis merged by scikit-build-core.
Two non-blocking notes inline: the new switch and API are tested on one side only, and the tightened float64 test has not run anywhere on this branch.
CI: 58 checks, 53 pass, 4 skipping (both CUDA test jobs, two release jobs), 1 fail (readthedocs, a toolchain error shared by other open PRs). The CPU run of the touched test file skips all 21 cases.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Reviewed delta c10b237 → 56d0764, including master dbca0b1 integration. No new correctness blocker found. All44 incoming master files preserve their exact master blobs and do not overlap the28 PR paths; new PR-specific changes are10 workflow lines and20 test lines.
The native rejection test disables autocast and checks the actual diagnostic: FP16 in both configurations, FP64 only when optional kernels are absent. Workflow steps assert a GPU, loaded native operators and the expected FP64 capability. The two earlier non-blocking review threads are addressed.
Independent validation is source/CI inspection, not a local CUDA build. The author reports fresh default-OFF21pass/2skip and FP64-ON22pass/1skip with numerical errors below unchanged1e-12. Actual CUDA run37717749700 now executes both GPU jobs: installation/capability and default dtype-contract steps succeeded, while full suites remain running. Earlier same-head runs either cancelled or skipped the GPU jobs and must not be counted as numerical passes. No CI activation or repository mutation was performed.
|
Agent: dot Actual CUDA validation has now completed successfully on unchanged head56d0764e9fec33141bf708042977f52442365903: https://github.com/deepmodeling/deepmd-kit/actions/runs/37717749700 . Both GPU jobs executed; this is not a skipped aggregate. Default-OFF C++ job113118072059 explicitly verified a GPU, loaded native operators and float64_available() == False. Both dtype-rejection tests passed; all4 CTests passed, followed by282 LAMMPS tests /18 skipped and8 further tests. FP64-ON Python job113118072251 explicitly verified the GPU/native operators and float64_available() == True. Its full suite finished with13,644 passed /8,696 skipped. The changed test_so2_value_train.py file executed22 passing cases and1 skip. The aggregate completed successfully at07:55:10 UTC. This supplies the previously missing hosted default-OFF/FP64-ON CUDA execution evidence. Skipped cases remain outside that coverage; no local GPU rerun or broader architecture/performance claim is implied. No CI activation, rerun or merge action was taken by this review. |
64d3266
Summary
Reduce the CUDA compilation cost of
libdeepmd_op_ptby removing unused DPA4 training instantiations, keeping framework headers out of kernel shards, and narrowing the headers included by operator hosts. The optional float64 validation path also gains double-precision accumulation throughout its forward and backward kernels.The supported production training paths remain float32 and bfloat16 autocast. This PR does not change model configuration, checkpoint formats, or the separate FP16x3 inference path.
Motivation
The PyTorch operator library dominates the installation cost of a CUDA build. Before this change, the DPA4 training operators alone generated 114 translation units:
{float32, float64, bfloat16};{float32, float64, float16, bfloat16}.Each shard parsed
<ATen/ATen.h>, although its device code did not need the framework API. Operator hosts also included the complete<torch/torch.h>C++ frontend. Those repeated parsing costs, together with unused dtype instantiations, made parallel compilation unnecessarily expensive.Implementation
Compile the training dtypes that are actually used
DEEPMD_ENABLE_DPA4_FP64, defaulting toOFF, to include float64 instantiations for numerical validation. The model binds the fused training path only for float32 parameters; autocast uses bfloat16.DPA4_SEZM_DISPATCH_TYPESconsistently across the training entries. Unsupported working dtypes are rejected by the value-path entry with an actionable error.deepmd::sezm_so2_value_supports_float64()and the Python helperfloat64_available()so validation can query the loaded library rather than infer its capabilities from the environment.The float64 consistency test skips when those kernels are absent. The CUDA Python CI job explicitly enables them, and the source-installation documentation describes the option. For a pip build:
The option is cached by CMake; an existing build tree configured with it enabled needs an explicit
-DDEEPMD_ENABLE_DPA4_FP64=OFFto disable it again.Use double accumulation in the optional float64 path
The training kernels use
acc_type<scalar_t>: float accumulation for float32/bfloat16 inputs and double accumulation for float64 inputs. This includes rotation and segment-sum reductions, the mixing traversal's pointwise kernels, and the resident forward/competition head's scalar parameters such as inverse temperature, label smoothing, and normalization epsilon.The float64 path therefore no longer contains the float32 arithmetic stages that previously limited its agreement with the eager reference. The source-branch measurements report relative discrepancies of approximately
1e-15to3e-11across the broader parity shapes, with the committed float64 test shapes around1e-14. The test threshold is tightened from5e-6to1e-12; a float32-rounded label-smoothing parameter alone previously produced approximately5e-10disagreement.The 384-thread, top-degree first-order rotation backward uses a static shared-memory-footprint condition to decide whether to stage edge rows. Its double-precision reduction scratch alone occupies 46,560 bytes, so staging those rows would exceed the 48 KiB static limit. Float32 and bfloat16 retain their existing staging.
Function-by-function SASS comparisons recorded on the source branch found unchanged code for representative float32/bfloat16 rotation and value-path shards and the three training hosts, except for
mixing_entry_bwd_kernel: it reuses the competition weight already loaded into a register instead of loading it again, with identical results.Keep kernel shards independent of ATen
rotate_mix_train/kernels.cuhandso2_conv_train/kernels.cuhcontain device code and launchers, without ATen includes.launch_so2_value_fwdreturns acudaError_t; the host reports shared-memory configuration errors and checks launch errors.sezm_types.cuhheader. Instantiation shards need only the lightweight scalar-type and CUDA headers, including<c10/util/BFloat16.h>.In the recorded measurements, including
<ATen/ATen.h>alone added approximately 33 seconds per translation unit.Narrow host headers and parallelize device passes
DPA4, DPA4C, DPA1, graph-fitting, and edge force/virial operator hosts and shared headers use
<torch/types.h>and<torch/library.h>where needed instead of<torch/torch.h>. This avoids parsing the unrelatednn,optim, data-loading, and serialization frontend.dpa1_graph_descriptor.cualso usesnvcc --threads=2, as the fused DPA4 convolution units already do, allowing its native-code and portable-PTX device passes to run concurrently.Build measurements
The following measurements were recorded on the source branch on a 192-core host, using CUDA 13.3 with
sm_120native code and acompute_75PTX fallback. Individual translation units were compiled in isolation:dpa1_graph_descriptor.cuThe narrowed framework-header set compiled in 34–38 seconds rather than 50–55 seconds, with peak compiler memory decreasing from approximately 2.6 GB to 1.5 GB.
A full default operator-library rebuild with 64 parallel jobs recorded 3,999 CPU-seconds across 134 translation units. The training-grid contributions were:
These are cumulative CPU times, not end-to-end elapsed times. The preceding full build used up to 94 concurrent jobs, so the aggregate figures are not a controlled wall-clock speedup comparison. The recorded critical path after this change is the fused DPA4 inference convolution units, at up to approximately 201 seconds per unit. No training or inference throughput improvement is claimed by these build measurements.
Validation
Validation completed on the source branch, as recorded in the source commit:
test_so2_value_train.pypasses with its two float64 cases skipped; grid-pair and segment-softmax kernel tests pass; DPA4/SeZM descriptor training-path, CUDA, Triton, and vacuum-reference tests pass; DPA4C, DPA1, and edge force/virial suites pass.DEEPMD_ENABLE_DPA4_FP64=ON:test_so2_value_train.pypasses in full, including the float64 cases at the tightened1e-12threshold.This PR cherry-picks that commit onto the latest
upstream/master. The suites and compilation benchmarks were not rerun during PR preparation; preparation checks cover patch equivalence, Ruff, formatting checks, andgit diff --check.Scope and compatibility
Summary by CodeRabbit
DEEPMD_ENABLE_DPA4_FP64=ON; they are disabled by default and increase build time.