Stabilize pipeline RL throughput controls

src/art/megatron/train.py

@@ -108,6 +108,7 @@
     _zero_contribution_inputs,
     _zero_contribution_sft_inputs,
     build_micro_sample_indices,
+    build_micro_sample_indices_by_dp_rank,
     build_rl_hybridep_token_counts,
     build_sft_hybridep_token_counts,
     resolve_global_grad_accumulation_sequences,
@@ -649,7 +650,12 @@ def run_megatron_rl_job(
                 hybridep_token_counts=hybridep_token_counts,
             )
             train_step_s = time.perf_counter() - train_step_started
-            global_packed_train_tokens = _global_packed_train_tokens(micro_inputs)
+            global_packed_train_tokens = _global_packed_train_tokens(
+                packed_tensors,
+                step_index=step_index,
+                num_sequences=num_sequences,
+                global_grad_accumulation_sequences=global_grad_accumulation_sequences,
+            )
             print0(
                 runtime.rank,
                 "Correlation between old and new probabilities:",
@@ -1108,6 +1114,7 @@ def _log_rl_step_result(
             "loss/grad_norm": step_result.grad_norm,
             "loss/probs_corr": step_result.probs_corr,
             TRAIN_GRADIENT_STEPS_KEY: num_gradient_steps,
+            "data/step_executed_packed_train_tokens": packed_train_tokens,
             "throughput/train_packed_tok_per_s": train_packed_tok_per_s,
         }
         if step_result.kl_policy_ref is not None:
@@ -1118,9 +1125,30 @@ def _log_rl_step_result(
         log_file.write(log_msg + "\n")
 
 
-def _global_packed_train_tokens(micro_inputs: list[PackedTensors]) -> int:
-    local_tokens = sum(int(micro["tokens"].numel()) for micro in micro_inputs)
-    return local_tokens * ps.get_data_parallel_world_size(with_context_parallel=False)
+def _global_packed_train_tokens(
+    packed_tensors: PackedTensors,
+    *,
+    step_index: int,
+    num_sequences: int,
+    global_grad_accumulation_sequences: int | None,
+) -> int:
+    sample_rows = build_micro_sample_indices_by_dp_rank(
+        step_index=step_index,
+        num_sequences=num_sequences,
+        global_grad_accumulation_sequences=global_grad_accumulation_sequences,
+    )
+    sequence_length = int(packed_tensors["tokens"].shape[1])
+    if ps.get_context_parallel_world_size() <= 1:
+        return sum(len(row) for row in sample_rows) * sequence_length
+    return sum(
+        int(
+            (packed_tensors["group_ids"][0 if index is None else index] != -1)
+            .sum()
+            .item()
+        )
+        for row in sample_rows
+        for index in row
+    )
 
 
 def _save_optimizer(

src/art/metrics.py

@@ -113,11 +113,23 @@ class MetricDefinition(pydantic.BaseModel):
         higher_is_better=False,
         dashboard_default=True,
     ),
+    MetricDefinition(
+        key="data/step_executed_packed_train_tokens",
+        title="Megatron executed packed train tokens",
+        description=(
+            "packed token rows included in Megatron throughput; CP excludes "
+            "configured packed-row padding that is not dispatched"
+        ),
+        kind="counter",
+        unit="tokens",
+        higher_is_better=None,
+    ),
     MetricDefinition(
         key="throughput/train_packed_tok_per_s",
         title="Megatron packed train tokens per second",
         description=(
-            "physical packed training-token throughput reported by the Megatron worker"
+            "physical training-token throughput reported by the Megatron worker; "
+            "CP excludes configured packed-row padding that is not dispatched"
         ),
         kind="rate",
         unit="tok/s",
@@ -276,6 +288,30 @@ class MetricDefinition(pydantic.BaseModel):
         kind="ratio",
         higher_is_better=False,
     ),
+    MetricDefinition(
+        key="queue/actual_stale_fraction",
+        title="Actual stale dequeue fraction",
+        description="stale groups divided by all groups dequeued for this train step",
+        kind="ratio",
+        higher_is_better=False,
+    ),
+    MetricDefinition(
+        key="queue/put_wait_frac",
+        title="Completed queue wait fraction",
+        description=(
+            "worker queue-put wait divided by queue-put wait plus rollout time"
+        ),
+        kind="ratio",
+        higher_is_better=False,
+    ),
+    MetricDefinition(
+        key="queue/put_wait_s",
+        title="Completed queue wait",
+        description="worker-seconds spent waiting to enqueue completed rollout groups",
+        kind="duration",
+        unit="seconds",
+        higher_is_better=False,
+    ),
 )
 
 PIPELINE_RL_DASHBOARD_DEFAULT_METRICS = tuple(

src/art/model.py

@@ -340,6 +340,7 @@ def __getattr__(self, name: str) -> Any:
         "offpolicy/token_weighted_policy_age_p95_steps",
         "throughput/accepted_train_tok_per_s",
         "throughput/train_packed_tok_per_s",
+        "data/step_executed_packed_train_tokens",
         "data/step_trainable_assistant_tokens",
         "data/step_non_padding_train_tokens",
         "data/step_padding_ratio",
@@ -348,6 +349,8 @@ def __getattr__(self, name: str) -> Any:
         "data/cum/num_gradient_steps",
         "discarded/cum/stale_groups",
         "discarded/cum/zero_variance_groups",
+        "discarded/step/stale_groups",
+        "discarded/step/zero_variance_groups",
         "discarded/rate/stale_groups",
         "discarded/rate/zero_variance_groups",
         "time/step_wall_s",
@@ -360,6 +363,9 @@ def __getattr__(self, name: str) -> Any:
         "pipeline_settings/max_batch_size",
         "pipeline_settings/target_groups_per_step",
         "pipeline_settings/queue_maxsize",
+        "queue/actual_stale_fraction",
+        "queue/put_wait_frac",
+        "queue/put_wait_s",
         "loss/train",
         "loss/entropy",
         "loss/kl_div",