Fix CI on torch 2.13: raw-process-group collectives + flaky GSPO metrics tolerance

[jlamypoirier]

Claude Opus 4.8:

Two independent fixes needed to get CI green on main again. Both surfaced when CI auto-upgraded to torch 2.13.0 (setup.cfg pins torch>=2.9.0, unbounded), so they land together.

Fix 1 — all_gather/reduce_scatter on raw process groups (torch 2.13)

CI was crashing on every tests/models/test_model.py distributed test (cascading through pytest-depends) with:

AttributeError: 'torch._C._distributed_c10d.ProcessGroupGloo' object has no attribute 'all_gather_single'

In torch 2.13, the deprecated torch.distributed.all_gather_into_tensor/reduce_scatter_tensor now dispatch through new all_gather_single/reduce_scatter_single methods. Those are only bound on the generic ProcessGroup wrapper from torch.distributed.new_group() — not on the raw Backend-derived objects (ProcessGroupGloo, ProcessGroupNCCL) that Fast-LLM constructs directly (fast_llm/engine/distributed/distributed.py deliberately bypasses init_process_group). Confirmed against PyTorch's C++ pybind bindings at the v2.13.0 tag.

fast_llm/core/distributed.py: reimplement both to call group._allgather_base(...)/group._reduce_scatter_base(...) directly — the actual underlying virtual methods, with an identical signature on Backend across torch 2.9.0, 2.13.0, and current main, so no version check is needed. fast_llm/engine/multi_stage/fsdp.py imported reduce_scatter_tensor straight from torch.distributed (same bug, FSDP gradient reduce-scatter path) — now routed through fast_llm.core.distributed like everything else.

Fix 2 — flaky GSPO/GRPO metrics tolerance (regression from #555)

Once Fix 1 let the distributed suite run, test_lm_loss_distributed[...gspo_metrics...] failed:

AssertionError: Rms diff too big (2.38e-07 > 2.21e-07, scale = 4.43e-03)

_check_policy_metrics (added in #555) compared each metric with a relative-only threshold (min_threshold=0). A near-zero-scale metric (e.g. kl_new_old ≈ 4e-3) got an effective bound of 5e-5 × scale ≈ 2.2e-7, below the float32 accumulation noise floor. The reference (naive loop) and implementation (fused kernel + tensor-parallel all-reduce) sum in different orders and differ by a few e-7 absolute; torch 2.13's rounding shifted it just over the line. Add a 1e-6 absolute floor, matching every other rms_close_relative call in the file (loss + log-prob comparisons). The floor only engages when 5e-5 × scale < 1e-6 (scale < 2e-2), so O(1)-scale metrics keep the full relative tolerance.

Test plan

• tests/models/test_model.py --models gpt_2 (30 tests, incl. dp2/tp2/sdp2 on Gloo) passes locally on torch 2.11.0 — no regression on the previously-working version.
• tests/layers/test_lm_losses.py (582 passed / 21 skipped, incl. the full distributed GSPO/GRPO metrics chain) passes locally on torch 2.11.0.
• CI on this PR runs torch 2.13.0 — the real test of both fixes (torch 2.13.0 isn't in the local dev venv).