Log weight-sync time as a streaming-trainer metric

[jlamypoirier]

What

The streaming trainer callback broadcasts the full model to the inference actor on every successful step — synchronously, inside the training loop (so it's part of step_time_ms) — but nothing measured how long that takes. This logs it as weight_sync_time_ms, so the share of step time spent shipping weights is visible in the training logs and W&B.

Part of the RL-diagnostics initiative (making the Fast-LLM trainer the single source of truth for RL metrics); shipped as an independent PR since it's self-contained and touches no wire schema.

How

• Timing: bracket the broadcast (iter_checkpoint ZeRO-gather + NCCL push) with CUDA events on the broadcasting rank. The result is read one step later (_weight_sync_pending + Event.query()), so the always-on timer never forces a CPU stall — the broadcast otherwise overlaps the next step's compute, and reading elapsed_time immediately would serialize them. broadcast() runs with async_op=False (work.wait()), so the current stream is ordered after the transfer and the end event captures true completion. The CPU/gloo broadcast path (non-NCCL) falls back to a plain perf_counter measurement, which needs no deferral.
• Plumbing: TrainerCallback.step_end may now return a dict[str, Any] | None of scalar metrics; the trainer merges any returned dicts into the step's training metrics alongside **train_metrics. Base callback returns None.
• Scope is trainer-side broadcast only — actor receive/load time is out of scope.

Cost

Zero on the disabled path (non-streaming runs never construct the callback). On the streaming path: two reusable CUDA events and one deferred elapsed_time read per logging step — negligible against a full-model gather+broadcast, and no added GPU sync.

Testing

No unit test: the path requires a live Redis + external-world rendezvous + CUDA, there's no streaming-callback test harness, and the only logic is the CUDA-event deferral state machine. Verified imports + black locally; the metric is validated on a live streaming RL run (the new key appears under training/*).

🤖 Generated with Claude Code