Aggregate metrics at log time only

[qgallouedec]

Branch sft-metrics-single-reduce (main merged in at 67dfbe2). Alternative to #6678 (@michaelbenayoun), which first moved SFT's metric collectives to log time; this takes the same idea with running sums instead of per-step buffers, and applies it to all main-code trainers (SFT, Distillation, GRPO, RLOO, DPO, KTO, Reward).

How it works

Same (total, count) pair that RLOOTrainer and GRPOTrainer already reduce today, kept as a running sum instead of being reduced every step:

# compute_loss: local only, no collective, no host sync
self._metric_stats[mode]["entropy"] += torch.stack([entropy_sum, total_tokens.float()])

# log(): one collective for the whole window
values = torch.stack([stats[key].double() for key in keys])   # keys sorted, so ranks agree on the order
totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}

• Extrema (GRPO, RLOO, Reward): _metric_mins / _metric_maxs, one reduce("max") each (min as -max(-x)); ±inf sentinels keep key sets rank-symmetric; empty window logs None.
• _metrics float lists remain only for values already global for free: stats derived from gathers the training needs anyway (completion lengths → num_input_tokens_seen, rewards → advantages) and Python constants (entropy_coef, step_time).
• Untouched: communication the training itself requires (num_items_in_batch, KTO's loss KL, ref-logp precompute).

Communication removed

Trainer	before / micro-batch	after / micro-batch	after / log window
SFT	up to 6	0	1
Distillation	2	0	1
GRPO	~8 (+5 sampling, +N log_metric)	0	≤3
RLOO	~7	0	≤3
DPO	8–13	0	1
KTO	~10	1 (loss KL, required)	1
Reward	4	0	≤3

Plus the matching .item() host syncs. Under gradient accumulation these gathers were the only communication on K−1 of K micro-batches (grad all-reduce fires once, under no_sync); in eval there is no grad sync at all.

Measured on 8×H100 (SmolLM2-135M, seq ≤256, bs 2, gas=64; two interleaved main/branch runs each):

	main	this PR
optimizer step (median)	12.38 / 12.59 s	11.90 / 11.80 s	−5.1%
eval pass (2000 samples)	4.59 / 4.66 s	4.51 / 4.53 s	−2.5%

The win grows with rank count, accumulation, and cheaper micro-batches; expect more on multi-node interconnects (unmeasured here).

Benchmark script

Run once on main and once on this branch:

torchrun --nproc_per_node 8 bench_metrics.py main   # then: bench_metrics.py branch

# bench_metrics.py
import json
import random
import statistics
import sys
import time

import torch
from datasets import Dataset
from transformers import TrainerCallback

from trl import SFTConfig, SFTTrainer


tag = sys.argv[1]


class StepTimer(TrainerCallback):
    def __init__(self):
        self.times = []
        self._t = None

    def on_step_begin(self, args, state, control, **kwargs):
        torch.cuda.synchronize()
        self._t = time.perf_counter()

    def on_step_end(self, args, state, control, **kwargs):
        torch.cuda.synchronize()
        self.times.append(time.perf_counter() - self._t)


torch.manual_seed(0)
rng = random.Random(0)
words = "the quick brown fox jumps over the lazy dog".split()
lengths = [max(10, min(350, int(rng.lognormvariate(4.3, 0.8)))) for _ in range(12000)]
texts = [" ".join(rng.choices(words, k=k)) for k in lengths]
train_dataset = Dataset.from_dict({"text": texts[:10000]})
eval_dataset = Dataset.from_dict({"text": texts[10000:]})

training_args = SFTConfig(
    output_dir=f"bench_out_{tag}",
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=64,
    max_steps=6,
    logging_steps=1000,
    save_strategy="no",
    report_to="none",
    seed=42,
    bf16=True,
    max_length=256,
)
timer = StepTimer()
trainer = SFTTrainer(
    model="HuggingFaceTB/SmolLM2-135M-Instruct",
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[timer],
)
trainer.train()

# Eval loop: time the second call (first warms up the eval dataloader)
eval_times = []
for _ in range(2):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    trainer.evaluate()
    torch.cuda.synchronize()
    eval_times.append(time.perf_counter() - t0)

if trainer.accelerator.is_main_process:
    steady = timer.times[2:]
    result = {"tag": tag, "train_median_ms": statistics.median(steady) * 1000, "eval_s": eval_times}
    print("RESULT", json.dumps(result))

Deliberate value changes

• Ratios are window-weighted (Σsum/Σcount), not means of per-step ratios → identical at logging_steps=1, and no longer depend on logging_steps.
• No more gather_for_metrics flat truncation at uneven last batches (it silently dropped whole ranks' scalars).
• Extrema are window extrema, not means of per-step extrema.
• KTO: unpaired-window rewards/* log 0.0 instead of omitting the key; margins still need both sides.
• num_tokens advances at train-time log(); an eval log in between can lag by ≤1 window.

Verification

• 2-proc gloo vs main at logging_steps=1, where the old and new semantics coincide: Reward exact (including the cross-rank min/max path), DPO 4e-6, GRPO and RLOO exact apart from step_time, SFT and KTO exact apart from epoch-boundary steps, which is the truncation fix above.
• test_log_averages_over_the_window_weighted_by_count pins the weighting: two steps of 1 token at entropy 3.0 and 9 tokens at entropy 1.0 log 1.2, not 2.0.
• Suites: SFT, Distillation, Reward, DPO, KTO, GRPO, RLOO all pass (GRPO has one pre-existing vLLM/flex-attention failure that fails identically on main).

Follow-up (not here, to bound review)

Experimental copies keep the old pattern: async_grpo, async_distillation, gmpo, grpo_with_replay_buffer (still work: GRPO keeps _metrics), and the DPO-family forks (cpo, orpo, bco, tpo, online_dpo, xpo, nash_md, ppo).