feat: log generation ISL/OSL histogram to wandb

nemo_rl/algorithms/grpo.py

@@ -40,6 +40,7 @@
 from nemo_rl.algorithms.utils import (
     calculate_baseline_and_std_per_prompt,
     log_generation_metrics_to_wandb,
+    log_histogram_metrics_to_wandb,
     print_performance_metrics,
     set_seed,
 )
@@ -1566,6 +1567,23 @@ def grpo_train(
                     logger,
                 )
 
+            # Plot ISL/OSL/ISL+OSL histograms to wandb
+            try:
+                for hist_metrics in [
+                    "gen_tokens_lengths",
+                    "input_tokens_lengths",
+                    "total_tokens_lengths",
+                ]:
+                    log_histogram_metrics_to_wandb(
+                        f"generation_metrics/{hist_metrics}",
+                        metrics[hist_metrics],
+                        total_steps + 1,
+                        logger,
+                    )
+            except Exception as e:
+                print(f"❌ Error plotting histograms to wandb: {e}")
+                pass
+
             print("\n📊 Training Results:")
 
             print(f"  • Loss: {metrics['loss']:.4f}")
@@ -2489,6 +2507,23 @@ def async_grpo_train(
                     logger,
                 )
 
+            # Plot ISL/OSL/ISL+OSL histograms to wandb
+            try:
+                for hist_metrics in [
+                    "gen_tokens_lengths",
+                    "input_tokens_lengths",
+                    "total_tokens_lengths",
+                ]:
+                    log_histogram_metrics_to_wandb(
+                        f"generation_metrics/{hist_metrics}",
+                        metrics[hist_metrics],
+                        step + 1,
+                        logger,
+                    )
+            except Exception as e:
+                print(f"❌ Error plotting histograms to wandb: {e}")
+                pass
+
             print("\n📊 Training Results:")
             print(f"  • Loss: {metrics['loss']:.4f}")
             print(f"  • Generation KL Error: {metrics['gen_kl_error']:.4f}")

nemo_rl/algorithms/utils.py

@@ -748,24 +748,46 @@ def visualize_per_worker_timeline(
 
 
 def log_generation_metrics_to_wandb(
-    vllm_logger_metrics: dict[str, dict[int, list[Any]]],
+    generation_logger_metrics: dict[str, dict[int, list[Any]]],
     step: int,
     timeline_interval: float,
     logger: Logger,
 ) -> None:
-    """Log vLLM metrics to wandb.
+    """Log generation metrics to wandb.
 
     Args:
-        vllm_logger_metrics: Dictionary of vLLM logger metrics
+        generation_logger_metrics: Dictionary of generation logger metrics
         step: Global step value
         timeline_interval: Interval between timeline points (in seconds)
         logger: Logger instance
     """
-    for vllm_metric in vllm_logger_metrics.keys():
+    for generation_metric in generation_logger_metrics.keys():
         logger.log_plot_per_worker_timeline_metrics(
-            vllm_logger_metrics[vllm_metric],
+            generation_logger_metrics[generation_metric],
             step=step,
-            prefix="vllm_metrics",
-            name=vllm_metric,
+            prefix="generation_metrics",
+            name=generation_metric,
             timeline_interval=timeline_interval,
         )
+
+
+def log_histogram_metrics_to_wandb(
+    metric_name: str,
+    metric_values: list[Any],
+    step: int,
+    logger: Logger,
+) -> None:
+    """Log histogram metrics to wandb.
+
+    Args:
+        metric_name: Name of the metric
+        metric_values: List of metric values
+        step: Global step value
+        logger: Logger instance
+    """
+    if logger.wandb_logger:
+        import wandb  # pyright: ignore[reportMissingImports]
+
+        logger.wandb_logger.run.log(
+            {metric_name: wandb.Histogram(metric_values)}, step=step
+        )

nemo_rl/experience/rollouts.py

@@ -21,6 +21,7 @@
 import statistics
 from collections import defaultdict
 from dataclasses import dataclass
+from functools import reduce
 from typing import Any, Optional
 
 import ray
@@ -654,6 +655,8 @@ async def run_sample_multi_turn_rollout(
 
     # Track per-turn metrics
     turn_gen_tokens = []
+    turn_input_tokens = []
+    turn_total_tokens = []
     # Track per-turn per-worker token accounting if available
     per_worker_token_counts = {}  # worker_idx -> token_count
 
@@ -685,6 +688,8 @@ async def run_sample_multi_turn_rollout(
             assistant_token_count += gen_token_count
             token_count += gen_token_count
             turn_gen_tokens.append(gen_token_count)
+            turn_input_tokens.append(int(input_lengths))
+            turn_total_tokens.append(int(input_lengths) + gen_token_count)
             # Per-worker load accounting
             if "gen_leader_worker_idx" in gen_metrics:
                 worker_idx = int(gen_metrics["gen_leader_worker_idx"])
@@ -770,6 +775,8 @@ async def run_sample_multi_turn_rollout(
         "max_turns_reached": max_turns_reached,
         "total_reward": total_reward,
         "turn_gen_tokens": turn_gen_tokens,
+        "turn_input_tokens": turn_input_tokens,
+        "turn_total_tokens": turn_total_tokens,
         # Pass-through per-worker per-turn accounting for aggregation at batch level
         "per_worker_token_counts": per_worker_token_counts,
     }
@@ -930,6 +937,17 @@ async def run_single_sample_with_error_handling(i, sample_state):
                     per_worker_token_counts[k] = per_worker_token_counts.get(k, 0) + v
             rollout_metrics["per_worker_token_counts"] = per_worker_token_counts
 
+        # Collect ISL, OSL, and ISL+OSL metrics for all samples
+        rollout_metrics["gen_tokens_lengths"] = reduce(
+            lambda x, y: x + y, [m["turn_gen_tokens"] for m in all_sample_metrics]
+        )
+        rollout_metrics["input_tokens_lengths"] = reduce(
+            lambda x, y: x + y, [m["turn_input_tokens"] for m in all_sample_metrics]
+        )
+        rollout_metrics["total_tokens_lengths"] = reduce(
+            lambda x, y: x + y, [m["turn_total_tokens"] for m in all_sample_metrics]
+        )
+
         return final_batch, rollout_metrics
 
     return asyncio.run(_async_rollout_implementation())