deepmd.pt_expt.utils.serialization

deepmd.pt_expt.utils.serialization#

Attributes#

`PT2_EXTRA_PREFIX`
`_TRACE_SENDLIST_KEEPALIVE`

Functions#

`_strip_shape_assertions`(→ None)	Neutralise shape-guard assertion nodes in a spin model's exported graph.
`_numpy_to_json_serializable`(→ dict)	Convert numpy arrays in a model dict to JSON-serializable lists.
`_json_to_numpy`(→ dict)	Convert JSON-serialized numpy arrays back to np.ndarray.
`_metadata_value_to_json`(→ Any)
`_needs_with_comm_artifact`(→ bool)	Return `True` if the model needs a "with-comm" AOTI artifact compiled.
`_make_comm_sample_inputs`(→ tuple[torch.Tensor, Ellipsis])	Build trivial-but-valid comm tensors for tracing the with-comm variant.
`_make_sample_inputs`(→ tuple[torch.Tensor, Ellipsis])	Create sample inputs for tracing forward_lower.
`_build_dynamic_shapes`(→ tuple)	Build dynamic shape specifications for torch.export.
`_collect_metadata`(→ dict)	Collect metadata from the model for C++ inference.
`serialize_from_file`(→ dict)	Serialize a .pte or .pt2 model file to a dictionary.
`_serialize_from_file_pte`(→ dict)	Serialize a .pte model file to a dictionary.
`_serialize_from_file_pt2`(→ dict)	Serialize a .pt2 model file to a dictionary.
`deserialize_to_file`(→ None)	Deserialize a dictionary to a .pte or .pt2 model file.
`_trace_and_export`(→ tuple)	Common logic: build model, trace, export.
`_deserialize_to_file_pte`(→ None)	Deserialize a dictionary to a .pte model file.
`_deserialize_to_file_pt2`(→ None)	Deserialize a dictionary to a .pt2 model file (AOTInductor).

Module Contents#

deepmd.pt_expt.utils.serialization.PT2_EXTRA_PREFIX = 'model/extra/'[source]#

deepmd.pt_expt.utils.serialization._strip_shape_assertions(graph_module: torch.nn.Module) → None[source]#

Neutralise shape-guard assertion nodes in a spin model’s exported graph.

torch.export inserts aten._assert_scalar nodes for symbolic shape relationships discovered during tracing. For the spin model, the atom- doubling logic creates slice patterns that depend on (nall - nloc), producing guards like Ne(nall, nloc). These guards are spurious: the model computes correct results even when nall == nloc (NoPBC, no ghost atoms).

This function is only called for spin models (guarded by if is_spin in _trace_and_export). The assertion messages use opaque symbolic variable names (e.g. Ne(s22, s96)) rather than human-readable names, so filtering by message content is not reliable. Since prefer_deferred_runtime_asserts_over_guards=True converts all shape guards into these deferred assertions, and the only shape relationships in the spin model involve nall/nloc, neutralising all of them is safe in this context.

We replace each assertion’s condition with True rather than erasing the node; erasing nodes can disturb the FX graph structure and produce NaN gradients on some Python/torch versions.

deepmd.pt_expt.utils.serialization._numpy_to_json_serializable(model_obj: dict) → dict[source]#: Convert numpy arrays in a model dict to JSON-serializable lists.

deepmd.pt_expt.utils.serialization._json_to_numpy(model_obj: dict) → dict[source]#: Convert JSON-serialized numpy arrays back to np.ndarray.

deepmd.pt_expt.utils.serialization._metadata_value_to_json(value: Any) → Any[source]#

deepmd.pt_expt.utils.serialization._needs_with_comm_artifact(model: torch.nn.Module) → bool[source]#

Return True if the model needs a “with-comm” AOTI artifact compiled.

The with-comm artifact carries the per-layer deepmd_export::border_op calls that exchange node-embedding tensors across MPI ranks. Multi-rank LAMMPS dispatches to it when the descriptor’s message passing extends across rank boundaries (i.e. layers consume neighbour features that live on a different rank). Non-GNN descriptors and GNN descriptors with use_loc_mapping=True keep all per-layer messaging local to each rank’s owned atoms; they need only the regular artifact.

Delegates to descriptor.has_message_passing_across_ranks(), which descriptor classes implement explicitly. Returns False defensively when the model has no single descriptor (linear/zbl/frozen) or when the method is somehow missing or raises.

deepmd.pt_expt.utils.serialization._TRACE_SENDLIST_KEEPALIVE: list[numpy.ndarray] = [][source]#

deepmd.pt_expt.utils.serialization._make_comm_sample_inputs(nloc: int, nghost: int, device: torch.device) → tuple[torch.Tensor, Ellipsis][source]#

Build trivial-but-valid comm tensors for tracing the with-comm variant.

Phase 0 finding: tracing with nswap == 0 causes the dim to specialize, so we must use nswap >= 1. We use nswap == 1 with a single self-send swap whose sendlist points to nghost local atoms (the actual indices don’t matter for the trace — only the validity of the pointer matters; border_op is opaque to torch.export via the deepmd_export::border_op wrapper).

Returns (send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal_ts, nghost_ts) — 8 tensors, matching the canonical positional order of forward_common_lower_exportable_with_comm.

deepmd.pt_expt.utils.serialization._make_sample_inputs(model: torch.nn.Module, nframes: int = 1, nloc: int = 7, has_spin: bool = False) → tuple[torch.Tensor, Ellipsis][source]#

Create sample inputs for tracing forward_lower.

Parameters:

modeltorch.nn.Module: The pt_expt model (must have get_rcut, get_sel, get_type_map, etc.).
nframesint: Number of frames.
nlocint: Number of local atoms.
has_spinbool: If True, create an extended spin tensor and return 7 tensors.

Returns:

tuple: (ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin) or (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam, charge_spin) when has_spin.

deepmd.pt_expt.utils.serialization._build_dynamic_shapes(*sample_inputs: torch.Tensor | None, has_spin: bool = False, with_comm_dict: bool = False, model_nnei: int = 1) → tuple[source]#

Build dynamic shape specifications for torch.export.

Marks nframes, nloc, nall and nnei as dynamic dimensions so the exported program handles arbitrary frame, atom and neighbor counts.

When with_comm_dict is True, 8 additional comm tensors are appended to the returned tuple — matching the positional order of forward_common_lower_exportable_with_comm. nswap is the only dynamic dim among them; the rest are scalar or fixed-size.

Parameters:

*sample_inputstorch.Tensor | None: Sample inputs: 6 tensors (non-spin) or 7 (spin), optionally followed by 8 comm tensors when with_comm_dict.
has_spinbool: Whether the inputs include an extended_spin tensor.
with_comm_dictbool: Whether the inputs include the 8 comm tensors.
model_nneiint: The model’s sum(sel). Used as the min for the dynamic nnei dim.
Returns a tuple (not dict) to match positional args of the make_fx
traced module, whose arg names may have suffixes like ``_1``.

deepmd.pt_expt.utils.serialization._collect_metadata(model: torch.nn.Module, is_spin: bool = False) → dict[source]#

Collect metadata from the model for C++ inference.

This metadata is stored as metadata.json in both .pt2 and .pte archives. Training config is stored separately in model_def_script.json. C++ reads flat JSON fields because compiling model API methods as AOTInductor entry points is impractical (~12 s per trivial function) and string outputs (get_type_map) cannot be expressed as tensor I/O.

The fitting_output_defs list is also included so that ModelOutputDef can be reconstructed without loading the full model.

deepmd.pt_expt.utils.serialization.serialize_from_file(model_file: str) → dict[source]#

Serialize a .pte or .pt2 model file to a dictionary.

Reads the model dict stored in the model archive.

Parameters:

model_filestr: The model file to be serialized (.pte or .pt2).

Returns:

dict: The serialized model data. If the archive contains model_def_script.json (training config), it is included under the "model_def_script" key.

deepmd.pt_expt.utils.serialization._serialize_from_file_pte(model_file: str) → dict[source]#: Serialize a .pte model file to a dictionary.

deepmd.pt_expt.utils.serialization._serialize_from_file_pt2(model_file: str) → dict[source]#

Serialize a .pt2 model file to a dictionary.

Reads the model dict stored in the model/extra/ directory of the .pt2 ZIP archive.

deepmd.pt_expt.utils.serialization.deserialize_to_file(model_file: str, data: dict, model_json_override: dict | None = None, do_atomic_virial: bool = False) → None[source]#

Deserialize a dictionary to a .pte or .pt2 model file.

Builds a pt_expt model from the dict, traces it via make_fx, exports with dynamic shapes, and saves.

Parameters:

model_filestr: The model file to be saved (.pte or .pt2).
datadict: The dictionary to be deserialized (same format as dpmodel’s serialize output, with “model” and optionally “model_def_script” keys). If data["model_def_script"] is present, it is embedded in the output so that --use-pretrain-script can extract descriptor/fitting params at finetune time.
model_json_overridedict or None: If provided, this dict is stored in model.json instead of data. Used by dp compress to store the compressed model dict while tracing the uncompressed model (make_fx cannot trace custom ops).
do_atomic_virialbool: If True, export with per-atom virial correction (3 extra backward passes, ~2.5x slower). Default False for best performance.

deepmd.pt_expt.utils.serialization._trace_and_export(data: dict, model_json_override: dict | None = None, with_comm_dict: bool = False, do_atomic_virial: bool = False) → tuple[source]#

Common logic: build model, trace, export.

Parameters:

data: Serialized model dict (with “model” and optionally “model_def_script” keys).
model_json_override: Optional alternate dict to embed as model.json (used by dp compress to store the compressed model dict while tracing the uncompressed one).
with_comm_dict: If True, trace forward_common_lower_exportable_with_comm instead of the regular variant. The resulting exported program accepts 8 additional positional comm tensors (send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost) used by the pt_expt Repflow/Repformer override to drive MPI ghost-atom exchange. Only valid for models that need cross-rank ghost-feature exchange (see _needs_with_comm_artifact).
do_atomic_virial: If True, the traced graph computes per-atom virial (extra autograd.grad backward passes); off by default to keep .pt2 inference fast. Mirrors PR #5407 in upstream master.

Returns:

tuple: (exported, metadata, data_for_json, output_keys).

deepmd.pt_expt.utils.serialization._deserialize_to_file_pte(model_file: str, data: dict, model_json_override: dict | None = None, do_atomic_virial: bool = False) → None[source]#: Deserialize a dictionary to a .pte model file.

deepmd.pt_expt.utils.serialization._deserialize_to_file_pt2(model_file: str, data: dict, model_json_override: dict | None = None, do_atomic_virial: bool = False) → None[source]#

Deserialize a dictionary to a .pt2 model file (AOTInductor).

Uses torch._inductor.aoti_compile_and_package to compile the exported program into a .pt2 package (ZIP archive with compiled shared libraries), then embeds metadata into the archive.

For models whose descriptor reports has_message_passing_across_ranks() == True (DPA2, DPA3 with use_loc_mapping=False, or hybrids wrapping such children), compiles a SECOND with-comm artifact and packs it alongside the regular one. The with-comm variant accepts comm-dict tensors as additional positional inputs and drives MPI ghost-atom exchange via deepmd_export::border_op. The C++ DeepPotPTExpt loader picks the artifact based on the LAMMPS rank count at runtime.

Layout inside the .pt2 ZIP (PyTorch 2.11 strict layout):: regular → artifact at model/ (AOTInductor’s own layout) with-comm → model/extra/forward_lower_with_comm.pt2 (nested ZIP) metadata → model/extra/metadata.json with

has_comm_artifact flag. The C++ reader matches by /-delimited suffix so the legacy root-level extra/ layout still loads.

Old .pt2 files (pre-this-change) lack has_comm_artifact so the C++ loader must default to False when the field is missing.