refactor(fsdp)(1/5): metrics buffer and train metrics integration

[Rockdu]

What

• FSDP train metrics now go through a MetricBuffer whose schema declares each metric's reduction op (MEAN / MAX / REPLICATED) once, replacing the defaultdict of per-microbatch tensors that was uniformly averaged.
• MEAN metrics record a local sum plus the item count behind it; the flush divides globally summed numerators by globally summed counts.
• MAX metrics (model_output_max_abs_diff, model_output_rel_max) are reduced with an actual ReduceOp.MAX.
• The reduction moved from a gloo gather_object over the whole world to two device all_reduces over dp_group, accumulating in float64.

Why

The old path had three defects:

1. Every metric was averaged, including the two named *_max_*. A rank that saw an outlier had it diluted by 1/world_size, which defeats the purpose of a max metric — it hid exactly the train-vs-rollout divergence those metrics exist to catch.
2. The mean was a mean of per-rank means, so a rank holding 1 sample counted as much as a rank holding 8. Under uneven microbatch splits the reported loss was not the batch loss.
3. gather_object pickled a dict per rank through gloo on every optimizer step and reduced on the Python side.

On the group choice. The buffer reduces over dp_group, not the world. SP peers in Miles-diffusion share samples (rollout data is indexed by dp_rank) and the sequence-parallel plan all-gathers the model output at the boundary hook before the loss, so SP peers hold identical values. To be precise, the reported numbers would be unchanged over the world group — a duplicate cancels out of both a sum/count and a max — so this is about not dragging sp_sizex the ranks into every collective, not about correctness.

Validation

Run on an H200 devbox using the project image rockdu/miles_diffusion:dev-cu129-sglang-main-20260719, CPU-only (CUDA_VISIBLE_DEVICES="").

Added the 4-rank test covers what single-rank tests cannot: that a mean weights items rather than ranks, that a max survives when only one rank sees the outlier, that a metric only one rank records still reduces correctly, and that a subgroup ignores ranks outside it (the way DP ignores SP peers).

Checklist

• pre-commit run --all-files passes
• Added/updated tests for new behaviour
• pytest -x is green — tests/fast is 125 passed with one pre-existing unrelated failure (see Validation)
• If launch flags changed, python3 train.py --help still parses — n/a, no flags changed
• If a public flag was added, it appears in the CLI reference docs — n/a, no flag added
• If an example was added, it has a real walkthrough — n/a, no example added
• e2e standard fixtures re-recorded — not done, needs the GPU e2e runners (see Not in this PR)

ci-sglang-repo: sgl-project/sglang
ci-sglang-pr: #32420

[zhihengy]

LGTM