[tunix] Accumulate weighted loss gradients globally

[ethannnnnn]

Motivation

Normalizing each microbatch independently biases gradients when microbatches contain different numbers or fractional weights of valid tokens. Diffusion objectives require one global numerator and denominator.

Scope

• Accumulate LossOutput gradients as unreduced weighted sums.
• Normalize once by the total denominator across all microbatches.
• Aggregate primary and auxiliary weighted metrics with the same rule.
• Select auxiliary-metric reducers by value type, preserving ordinary scalar
  diagnostics alongside weighted objectives.
• Preserve scalar-loss behavior and define finite all-zero handling.
• Preserve each optimizer's state dtypes across jitted update/skip branches.

Design

For each microbatch, the trainer differentiates the loss numerator and accumulates its denominator. The final gradient is the sum of numerator gradients divided by the sum of denominators. A zero total denominator returns zero gradients without division by zero. Scalar/legacy loss functions use denominator one and retain their existing mean-of-microbatches behavior.

Before an optimizer update, the trainer records the optimizer-state dtype tree and restores floating leaves to that contract after Optax returns. This keeps nnx.cond update and skip branches type-identical for bfloat16 models without globally promoting optimizer state at construction time.

Buffered WeightedMetric values are combined as one numerator and denominator
in both training and evaluation. Plain array metrics continue to use their
existing arithmetic mean. The reducer is selected from the metric value
itself, so a weighted primary loss does not incorrectly force unrelated RL
diagnostics through the weighted metric API. A metric key cannot change kind
within one buffer, and combined weighted metrics must agree on eps and
min_denom; inconsistent inputs fail before producing misleading telemetry.

Compatibility

Ordinary scalar losses and plain auxiliary diagnostics are unchanged.
LossOutput callers receive mathematically correct weighted reduction for
weighted training and evaluation metrics. Direct and injected Optax
transformations retain their configured state dtype.

Extensibility

The reduction is objective agnostic and supports sparse masks, fractional weights, variable sequence lengths, and future weighted objectives beyond diffusion.

Tests

• Targeted PeftTrainer regressions cover unequal and fractional denominators, all-zero weights, gradients, and metric aggregation.
• tests/sft/peft_trainer_test.py: 56 passed, covering mixed weighted/plain
  train and evaluation metrics, per-key kind consistency,
  denominator-bound consistency, unequal and fractional denominators,
  all-zero weights, and optimizer-state dtype regressions.
• Jitted bfloat16 coverage exercises direct and injected AdamW with one- and two-microbatch updates, asserts parameter updates, and verifies optimizer-state dtype stability.
• The exact final-head CPU gate completed 119 tests in 107.08 seconds: 56
  PeftTrainer tests and 63 focused RL/role-sharding tests. Six optional
  vLLM/SGL-JAX integration tests were deselected.
• A downstream 16-worker integration run completed all three configured actor
  updates after the fix, with finite losses and rewards, nonzero gradient
  norms, a finalized step checkpoint, and a verified model export. This is
  correctness smoke evidence, not a throughput or model-quality claim.
• Ruff, Pyrefly, scoped Pylint, Python compilation, and git diff --check
  passed.

Known limitations

This PR does not add a diffusion loss by itself; it provides the denominator-aware trainer primitive used by later adapters.

Stack

Depends on the preceding upstream PR: #1743

Tunix block-diffusion design document

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.