fix

Author: OutisLi

deepmd/pd/entrypoints/main.py

@@ -64,6 +64,7 @@
     normalize,
 )
 from deepmd.utils.compat import (
+    convert_optimizer_to_new_format,
     update_deepmd_input,
 )
 from deepmd.utils.data_system import (
@@ -292,6 +293,8 @@ def train(
 
     # argcheck
     config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
+    # Backward compatibility: convert old optimizer format
+    config = convert_optimizer_to_new_format(config)
     config = normalize(config, multi_task=multi_task)
 
     # do neighbor stat

deepmd/pd/train/training.py

@@ -72,9 +72,6 @@
     nvprof_context,
     to_numpy_array,
 )
-from deepmd.utils.compat import (
-    convert_optimizer_to_new_format,
-)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -116,8 +113,6 @@ def __init__(
             resume_model = None
         resuming = resume_model is not None
         self.restart_training = restart_model is not None
-        # Backward compatibility: convert old optimizer format
-        config = convert_optimizer_to_new_format(config)
         model_params = config["model"]
         training_params = config["training"]
         optimizer_params = config.get("optimizer", {})

deepmd/pt/entrypoints/main.py

@@ -83,6 +83,7 @@
     normalize,
 )
 from deepmd.utils.compat import (
+    convert_optimizer_to_new_format,
     update_deepmd_input,
 )
 from deepmd.utils.data_system import (
@@ -325,6 +326,8 @@ def train(
 
     # argcheck
     config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
+    # Backward compatibility: convert old optimizer format
+    config = convert_optimizer_to_new_format(config)
     config = normalize(config, multi_task=multi_task)
 
     # do neighbor stat

deepmd/pt/train/training.py

@@ -73,9 +73,6 @@
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
-from deepmd.utils.compat import (
-    convert_optimizer_to_new_format,
-)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -126,8 +123,6 @@ def __init__(
             resume_model = None
         resuming = resume_model is not None
         self.restart_training = restart_model is not None
-        # Backward compatibility: convert old optimizer format
-        config = convert_optimizer_to_new_format(config)
         model_params = config["model"]
         training_params = config["training"]
         optimizer_params = config.get("optimizer", {})

deepmd/tf/entrypoints/train.py

@@ -37,6 +37,7 @@
     normalize,
 )
 from deepmd.tf.utils.compat import (
+    convert_optimizer_to_new_format,
     update_deepmd_input,
 )
 from deepmd.tf.utils.finetune import (
@@ -162,6 +163,8 @@ def train(
         jdata["model"] = json.loads(t_training_script)["model"]
 
     jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
+    # Backward compatibility: convert old optimizer format
+    jdata = convert_optimizer_to_new_format(jdata)
 
     jdata = normalize(jdata)
 

deepmd/tf/train/trainer.py

@@ -57,9 +57,6 @@
 from deepmd.tf.utils.sess import (
     run_sess,
 )
-from deepmd.utils.compat import (
-    convert_optimizer_to_new_format,
-)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -124,8 +121,6 @@ def get_lr_and_coef(lr_param):
         lr_param = jdata["learning_rate"]
         self.lr, self.scale_lr_coef = get_lr_and_coef(lr_param)
         # optimizer
-        # Backward compatibility: convert old optimizer format
-        jdata = convert_optimizer_to_new_format(jdata)
         # Note: Default values are already filled by argcheck.normalize()
         optimizer_param = jdata.get("optimizer", {})
         self.optimizer_type = optimizer_param.get("type", "Adam")

deepmd/tf/utils/compat.py

@@ -4,13 +4,15 @@
 from deepmd.utils.compat import (
     convert_input_v0_v1,
     convert_input_v1_v2,
+    convert_optimizer_to_new_format,
     deprecate_numb_test,
     update_deepmd_input,
 )
 
 __all__ = [
     "convert_input_v0_v1",
     "convert_input_v1_v2",
+    "convert_optimizer_to_new_format",
     "deprecate_numb_test",
     "update_deepmd_input",
 ]

deepmd/utils/compat.py

@@ -412,53 +412,71 @@ def convert_optimizer_to_new_format(
     dict[str, Any]
         converted output with optimizer section
     """
+    # Default optimizer values (must match argcheck.py defaults)
+    default_optimizer = {
+        "type": "Adam",
+        "adam_beta1": 0.9,
+        "adam_beta2": 0.999,
+        "weight_decay": 0.0,
+    }
+
     training_cfg = jdata.get("training", {})
-    if "opt_type" not in training_cfg:
-        # No conversion needed
-        return jdata
-
-    # Optimizer parameters that may be in the training section
-    optimizer_keys = [
-        "opt_type",
-        "kf_blocksize",
-        "kf_start_pref_e",
-        "kf_limit_pref_e",
-        "kf_start_pref_f",
-        "kf_limit_pref_f",
-        "weight_decay",
-        "momentum",
-        "muon_momentum",
-        "adam_beta1",
-        "adam_beta2",
-        "lr_adjust",
-        "lr_adjust_coeff",
-        "muon_2d_only",
-        "min_2d_dim",
-    ]
-
-    # Extract optimizer parameters from training section
-    optimizer_cfg = {}
-    for key in optimizer_keys:
-        if key in training_cfg:
-            optimizer_cfg[key] = training_cfg.pop(key)
-
-    # Convert opt_type to type for new format
-    if "opt_type" in optimizer_cfg:
-        optimizer_cfg["type"] = optimizer_cfg.pop("opt_type")
-
-    # Set the optimizer section if not already present
-    if "optimizer" not in jdata:
-        jdata["optimizer"] = optimizer_cfg
-    else:
-        # Merge with existing optimizer config (new config from conversion takes precedence)
-        jdata["optimizer"].update(optimizer_cfg)
+    optimizer_cfg = jdata.get("optimizer", {})
+
+    # Case 1: Old format - optimizer params in training section
+    if "opt_type" in training_cfg:
+        # Optimizer parameters that may be in the training section
+        optimizer_keys = [
+            "opt_type",
+            "kf_blocksize",
+            "kf_start_pref_e",
+            "kf_limit_pref_e",
+            "kf_start_pref_f",
+            "kf_limit_pref_f",
+            "weight_decay",
+            "momentum",
+            "muon_momentum",
+            "adam_beta1",
+            "adam_beta2",
+            "lr_adjust",
+            "lr_adjust_coeff",
+            "muon_2d_only",
+            "min_2d_dim",
+        ]
+
+        # Extract optimizer parameters from training section
+        extracted_cfg = {}
+        for key in optimizer_keys:
+            if key in training_cfg:
+                extracted_cfg[key] = training_cfg.pop(key)
+
+        # Convert opt_type to type for new format
+        if "opt_type" in extracted_cfg:
+            extracted_cfg["type"] = extracted_cfg.pop("opt_type")
+
+        # Merge with existing optimizer config (conversion takes precedence)
+        optimizer_cfg = {**optimizer_cfg, **extracted_cfg}
 
-    if warning:
-        warnings.warn(
-            "Placing optimizer parameters (opt_type, kf_blocksize, etc.) in the training section "
-            "is deprecated. Use a separate 'optimizer' section with 'type' field instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
+        if warning:
+            warnings.warn(
+                "Placing optimizer parameters (opt_type, kf_blocksize, etc.) in the training section "
+                "is deprecated. Use a separate 'optimizer' section with 'type' field instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+    # Case 2: Fill in missing defaults
+    # If type is not specified, default to Adam
+    if "type" not in optimizer_cfg:
+        optimizer_cfg["type"] = default_optimizer["type"]
+
+    # Fill in defaults for Adam optimizer type
+    if optimizer_cfg["type"] in ("Adam", "AdamW"):
+        for key, value in default_optimizer.items():
+            if key not in optimizer_cfg:
+                optimizer_cfg[key] = value
+
+    # Set/update the optimizer section
+    jdata["optimizer"] = optimizer_cfg
 
     return jdata

source/tests/pd/test_training.py

@@ -21,6 +21,9 @@
 from deepmd.pd.utils.finetune import (
     get_finetune_rules,
 )
+from deepmd.utils.compat import (
+    convert_optimizer_to_new_format,
+)
 
 from .model.test_permutation import (
     model_dpa1,
@@ -151,6 +154,7 @@ def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        self.config = convert_optimizer_to_new_format(self.config, warning=False)
         data_file = [str(Path(__file__).parent / "water/data/data_0")]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
@@ -168,6 +172,7 @@ def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        self.config = convert_optimizer_to_new_format(self.config, warning=False)
         data_file = [str(Path(__file__).parent / "water/data/data_0")]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
@@ -188,6 +193,7 @@ def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        self.config = convert_optimizer_to_new_format(self.config, warning=False)
         data_file = [str(Path(__file__).parent / "water/data/data_0")]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
@@ -209,6 +215,7 @@ def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        self.config = convert_optimizer_to_new_format(self.config, warning=False)
         data_file = [str(Path(__file__).parent / "water/data/data_0")]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
@@ -225,6 +232,7 @@ def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        self.config = convert_optimizer_to_new_format(self.config, warning=False)
         data_file = [str(Path(__file__).parent / "water/data/data_0")]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file