From d1207d5070d1113553d63a8f7c485151ea45d015 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Sat, 7 Nov 2026 19:51:04 +0800 Subject: [PATCH 1/8] fix(pt_expt): skip neighbor statistics for graph-native DPA4 models --- deepmd/pt_expt/entrypoints/main.py | 43 +++++++- deepmd/pt_expt/model/graph_lower.py | 73 +++++++++++++ deepmd/utils/argcheck.py | 11 +- doc/train/training-advanced.md | 19 ++++ .../pt_expt/model/test_graph_lower_config.py | 97 +++++++++++++++++ source/tests/pt_expt/test_entrypoint.py | 102 ++++++++++++++++++ 6 files changed, 340 insertions(+), 5 deletions(-) create mode 100644 source/tests/pt_expt/model/test_graph_lower_config.py diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 72de9277da..c5a731471e 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/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 diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index 5fa111661a..e09798934a 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -71,3 +71,76 @@ 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, + ) + from deepmd.utils.bridging import ( + expand_bridging_method, + ) + + config = copy.deepcopy(expand_bridging_method(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 diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index c7929a0334..16b6ddb84a 100644 --- a/deepmd/utils/argcheck.py +++ b/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", diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 353eda9bb5..0753edfd25 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -266,6 +266,25 @@ 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. + +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. diff --git a/source/tests/pt_expt/model/test_graph_lower_config.py b/source/tests/pt_expt/model/test_graph_lower_config.py new file mode 100644 index 0000000000..4c34d71e50 --- /dev/null +++ b/source/tests/pt_expt/model/test_graph_lower_config.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Pre-construction graph capability agrees with runtime model routing.""" + +import pytest + +from deepmd.pt_expt.model.graph_lower import ( + model_config_skips_neighbor_stat, + model_uses_graph_lower, +) +from deepmd.pt_expt.model.get_model import get_model + + +@pytest.mark.parametrize("descriptor_type", ["dpa4", "dpa4c"]) +def test_graph_config_matches_constructed_model(descriptor_type): + descriptor = { + "type": descriptor_type, + "rcut": 3.0, + "channels": 8, + "n_radial": 4, + "lmax": 2, + "seed": 1, + } + if descriptor_type == "dpa4": + descriptor.update(sel=20, n_blocks=1, mmax=1) + config = { + "type_map": ["H"], + "descriptor": descriptor, + "fitting_net": {"neuron": [8]}, + } + assert model_config_skips_neighbor_stat(config) + assert model_uses_graph_lower(get_model(config)) + + +@pytest.mark.parametrize( + "model,expected", + [ + ({"descriptor": {"type": "dpa1", "sel": "auto:1.2"}}, False), + ( + {"descriptor": {"type": "dpa1", "sel": 8}, "fitting_net": {"type": "dos"}}, + False, + ), + ( + {"descriptor": {"type": "dpa1", "sel": 8}, "spin": {"use_spin": [True]}}, + False, + ), + ({"descriptor": {"type": "dpa4"}, "spin": {"scheme": "native"}}, True), + ({"descriptor": {"type": "dpa4c"}}, True), + ({"type": "dpa4"}, True), + ({"descriptor": {"type": "dpa1", "sel": 8}}, False), + ({"descriptor": {"type": "dpa2"}}, False), + ({"descriptor": {"type": "dpa4", "sel": "auto"}}, True), + ({"descriptor": {"type": "dpa4"}, "fitting_net": {"type": "dos"}}, False), + ({"type": "frozen"}, False), + ({"descriptor": {"type": "dpa1", "sel": 8}, "use_srtab": "table.txt"}, False), + ( + { + "type": "linear_ener", + "models": [ + {"descriptor": {"type": "dpa4"}}, + {"type": "inner_potential"}, + ], + }, + True, + ), + ( + { + "type": "linear_ener", + "models": [{"descriptor": {"type": "dpa4"}}, {"type": "pairtab"}], + }, + False, + ), + ({"type": "linear_ener", "models": []}, False), + ], +) +def test_graph_config_conservative_boundaries(model, expected): + assert model_config_skips_neighbor_stat(model) is expected + + +@pytest.mark.parametrize("sel", ["auto", "auto:1.5", 32, [16, 16]]) +def test_stat_free_selection_uses_schema_default_without_mutation(sel): + from deepmd.pt_expt.model.graph_lower import resolve_stat_free_selections + from deepmd.utils.argcheck import DEFAULT_DPA4_SEL + + config = {"descriptor": {"type": "dpa4", "sel": sel}} + actual = resolve_stat_free_selections(config) + assert actual["descriptor"]["sel"] == ( + DEFAULT_DPA4_SEL if isinstance(sel, str) else sel + ) + assert config["descriptor"]["sel"] == sel + + +@pytest.mark.parametrize("sel", ["invalid", "auto:bad", "auto:1:2"]) +def test_stat_free_selection_rejects_invalid_auto(sel): + from deepmd.pt_expt.model.graph_lower import resolve_stat_free_selections + + with pytest.raises((ValueError, RuntimeError)): + resolve_stat_free_selections({"descriptor": {"type": "dpa4", "sel": sel}}) diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index 24bd98efaf..330dc52452 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -148,3 +148,105 @@ def test_pt_expt_entrypoint_rejects_random_model_key( {"model": {"model_dict": {"RANDOM": {}}}}, TrainEntrypointOptions(input_file="input.json"), ) + + +@pytest.mark.parametrize("restore", [None, "restart", "init_model", "finetune"]) +def test_graph_neighbor_statistics_skip_is_per_task(monkeypatch, restore): + """Graph tasks must not even open the full neighbor-stat dataset.""" + from deepmd.pt_expt.model import BaseModel + + calls = [] + monkeypatch.setattr( + "deepmd.pt_expt.entrypoints.main.get_neighbor_stat_data", + lambda params, type_map: calls.append(params["systems"]), + ) + monkeypatch.setattr(BaseModel, "update_sel", lambda *args: (args[-1], 0.5)) + models = { + "graph": {"descriptor": {"type": "dpa4", "sel": 120}}, + "dense": {"descriptor": {"type": "se_e2_a", "sel": [20, 20]}}, + "auto": {"descriptor": {"type": "dpa1", "sel": "auto"}}, + } + config = { + "model": {"model_dict": models}, + "training": { + "data_dict": {key: {"training_data": {"systems": key}} for key in models} + }, + } + options = TrainEntrypointOptions(input_file="unused.json") + if restore: + setattr(options, restore, "checkpoint.pt") + _, distances = PTExptTrainEntrypoint().update_neighbor_stat( + config, options, multi_task=True + ) + assert calls == (["graph", "dense", "auto"] if restore else ["dense", "auto"]) + assert distances == {"graph": 0.5 if restore else None, "dense": 0.5, "auto": 0.5} + assert config["model"]["model_dict"]["graph"] == models["graph"] + + +def test_graph_single_task_skips_neighbor_statistics(monkeypatch): + def unexpected(*args): + raise AssertionError("graph-only model opened neighbor-stat data") + + monkeypatch.setattr( + "deepmd.pt_expt.entrypoints.main.get_neighbor_stat_data", unexpected + ) + config = { + "model": {"descriptor": {"type": "dpa4c"}}, + "training": {"training_data": {}}, + } + actual, distance = PTExptTrainEntrypoint().update_neighbor_stat( + config, TrainEntrypointOptions(input_file="unused.json"), multi_task=False + ) + assert actual is config + assert distance is None + + +@pytest.mark.parametrize( + "model", + [ + {"type": "dpa4", "descriptor": {"sel": "auto:1.5"}, "fitting_net": {}}, + {"preset": "dpa4c-nano-v20260911"}, + ], +) +def test_graph_default_skip_after_normalization(monkeypatch, tmp_path, model): + import json + + from deepmd.utils.argcheck import DEFAULT_DPA4_SEL + + model["type_map"] = ["H"] + config = { + "model": model, + "learning_rate": {"type": "exp", "start_lr": 0.001, "stop_lr": 0.0001}, + "loss": {"type": "ener"}, + "training": { + "training_data": {"systems": "must-not-open"}, + "numb_steps": 1, + }, + } + input_file = tmp_path / "input.json" + output = tmp_path / "out.json" + input_file.write_text(json.dumps(config)) + entrypoint = PTExptTrainEntrypoint() + observed = [] + monkeypatch.setattr( + entrypoint, "run_training", lambda c, o, n: observed.append((c, n)) + ) + monkeypatch.setattr(entrypoint, "print_summary", lambda: None) + monkeypatch.setattr(entrypoint, "setup_run", lambda *args: None) + monkeypatch.setattr(entrypoint, "teardown_run", lambda *args: None) + + def unexpected(*args): + raise AssertionError("default graph training must not scan neighbor statistics") + + monkeypatch.setattr( + "deepmd.pt_expt.entrypoints.main.get_neighbor_stat_data", unexpected + ) + entrypoint.run( + TrainEntrypointOptions(input_file=str(input_file), output=str(output)) + ) + normalized, distance = observed[0] + assert distance is None + descriptor = normalized["model"]["descriptor"] + if descriptor["type"] == "dpa4": + assert descriptor["sel"] == DEFAULT_DPA4_SEL + assert json.loads(output.read_text()) == normalized From 0f6da84097e59ac1a8110227cbc44da8f6615e77 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Sat, 7 Nov 2026 20:53:54 +0800 Subject: [PATCH 2/8] test(pt_expt): cover automatic selection in mixed training tasks --- .../pt_expt/model/test_graph_lower_config.py | 20 +++++++++--- source/tests/pt_expt/test_entrypoint.py | 32 ++++++++++++++++--- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/source/tests/pt_expt/model/test_graph_lower_config.py b/source/tests/pt_expt/model/test_graph_lower_config.py index 4c34d71e50..2f4db67d12 100644 --- a/source/tests/pt_expt/model/test_graph_lower_config.py +++ b/source/tests/pt_expt/model/test_graph_lower_config.py @@ -3,15 +3,18 @@ import pytest +from deepmd.pt_expt.model.get_model import ( + get_model, +) from deepmd.pt_expt.model.graph_lower import ( model_config_skips_neighbor_stat, model_uses_graph_lower, ) -from deepmd.pt_expt.model.get_model import get_model @pytest.mark.parametrize("descriptor_type", ["dpa4", "dpa4c"]) def test_graph_config_matches_constructed_model(descriptor_type): + """The allowlisted families actually construct graph-lower models.""" descriptor = { "type": descriptor_type, "rcut": 3.0, @@ -73,13 +76,19 @@ def test_graph_config_matches_constructed_model(descriptor_type): ], ) def test_graph_config_conservative_boundaries(model, expected): + """Only supported energy compositions opt out of the scan.""" assert model_config_skips_neighbor_stat(model) is expected @pytest.mark.parametrize("sel", ["auto", "auto:1.5", 32, [16, 16]]) def test_stat_free_selection_uses_schema_default_without_mutation(sel): - from deepmd.pt_expt.model.graph_lower import resolve_stat_free_selections - from deepmd.utils.argcheck import DEFAULT_DPA4_SEL + """Auto uses the schema default while numeric capacities remain unchanged.""" + from deepmd.pt_expt.model.graph_lower import ( + resolve_stat_free_selections, + ) + from deepmd.utils.argcheck import ( + DEFAULT_DPA4_SEL, + ) config = {"descriptor": {"type": "dpa4", "sel": sel}} actual = resolve_stat_free_selections(config) @@ -91,7 +100,10 @@ def test_stat_free_selection_uses_schema_default_without_mutation(sel): @pytest.mark.parametrize("sel", ["invalid", "auto:bad", "auto:1:2"]) def test_stat_free_selection_rejects_invalid_auto(sel): - from deepmd.pt_expt.model.graph_lower import resolve_stat_free_selections + """Skipping the scan does not bypass automatic-selection syntax validation.""" + from deepmd.pt_expt.model.graph_lower import ( + resolve_stat_free_selections, + ) with pytest.raises((ValueError, RuntimeError)): resolve_stat_free_selections({"descriptor": {"type": "dpa4", "sel": sel}}) diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index 330dc52452..db48840a91 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -153,7 +153,12 @@ def test_pt_expt_entrypoint_rejects_random_model_key( @pytest.mark.parametrize("restore", [None, "restart", "init_model", "finetune"]) def test_graph_neighbor_statistics_skip_is_per_task(monkeypatch, restore): """Graph tasks must not even open the full neighbor-stat dataset.""" - from deepmd.pt_expt.model import BaseModel + from deepmd.pt_expt.model import ( + BaseModel, + ) + from deepmd.utils.argcheck import ( + DEFAULT_DPA4_SEL, + ) calls = [] monkeypatch.setattr( @@ -163,6 +168,7 @@ def test_graph_neighbor_statistics_skip_is_per_task(monkeypatch, restore): monkeypatch.setattr(BaseModel, "update_sel", lambda *args: (args[-1], 0.5)) models = { "graph": {"descriptor": {"type": "dpa4", "sel": 120}}, + "graph_auto": {"descriptor": {"type": "dpa4", "sel": "auto"}}, "dense": {"descriptor": {"type": "se_e2_a", "sel": [20, 20]}}, "auto": {"descriptor": {"type": "dpa1", "sel": "auto"}}, } @@ -178,12 +184,25 @@ def test_graph_neighbor_statistics_skip_is_per_task(monkeypatch, restore): _, distances = PTExptTrainEntrypoint().update_neighbor_stat( config, options, multi_task=True ) - assert calls == (["graph", "dense", "auto"] if restore else ["dense", "auto"]) - assert distances == {"graph": 0.5 if restore else None, "dense": 0.5, "auto": 0.5} - assert config["model"]["model_dict"]["graph"] == models["graph"] + assert calls == ( + ["graph", "graph_auto", "dense", "auto"] if restore else ["dense", "auto"] + ) + assert distances == { + "graph": 0.5 if restore else None, + "graph_auto": 0.5 if restore else None, + "dense": 0.5, + "auto": 0.5, + } + tasks = config["model"]["model_dict"] + assert tasks["graph"]["descriptor"]["sel"] == 120 + assert tasks["graph_auto"]["descriptor"]["sel"] == ( + "auto" if restore else DEFAULT_DPA4_SEL + ) def test_graph_single_task_skips_neighbor_statistics(monkeypatch): + """A graph-only task never opens data for neighbor statistics.""" + def unexpected(*args): raise AssertionError("graph-only model opened neighbor-stat data") @@ -209,9 +228,12 @@ def unexpected(*args): ], ) def test_graph_default_skip_after_normalization(monkeypatch, tmp_path, model): + """Preset expansion and auto normalization precede the trainer and output dump.""" import json - from deepmd.utils.argcheck import DEFAULT_DPA4_SEL + from deepmd.utils.argcheck import ( + DEFAULT_DPA4_SEL, + ) model["type_map"] = ["H"] config = { From 2b7e016bcb924c9c4258fca577eead4353e4727e Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Mon, 9 Nov 2026 11:31:10 +0800 Subject: [PATCH 3/8] fix(pt_expt): preserve reloadable configs when skipping neighbor stats --- deepmd/pt_expt/model/graph_lower.py | 7 +- .../model/test_stat_free_graph_equivalence.py | 133 ++++++++++++++++++ source/tests/pt_expt/test_entrypoint.py | 84 +++++++++++ 3 files changed, 220 insertions(+), 4 deletions(-) create mode 100644 source/tests/pt_expt/model/test_stat_free_graph_equivalence.py diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index e09798934a..a7aad0b35f 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -130,11 +130,10 @@ def resolve_stat_free_selections(config: dict) -> dict: from deepmd.utils.argcheck import ( DEFAULT_DPA4_SEL, ) - from deepmd.utils.bridging import ( - expand_bridging_method, - ) - config = copy.deepcopy(expand_bridging_method(config)) + # 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", {}) diff --git a/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py b/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py new file mode 100644 index 0000000000..b7c919a33d --- /dev/null +++ b/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Guard that resolving unused DPA4 capacities cannot change graph training.""" + +import copy + +import numpy as np +import torch + +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.model.graph_lower import ( + model_config_skips_neighbor_stat, + resolve_stat_free_selections, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) + + +def test_stat_free_dpa4_graph_training_equivalence(monkeypatch): + """A binding sel, a large sel and resolved auto have identical training.""" + config = { + "type": "dpa4", + "type_map": ["H", "O"], + "descriptor": { + "type": "dpa4", + "sel": "auto", + "rcut": 4.0, + "channels": 8, + "n_radial": 4, + "lmax": 2, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8], + "precision": "float64", + "seed": 1, + }, + } + assert model_config_skips_neighbor_stat(config) + resolved = resolve_stat_free_selections(config) + assert config["descriptor"]["sel"] == "auto" + configs = [copy.deepcopy(config), copy.deepcopy(config), resolved] + configs[0]["descriptor"]["sel"] = 1 + configs[1]["descriptor"]["sel"] = 256 + models = [get_model(c).to(env.DEVICE) for c in configs] + # Compare independently initialized weights/buffers, not loaded clones. + states = [m.state_dict() for m in models] + for state in states[1:]: + assert state.keys() == states[0].keys() + for key in state: + torch.testing.assert_close(state[key], states[0][key], rtol=0, atol=0) + # Fresh DPA4 residuals are zero: make the model sensitive to geometry. + for model in models: + data = model.atomic_model.descriptor.serialize() + data = jitter_zero_arrays(data, np.random.default_rng(99)) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data).to(env.DEVICE) + model.train() + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.9, 0.1, 0.0], [0.2, 1.1, 0.3], [1.0, 1.0, 0.7]]], + dtype=torch.float64, + device=env.DEVICE, + ) + atype = torch.tensor([[0, 1, 0, 1]], device=env.DEVICE) + box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 12).reshape(1, 9) + sample = [{"coord": coord, "box": box, "atype": atype}] + for model in models: + model.atomic_model.descriptor.compute_input_stats(sample) + stat_states = [m.state_dict() for m in models] + for state in stat_states[1:]: + assert state.keys() == stat_states[0].keys() + for key in state: + torch.testing.assert_close(state[key], stat_states[0][key], rtol=0, atol=0) + calls = [] + original = DescrptDPA4.call_graph + + def spy(self, graph, *args, **kwargs): + calls.append(int(graph.edge_mask.sum().item())) + return original(self, graph, *args, **kwargs) + + monkeypatch.setattr(DescrptDPA4, "call_graph", spy) + outputs, grads, updated = [], [], [] + for model in models: + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) + out = model.forward_common(coord.clone().requires_grad_(True), atype, box) + outputs.append( + { + k: out[k].detach().clone() + for k in ("energy_redu", "energy_derv_r", "energy_derv_c_redu") + } + ) + assert out["energy_derv_r"].abs().max().item() > 1e-8 + loss = out["energy_redu"].square().sum() + out["energy_derv_r"].square().sum() + loss.backward() + grads.append( + { + n: None if p.grad is None else p.grad.detach().clone() + for n, p in model.named_parameters() + } + ) + optimizer.step() + updated.append(copy.deepcopy(model.state_dict())) + # Four mutually neighbouring atoms have 12 directed edges; sel=1 must + # not truncate this to four edges on the actual model forward route. + assert len(calls) >= 3 + assert all(count == 12 for count in calls) + for records in (outputs, grads, updated): + for record in records[1:]: + assert record.keys() == records[0].keys() + for key, value in record.items(): + reference = records[0][key] + if reference is None: + assert value is None + else: + torch.testing.assert_close(value, reference, rtol=1e-10, atol=1e-12) + perturbed = coord.clone() + perturbed[0, 0, 0] += 0.1 + model = models[0] + baseline = model.forward_common(coord.clone(), atype, box)["energy_redu"] + changed = model.forward_common(perturbed, atype, box)["energy_redu"] + assert (changed - baseline).abs().max().item() > 1e-8 diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index db48840a91..69c31d6b4e 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -272,3 +272,87 @@ def unexpected(*args): if descriptor["type"] == "dpa4": assert descriptor["sel"] == DEFAULT_DPA4_SEL assert json.loads(output.read_text()) == normalized + + +@pytest.mark.parametrize("restore", ["json", "init_model", "finetune"]) +def test_bridging_auto_selection_preserves_reloadable_model_params( + monkeypatch, tmp_path, restore +): + """Do not persist constructor-only bridging expansion in training artifacts.""" + import copy + import json + + import torch + + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + from deepmd.utils.argcheck import ( + normalize, + ) + + config = { + "model": { + "type": "dpa4", + "type_map": ["H"], + "descriptor": { + "type": "dpa4", + "sel": "auto", + "channels": 8, + "n_radial": 4, + "lmax": 2, + "mmax": 1, + "n_blocks": 1, + "seed": 1, + }, + "fitting_net": {"neuron": [8], "seed": 1}, + "bridging_method": "zbl", + "data_stat_nbatch": 3, + "data_stat_protect": 0.02, + }, + "learning_rate": {"type": "exp", "start_lr": 0.001, "stop_lr": 0.0001}, + "loss": {"type": "ener"}, + "training": {"training_data": {"systems": "must-not-open"}, "numb_steps": 1}, + } + original = normalize(config) + entrypoint = PTExptTrainEntrypoint() + + def unexpected(*args): + raise AssertionError( + "eligible bridging model must not scan neighbor statistics" + ) + + monkeypatch.setattr( + "deepmd.pt_expt.entrypoints.main.get_neighbor_stat_data", unexpected + ) + updated, distance = entrypoint.update_neighbor_stat( + copy.deepcopy(original), + TrainEntrypointOptions(input_file="unused.json"), + multi_task=False, + ) + expected = copy.deepcopy(original["model"]) + expected["descriptor"]["sel"] = 256 + assert updated["model"] == expected + assert distance is None + assert original["model"]["descriptor"]["sel"] == "auto" + assert updated["model"]["data_stat_nbatch"] == 3 + assert updated["model"]["data_stat_protect"] == 0.02 + + if restore == "json": + output = tmp_path / "out.json" + entrypoint.dump_config(updated, str(output)) + reloaded = json.loads(output.read_text()) + else: + model = get_model(updated["model"]) + wrapper = ModelWrapper(model, model_params=updated["model"]) + checkpoint = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, checkpoint) + options = TrainEntrypointOptions( + input_file="unused.json", use_pretrain_script=True + ) + setattr(options, restore, str(checkpoint)) + reloaded = entrypoint.preprocess_config(copy.deepcopy(original), options) + assert normalize(reloaded)["model"] == expected From 6adfbd08617473837506557174712b7b62ff8493 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Tue, 10 Nov 2026 13:51:10 +0800 Subject: [PATCH 4/8] fix: reject legacy ZBL for DPA4 and guard skip-path statistics --- deepmd/dpmodel/model/model_factory.py | 3 ++ deepmd/utils/bridging.py | 21 +++++++- doc/train/training-advanced.md | 3 ++ source/tests/common/test_bridging.py | 30 +++++++++++ .../model/test_stat_free_graph_equivalence.py | 8 --- source/tests/pt_expt/test_training.py | 50 +++++++++++++++++-- 6 files changed, 101 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 118b09a75f..5fdfff9c94 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/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 " @@ -552,6 +554,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: diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 2fa225cc64..204b27563b 100644 --- a/deepmd/utils/bridging.py +++ b/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 diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 0753edfd25..aa6d8e8b06 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -274,6 +274,9 @@ unnecessary. The decision is made separately for each training task; composition 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. + 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 diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 6432d46955..918cf7790f 100644 --- a/source/tests/common/test_bridging.py +++ b/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 diff --git a/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py b/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py index b7c919a33d..7ebe462590 100644 --- a/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py +++ b/source/tests/pt_expt/model/test_stat_free_graph_equivalence.py @@ -75,14 +75,6 @@ def test_stat_free_dpa4_graph_training_equivalence(monkeypatch): ) atype = torch.tensor([[0, 1, 0, 1]], device=env.DEVICE) box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 12).reshape(1, 9) - sample = [{"coord": coord, "box": box, "atype": atype}] - for model in models: - model.atomic_model.descriptor.compute_input_stats(sample) - stat_states = [m.state_dict() for m in models] - for state in stat_states[1:]: - assert state.keys() == stat_states[0].keys() - for key in state: - torch.testing.assert_close(state[key], stat_states[0][key], rtol=0, atol=0) calls = [] original = DescrptDPA4.call_graph diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index fd87f8a0ea..bd73dc46d1 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -9,6 +9,7 @@ """ import copy +import json import math import os import platform @@ -30,6 +31,9 @@ import pytest import torch +from deepmd.dpmodel.train import ( + TrainEntrypointOptions, +) from deepmd.dpmodel.utils.dist_check import ( compute_min_pair_margin_batch, pair_half_thresholds, @@ -38,6 +42,7 @@ HybridMuonOptimizer, ) from deepmd.pt_expt.entrypoints.main import ( + PTExptTrainEntrypoint, get_trainer, ) from deepmd.pt_expt.model import ( @@ -611,13 +616,48 @@ def test_ema_full_validation_is_ignored_without_ema(self) -> None: self.assertIsNone(trainer.ema_full_validator) def test_training_loop_dpa4(self) -> None: - """Run a few DPA4/SeZM training steps (model type "dpa4" dispatch).""" + """The auto-skip entrypoint still initializes stats and trains normally.""" config = _make_config(self.data_dir, numb_steps=5) config["model"] = copy.deepcopy(_MODEL_DPA4) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - self.assertEqual(config["model"]["type"], "dpa4") - self._run_training(config) + config["model"]["descriptor"]["sel"] = "auto" + config["model"]["data_stat_protect"] = 0.02 + + def checked_trainer(*args, **kwargs): + trainer = get_trainer(*args, **kwargs) + model_config = args[0]["model"] + self.assertEqual(model_config["data_stat_nbatch"], 1) + self.assertEqual(model_config["data_stat_protect"], 0.02) + # Water labels are nonzero. A skipped statistics initialization + # leaves this buffer at zero, even though training can still run. + bias = trainer.model.atomic_model.out_bias + self.assertTrue(torch.isfinite(bias).all()) + self.assertGreater(bias.abs().max().item(), 1.0) + return trainer + + with tempfile.TemporaryDirectory() as tmpdir: + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + with open("input.json", "w") as f: + json.dump(config, f) + with ( + patch( + "deepmd.pt_expt.entrypoints.main.get_neighbor_stat_data", + side_effect=AssertionError("unexpected neighbor-stat scan"), + ), + patch( + "deepmd.pt_expt.entrypoints.main.get_trainer", + side_effect=checked_trainer, + ) as factory, + ): + PTExptTrainEntrypoint().run( + TrainEntrypointOptions(input_file="input.json") + ) + factory.assert_called_once() + self.assertTrue(Path("lcurve.out").is_file()) + self.assertTrue(list(Path().glob("*.pt"))) + finally: + os.chdir(old_cwd) @REQUIRES_SUPPORTED_COMPILE def test_training_loop_compiled(self) -> None: From 293c14f3749218dcb55db9f61dce00814cefcec1 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Tue, 10 Nov 2026 16:03:35 +0800 Subject: [PATCH 5/8] fix: enforce DPA4 tabulated ZBL rejection across builders --- deepmd/dpmodel/model/model_factory.py | 1 + deepmd/pt/model/model/__init__.py | 3 +++ doc/train/training-advanced.md | 5 +++++ .../tests/pt/model/test_get_model_bridging.py | 20 +++++++++++++++++++ .../pt_expt/model/test_get_model_bridging.py | 16 +++++++++++++++ 5 files changed, 45 insertions(+) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 5fdfff9c94..cbb1f5284e 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -406,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) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index e3ea613b44..26d4de4dc2 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/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, @@ -362,6 +363,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 +385,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"]) diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index aa6d8e8b06..af5da1b1b9 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -276,6 +276,11 @@ 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 diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index 0089e69972..bfbfade728 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -434,3 +434,23 @@ def test_spin_requires_native_scheme(canonical: bool, scheme: str | None) -> Non cfg["spin"]["scheme"] = scheme with pytest.raises(ValueError, match=r"require model\.spin\.scheme='native'"): get_model(cfg) + + +@pytest.mark.parametrize("location", ["composition", "child"]) +def test_canonical_rejects_tabulated_zbl(location): + config = _canonical_config() + target = config if location == "composition" else config["models"][0] + target["use_srtab"] = "unused-table.txt" + with pytest.raises(ValueError, match="legacy tabulated ZBL"): + get_model(config) + + +def test_training_dispatch_rejects_tabulated_zbl(): + from deepmd.pt.train.training import get_single_model + + config = _sugar_config() + config.pop("type") + config.pop("bridging_method") + config["use_srtab"] = "unused-table.txt" + with pytest.raises(ValueError, match="legacy tabulated ZBL"): + get_single_model(config) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index a73674e48a..9340d8a807 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -399,3 +399,19 @@ def test_nonempty_shared_dict_raises() -> None: cfg["shared_dict"] = {"my_descriptor": "descriptor"} with pytest.raises(NotImplementedError, match="shared_dict"): get_model(cfg) + + +@pytest.mark.parametrize("location", ["composition", "child"]) +@pytest.mark.parametrize("backend", ["dpmodel", "pt_expt"]) +def test_canonical_rejects_tabulated_zbl(location, backend): + from deepmd.utils.bridging import expand_bridging_method + + config = expand_bridging_method(_bridged(_dpa4_standard_config())) + target = config if location == "composition" else config["models"][0] + target["use_srtab"] = "unused-table.txt" + if backend == "dpmodel": + from deepmd.dpmodel.model.model import get_model as build + else: + build = get_model + with pytest.raises(ValueError, match="legacy tabulated ZBL"): + build(config) From 8197f4c1e974f6397dcf1ef1222c587855b7c6fd Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Tue, 10 Nov 2026 16:35:53 +0800 Subject: [PATCH 6/8] test: separate DPA4 ZBL and long-range rejection contracts --- source/tests/pt_expt/model/test_long_range.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 0581dd8626..cd286dd560 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -1172,7 +1172,17 @@ def test_unsupported_combinations_fail() -> None: config = _config("les") config["use_srtab"] = "table.txt" config["sw_rmin"], config["sw_rmax"] = 0.5, 1.0 - with pytest.raises(NotImplementedError, match="ZBL"): + with pytest.raises(ValueError, match="legacy tabulated ZBL"): + get_model(config) + # DPA4C rejects tabulated ZBL regardless of long range. Use an otherwise + # supported descriptor to exercise the separate long-range/ZBL boundary. + config["descriptor"] = { + "type": "se_e2_a", + "sel": [4, 4], + "rcut": 3.0, + "rcut_smth": 2.5, + } + with pytest.raises(NotImplementedError, match=r"long-range.*tabulated ZBL"): get_model(config) config = _config("les") config["type_map"] = ["O", "H"] From 3ffab455815ff00eaea11d85868388074c33bb87 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:02:15 +0800 Subject: [PATCH 7/8] fix: validate unbridged DPA4 children and format test imports --- deepmd/pt/model/model/__init__.py | 1 + source/tests/pt/model/test_get_model_bridging.py | 12 +++++++++++- .../tests/pt_expt/model/test_get_model_bridging.py | 4 +++- 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 26d4de4dc2..e69fdae74d 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -252,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 diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index bfbfade728..ad793b9013 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -446,7 +446,9 @@ def test_canonical_rejects_tabulated_zbl(location): def test_training_dispatch_rejects_tabulated_zbl(): - from deepmd.pt.train.training import get_single_model + from deepmd.pt.train.training import ( + get_single_model, + ) config = _sugar_config() config.pop("type") @@ -454,3 +456,11 @@ def test_training_dispatch_rejects_tabulated_zbl(): config["use_srtab"] = "unused-table.txt" with pytest.raises(ValueError, match="legacy tabulated ZBL"): get_single_model(config) + + +def test_unbridged_linear_rejects_tabulated_zbl_child(): + config = _canonical_config() + config["models"].pop() + config["models"][0]["use_srtab"] = "unused-table.txt" + with pytest.raises(ValueError, match="legacy tabulated ZBL"): + get_model(config) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 9340d8a807..577fc7df57 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -404,7 +404,9 @@ def test_nonempty_shared_dict_raises() -> None: @pytest.mark.parametrize("location", ["composition", "child"]) @pytest.mark.parametrize("backend", ["dpmodel", "pt_expt"]) def test_canonical_rejects_tabulated_zbl(location, backend): - from deepmd.utils.bridging import expand_bridging_method + from deepmd.utils.bridging import ( + expand_bridging_method, + ) config = expand_bridging_method(_bridged(_dpa4_standard_config())) target = config if location == "composition" else config["models"][0] From 95d2cd61d1fa5f97e50276156d95b5a11da8ac91 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:03:18 +0800 Subject: [PATCH 8/8] ci: install pinned Paddle wheels without the nightly index --- .github/workflows/test_python.yml | 4 +- source/install/uv_with_retry.sh | 42 ++++--- source/tests/common/test_uv_with_retry.py | 146 ++++++++++++++++++++++ 3 files changed, 171 insertions(+), 21 deletions(-) create mode 100644 source/tests/common/test_uv_with_retry.py diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index a98215f30d..2bd0211284 100644 --- a/.github/workflows/test_python.yml +++ b/.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 diff --git a/source/install/uv_with_retry.sh b/source/install/uv_with_retry.sh index deb0228d82..80ced70cc8 100755 --- a/source/install/uv_with_retry.sh +++ b/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 diff --git a/source/tests/common/test_uv_with_retry.py b/source/tests/common/test_uv_with_retry.py new file mode 100644 index 0000000000..8ce16edc0a --- /dev/null +++ b/source/tests/common/test_uv_with_retry.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import shlex +import shutil +import subprocess +import tempfile +import unittest +from pathlib import ( + Path, +) + +SCRIPT = Path(__file__).resolve().parents[2] / "install" / "uv_with_retry.sh" + + +@unittest.skipUnless(shutil.which("bash"), "requires bash") +class TestUvWithRetry(unittest.TestCase): + def run_uv(self, outcomes, delayed_tee=False): + """Run the wrapper with fake uv and sleep, without network access.""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + temporary = root / "tmp" + temporary.mkdir() + cases = [] + for attempt, (code, message) in enumerate(outcomes, start=1): + cases.append( + f"{attempt}) printf '%s\\n' {shlex.quote(message)} >&2; " + f"exit {code} ;;" + ) + commands = { + "uv": """#!/bin/bash +attempt=0 +if [ -f "$STATE/attempt" ]; then + read -r attempt < "$STATE/attempt" +fi +attempt=$((attempt + 1)) +printf '%s\\n' "$attempt" > "$STATE/attempt" +printf '%s\\n' "$@" > "$STATE/args" +printf 'stdout %s\\n' "$attempt" +case "$attempt" in +""" + + "\n".join(cases) + + "\n*) exit 99 ;;\nesac\n", + "sleep": '#!/bin/bash\nprintf "%s\\n" "$@" >> "$STATE/sleeps"\n', + } + if delayed_tee: + commands["tee"] = ( + "#!/bin/bash\n" + f"{shlex.quote(shutil.which('sleep'))} 0.1\n" + f'exec {shlex.quote(shutil.which("tee"))} "$@"\n' + ) + for name, content in commands.items(): + executable = root / name + executable.write_text(content, encoding="utf-8") + executable.chmod(0o755) + arguments = ["pip", "install", "package with spaces", "--find-links", "url"] + result = subprocess.run( + ["bash", str(SCRIPT), *arguments], + env={ + **os.environ, + "PATH": f"{root}{os.pathsep}{os.environ['PATH']}", + "STATE": str(root), + "TMPDIR": str(temporary), + }, + capture_output=True, + text=True, + check=False, + timeout=10, + ) + self.assertEqual((root / "args").read_text().splitlines(), arguments) + self.assertEqual(list(temporary.iterdir()), []) + attempts = int((root / "attempt").read_text()) + sleeps = root / "sleeps" + delays = sleeps.read_text().splitlines() if sleeps.exists() else [] + self.assertEqual( + result.stderr, + "".join(f"{message}\n" for _, message in outcomes[:attempts]), + ) + for attempt in range(1, attempts + 1): + self.assertIn(f"stdout {attempt}\n", result.stdout) + return result, attempts, delays + + def test_success(self): + result, attempts, delays = self.run_uv([(0, "installed")]) + self.assertEqual((result.returncode, attempts, delays), (0, 1, [])) + + def test_transient_errors(self): + messages = [ + "error decoding response body", + "network timeout", + "I/O operation failed during extraction", + "Failed to download distribution", + "HTTP status client error (429 Too Many Requests) for url (https://example.com)", + ] + for code in (500, 502, 503, 504, 599): + messages.append( + "error: Failed to read `--find-links` URL: " + "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/\n" + " Caused by: Request failed after 3 retries\n" + f" Caused by: HTTP status server error ({code} Server Error) for url " + "(https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/)" + ) + for message in messages: + with self.subTest(message=message): + result, attempts, delays = self.run_uv([(2, message), (0, "installed")]) + self.assertEqual((result.returncode, attempts, delays), (0, 2, ["5"])) + + def test_permanent_errors(self): + for message in ( + "No solution found when resolving dependencies", + "HTTP status client error (401 Unauthorized)", + "HTTP status client error (403 Forbidden)", + "HTTP status client error (404 Not Found)", + "invalid argument: 502", + ): + with self.subTest(message=message): + result, attempts, delays = self.run_uv([(2, message)]) + self.assertEqual((result.returncode, attempts, delays), (2, 1, [])) + + def test_permanent_error_after_transient_error(self): + result, attempts, delays = self.run_uv( + [(1, "network timeout"), (2, "No solution found")] + ) + self.assertEqual((result.returncode, attempts, delays), (2, 2, ["5"])) + + def test_success_on_final_attempt(self): + result, attempts, delays = self.run_uv( + [(1, "network timeout"), (1, "network timeout"), (0, "installed")] + ) + self.assertEqual((result.returncode, attempts, delays), (0, 3, ["5", "10"])) + + def test_retry_limit_preserves_failure(self): + result, attempts, delays = self.run_uv([(7, "network timeout")] * 3) + self.assertEqual((result.returncode, attempts, delays), (7, 3, ["5", "10"])) + self.assertIn("Max retry reached", result.stdout) + self.assertEqual(result.stdout.count("Retrying uv"), 2) + + def test_waits_for_stderr(self): + result, attempts, delays = self.run_uv( + [(1, "HTTP status server error (502 Bad Gateway)"), (0, "installed")], + delayed_tee=True, + ) + self.assertEqual((result.returncode, attempts, delays), (0, 2, ["5"])) + + +if __name__ == "__main__": + unittest.main()