fix: swanlab logger error caused by define_metric (#1615)

Zeyi-Lin · clumsy · azzhipa · web-flow · commit b33e440b80b7 · 2025-12-13T03:40:47.000Z
Signed-off-by: ZeYi Lin &lt;944270057@qq.com&gt;
Signed-off-by: Alexander Zhipa &lt;azzhipa@amazon.com&gt;
Co-authored-by: Alexander Zhipa &lt;alex.zhipa@proton.me&gt;
Co-authored-by: Alexander Zhipa &lt;azzhipa@amazon.com&gt;
Co-authored-by: Terry Kong &lt;terrycurtiskong@gmail.com&gt;
diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
@@ -226,6 +226,9 @@ logger:
     wandb:
         project: "nemo-distillation"
         name: "distillation-${data.dataset_name}-${teacher.model_name}-${policy.model_name}-${loss_fn.kl_type}-${distillation.topk_logits_k}"
+    swanlab:
+        project: "nemo-distillation"
+        name: "distillation-${data.dataset_name}-${teacher.model_name}-${policy.model_name}-${loss_fn.kl_type}-${distillation.topk_logits_k}"
     tensorboard:
         log_dir: "tb_logs-distillation-${data.dataset_name}"
     mlflow:
diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml
@@ -207,6 +207,9 @@ logger:
   wandb:
     project: "dpo-dev"
     name: "dpo"
+  swanlab:
+    project: "dpo-dev"
+    name: "dpo"
   tensorboard:
     log_dir: "tb_logs-dpo-dev"
   mlflow:
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
@@ -283,6 +283,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "grpo-dev-logger"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-dev-logger"
   tensorboard: {}
   mlflow:
     experiment_name: "grpo-dev"
diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml
@@ -175,6 +175,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "sj_megatron_1B"
+  swanlab:
+    project: "grpo-dev"
+    name: "sj_megatron_1B"
   tensorboard: {}
   mlflow:
     experiment_name: "grpo-dev"
diff --git a/examples/configs/grpo_sliding_puzzle.yaml b/examples/configs/grpo_sliding_puzzle.yaml
@@ -68,6 +68,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "grpo-dev-sliding_puzzle"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-dev-sliding_puzzle"
   tensorboard: {}
   mlflow:
     experiment_name: "grpo-dev"
diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml
@@ -171,6 +171,9 @@ logger:
   wandb:
     project: "rm-dev"
     name: "rm-dev-${data.dataset_name}"
+  swanlab:
+    project: "rm-dev"
+    name: "rm-dev-${data.dataset_name}"
   tensorboard:
     log_dir: "tb_logs-rm-dev-${data.dataset_name}"
   gpu_monitoring:
diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml
@@ -203,6 +203,9 @@ logger:
   wandb:
     project: "sft-dev"
     name: "sft-dev-${data.dataset_name}"
+  swanlab:
+    project: "sft-dev"
+    name: "sft-dev-${data.dataset_name}"
   tensorboard:
     log_dir: "tb_logs-sft-dev-${data.dataset_name}"
   mlflow:
diff --git a/examples/configs/sft_openmathinstruct2.yaml b/examples/configs/sft_openmathinstruct2.yaml
@@ -88,6 +88,9 @@ logger:
   wandb:
     project: "sft-dev"
     name: "openmathinstruct-nemorl-1M_train"
+  swanlab:
+    project: "sft-dev"
+    name: "openmathinstruct-nemorl-1M_train"
   tensorboard:
     log_dir: "tb_logs-openmathinstruct-nemorl-1M_train"
   mlflow:
diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml
@@ -271,6 +271,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "grpo-dev-logger"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-dev-logger"
   tensorboard: {}
   gpu_monitoring:
     collection_interval: 10  # How often to collect GPU usage metrics (in seconds)
diff --git a/examples/configs/vlm_grpo_3B_megatron.yaml b/examples/configs/vlm_grpo_3B_megatron.yaml
@@ -221,6 +221,9 @@ logger:
   wandb:
     project: grpo-dev
     name: vlm-grpo-3b-megatron
+  swanlab:
+    project: grpo-dev
+    name: vlm-grpo-3b-megatron
   tensorboard: {}
   gpu_monitoring:
     collection_interval: 10
diff --git a/examples/nemo_gym/grpo_dapo17k_bytedtsinghua_qwen3_4binstruct_nf.yaml b/examples/nemo_gym/grpo_dapo17k_bytedtsinghua_qwen3_4binstruct_nf.yaml
@@ -262,6 +262,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "grpo-dev-logger"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-dev-logger"
   tensorboard: {}
   mlflow:
     experiment_name: "grpo-dev"
diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py
@@ -374,58 +374,45 @@ class SwanlabLogger(LoggerInterface):
     """SwanLab logger backend."""
 
     def __init__(self, cfg: SwanlabConfig, log_dir: Optional[str] = None):
+        """Initialize the SwanlabLogger by starting a Swanlab run and storing the resulting run on self.run.
+
+        Parameters:
+            cfg (SwanlabConfig): Configuration for the Swanlab run (e.g., project and name).
+            log_dir (Optional[str]): Optional offline log directory passed to Swanlab's init.
+        """
         self.run = swanlab.init(**cfg, logdir=log_dir)
         print(
             f"Initialized SwanlabLogger for project {cfg.get('project')}, run {cfg.get('name')} (with offline logdir={log_dir})"
         )
 
-    def define_metric(
-        self,
-        name: str,
-        step_metric: Optional[str] = None,
-    ) -> None:
-        """Define a metric with custom step metric.
-
-        Args:
-            name: Name of the metric or pattern (e.g. 'ray/*')
-            step_metric: Optional name of the step metric to use
-        """
-        self.run.define_metric(name, step_metric=step_metric)
-
     def log_metrics(
         self,
         metrics: dict[str, Any],
         step: int,
         prefix: Optional[str] = "",
         step_metric: Optional[str] = None,
     ) -> None:
-        """Log metrics to swanlab.
+        """Log metrics to the associated Swanlab run.
 
-        Args:
-            metrics: Dict of metrics to log
-            step: Global step value
-            prefix: Optional prefix for metric names
-            step_metric: Optional name of a field in metrics to use as step instead
-                         of the provided step value
+        Parameters:
+            metrics (dict[str, Any]): Mapping of metric names to metric values.
+            step (int): Global step value to associate with all logged metrics.
+            prefix (Optional[str]): Optional prefix applied to metric names; metric names equal to `step_metric` are not prefixed.
+            step_metric (Optional[str]): Name of a metric that should be excluded from prefixing.
         """
         if prefix:
             metrics = {
                 f"{prefix}/{k}" if k != step_metric else k: v
                 for k, v in metrics.items()
             }
 
-        # If step_metric is provided, use the corresponding value from metrics as step
-        if step_metric and step_metric in metrics:
-            # commit=False so the step does not get incremented
-            self.run.log(metrics, commit=False)
-        else:
-            self.run.log(metrics, step=step)
+        self.run.log(metrics, step=step)
 
     def log_hyperparams(self, params: Mapping[str, Any]) -> None:
-        """Log hyperparameters to swanlab.
+        """Update the Swanlab run configuration with the provided hyperparameters.
 
-        Args:
-            params: Dict of hyperparameters to log
+        Parameters:
+            params (Mapping[str, Any]): Mapping of hyperparameter names to values to store in the run configuration.
         """
         self.run.config.update(params)
 
@@ -436,7 +423,7 @@ def log_plot(self, figure: plt.Figure, step: int, name: str) -> None:
             figure: Matplotlib figure to log
             step: Global step value
         """
-        self.run.log({name: figure}, step=step)
+        self.run.log({name: swanlab.Image(figure)}, step=step)
 
     def log_histogram(self, histogram: list[Any], step: int, name: str) -> None:
         """Log histogram metrics to swanlab."""
@@ -833,19 +820,15 @@ class Logger(LoggerInterface):
     """Main logger class that delegates to multiple backend loggers."""
 
     def __init__(self, cfg: LoggerConfig):
-        """Initialize the logger.
-
-        Args:
-            cfg: Config dict with the following keys:
-                - wandb_enabled
-                - tensorboard_enabled
-                - mlflow_enabled
-                - wandb
-                - tensorboard
-                - mlflow
-                - monitor_gpus
-                - gpu_collection_interval
-                - gpu_flush_interval
+        """Create and configure enabled logging backends and optionally start GPU monitoring.
+
+        Parameters:
+            cfg (LoggerConfig): Configuration mapping. Expected keys include:
+                - "log_dir": base directory for backend logs.
+                - "wandb_enabled", "swanlab_enabled", "tensorboard_enabled", "mlflow_enabled": booleans to enable backends.
+                - "wandb", "swanlab", "tensorboard", "mlflow": per-backend configuration dicts.
+                - "monitor_gpus": boolean to enable Ray GPU monitoring.
+                - "gpu_monitoring": dict with "collection_interval" and "flush_interval" when GPU monitoring is enabled.
         """
         self.loggers: list[LoggerInterface] = []
         self.wandb_logger = None
@@ -892,11 +875,6 @@ def __init__(self, cfg: LoggerConfig):
                     f"{metric_prefix}/*", step_metric=step_metric
                 )
 
-            if cfg["swanlab_enabled"] and self.swanlab_logger:
-                self.swanlab_logger.define_metric(
-                    f"{metric_prefix}/*", step_metric=step_metric
-                )
-
             self.gpu_monitor = RayGpuMonitorLogger(
                 collection_interval=cfg["gpu_monitoring"]["collection_interval"],
                 flush_interval=cfg["gpu_monitoring"]["flush_interval"],
diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml
@@ -242,6 +242,9 @@ logger:
   wandb:
     project: "grpo-dev"
     name: "grpo-dev-logger"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-dev-logger"
   tensorboard: {}
   mlflow:
     experiment_name: "grpo-dev"
diff --git a/tests/unit/utils/test_logger.py b/tests/unit/utils/test_logger.py
@@ -341,10 +341,9 @@ def test_log_metrics_with_step_metric(self, mock_swanlab):
 
         logger.log_metrics(metrics, step, step_metric=step_metric)
 
-        # Check that log was called with metrics and commit=False
-        # When using step_metric, step should be ignored and commit=False should be used
+        # Check that log was called with metrics
         mock_run = mock_swanlab.init.return_value
-        mock_run.log.assert_called_once_with(metrics, commit=False)
+        mock_run.log.assert_called_once_with(metrics, step=step)
 
     @patch("nemo_rl.utils.logger.swanlab")
     def test_log_metrics_with_prefix_and_step_metric(self, mock_swanlab):
@@ -362,30 +361,14 @@ def test_log_metrics_with_prefix_and_step_metric(self, mock_swanlab):
 
         logger.log_metrics(metrics, step, prefix=prefix, step_metric=step_metric)
 
-        # Check that log was called with prefixed metrics and commit=False
-        # The step_metric key gets prefixed based on the current implementation
+        # Check that log was called with prefixed metrics
         mock_run = mock_swanlab.init.return_value
         expected_metrics = {
             "train/loss": 0.5,
             "train/accuracy": 0.8,
             "train/iteration": 15,
         }
-        mock_run.log.assert_called_once_with(expected_metrics, commit=False)
-
-    @patch("nemo_rl.utils.logger.swanlab")
-    def test_define_metric(self, mock_swanlab):
-        """Test defining a metric with a custom step metric."""
-        cfg = {}
-        logger = SwanlabLogger(cfg)
-
-        # Define metric pattern and step metric
-        logger.define_metric("ray/*", step_metric="ray/ray_step")
-
-        # Check that define_metric was called
-        mock_run = mock_swanlab.init.return_value
-        mock_run.define_metric.assert_called_once_with(
-            "ray/*", step_metric="ray/ray_step"
-        )
+        mock_run.log.assert_called_once_with(expected_metrics, step=step)
 
     @patch("nemo_rl.utils.logger.swanlab")
     def test_log_hyperparams(self, mock_swanlab):
