compat

Author: OutisLi

deepmd/utils/compat.py

@@ -370,6 +370,69 @@ def deprecate_numb_test(
     return jdata
 
 
+def migrate_training_warmup(
+    jdata: dict[str, Any], warning: bool = True
+) -> dict[str, Any]:
+    """
+    Migrate legacy warmup settings from training to learning_rate.
+
+    Parameters
+    ----------
+    jdata : dict[str, Any]
+        Input configuration dictionary.
+    warning : bool, optional
+        Whether to show a deprecation warning, by default True.
+
+    Returns
+    -------
+    dict[str, Any]
+        Updated configuration dictionary.
+    """
+    training = jdata.get("training")
+    if not isinstance(training, dict):
+        return jdata
+
+    warmup_keys = ("warmup_steps", "warmup_ratio", "warmup_start_factor")
+    legacy_keys = [key for key in warmup_keys if key in training]
+    if not legacy_keys:
+        return jdata
+
+    lr = jdata.get("learning_rate")
+    if not isinstance(lr, dict):
+        for key in legacy_keys:
+            training.pop(key)
+        if warning:
+            warnings.warn(
+                "Found legacy warmup settings under training, but learning_rate "
+                "is missing or invalid. The warmup keys were removed from training."
+            )
+        return jdata
+
+    moved_keys = []
+    skipped_keys = []
+    # === Step 1. Move legacy warmup keys ===
+    for key in legacy_keys:
+        value = training.pop(key)
+        if key in lr:
+            skipped_keys.append(key)
+            continue
+        lr[key] = value
+        moved_keys.append(key)
+
+    if warning:
+        if skipped_keys:
+            warnings.warn(
+                "Legacy warmup settings under training were ignored because "
+                f"learning_rate already defines them: {', '.join(skipped_keys)}."
+            )
+        else:
+            warnings.warn(
+                "Legacy warmup settings under training were moved to learning_rate: "
+                f"{', '.join(moved_keys)}."
+            )
+    return jdata
+
+
 def update_deepmd_input(
     jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None
 ) -> dict[str, Any]:
@@ -389,4 +452,5 @@ def is_deepmd_v1_input(jdata: dict[str, Any]) -> bool:
     else:
         jdata = deprecate_numb_test(jdata, warning, dump)
 
+    jdata = migrate_training_warmup(jdata, warning=warning)
     return jdata