[Seq Packing] KK load balancing: balance packs across DP shards (P5)

[yuxuandexter]

What

Load balancing for sequence packing (P5), stacked on the loss-normalization PR (#1662). Replaces the greedy first-fit packer with Karmarkar-Karp balanced packing so the [num_packs, T] rows — one per DP shard — carry near-equal work.

Why stacked on P3 (not parallel)

KK changes which sequences land in which pack / how many sequences per microbatch, which changes each microbatch's sequence count and denominator. P3's Layer 2 (global gradient weighting, Σgrads / Σdenom) is exactly what keeps the loss correct when microbatch denominators vary. So KK's regrouping is only end-to-end correct on top of P3 — hence the sequential stack.

Change

• utils.py:
  • karmarkar_karp(vals, k) — largest-differencing heuristic (pure python + heapq): each sequence starts as its own arrangement, repeatedly merge the two most imbalanced arrangements (pairing each one's largest bucket with the other's smallest) until one k-bucket arrangement remains.
  • group_by_version(valid) — a pack never mixes policy_versions (single group under synchronous training).
  • balanced_pack(...) — balances by token count; picks k = smallest multiple of num_packs whose average pack fits the budget, then bumps k += num_packs until every pack fits T (a static-shape capacity guard KK itself does not enforce). Balances by workload, capacity-checks by tokens — kept separate for a future attention-aware (L^2) workload; see TODO.
  • pack_sequences: new packing_strategy arg — "kk" (default) or "first_fit". The first-fit block is unchanged, kept as a fallback.

What KK buys (and what it doesn't)

• Straggler removal (main win): balanced per-shard token counts → no shard lags the step.
• Dummy ratio: improves only in quantized jumps (the microbatch count is rounded to a num_packs multiple), so the gain is real but lumpy; balancing by token count also leaves attention O(L^2) imbalanced (TODO).

Tests (tests/rl/rl_utils_test.py)

• karmarkar_karp known answers: [8,7,6,5,4] k=2 → spread 2 (14/16); [8..1] k=4 → all 9. No drop/dup; k > n leaves empty groups.
• pack_sequences KK: preserves all sequences (no drop/dup); capacity guard keeps every pack <= T; invalid strategy raises. The existing first-fit layout test is pinned to packing_strategy="first_fit".

Notes

Stacked on yuxzhang/seqpack_loss_norm (#1662). Re-target to main once the parent PRs merge.