[BUG] Property models trained before #5322 give shifted predictions after upgrade: distinguish_types defaults to True when a .pt checkpoint is loaded #6082

Summary

A property-fitting model (intensive=true) saved before PR #5322 (merge commit 7cd0d191b80e) has no distinguish_types key in its model_params. On newer versions it is loaded with distinguish_types=True, so out_bias is applied per atom type instead of out_bias[0] for every atom. Whenever the stored per-type biases differ, each prediction is offset by a composition-dependent constant:

pred_new - pred_old = mean_over_atoms(out_bias[atype]) - out_bias[0]

The network weights and stored statistics are identical. Only how out_bias is applied has changed. .pt checkpoints are rebuilt through the constructor path, where distinguish_types defaults to True. The deserialize() path falls back to False instead.

Versions

Minimal reproduction (no external data or checkpoint)

Everything below runs on CPU in about a minute. Steps 1-3 use the old version, step 4 runs in both.

1. Synthetic 3-type data (mkdata.py):

import numpy as np, os
rng = np.random.default_rng(0)
for name, n in [("train", 8), ("test", 4)]:
    d = f"{name}/set.000"; os.makedirs(d, exist_ok=True)
    atype = np.array([0, 0, 0, 0, 1, 1, 2, 2])
    np.savetxt(f"{name}/type.raw", atype, fmt="%d")
    open(f"{name}/type_map.raw", "w").write("A\nB\nC\n")
    L = 6.0
    np.save(f"{d}/box.npy", np.tile(np.eye(3).ravel() * L, (n, 1)))
    np.save(f"{d}/coord.npy", rng.random((n, 24)) * L)
    np.save(f"{d}/property.npy", rng.normal(5, 1, (n, 1)))

2. Train a tiny property model for 2 steps with the old version: dp --pt train input.json

{"model":{"type_map":["A","B","C"],
 "descriptor":{"type":"se_e2_a","sel":[20,20,20],"rcut":4.0,"rcut_smth":3.5,"neuron":[4,8],"axis_neuron":4,"seed":1},
 "fitting_net":{"type":"property","property_name":"property","task_dim":1,"intensive":true,"neuron":[8],"seed":1}},
 "learning_rate":{"type":"exp","start_lr":1e-3,"stop_lr":1e-4,"decay_steps":10},
 "loss":{"type":"property"},
 "training":{"training_data":{"systems":["train"],"batch_size":2},"numb_steps":2,"disp_freq":1,"save_freq":2}}

3. Give the checkpoint unequal per-type biases, as an old checkpoint can carry (prep.py). A copy with the workaround key is also written:

import torch
sd = torch.load("model.ckpt.pt", map_location="cpu", weights_only=False)
m = sd["model"]
fp = m["_extra_state"]["model_params"]["fitting_net"]
print("distinguish_types in params:", "distinguish_types" in fp)   # False
k = "model.Default.atomic_model.out_bias"
m[k] = torch.tensor([-1.0, 2.0, 3.0], dtype=m[k].dtype).reshape(m[k].shape)
torch.save(sd, "model_bias.pt")
fp["distinguish_types"] = False
torch.save(sd, "model_bias_fix.pt")

4. Evaluate in both versions (check.py <ckpt>):

import sys, torch, numpy as np, deepmd
from deepmd.pt.infer.deep_eval import DeepProperty
ckpt = sys.argv[1]
sd = torch.load(ckpt, map_location="cpu", weights_only=False)["model"]
ob = sd["model.Default.atomic_model.out_bias"][0, :, 0].double()
coord = np.load("test/set.000/coord.npy"); cell = np.load("test/set.000/box.npy")
atype = np.loadtxt("test/type.raw", dtype=int)
pred = DeepProperty(ckpt).eval(coord, cell, atype)[0].ravel()
t, c = np.unique(atype, return_counts=True)
shift = float((ob[t] * torch.tensor(c)).sum() / c.sum() - ob[0])
print(deepmd.__version__, ckpt, np.round(pred, 6).tolist(), "expected shift:", round(shift, 6))

Observed output

Every frame shifts by exactly +1.750000, which equals the expected 0.5*(-1.0) + 0.25*2.0 + 0.25*3.0 - (-1.0) = 1.75. With the key set to false the new version matches the old version to six decimals.

Expected behavior

A checkpoint saved before #5322 should give the same predictions after upgrading, i.e. it should load with distinguish_types=False.

Root cause

• deepmd/pt/model/atomic_model/property_atomic_model.py, apply_out_stat:
  • Before feat(property): Support element type-wise bias in property fitting #5322 (b494a0d): ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]
  • After feat(property): Support element type-wise bias in property fitting #5322 (27a18b604): if self.get_compute_stats_distinguish_types(), ret[kk] = ret[kk] * out_std[kk][atype] + out_bias[kk][atype], otherwise the old rule.
• deepmd/pt/model/task/property.py:
  • L96: PropertyFittingNet.__init__(..., distinguish_types: bool = True)
  • L148: deserialize() uses data.setdefault("distinguish_types", False) for backward compatibility.
• Loading a .pt checkpoint rebuilds the model from model_params via get_model and the constructor, so the deserialize() fallback never runs and old checkpoints silently get True. The argcheck default (True) may have the same effect when params are normalized.

Suggested fix

1. When distinguish_types is missing from a loaded checkpoint's model_params, default it to False, matching deserialize().
2. Warn when an old property checkpoint without the key is loaded, naming the assumed value.
3. Add a regression test: a property checkpoint saved without distinguish_types must reproduce its pre-feat(property): Support element type-wise bias in property fitting #5322 predictions.

Workaround

Add "distinguish_types": false to model["_extra_state"]["model_params"]["fitting_net"] in a copy of the checkpoint (step 3 above), or pin a version before #5322.