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

[Zmcmc123]

Question

Training of a DPA4 model finishes successfully, but the freeze step fails. During dp --pt freeze, torch.jit.script(model) raises RuntimeError: Unknown type name 'EdgeFeatureCache' while compiling deepmd/pt/model/descriptor/sezm_nn/embedding.py.

DeePMD-kit Version

3.2.0b1.dev9+gbe6af6740

Backend and its version

PyTorch v2.8.0+cu128-ga1cb3cc05d4

Python Version, CUDA Version, GCC Version, LAMMPS Version, etc

Python 3.10; CUDA 12.8; Installed via conda env dp-master, built from source.

Details

I trained DPA4 model via dpgen and deepmd-kit. The model trains to completion without any issue. The error only occurs during freezing. Then I tried to freeze the model manually, and got the same error.

Via dpgen, which calls:
dp --pt freeze

Manually submitting a freeze job:
dp --pt freeze -c model.ckpt.pt -o frozen_model

Error log (identical for both):

[2026-06-19 14:41:47,279] DEEPMD INFO DeePMD version: 3.2.0b1.dev9+gbe6af6740
Traceback (most recent call last):
File ".../bin/dp", line 6, in
sys.exit(main())
File ".../deepmd/main.py", line 1061, in main
deepmd_main(args)
File ".../torch/distributed/elastic/multiprocessing/errors/init.py", line 357, in wrapper
return f(*args, **kwargs)
File ".../deepmd/pt/entrypoints/main.py", line 670, in main
freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
File ".../deepmd/pt/entrypoints/main.py", line 480, in freeze
model = torch.jit.script(model)
File ".../torch/jit/_script.py", line 1443, in script
ret = _script_impl(
...
File ".../torch/jit/_recursive.py", line 466, in create_methods_and_properties_from_stubs
concrete_type._create_methods_and_properties(
RuntimeError:
Unknown type name 'EdgeFeatureCache':
File ".../deepmd/pt/model/descriptor/sezm_nn/embedding.py", line 463
self,
*,
edge_cache: EdgeFeatureCache,
~~~~~~~~~~~~~~~~ <--- HERE
atype_flat: torch.Tensor,
n_nodes: int,

Reproducible Example, Input Files, and Commands

input.json
train.log

Further Information, Files, and Links

No response

[krotname]

This looks like a TorchScript/source bug in the current dev tree, not a problem with your input file or the trained checkpoint.

In deepmd/pt/model/descriptor/sezm_nn/embedding.py, the module has:

from __future__ import annotations
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from .edge_cache import EdgeFeatureCache

but the scripted forward() signatures use EdgeFeatureCache at runtime, for example:

edge_cache: EdgeFeatureCache

During normal eager training this can work because the annotation is mainly for type checkers. During dp --pt freeze, torch.jit.script(model) compiles the model and tries to resolve that type name. Since EdgeFeatureCache was imported only under TYPE_CHECKING, it is not in the runtime namespace, so TorchScript raises exactly:

RuntimeError: Unknown type name 'EdgeFeatureCache'

A maintainer-side fix is likely to make EdgeFeatureCache available at runtime in embedding.py, e.g. move the import out of the TYPE_CHECKING block:

from .edge_cache import EdgeFeatureCache

If that creates a circular import, the alternative is to remove/loosen that annotation on scripted paths, but the current state is definitely incompatible with TorchScript freeze.

For your local source build, I would test the minimal runtime import first, then rerun only:

dp --pt freeze -c model.ckpt.pt -o frozen_model

If it then moves to a different TorchScript error, that would be the next blocker, but this first error is caused by the annotation being hidden behind TYPE_CHECKING.

[Zmcmc123]

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.

[njzjz-bot]

Thanks for the detailed follow-up. This is very useful, and I agree with your conclusion: please do not keep patching exp_l4 locally just to make torch.jit.script pass, because at that point we may accidentally produce a frozen model whose DPA4/SO(3) path is not semantically correct.

From the traceback sequence, this is a DeePMD-kit TorchScript compatibility issue in the current master DPA4/sezm_nn implementation, not a problem with your input data or checkpoint. The first fixes you tried (runtime import of EdgeFeatureCache, TorchScript-friendly NamedTuple annotations, static view shape, and replacing/guarding the NVTX context manager during scripting) are all in the right category. The exp_l4 error needs a maintainer-side fix in WignerSmallOrderCoefficients/WignerDCalculator, e.g. by making the registered-buffer set TorchScript-stable or by splitting/marking unused paths in a way that preserves correctness for lower lmax.

Also, you are right that v3.2.0b0 is not a good workaround for a checkpoint produced by current master; the DPA4/SO3 code has changed, so freezing a master checkpoint with the older release is not expected to work.

I will treat this as a DPA4 freeze bug in the current development branch. For now, the safe recommendation is:

• use the same source tree for training and freezing;
• avoid relying on a locally monkey-patched frozen model until we have a tested upstream fix;
• if you only need to continue validation in Python, use the checkpoint/eager path rather than dp --pt freeze.

Thanks again for reducing the issue to the exact TorchScript blockers. That saves us a lot of debugging time, and sorry for the frustration here.

— OpenClaw 2026.6.8 (844f405), model: custom-chat-jinzhezeng-group/gpt-5.5

[Zmcmc123]

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.