deepmodeling / deepmodeling/deepmd-kit
pt_expt: freeze-time opt-out for the with-comm AOTI artifact (halve .pt2 compile time for MP models)
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
Every `.pt2` freeze of a message-passing model (DPA-2, DPA-3 with `use_loc_mapping=False`, hybrids wrapping such children) compiles **two** AOTInductor artifacts: the regular lower plus the nested `model/extra/forward_lower_with_comm.pt2` (`deepmd/pt_expt/utils/serialization.py`, `_deserialize_to_file_pt2`). For DPA-2 the second compile is unconditional — `DescrptDPA2.has_message_passing_across_ranks()` is always `True` (repformers pass `g1` in `nall` layout with no block-level `use_loc_mapping` opt-out) — so single-rank users and virtually every unit test pay roughly 2× the AOTI compile time for an artifact that is only ever loaded by multi-rank LAMMPS.
## Proposal
Add a freeze-time opt-out for the with-comm artifact:
- `deserialize_to_file(..., with_comm="auto" | "never" | "always")` (default `"auto"` = current behavior, compile when `_needs_with_comm_artifact(model)`).
- Surface as `dp freeze --no-comm-artifact` (or equivalent).
- `"never"`: skip the second trace+compile, set `has_comm_artifact: false` in `metadata.json`.
This is safe by design: the C++ `DeepPotPTExpt` loader already fail-fasts with a clear "re-freeze" error when a multi-rank run meets a `.pt2` lacking the artifact, so an opted-out model degrades loudly, not silently.
Applies to both with-comm flavors (dense and graph lower); the knob belongs to the shared freeze machinery, not to any one descriptor.
## Tests (both branches)
- `"never"`: zip entry absent, `has_comm_artifact` false, multi-rank loader fail-fast message fires.
- `"always"` / `"auto"`: entry present and loadable (existing with-comm tests keep covering the artifact itself).
Coverage note: intentional with-comm coverage lives in the dedicated tests that load and run the artifact (`test_graph_with_comm_export.py`, C++ multi-rank gtests, LAMMPS MPI runs) — those keep compiling it. The incidental compiles being removed only verified "the with-comm trace didn't crash for this config"; keeping one with-comm compile per trace-distinct config family (plain / fparam+aparam / spin) preserves that smoke signal. The dedicated with-comm tests must stay in the per-PR CI lane.
## Follow-up (separate change)
Test-suite compile-time cleanup that stacks with this knob:
- Flip incidental dpa2/dpa3 test freezes to `"never"`.
- `compile` / `aot` / `ddp` pytest markers + dedicated CI lane, or `duration_based_chunks` splitting so pytest-split stops scattering class methods across shards and re-running expensive AOT `setUpClass`.
- Persist `TORCHINDUCTOR_CACHE_DIR` across CI runs (measured ~44 s cold vs ~12 s warm per compile).
- Session-scoped frozen-artifact cache keyed by model-config hash; `SetUpTestSuite()` for C++ per-test `.pt2` loads.
## Sequencing
Blocked on #5779 (reworks the same `_deserialize_to_file_pt2` / `_trace_and_export` region); branch off master after it merges.
Contributor guide
Assessment
This issue has not been assessed yet.