perf(rollout): group the synchronous trainable-subset R2R broadcast

[karkuspeter]

Summary

The synchronous R2R path offers two shapes and the cheap one is neither of them:

• the per-parameter branch sends only the trainable subset, but pays one collective per
  tensor — each a handshake across every rollout replica;
• the grouped branch (broadcast_all_params=true) pays one collective, but walks the whole
  state dict and so ships the frozen parameters on every sync.

This groups the per-parameter loop, so both savings apply at once and broadcast_all_params
means only what its name says: whether frozen parameters travel too.

Measurements

From an RL run with twelve single-rank rollout replicas on two nodes, 3.30 GiB of trainable
parameters across 321 tensors:

Broadcast shape	Bytes per sync	Collectives	Transfer, median	Effective rate
per-parameter, trainable subset	3.30 GiB	321	0.34 s	9.7 GiB/s
grouped, whole state dict	7.88 GiB	1	0.16 s	49 GiB/s
grouped, trainable subset — this PR	3.30 GiB	1	0.05 s	66 GiB/s

At equal bytes grouping is worth about 7x here, and it costs no memory: a contiguous
parameter is still broadcast in place, and model VRAM measured identical at 7.5 GiB per rank
across all three shapes.

Implementation notes

A non-contiguous parameter's receive buffer only holds its data once the group closes, so those
copies back are deferred to after nccl_group_end — the same way do_nccl_broadcast_grouped
already handles them.

The obvious alternative, giving do_nccl_broadcast_grouped a trainable_only flag and calling
it from this branch, would change what gets sent: that function iterates _buffer_state_dict or
a raw state_dict(), whereas this branch iterates rollout.model_param_map(weight_mapper),
which is HF-name-mapped and includes get_quantized_tensors. Names and membership both differ,
and trainable_params is keyed to the latter. Grouping in place keeps the set of tensors and
their order exactly as they were.

Testing

tests/test_r2r_broadcast_grouping.py, CPU-only, with the transport stubbed by a fake that
records the call order and, like a real group, only fills a receive buffer at group end. Both
the grouping and the deferred copy-back fail against the pre-change code, so they are real
regression checks rather than restatements. Added to the pytest line in tests/run_test.sh.

Validation

• python -m pytest -q tests/test_r2r_broadcast_grouping.py tests/test_weight_sync.py tests/test_ranked_rollout_end_and_wst_fence.py — 49 passed
• uvx ruff@0.12.7 format --check and uvx ruff@0.12.7 check on the changed files — clean
• Verified the new assertions fail with the change reverted