Add RL token throughput and packing metrics

megatron/rl/rl_utils.py

@@ -57,6 +57,10 @@
     get_sequence_packing_tensorboard_metrics,
     get_sequence_packing_log_info,
     get_default_packed_seq_params,
+    get_packing_actual_tokens,
+    get_packing_compute_tokens,
+    get_packing_efficiency,
+    get_packing_avg_seq_length,
     update_microbatch_calculator,
 )
 from megatron.rl.agent.api import (
@@ -301,11 +305,26 @@ def __init__(self):
         self.last_collection_iteration = 0
         self.sequences_this_iteration_on_rank = 0
         self.latest_batch_num_sequences = 0
+        # Derived throughput metrics (set by log_rl_throughput_metrics, read by RLProfiler).
+        # Per-GPU variants are available via methods that divide by world_size.
+        self.world_size = None
+        # batch_size * seq_length / time: nominal throughput based on batch configuration
+        self.tokens_per_sec = None
+        # Total tokens in packed bins across all DP ranks / time: what the GPU actually processes.
+        self.compute_tokens_per_sec = None
+        # Real non-padding tokens across all DP ranks / time: true useful throughput.
+        self.actual_tokens_per_sec = None
+        # Fraction of bin capacity filled with real tokens (actual / total capacity)
+        self.packing_efficiency = None
 
     def reset_iteration_counters(self, iteration):
         """Reset per-iteration counters."""
         self.sequences_this_iteration_on_rank = 0
         self.last_collection_iteration = iteration
+        self.tokens_per_sec = None
+        self.compute_tokens_per_sec = None
+        self.actual_tokens_per_sec = None
+        self.packing_efficiency = None
 
     def increment_sequences(self, count):
         """Increment the sequence counter."""
@@ -322,6 +341,99 @@ def get_rl_runtime_state():
     return _rl_runtime_state
 
 
+def log_rl_throughput_metrics(args, batch_size, elapsed_time_per_iteration, iteration, wandb_writer):
+    """Compute, log, and store RL token throughput metrics.
+
+    Returns a string fragment to append to the training log line.
+    Also logs metrics to wandb and stores them on RLRuntimeState for
+    downstream consumers (e.g. RLProfiler).
+    """
+    log_string = ''
+    tokens_per_sec = None
+    tokens_per_sec_per_gpu = None
+    compute_tokens_per_sec = None
+    compute_tokens_per_sec_per_gpu = None
+    actual_tokens_per_sec = None
+    actual_tokens_per_sec_per_gpu = None
+    packing_efficiency = None
+
+    if args.seq_length > 0:
+        tokens_per_iteration = batch_size * args.seq_length
+        tokens_per_sec = tokens_per_iteration / elapsed_time_per_iteration
+        tokens_per_sec_per_gpu = tokens_per_sec / args.world_size
+
+        # For sequence packing, break down into compute vs actual tokens
+        if args.rl_use_sequence_packing:
+            runtime_state = get_rl_runtime_state()
+            if runtime_state.packing_context is not None:
+                dp_world_size = mpu.get_data_parallel_world_size()
+
+                compute_tokens = get_packing_compute_tokens(runtime_state.packing_context)
+                all_ranks_compute_tokens = compute_tokens * dp_world_size
+                compute_tokens_per_sec = all_ranks_compute_tokens / elapsed_time_per_iteration
+                compute_tokens_per_sec_per_gpu = compute_tokens_per_sec / args.world_size
+
+                actual_tokens = get_packing_actual_tokens(runtime_state.packing_context)
+                all_ranks_actual_tokens = actual_tokens * dp_world_size
+                actual_tokens_per_sec = all_ranks_actual_tokens / elapsed_time_per_iteration
+                actual_tokens_per_sec_per_gpu = actual_tokens_per_sec / args.world_size
+
+                packing_efficiency = get_packing_efficiency(runtime_state.packing_context)
+
+        # Add tokens/sec to log string
+        log_string += f' toks/s: {tokens_per_sec:.0f} |'
+        log_string += f' toks/s/gpu: {tokens_per_sec_per_gpu:.0f} |'
+        if compute_tokens_per_sec is not None:
+            log_string += f' compute_toks/s: {compute_tokens_per_sec:.0f} |'
+            log_string += f' compute_toks/s/gpu: {compute_tokens_per_sec_per_gpu:.0f} |'
+        if actual_tokens_per_sec is not None:
+            log_string += f' actual_toks/s: {actual_tokens_per_sec:.0f} |'
+            log_string += f' actual_toks/s/gpu: {actual_tokens_per_sec_per_gpu:.0f} |'
+            log_string += f' packing_eff: {packing_efficiency:.1%} |'
+
+    # Log throughput metrics to wandb
+    if wandb_writer is not None:
+        if tokens_per_sec is not None:
+            wandb_writer.log({
+                'throughput/tokens_per_sec': tokens_per_sec,
+                'throughput/tokens_per_sec_per_gpu': tokens_per_sec_per_gpu,
+            }, iteration)
+        if compute_tokens_per_sec is not None:
+            wandb_writer.log({
+                'throughput/compute_tokens_per_sec': compute_tokens_per_sec,
+                'throughput/compute_tokens_per_sec_per_gpu': compute_tokens_per_sec_per_gpu,
+            }, iteration)
+        if actual_tokens_per_sec is not None:
+            wandb_writer.log({
+                'throughput/actual_tokens_per_sec': actual_tokens_per_sec,
+                'throughput/actual_tokens_per_sec_per_gpu': actual_tokens_per_sec_per_gpu,
+                'throughput/packing_efficiency': packing_efficiency,
+            }, iteration)
+
+    # Store derived throughput metrics on RLRuntimeState so that
+    # downstream consumers (e.g. RLProfiler) can read them.
+    # Per-GPU values are derived via methods on RLRuntimeState.
+    runtime_state = get_rl_runtime_state()
+    runtime_state.world_size = args.world_size
+    runtime_state.tokens_per_sec = tokens_per_sec
+    runtime_state.compute_tokens_per_sec = compute_tokens_per_sec
+    runtime_state.actual_tokens_per_sec = actual_tokens_per_sec
+    runtime_state.packing_efficiency = packing_efficiency
+
+    # Log average sequence length. With packing this shows real sequence
+    # lengths; without packing it equals seq_length as a baseline.
+    packing_ctx = runtime_state.packing_context
+    if args.rl_use_sequence_packing and packing_ctx is not None:
+        avg_seq_length = get_packing_avg_seq_length(packing_ctx)
+        log_string += f' avg_seq_len: {avg_seq_length:.1f} |'
+        if wandb_writer is not None:
+            wandb_writer.log({'throughput/avg_seq_length': avg_seq_length}, iteration)
+    elif args.log_throughput:
+        log_string += f' avg_seq_len: {args.seq_length} |'
+
+    return log_string
+
+
 def update_inference_logprobs_group_stats(
     old_logprobs: torch.Tensor,
     inference_logprobs: torch.Tensor,

megatron/rl/sequence_packing_utils.py

@@ -1117,3 +1117,61 @@ def get_sequence_packing_tensorboard_metrics(args):
         metrics['bin-batch-size'] = bin_batch_size
         metrics['consumed-bins'] = args.consumed_train_bins
     return metrics
+
+
+def get_packing_actual_tokens(packing_context: PackingContext) -> int:
+    """Get the actual number of tokens (non-padding) in the packed sequences for this rank.
+
+    Args:
+        packing_context: The PackingContext containing packing information.
+
+    Returns:
+        Total number of actual tokens across all bins on this rank.
+    """
+    return sum(
+        packing_context.packing_info.seq_lengths[idx]
+        for indices in packing_context.packing_info.bin_seq_indices
+        for idx in indices
+    )
+
+
+def get_packing_compute_tokens(packing_context: PackingContext) -> int:
+    """Get the total compute tokens (including padding) for packed sequences on this rank.
+
+    Args:
+        packing_context: The PackingContext containing packing information.
+
+    Returns:
+        Total compute tokens (num_bins * bin_size) on this rank.
+    """
+    return packing_context.packed_trajs.shape[0] * packing_context.packed_trajs.shape[1]
+
+
+def get_packing_efficiency(packing_context: PackingContext) -> float:
+    """Get the packing efficiency (actual_tokens / total_capacity) across all DP ranks.
+
+    Args:
+        packing_context: The PackingContext containing packing information.
+
+    Returns:
+        Packing efficiency as a float between 0 and 1.
+    """
+    total_actual_tokens = sum(packing_context.packing_info.seq_lengths)
+    num_ranks = mpu.get_data_parallel_world_size()
+    bins_per_rank = packing_context.packed_trajs.shape[0]
+    bin_size = packing_context.packed_trajs.shape[1]
+    total_capacity = bins_per_rank * bin_size * num_ranks
+
+    if total_capacity == 0:
+        return 0.0
+
+    return total_actual_tokens / total_capacity
+
+
+def get_packing_avg_seq_length(packing_context: PackingContext) -> float:
+    """Get the average sequence length across all sequences in the packing context."""
+    seq_lengths = packing_context.packing_info.seq_lengths
+    if not seq_lengths:
+        return 0.0
+
+    return sum(seq_lengths) / len(seq_lengths)

megatron/training/training.py

@@ -2293,6 +2293,13 @@ def training_log(
             total_loss_dict[skipped_iters_key]
         )
         log_string += ' number of nan iterations: {:3d} |'.format(total_loss_dict[nan_iters_key])
+
+        # RL token throughput metrics.
+        if args.perform_rl_step:
+            log_string += rl_utils.log_rl_throughput_metrics(
+                args, batch_size, elapsed_time_per_iteration, iteration, wandb_writer,
+            )
+
         if should_reset:
             total_loss_dict[advanced_iters_key] = 0
             total_loss_dict[skipped_iters_key] = 0

tests/unit_tests/rl/test_sequence_packing_utils.py

@@ -409,6 +409,50 @@ def test_compute_packed_inference_logprobs_stats_shape_mismatch():
     assert group_stats.mean_piold_to_inf_prob is None
 
 
+def test_packing_observability_metrics():
+    """Test various observability metrics related to sequence packing."""
+
+    # 4 sequences with known lengths packed into 2 bins of size 16.
+    # Bin 0 holds seqs 0 (len 5) and 1 (len 3) → 8 actual tokens
+    # Bin 1 holds seqs 2 (len 10) and 3 (len 4) → 14 actual tokens
+    seq_lengths = [5, 3, 10, 4]
+    packing_info = sequence_packing_utils.PackingInfo(
+        bin_seq_indices=[[0, 1], [2, 3]],
+        seq_starts={0: [0, 5], 1: [0, 10]},
+        seq_lengths=seq_lengths,
+        seq_to_bin_idx=[0, 0, 1, 1],
+        packing_algo='fifo',
+    )
+
+    num_bins, bin_size = 2, 16
+    packed_trajs = torch.zeros(num_bins, bin_size, dtype=torch.long)
+    ctx = sequence_packing_utils.PackingContext(
+        bin_size=bin_size,
+        packer=None,
+        packing_info=packing_info,
+        original_generation_masks=None,
+        original_trajs=None,
+        packed_trajs=packed_trajs,
+        packed_position_ids=None,
+        packed_loss_mask=None,
+    )
+
+    # actual tokens = sum of all seq_lengths referenced by bin_seq_indices
+    assert sequence_packing_utils.get_packing_actual_tokens(ctx) == 5 + 3 + 10 + 4
+
+    # compute tokens = num_bins * bin_size
+    assert sequence_packing_utils.get_packing_compute_tokens(ctx) == 2 * 16
+
+    # avg seq length = mean of seq_lengths
+    assert sequence_packing_utils.get_packing_avg_seq_length(ctx) == pytest.approx(22 / 4)
+
+    # efficiency = total_actual / (bins_per_rank * bin_size * num_ranks)
+    with patch('megatron.core.mpu.get_data_parallel_world_size', return_value=4):
+        eff = sequence_packing_utils.get_packing_efficiency(ctx)
+        # total_actual = sum(seq_lengths) = 22, capacity = 2 * 16 * 4 = 128
+        assert eff == pytest.approx(22 / 128)
+
+
 @pytest.mark.parametrize("num_sequences", [1, 10, 48, 49, 50])
 def test_cu_seqlens_size(num_sequences):
     """Test that cu_seqlens always has a fixed size regardless of how many sequences are packed."""