Skip to content

Commit 5c22e17

Browse files
authored
feat(pt): add full validation (#5336)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Full-training validation: configurable start/frequency/metric, per-system energy/force/virial MAE & RMSE, logging to val.log, and optional top‑K best‑checkpoint saving/rotation. * New evaluation utilities producing standardized metric outputs for energy/force/virial/spin. * Auto-batch-size gain a `silent` option to suppress informational logs. * **Changes** * Read-only access to underlying data system from the dataset loader. * test report function now returns structured metric dictionaries. * **Tests** * New unit and integration tests covering full-validation, config validation, and checkpoint rotation. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent dadb45c commit 5c22e17

14 files changed

Lines changed: 2089 additions & 244 deletions

File tree

deepmd/entrypoints/test.py

Lines changed: 390 additions & 214 deletions
Large diffs are not rendered by default.

deepmd/jax/utils/auto_batch_size.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,13 @@ def __init__(
2424
self,
2525
initial_batch_size: int = 1024,
2626
factor: float = 2.0,
27+
*,
28+
silent: bool = False,
2729
) -> None:
2830
super().__init__(
2931
initial_batch_size=initial_batch_size,
3032
factor=factor,
33+
silent=silent,
3134
)
3235

3336
def is_gpu_available(self) -> bool:

deepmd/pd/utils/auto_batch_size.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,13 @@ def __init__(
2222
self,
2323
initial_batch_size: int = 1024,
2424
factor: float = 2.0,
25+
*,
26+
silent: bool = False,
2527
) -> None:
2628
super().__init__(
2729
initial_batch_size=initial_batch_size,
2830
factor=factor,
31+
silent=silent,
2932
)
3033

3134
def is_gpu_available(self) -> bool:

deepmd/pt/train/training.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@
5454
KFOptimizerWrapper,
5555
LKFOptimizer,
5656
)
57+
from deepmd.pt.train.validation import (
58+
FullValidator,
59+
resolve_full_validation_start_step,
60+
)
5761
from deepmd.pt.train.wrapper import (
5862
ModelWrapper,
5963
)
@@ -942,11 +946,89 @@ def single_model_finetune(
942946
self.enable_profiler = training_params.get("enable_profiler", False)
943947
self.profiling = training_params.get("profiling", False)
944948
self.profiling_file = training_params.get("profiling_file", "timeline.json")
949+
validating_params = config.get("validating") or {}
950+
self.full_validator = self._create_full_validator(
951+
validating_params=validating_params,
952+
validation_data=validation_data,
953+
)
945954

946955
# Log model parameter count
947956
if self.rank == 0:
948957
self._log_parameter_count()
949958

959+
def _create_full_validator(
960+
self,
961+
*,
962+
validating_params: dict[str, Any],
963+
validation_data: DpLoaderSet | None,
964+
) -> FullValidator | None:
965+
"""Create the runtime full validator when it is active."""
966+
if not self._is_full_validation_requested(validating_params):
967+
return None
968+
self._raise_if_full_validation_unsupported(validation_data)
969+
if validation_data is None:
970+
raise RuntimeError(
971+
"validation_data must be available after full validation checks."
972+
)
973+
return FullValidator(
974+
validating_params=validating_params,
975+
validation_data=validation_data,
976+
model=self.model,
977+
train_infos=self._get_inner_module().train_infos,
978+
num_steps=self.num_steps,
979+
rank=self.rank,
980+
zero_stage=self.zero_stage,
981+
restart_training=self.restart_training,
982+
checkpoint_dir=Path(self.save_ckpt).parent,
983+
)
984+
985+
def _is_full_validation_requested(self, validating_params: dict[str, Any]) -> bool:
986+
"""Check whether full validation can trigger during this training run."""
987+
if not validating_params.get("full_validation", False):
988+
return False
989+
start_step = resolve_full_validation_start_step(
990+
validating_params.get("full_val_start", 0.5),
991+
self.num_steps,
992+
)
993+
return start_step is not None and start_step <= self.num_steps
994+
995+
def _raise_if_full_validation_unsupported(
996+
self,
997+
validation_data: DpLoaderSet | None,
998+
) -> None:
999+
"""Validate runtime full validation constraints."""
1000+
if self.multi_task:
1001+
raise ValueError(
1002+
"validating.full_validation only supports single-task energy "
1003+
"training; multi-task training is not supported."
1004+
)
1005+
1006+
has_spin = getattr(self.model, "has_spin", False)
1007+
if callable(has_spin):
1008+
has_spin = has_spin()
1009+
if has_spin or isinstance(self.loss, EnergySpinLoss):
1010+
raise ValueError(
1011+
"validating.full_validation only supports single-task energy "
1012+
"training; spin-energy training is not supported."
1013+
)
1014+
1015+
if not isinstance(self.loss, EnergyStdLoss):
1016+
raise ValueError(
1017+
"validating.full_validation only supports single-task energy training."
1018+
)
1019+
1020+
if validation_data is None:
1021+
raise ValueError(
1022+
"validating.full_validation requires `training.validation_data` "
1023+
"to be configured."
1024+
)
1025+
1026+
if self.zero_stage >= 2:
1027+
raise ValueError(
1028+
"validating.full_validation only supports single-task energy "
1029+
"training with training.zero_stage < 2."
1030+
)
1031+
9501032
@staticmethod
9511033
def _count_parameters(model: torch.nn.Module) -> tuple[int, int]:
9521034
"""
@@ -1448,6 +1530,14 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
14481530
fout, display_step_id, cur_lr, train_results, valid_results
14491531
)
14501532

1533+
if self.full_validator is not None:
1534+
self.full_validator.run(
1535+
step_id=_step_id,
1536+
display_step=display_step_id,
1537+
lr=cur_lr,
1538+
save_checkpoint=self.save_model,
1539+
)
1540+
14511541
if (
14521542
(
14531543
(display_step_id) % self.save_freq == 0

0 commit comments

Comments
 (0)