fix(pt_expt): skip neighbor statistics for graph-native DPA4 models

.github/workflows/test_python.yml

@@ -34,7 +34,9 @@ jobs:
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           source/install/uv_with_retry.sh pip install --system -e .[test,jax,torch] mpi4py --group pin_jax_cpu
-          source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn paddlepaddle==3.4.0.dev20260310
+          # Fetch the pinned Linux x86_64 wheel directly: the nightly index can return HTTP 502.
+          paddle_python_tag=$(python -c 'import sys; print(f"cp{sys.version_info.major}{sys.version_info.minor}")')
+          source/install/uv_with_retry.sh pip install --system --index-url https://pypi.org/simple "https://paddle-whl.bj.bcebos.com/nightly/cpu/paddlepaddle/paddlepaddle-3.4.0.dev20260310-${paddle_python_tag}-${paddle_python_tag}-linux_x86_64.whl"
         env:
           # Please note that uv has some issues with finding
           # existing TensorFlow package. Currently, it uses

deepmd/dpmodel/model/model_factory.py

@@ -13,6 +13,7 @@
 from deepmd.utils.bridging import (
     resolve_bridging_window,
     route_canonical_learned_options,
+    validate_tabulated_zbl,
 )
 from deepmd.utils.spin import (
     Spin,
@@ -188,6 +189,7 @@ def get_zbl_model(
     backend_name: str,
 ) -> Any:
     """Construct a ZBL model from backend-native atomic model classes."""
+    validate_tabulated_zbl(data)
     if data.get("long_range") is not None:
         raise NotImplementedError(
             "a long-range module is not supported with the tabulated ZBL "
@@ -404,6 +406,7 @@ def get_linear_atomic_model(
                 "child's type_map to match the composition type_map."
             )
         if "descriptor" in sub:
+            validate_tabulated_zbl(sub)
             child = None
             if descriptor_child_builder is not None:
                 child = descriptor_child_builder(sub)
@@ -552,6 +555,7 @@ def get_model(
     Any
         The backend-native model instance.
     """
+    validate_tabulated_zbl(data)
     model_type = data.get("type", "standard")
     if model_type == "standard":
         if "spin" in data:

deepmd/pt/model/model/__init__.py

@@ -43,6 +43,7 @@
     expand_bridging_method,
     resolve_bridging_window,
     route_canonical_learned_options,
+    validate_tabulated_zbl,
 )
 from deepmd.utils.spin import (
     Spin,
@@ -251,6 +252,7 @@ def get_linear_model(model_params: dict) -> BaseModel:
         if "type_map" not in sub_model_params:
             sub_model_params["type_map"] = model_params["type_map"]
         if "descriptor" in sub_model_params:
+            validate_tabulated_zbl(sub_model_params)
             # descriptor
             sub_ntypes = len(sub_model_params["type_map"])
             sub_model_params["descriptor"]["ntypes"] = sub_ntypes
@@ -362,6 +364,7 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel:
             "match the bridged linear_ener composition's `type_map`."
         )
     route_canonical_learned_options(model_params, learned)
+    validate_tabulated_zbl(learned)
     inner_cfg = inner_cfgs[0]
     learned["type"] = "dpa4"
     learned["type_map"] = copy.deepcopy(model_params["type_map"])
@@ -383,6 +386,7 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel:
 
 
 def get_zbl_model(model_params: dict) -> DPZBLModel:
+    validate_tabulated_zbl(model_params)
     _reject_long_range(model_params)
     model_params = copy.deepcopy(model_params)
     ntypes = len(model_params["type_map"])

deepmd/pt_expt/entrypoints/main.py

@@ -320,16 +320,51 @@ def update_neighbor_stat(
         multi_task: bool,
     ) -> tuple[dict[str, Any], float | dict[str, float | None] | None]:
         """Update pt_expt descriptor selections from neighbor statistics."""
-        log.info(
-            "Calculate neighbor statistics... "
-            "(add --skip-neighbor-stat to skip this step)"
-        )
         from deepmd.pt_expt.model import (
             BaseModel,
         )
+        from deepmd.pt_expt.model.graph_lower import (
+            model_config_skips_neighbor_stat,
+            resolve_stat_free_selections,
+        )
+
+        # Restored descriptors may explicitly disable the graph route; config
+        # alone cannot determine their effective checkpoint capability.
+        restoring = any(
+            (
+                options.restart,
+                options.init_model,
+                options.init_frz_model,
+                options.finetune,
+            )
+        )
 
         min_nbor_dist: dict[str, float | None] = {}
         for task_config in iter_training_task_configs(config):
+            if not restoring and model_config_skips_neighbor_stat(
+                task_config.model_params
+            ):
+                log.info(
+                    "Skip neighbor statistics for graph-native task %s: "
+                    "its forward carries every neighbor within the cutoff.",
+                    task_config.key,
+                )
+                updated_model_params = resolve_stat_free_selections(
+                    task_config.model_params
+                )
+                if multi_task:
+                    config["model"]["model_dict"][task_config.key] = (
+                        updated_model_params
+                    )
+                else:
+                    config["model"] = updated_model_params
+                min_nbor_dist[task_config.key] = None
+                continue
+            log.info(
+                "Calculate neighbor statistics for task %s... "
+                "(add --skip-neighbor-stat to skip this step)",
+                task_config.key,
+            )
             type_map = task_config.model_params.get("type_map")
             train_data = get_neighbor_stat_data(
                 task_config.training_data_params, type_map

deepmd/pt_expt/model/graph_lower.py

@@ -71,3 +71,75 @@ def model_uses_graph_lower(model: Any) -> bool:
     if atomic_model is None:
         return False
     return bool(atomic_model.uses_graph_lower())
+
+
+# Training families whose graph forward does not truncate neighbors. Keep the
+# allowlist explicit so new families can opt in after their training path is checked.
+_NEIGHBOR_STAT_FREE_DESCRIPTORS = frozenset({"dpa4", "dpa4c", "DPA4", "sezm", "SeZM"})
+
+
+def model_config_skips_neighbor_stat(config: dict) -> bool:
+    """Whether a fresh DPA4-family training model can omit neighbor statistics.
+
+    Unlike the runtime predicate, this operates before constructing a model.
+    Checkpoint restoration must be excluded by the caller, since restored buffers may
+    disable the graph route.
+    """
+    from deepmd.utils.bridging import (
+        expand_bridging_method,
+    )
+
+    config = expand_bridging_method(config)
+    model_type = config.get("type", "standard")
+    if model_type == "linear_ener":
+        children = config.get("models", [])
+        return bool(children) and all(
+            child.get("type") == "inner_potential"
+            or model_config_skips_neighbor_stat(child)
+            for child in children
+        )
+    dpa4_model = model_type in ("dpa4", "DPA4", "sezm", "SeZM")
+    if model_type != "standard" and not dpa4_model:
+        return False
+    if "spin" in config and config["spin"].get("scheme", "deepspin") != "native":
+        return False
+    if "use_srtab" in config:
+        return False
+    fitting_type = config.get("fitting_net", {}).get(
+        "type", "dpa4_ener" if dpa4_model else "ener"
+    )
+    if fitting_type not in ("ener", "direct_force_ener", "dpa4_ener", "sezm_ener"):
+        return False
+    descriptor = config.get("descriptor", {})
+    descriptor_type = descriptor.get("type", "dpa4" if dpa4_model else "se_e2_a")
+    return descriptor_type in _NEIGHBOR_STAT_FREE_DESCRIPTORS
+
+
+def resolve_stat_free_selections(config: dict) -> dict:
+    """Replace unused automatic capacities with their ordinary schema default.
+
+    Call only after ``model_config_skips_neighbor_stat`` succeeds. No model
+    weights or normalization statistics depend on this capacity on that route.
+    Keep numeric capacities for explicit dense-lower export compatibility.
+    """
+    import copy
+
+    from deepmd.dpmodel.utils.update_sel import (
+        UpdateSel,
+    )
+    from deepmd.utils.argcheck import (
+        DEFAULT_DPA4_SEL,
+    )
+
+    # Bridging expansion is a constructor detail, not a training config.
+    # Preserve concise models and trainer-owned keys for dump/checkpoint reload.
+    config = copy.deepcopy(config)
+    if config.get("type") == "linear_ener":
+        config["models"] = [resolve_stat_free_selections(m) for m in config["models"]]
+    descriptor = config.get("descriptor", {})
+    if isinstance(descriptor.get("sel"), str):
+        # Validate the existing auto syntax without opening training data.
+        UpdateSel().parse_auto_sel_ratio(descriptor["sel"])
+        # The capacity is unused in graph training, including any auto ratio.
+        descriptor["sel"] = DEFAULT_DPA4_SEL
+    return config

deepmd/utils/argcheck.py

@@ -615,6 +615,9 @@ def descrpt_dpa4c_args() -> list[Argument]:
     ]
 
 
+DEFAULT_DPA4_SEL = 256
+
+
 @descrpt_args_plugin.register(
     "dpa4",
     alias=["DPA4", "SeZM", "sezm"],
@@ -941,7 +944,13 @@ def descrpt_se_zm_args() -> list[Argument]:
     doc_trainable = "If the parameters in the descriptor are trainable."
     doc_seed = "Random seed for parameter initialization."
     return [
-        Argument("sel", [int, list[int], str], optional=True, default=256, doc=doc_sel),
+        Argument(
+            "sel",
+            [int, list[int], str],
+            optional=True,
+            default=DEFAULT_DPA4_SEL,
+            doc=doc_sel,
+        ),
         Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut),
         Argument(
             "env_exp",

deepmd/utils/bridging.py

@@ -34,6 +34,7 @@
     "is_bridged_sezm_config",
     "migrate_inner_clamp_keys",
     "resolve_bridging_window",
+    "validate_tabulated_zbl",
     "window_midpoint",
 ]
 
@@ -311,6 +312,22 @@ def concise_bridging_window(data: dict) -> dict:
     }
 
 
+def validate_tabulated_zbl(data: dict) -> None:
+    """Reject legacy tabulated switching for DPA4-family model configs."""
+    # Legacy tabulated switching is not the DPA4 analytical inner potential:
+    # reject it rather than silently ignoring it in the dedicated model builder.
+    model_type = str(data.get("type", "standard")).lower()
+    descriptor_type = str(data.get("descriptor", {}).get("type", "")).lower()
+    if "use_srtab" in data and (
+        model_type in _DPA4_FAMILY_TYPES
+        or descriptor_type in (*_DPA4_FAMILY_TYPES, "dpa4c")
+    ):
+        raise ValueError(
+            "DPA4/DPA4C do not support the legacy tabulated ZBL (`use_srtab`). "
+            "Use `bridging_method: zbl` for analytical inner-potential bridging."
+        )
+
+
 def expand_bridging_method(data: dict) -> dict:
     """Expand the ``bridging_method`` sugar into a ``linear_ener`` config.
 
@@ -342,9 +359,11 @@ def expand_bridging_method(data: dict) -> dict:
     ------
     ValueError
         If ``bridging_method`` is set on a model type that does not
-        support it, or if ``pair_exclude_types`` and
+        support it, if a DPA4-family model requests legacy ``use_srtab``,
+        or if ``pair_exclude_types`` and
         ``descriptor.exclude_types`` are both given and differ.
     """
+    validate_tabulated_zbl(data)
     method = str(data.get("bridging_method", "none"))
     if method.lower() in ("none", ""):
         return data

doc/train/training-advanced.md

@@ -266,6 +266,33 @@ training or of restoring every stochastic data transform.
 
 **`--skip-neighbor-stat`** will skip calculating neighbor statistics if one is concerned about performance. Some features will be disabled.
 
+The experimental PyTorch backend (`--pt-expt`) automatically skips this full-dataset
+scan for newly initialized graph-native DPA4 and DPA4C energy models
+(DPA4C does not have a selection parameter). Their graph forward
+includes every neighbor inside the cutoff, so checking the selection capacity is
+unnecessary. The decision is made separately for each training task; compositions
+skip the scan only when all components use the graph route. Input/output
+normalization statistics are still computed as usual.
+
+DPA4/DPA4C do not support legacy tabulated ZBL switching (`use_srtab`);
+use `bridging_method: "zbl"` for analytical inner-potential bridging instead.
+This restriction also applies to `type: "standard"` models using these descriptors
+and to learned children of compositions. Existing checkpoints configured with
+DPA4/DPA4C and `use_srtab` must be evaluated with the version that created them;
+changing the config to analytical bridging is not an equivalent conversion of
+the trained potential. Other descriptor families retain tabulated ZBL support.
+
+For these models, `sel: "auto"` (with or without a ratio) uses the ordinary
+DPA4 default capacity of 256 instead of scanning the dataset. This capacity is
+unused by the graph forward and its input statistics; numeric selections are
+kept unchanged. Explicit dense-lower exports still use a bounded neighbor list,
+so their capacity must cover the intended dataset. Restarts, initialization and
+fine-tuning from checkpoints retain the scan because restored state can disable
+graph lowering. Other descriptor families, backends, and dense/non-energy models
+are unchanged. `--skip-neighbor-stat` remains an
+explicit override. As with that flag, automatically skipped models do not record
+the minimum neighbor distance; compression can compute it later from training data.
+
 To maximize the performance, one should follow [FAQ: How to control the parallelism of a job](../troubleshooting/howtoset_num_nodes.md) to control the number of threads.
 See [Runtime environment variables](../env.md) for all runtime environment variables.
 

source/install/uv_with_retry.sh

@@ -5,28 +5,30 @@
 # https://github.com/astral-sh/uv/issues/3456
 # https://github.com/astral-sh/uv/issues/3514
 # https://github.com/astral-sh/uv/issues/4402
-tmpstderr=$(mktemp)
-max_retry=3
+tmpstderr=$(mktemp) || exit 1
+trap 'rm -f "${tmpstderr}"' EXIT
+max_attempts=3
+attempt=1
+delay=5
 while true; do
-	uv "$@" 2> >(tee -a "${tmpstderr}" >&2)
-	exit_code=$?
-	# exit if ok
-	if [ $exit_code -eq 0 ]; then
-		rm -f "${tmpstderr}"
+	# Keep stdout separate and wait for tee before inspecting this attempt's stderr.
+	{
+		uv "$@" 2>&1 >&3 | tee "${tmpstderr}" >&2
+		exit_code=${PIPESTATUS[0]}
+	} 3>&1
+	if [ "$exit_code" -eq 0 ]; then
 		exit 0
 	fi
-	# check if a retryable network error is in the stderr
-	if grep -qE "error decoding response body|network timeout|I/O operation failed during extraction|Failed to download distribution" "${tmpstderr}"; then
-		echo "Retrying uv in 5 s..."
-		max_retry=$((max_retry - 1))
-		if [ $max_retry -eq 0 ]; then
-			echo "Max retry reached, exiting..."
-			rm -f "${tmpstderr}"
-			exit 1
-		fi
-		sleep 5
-	else
-		rm -f "${tmpstderr}"
-		exit $exit_code
+	# Retry transient network failures, including errors fetching --find-links.
+	if ! grep -qE "error decoding response body|network timeout|I/O operation failed during extraction|Failed to download distribution|HTTP status (server|client) error \((5[0-9]{2}|429)([[:space:]]|\))" "${tmpstderr}"; then
+		exit "$exit_code"
 	fi
+	if [ "$attempt" -ge "$max_attempts" ]; then
+		echo "Max retry reached, exiting..."
+		exit "$exit_code"
+	fi
+	echo "Retrying uv in $delay s..."
+	sleep "$delay"
+	attempt=$((attempt + 1))
+	delay=$((delay * 2))
 done

source/tests/common/test_bridging.py

@@ -464,3 +464,33 @@ def test_an_unbridged_model_without_a_floor_registers_no_filter() -> None:
     assert requirement is not None
     assert requirement.default == 0.5
     assert requirement.length_scale == "absolute"
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        {"type": "dpa4"},
+        {"type": "DPA4"},
+        {"type": "SeZM"},
+        {"descriptor": {"type": "dpa4"}},
+        {"descriptor": {"type": "dpa4c"}},
+    ],
+)
+@pytest.mark.parametrize("method", ["none", "zbl"])
+def test_dpa4_rejects_legacy_tabulated_zbl(model, method):
+    """An inert or dense table must never masquerade as analytical bridging."""
+    model = {**model, "use_srtab": "unused-table.txt", "bridging_method": method}
+    with pytest.raises(ValueError, match="legacy tabulated ZBL"):
+        expand_bridging_method(model)
+    # Backends such as JAX dispatch directly through the shared factory.
+    from deepmd.dpmodel.model.model_factory import (
+        get_model,
+    )
+
+    with pytest.raises(ValueError, match="legacy tabulated ZBL"):
+        get_model(model, base_model=None, standard_model_factory=None)
+
+
+def test_other_descriptors_keep_legacy_tabulated_zbl():
+    model = {"descriptor": {"type": "se_e2_a"}, "use_srtab": "table.txt"}
+    assert expand_bridging_method(model) is model