[Q&A] dp --pt freeze fails for DPA4 with RuntimeError: Unknown type name 'EdgeFeatureCache' #5557

Thanks a lot for the detailed explanation and suggestion. Following your advice, I moved

from .edge_cache import EdgeFeatureCache

out of the if TYPE_CHECKING: block in deepmd/pt/model/descriptor/sezm_nn/embedding.py. That fixed the original Unknown type name 'EdgeFeatureCache' error.

However, after that, dp --pt freeze hit a chain of further TorchScript-compatibility problems in the sezm_nn (DPA4) module. I worked through them one by one with the help of Claude just to keep moving forward. Below is the full sequence of errors and the (temporary, local) changes I made. I think this is useful information for a proper maintainer-side fix.

All of this is with DeePMD-kit 3.2.0b1.dev9+gbe6af6740, installed from source from the master branch (PyTorch backend, Python 3.10, CUDA 12.8). Training completes fine; only freezing fails.

Blocker 1 — EdgeFeatureCache hidden behind TYPE_CHECKING (the original issue)

Error:

RuntimeError: 
Unknown type name 'EdgeFeatureCache':
 File ".../deepmd/pt/model/descriptor/sezm_nn/embedding.py", line 463
 edge_cache: EdgeFeatureCache,
 ~~~~~~~~~~~~~~~~ <--- HERE

Fix (as suggested): move the import out of TYPE_CHECKING in embedding.py:

from .edge_cache import EdgeFeatureCache

→ This error is resolved.

Blocker 2 — EdgeFeatureCache NamedTuple uses X | None (ForwardRef)

Next error:

ValueError: Unknown type annotation: 'ForwardRef('torch.Tensor | None')' in NamedTuple EdgeFeatureCache.
Likely due to partial support for ForwardRef parameters in NamedTuples, see #95858.
 File ".../deepmd/pt/model/descriptor/sezm_nn/embedding.py", line 463
 edge_cache: EdgeFeatureCache,

Cause: edge_cache.py has from __future__ import annotations, so the EdgeFeatureCache NamedTuple field annotations become ForwardRef strings, and TorchScript cannot resolve PEP-604 X | None / builtin dict[...] forms (pytorch/pytorch#95858).

Fix in edge_cache.py — switch the NamedTuple optional fields from X | None / dict[...] | None to Optional[...] / Optional[Dict[...]]:

from typing import (
 Dict,
 NamedTuple,
 Optional,
)
class EdgeFeatureCache(NamedTuple):
 ...
 D_full: Optional[torch.Tensor] = None
 Dt_full: Optional[torch.Tensor] = None
 D_to_m_cache: Optional[Dict[str, torch.Tensor]] = None
 Dt_from_m_cache: Optional[Dict[str, torch.Tensor]] = None
 edge_src_gate: Optional[torch.Tensor] = None
 edge_quat: Optional[torch.Tensor] = None

→ This error is resolved.

Blocker 3 — Variable-length list unpacking in view(...) (norm.py)

Next error:

RuntimeError: 
cannot statically infer the expected size of a list in this context:
 scale = self.adam_scale.view(*([1] * (x.ndim - 1)), self.channels)
 ~~~~~~~~~~~~~~~~~~~ <--- HERE

Cause: view(*([1] * (x.ndim - 1)), self.channels) unpacks a list whose length is not statically known, which TorchScript cannot handle.

Fix in deepmd/pt/model/descriptor/sezm_nn/norm.py — build an explicit list[int] shape instead of star-unpacking:

shape: list[int] = [1] * (x.ndim - 1) + [self.channels]
scale = self.adam_scale.view(shape)

→ This error is resolved.

Blocker 4 — nvtx_range is a @contextmanager generator (yield)

Next error:

torch.jit.frontend.UnsupportedNodeError: Yield aren't supported:
 File ".../contextlib.py", line 106
 with nvtx.range(name):
 yield

Cause: nvtx_range in deepmd/pt/model/descriptor/sezm_nn/utils.py is a @contextlib.contextmanager generator, and it is used as with nvtx_range(...): throughout the scripted forward paths. TorchScript cannot compile yield.

Blocker 5 / 6 — making nvtx_range a TorchScript-compatible no-op context manager

My first attempt (returning contextlib.nullcontext() while still decorated with @contextmanager) failed at runtime:

TypeError: 'nullcontext' object is not an iterator

and, after removing @contextmanager, returning nullcontext() failed under torch.jit.script:

RuntimeError: Expected <function nullcontext.__aenter__ ...> to be a function

(nullcontext carries __aenter__/__aexit__, which TorchScript chokes on.)

Then I replaced nvtx_range with a custom synchronous no-op context manager, which initially failed with:

RuntimeError: 
argument 1 of __exit__ must have Any type; TorchScript does not currently support
passing exception type, value, or traceback to the __exit__ function.:
 File ".../sezm_nn/wignerd.py", line 508
 with nvtx_range("WignerD/l1"):

Final working fix in utils.py — a custom context manager whose __exit__ arguments are annotated Any:

from typing import Any
class _NvtxNoOp:
 def __init__(self) -> None:
 pass
 def __enter__(self) -> None:
 return None
 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
 return None
def nvtx_range(name: str) -> "_NvtxNoOp":
 return _NvtxNoOp()

→ All with nvtx_range(...) sites now compile. (Note: this disables NVTX profiling, which is irrelevant for freeze/inference. A maintainer-side fix that guards profiling with torch.jit.is_scripting() would be cleaner.)

Blocker 7 — conditionally-registered attribute exp_l4 referenced unconditionally (current blocker)

Current error:

RuntimeError: 
Module 'WignerSmallOrderCoefficients' has no attribute 'exp_l4' :
 File ".../sezm_nn/wignerd.py", line 1102
 monomials = self._build_monomial_matrix(
 powers,
 self.small_order_kernels.exp_l4,
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
 )
'WignerDCalculator._compute_l3l4_blocks' is being compiled since it was called from 'WignerDCalculator.forward'
 File ".../sezm_nn/wignerd.py", line 518
 if self.lmax >= 4:
 with nvtx_range("WignerD/l3l4"):
 D_l3, D_l4 = self._compute_l3l4_blocks(edge_quaternion)

Cause: _compute_l3l4_blocks references self.small_order_kernels.exp_l4, which appears to be registered only when lmax >= 4. In eager mode the if self.lmax >= 4: guard means this branch is never entered for my model (lmax < 4), so it works. But torch.jit.script statically compiles the whole method body regardless of the runtime lmax, and requires the attribute to exist.

I have not patched this one, because a proper fix likely needs maintainer knowledge of how WignerSmallOrderCoefficients registers exp_l1..exp_l4 (e.g. always registering the buffers, or guarding _compute_l3l4_blocks with @torch.jit.unused). I'd rather not introduce a workaround that makes freeze "succeed" but produce an incorrect frozen model.

I don't know how the official DPA4 dp --pt freeze command works. I also tried the v3.2.0b0 version from the releases. This version doesn’t seem to include any SO3-related parameters, and it’s older than the master branch. So I can’t use it to freeze the model.ckpt.pt file generated by dp using the source code from the master branch. I’m feeling sad right now.
Anyway, thanks again for your reply. Have a good day.

If I don't freeze it, I won't be able to run MD in LAMMPS. I also won't be able to continue using dpgen for iterations. I hope the developers can fix this incompatibility as soon as possible. Thank you.