[bugfix] reset train metric when restoring a fine-tune checkpoint (#651)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Author: tiankongdeguiji

docs/source/models/evaluation_metrics.md

@@ -60,6 +60,8 @@ model_config {
   - decay_rate: 默认为0.9，历史指标的衰减率。
   - decay_step: 默认为100，历史指标每训练多少step会进行一次衰减，配置需要可以被train_config.log_step_count_steps整除。
 
+训练时指标会保存在checkpoint中，使用--continue_train续跑时会一并恢复。从fine_tune_checkpoint恢复时则会重置，避免继承上一个模型的指标，重置后在第一个decay_step之前打印为0。
+
 ______________________________________________________________________
 
 ## 指标详情

tzrec/main.py

@@ -348,6 +348,7 @@ def _train_and_evaluate(
     eval_result_filename: str = TRAIN_EVAL_RESULT_FILENAME,
     check_all_workers_data_status: bool = False,
     ignore_restore_optimizer: bool = False,
+    restore_from_model_dir: bool = False,
     dataloader_state: Optional[Dict[str, Any]] = None,
     delta_embedding_dumper: Optional[DeltaEmbeddingDumper] = None,
     pipeline_config_path: Optional[str] = None,
@@ -526,6 +527,10 @@ def run_eval(step: int, epoch: int) -> None:
                         train_config.fine_tune_ckpt_param_map,
                         dense_ema=dense_ema,
                     )
+                if not restore_from_model_dir:
+                    # a fine-tune checkpoint's train metric describes the source
+                    # model, not the model this job is training.
+                    _model.reset_train_metric()
                 if delta_embedding_dumper is not None:
                     delta_embedding_dumper.clear()
 
@@ -933,6 +938,7 @@ def train_and_evaluate(
         ckpt_path=ckpt_path,
         check_all_workers_data_status=check_all_workers_data_status,
         ignore_restore_optimizer=ignore_restore_optimizer,
+        restore_from_model_dir=restore_from_model_dir,
         dataloader_state=dataloader_state,
         delta_embedding_dumper=delta_embedding_dumper,
         pipeline_config_path=os.path.join(pipeline_config.model_dir, "pipeline.config"),

tzrec/metrics/train_metric_wrapper.py

@@ -26,18 +26,18 @@ class TrainMetricWrapper(nn.Module):
         decay_step (int): decay step for decay,
     """
 
+    _value: Tensor
+    _step_cnt: Tensor
+
     def __init__(
         self, metric_module: Metric, decay_rate: float = 0.5, decay_step: int = 100
     ) -> None:
         super().__init__()
         self._decay_rate = decay_rate
         self._decay_step = decay_step
         self._metric_module = metric_module
-        self._value = nn.Parameter(torch.tensor(0.0), requires_grad=False)
-        self._step_total_value = nn.Parameter(torch.tensor(0.0), requires_grad=False)
-        self._step_cnt = nn.Parameter(
-            torch.tensor(0, dtype=torch.int), requires_grad=False
-        )
+        self.register_buffer("_value", torch.tensor(0.0))
+        self.register_buffer("_step_cnt", torch.tensor(0, dtype=torch.int))
 
     def update(self, preds: Tensor, target: Tensor) -> None:
         """Update metric module."""
@@ -60,3 +60,9 @@ def update(self, preds: Tensor, target: Tensor) -> None:
     def compute(self) -> Tensor:
         """Get metric value."""
         return self._value.data
+
+    def reset(self) -> None:
+        """Reset metric state."""
+        self._metric_module.reset()
+        self._value.zero_()
+        self._step_cnt.zero_()

tzrec/metrics/train_metric_wrapper_test.py

@@ -48,6 +48,33 @@ def test_module_is_mean_absolute_error(self):
         value = metric_wrapper.compute()
         torch.testing.assert_close(value, torch.tensor(0.11))
 
+    def test_state_restore_and_reset(self):
+        metric_wrapper = TrainMetricWrapper(
+            torchmetrics.MeanAbsoluteError(), decay_rate=0.9, decay_step=1
+        )
+        preds = torch.tensor([0.1, 0.2])
+        metric_wrapper.update(preds, torch.tensor([0.6, 0.7]))
+        torch.testing.assert_close(metric_wrapper.compute(), torch.tensor(0.5))
+        self.assertEqual(
+            list(metric_wrapper.state_dict().keys()), ["_value", "_step_cnt"]
+        )
+        self.assertEqual(list(metric_wrapper.parameters()), [])
+
+        # continue train restores the running metric of the interrupted job.
+        restored = TrainMetricWrapper(
+            torchmetrics.MeanAbsoluteError(), decay_rate=0.9, decay_step=1
+        )
+        restored.load_state_dict(metric_wrapper.state_dict())
+        torch.testing.assert_close(restored.compute(), torch.tensor(0.5))
+        restored.update(preds, torch.tensor([0.2, 0.3]))
+        torch.testing.assert_close(restored.compute(), torch.tensor(0.46))
+
+        # fine tune resets it, so the next value is not blended with the old one.
+        restored.reset()
+        torch.testing.assert_close(restored.compute(), torch.tensor(0.0))
+        restored.update(preds, torch.tensor([0.2, 0.3]))
+        torch.testing.assert_close(restored.compute(), torch.tensor(0.1))
+
 
 if __name__ == "__main__":
     unittest.main()

tzrec/models/model.py

@@ -156,6 +156,11 @@ def compute_train_metric(self) -> Dict[str, torch.Tensor]:
             metric_results[metric_name] = metric.compute()
         return metric_results
 
+    def reset_train_metric(self) -> None:
+        """Reset train metric state."""
+        for metric in self._train_metric_modules.values():
+            metric.reset()
+
     def on_train_end(self) -> None:
         """Hook fired once after the train_eval loop exits.
 

tzrec/utils/export_util.py

@@ -639,7 +639,7 @@ def _rewrite_dense_serving_graph(
     dense_graph_config["sequence__ec"] = []
     for node in list(graph.nodes):
         if node.op == "call_function" and node.target == fx_mark_keyed_tensor:
-            name = node.args[0]
+            name = cast(str, node.args[0])
             if node.kwargs.get("is_dense", False):
                 continue
             node_kt = node.args[1]

tzrec/version.py

@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.4.0"
+__version__ = "1.4.1"