[Seq Packing] Loss normalization under packing (P3): unpack aggregation + global grad weighting

[yuxuandexter]

What

Loss normalization correctness under sequence packing (P3), stacked on the packed-logprob + pack-first PR (#1659). Two independent layers:

Layer 1 — per-sequence aggregation (unpack, veRL pad_input equivalent)

With packing, per_token_loss is [num_packs, T] where each row holds ~8 sequences, but aggregate_loss only knows completion_mask — it treats a whole pack row as one sequence, so every sequence-mean mode is wrong (only token-mean is packing-invariant).

Fix (mirrors veRL, which unpacks logprobs back to [batch, resp_len] before the loss so aggregation is transparent to packing):

• common.py: make_unpack_indices + aggregate_loss_packed. It scatters the packed [num_packs, T] tensor back to a rectangular [n_max, T] (row = sequence) via a segment_ids scatter, then runs the original aggregate_loss so all 5 modes reduce per-sequence. segment_ids=None delegates unchanged (non-packed path byte-identical).
• common.py: sequence-mean-token-mean / sequence-mean-token-scale denominators use the non-empty-row count so unpacked padding rows do not inflate them.
• algo_core.py: 4 aggregate_loss call sites (clipfrac / main loss / kl / entropy) → aggregate_loss_packed.
• agentic_grpo_learner.py: pass unpack_n_max = train_micro_batch_size * num_generations (a static config constant).

Layer 2 — cross-micro-batch gradient (b/491970038)

peft_trainer accumulated a mean-of-means (local 1/denom scale + denom=1). Since diff_fn differentiates the unreduced sum, grads = grad(sum). Now it accumulates grad(sum) with the loss's real denominator, so the optimizer sees Σ grads / Σ denom — token-weighted for token-mean, sequence-weighted for sequence-mean. This is exact per mode (each mode contributes its own denominator), whereas veRL uses a fixed sequence count.

Tests

• tests/rl/common_test.py::AggregateLossPackedTest — all 5 modes: aggregate_loss_packed(packed) == aggregate_loss(rectangular) (atol 1e-4); token-mean ground truth; decision-B empty-row protection; segment_ids=None matches the original (non-packed unchanged).
• tests/sft/peft_trainer_test.py::TrainStepGlobalWeightedTest — uneven denominators across micro-batches: get() == Σ grads / Σ denom (global), and differs from mean-of-means.

Notes

Stacked on yuxzhang/unified_seqpack (#1659); depends on its LossOutput / WeightedMetric infra. Re-target to main once #1659 merges.