fix: skip zero-variance GRPO updates

src/art/local/backend.py

@@ -1114,6 +1114,13 @@ def _get_packed_tensors(
             pack_results=self._supports_result_packing,
             include_moe_routing=include_moe_routing,
         )
+        # Position 0 has no preceding token, so it never carries loss.
+        if not bool(
+            (
+                packed_tensors["advantages"][:, 1:] * packed_tensors["weights"][:, 1:]
+            ).any()
+        ):
+            return None
         if (
             not allow_training_without_logprobs
             and np.isnan(packed_tensors["logprobs"]).all()

src/art/megatron/backend.py

@@ -567,6 +567,8 @@ async def _prepare_training_batch(
         *,
         include_moe_routing: bool,
     ) -> _PackedTrainingBatch | None:
+        if not trajectory_groups:
+            return None
         prepared = tuple(group._prepared_training_batch for group in trajectory_groups)
         collect_packing_shapes = any(
             group._collect_packing_shape for group in trajectory_groups

src/art/metrics_taxonomy.py

@@ -4,6 +4,7 @@
 from typing import Any, Iterable
 
 from .trajectories import TrajectoryGroup
+from .utils.grpo import rewards_have_zero_variance
 
 TRAIN_GRADIENT_STEPS_KEY = "data/step_num_gradient_steps"
 SFT_METRIC_PREFIX = "sft"
@@ -107,7 +108,7 @@ def build_training_summary_metrics(
 
 def _group_is_trainable(group: TrajectoryGroup) -> bool:
     rewards = [trajectory.reward for trajectory in group.trajectories]
-    return len(rewards) > 1 and len(set(rewards)) > 1
+    return not rewards_have_zero_variance(rewards)
 
 
 def _extract_scenario_id(group: TrajectoryGroup) -> str | None:

src/art/pipeline_trainer/trainer.py

@@ -51,6 +51,7 @@
     RolloutWorkerController,
 )
 from art.preprocessing.policy_spans import PolicyTokenSpan
+from art.utils.grpo import rewards_have_zero_variance
 
 from .checkpoint_retention import (
     CHECKPOINT_CREATED_AT_METRIC,
@@ -1827,11 +1828,7 @@ async def _log_zero_variance_groups(self, step: int) -> None:
 
     @staticmethod
     def _group_zero_variance(group: TrajectoryGroup) -> bool:
-        rewards = [t.reward for t in group.trajectories]
-        if len(rewards) <= 1:
-            return True
-        first = rewards[0]
-        return all(abs(r - first) <= 1e-12 for r in rewards[1:])
+        return rewards_have_zero_variance([t.reward for t in group.trajectories])
 
     def _group_initial_version(self, group: TrajectoryGroup) -> int | None:
         versions = [

src/art/preprocessing/pack.py

@@ -341,9 +341,9 @@ def prefix_tree_pack(
         assistant_mask_tensor, weights_tensor, torch.zeros_like(weights_tensor)
     )
     if bool(assistant_mask_tensor.any()):
-        weights_tensor[assistant_mask_tensor] /= weights_tensor[
-            assistant_mask_tensor
-        ].mean()
+        weights_mean = weights_tensor[assistant_mask_tensor].mean()
+        if weights_mean > 0:
+            weights_tensor[assistant_mask_tensor] /= weights_mean
     advantages_tensor = torch.from_numpy(advantages_np)
     advantages_tensor = torch.where(
         assistant_mask_tensor, advantages_tensor, torch.zeros_like(advantages_tensor)
@@ -361,10 +361,12 @@ def prefix_tree_pack(
             advantages_tensor * (1 + advantage_balance),
         )
     if bool(assistant_mask_tensor.any()):
-        advantages_tensor[assistant_mask_tensor] /= (
+        advantages_mean = (
             advantages_tensor[assistant_mask_tensor].abs()
             * weights_tensor[assistant_mask_tensor]
         ).mean()
+        if advantages_mean > 0:
+            advantages_tensor[assistant_mask_tensor] /= advantages_mean
 
     packed_tensors: PackedTensors = {
         "tokens": torch.from_numpy(tokens_np),

src/art/preprocessing/tokenize.py

@@ -36,6 +36,7 @@
     merge_chat_template_kwargs,
     normalize_tool_call_arguments_for_chat_template,
 )
+from ..utils.grpo import rewards_have_zero_variance
 from .dynamo_tokens import choice_completion_logprobs, has_completion_logprobs
 from .moe_routing import (
     MoeRouteArray,
@@ -758,15 +759,20 @@ def tokenize_trajectory_groups(
         if not group:
             continue
         results: list[TokenizedResult] = []
+        zero_variance = rewards_have_zero_variance(
+            [trajectory.reward for trajectory in group]
+        )
         # Calculate GRPO group mean and standard deviation
         reward_mean = sum(trajectory.reward for trajectory in group) / len(group)
-        reward_std = math.sqrt(
-            sum((trajectory.reward - reward_mean) ** 2 for trajectory in group)
-            / len(group)
-        )
+        reward_std = 0.0
+        if not zero_variance:
+            reward_std = math.sqrt(
+                sum((trajectory.reward - reward_mean) ** 2 for trajectory in group)
+                / len(group)
+            )
         for trajectory in group:
             # Calculate GRPO advantage for this trajectory
-            advantage = trajectory.reward - reward_mean
+            advantage = 0.0 if zero_variance else trajectory.reward - reward_mean
             if scale_rewards:
                 advantage /= reward_std + 1e-6
             if advantage == 0 and drop_zero_advantage_trajectories:

src/art/utils/grpo.py

@@ -0,0 +1,8 @@
+from collections.abc import Sequence
+
+
+def rewards_have_zero_variance(rewards: Sequence[float]) -> bool:
+    if len(rewards) <= 1:
+        return True
+    first = rewards[0]
+    return all(abs(reward - first) <= 1e-12 for reward in rewards[1:])