fix: Fix the logger error in non-colocated sync-grpo code path (#1355)

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Author: youngeunkwon0405

nemo_rl/algorithms/grpo.py

@@ -1798,9 +1798,6 @@ def async_grpo_train(
                 train_results, metrics, timing_metrics, master_config
             )
 
-            if "per_worker_token_counts" in metrics:
-                del metrics["per_worker_token_counts"]
-
             logger.log_metrics(performance_metrics, step + 1, prefix="performance")
             logger.log_metrics(metrics, step + 1, prefix="train")
             logger.log_metrics(timing_metrics, step + 1, prefix="timing/train")

nemo_rl/algorithms/utils.py

@@ -553,6 +553,15 @@ def visualize_per_worker_load(per_worker_token_counts: dict[int, int]) -> float:
                 total_tflops / theoretical_tflops
             )
 
+    # =====================================================
+    # Clean up metrics
+    # =====================================================
+
+    # Clean up metrics to avoid wandb logging errors
+    # Dict structures cannot be logged to wandb
+    if "per_worker_token_counts" in metrics:
+        del metrics["per_worker_token_counts"]
+
     # =====================================================
     # Logging
     # =====================================================