[fix][metrics] Reduce policy_entropy and logprob-diff metrics via sum/count

[ayucpp]

Fixes #1822

What does this PR do?

Updates metric aggregation for policy_entropy and minibatch_rollout_logprobs_abs_diff_* to compute the exact global mean instead of a mean-of-means.

Why

Workers reported per-micro-batch masked means, averaged across micro-batches, DP ranks, and mini-batches — over-weighting shards with few unmasked tokens.

Example: Shard A = 1 token (entropy 1.0), Shard B = 100 tokens (entropy 2.0):

• Old (mean-of-means): (1.0 + 2.0) / 2 = 1.5
• New (exact): 201 / 101 ≈ 1.99

Token-based batching (or DP workers getting mostly failed trajectories) makes shards imbalanced, so the old value wasn't the global-batch metric.

What changes

• Workers emit masked per-token sums (policy_entropy_sum, minibatch_rollout_logprobs_abs_diff_sum/_sq_sum) and counts (loss_mask_nnz, ..._nnz); reduce_metrics/all_reduce_metrics always sum these.
• Final derivations (entropy = sum/count, std = sqrt(E[x²] − E[x]²)) moved to the trainer, computed once.
• Workers still emit a final policy_entropy after the DP all-reduce, so Tinker _extract_metrics and worker-level GPU assertions are unchanged.
• Megatron packed-entropy helpers now return the CP-global sum/count, keeping CP invariance.

Testing

• Added imbalanced-shard exactness tests in test_worker_utils.py and test_skip_fwd_logprobs.py.
• CPU suite: 1393 passed (-m "not vllm").
• Pre-commit clean.

[gemini-code-assist]

Code Review

This pull request refactors the calculation of policy entropy and rollout logprob difference metrics to prevent "mean-of-means" inaccuracies when dealing with imbalanced shards. Instead of averaging means at the worker level, the code now tracks and reduces raw masked sums and non-zero counts (nnz) across micro-batches, DP ranks, and mini-batches, performing a single final division at the trainer level. A critical issue was identified in megatron_model_wrapper.py where entropy_sum and entropy_nnz are returned as PyTorch tensors in packed sequence branches but as Python floats in non-packed branches. This type mismatch can cause the metrics to be silently skipped during the reduction step, so calling .item() on the tensors is recommended to ensure they are properly logged.

[gemini-code-assist]

In the packed sequence branches (e.g., when calling vocab_parallel_entropy_packed_sequences or from_parallel_hidden_to_entropy_packed_sequences), entropy_sum and entropy_nnz are returned as torch.Tensors on GPU. However, in the non-packed branches (using compute_masked_sum_and_count), they are returned as Python floats.\n\nIf they are torch.Tensors, they will be added to the metrics dictionary as tensors. Later, when reduce_metrics is called, it checks isinstance(x, (int, float)) for each metric value. Since torch.Tensor is not an instance of int or float, these metrics will be silently skipped and omitted from the reduced metrics, preventing policy_entropy from being finalized and logged.\n\nTo fix this, we should call .item() on entropy_sum and entropy_nnz if they are PyTorch tensors.

                POLICY_ENTROPY_SUM_KEY: entropy_sum.item() if torch.is_tensor(entropy_sum) else entropy_sum,\n                LOSS_MASK_NNZ_KEY: entropy_nnz.item() if torch.is_tensor(entropy_nnz) else entropy_nnz,

[ayucpp]

Fixed in ac1318c — I normalized to Python floats right after the branch chain (entropy_sum = float(entropy_sum); entropy_nnz = float(entropy_nnz)), which covers both the packed branches (tensors) and compute_masked_sum_and_count (already floats) in one place. float(tensor) on a scalar tensor is equivalent to .item(), so the metrics now survive the isinstance(x, (int, float)) check in reduce_metrics. Full CPU suite: 1393 passed, no regressions.