From 5012675f4eecf2d9ef436102b345383e02601846 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 18:04:40 -0400 Subject: [PATCH 1/2] Fix all_gather/reduce_scatter for raw process groups on torch 2.13 torch 2.13 routes the deprecated all_gather_into_tensor/reduce_scatter_tensor through new all_gather_single/reduce_scatter_single methods that only exist on the generic ProcessGroup wrapper (torch.distributed.new_group()), not on the raw Backend objects (ProcessGroupGloo/ProcessGroupNCCL) Fast-LLM constructs directly, breaking every distributed CI test on the Gloo (CPU) backend. Call group._allgather_base/_reduce_scatter_base directly instead: these are the actual underlying implementations, present with a stable signature across torch 2.9-2.13+. Co-Authored-By: Claude Sonnet 5 --- fast_llm/core/distributed.py | 46 +++++++++++++++++++++++++++-- fast_llm/engine/multi_stage/fsdp.py | 3 +- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/fast_llm/core/distributed.py b/fast_llm/core/distributed.py index a71cbc306..11c19fee9 100644 --- a/fast_llm/core/distributed.py +++ b/fast_llm/core/distributed.py @@ -15,14 +15,13 @@ import torch import torch.monitor +from torch._C._distributed_c10d import AllgatherOptions, ReduceScatterOptions from torch.distributed import ( # noqa ProcessGroup, ReduceOp, all_gather, - all_gather_into_tensor, all_reduce, reduce_scatter, - reduce_scatter_tensor, ) logger = logging.getLogger(__name__) @@ -65,6 +64,47 @@ def broadcast( return None +def all_gather_into_tensor( + output_tensor: torch.Tensor, input_tensor: torch.Tensor, group: ProcessGroup, async_op: bool = False +) -> torch.distributed.Work | None: + """ + Same as torch.distributed.all_gather_into_tensor, but calling the group directly instead of going through + torch's deprecated public wrapper. As of torch 2.13, that wrapper only supports the generic `ProcessGroup` + returned by `torch.distributed.new_group` and raises an `AttributeError` on the raw `Backend` objects + (e.g. `ProcessGroupGloo`, `ProcessGroupNCCL`) used here. + """ + assert group is not None + opts = AllgatherOptions() + opts.asyncOp = async_op + work = group._allgather_base(output_tensor, input_tensor, opts) + if async_op: + return work + elif work is not None: + work.wait() + return None + + +def reduce_scatter_tensor( + output: torch.Tensor, + input_: torch.Tensor, + group: ProcessGroup, + op: ReduceOp.RedOpType = ReduceOp.SUM, + async_op: bool = False, +) -> torch.distributed.Work | None: + """Same as torch.distributed.reduce_scatter_tensor, bypassing torch's deprecated wrapper for the same reason + as `all_gather_into_tensor` above.""" + assert group is not None + opts = ReduceScatterOptions() + opts.reduceOp = op + opts.asyncOp = async_op + work = group._reduce_scatter_base(output, input_, opts) + if async_op: + return work + elif work is not None: + work.wait() + return None + + def check_parallel_match(tensor: torch.Tensor, group: ProcessGroup | None, name: str) -> None: # A utility function to check for tensor-parallel (or other) mismatches. all_tensors = tensor.new_empty((group.size(),) + tensor.shape) @@ -114,7 +154,7 @@ def all_gather_scalar( value = torch.full([1], value, dtype=dtype, device=_get_device(group)) output_tensor = value.new_empty((group.size(),)) with set_timeout(group, timeout): - torch.distributed.all_gather_into_tensor(output_tensor, value, group=group) + all_gather_into_tensor(output_tensor, value, group=group) return output_tensor.tolist() else: return value diff --git a/fast_llm/engine/multi_stage/fsdp.py b/fast_llm/engine/multi_stage/fsdp.py index 6c3600f4d..83516fa03 100644 --- a/fast_llm/engine/multi_stage/fsdp.py +++ b/fast_llm/engine/multi_stage/fsdp.py @@ -4,9 +4,8 @@ import typing import torch -from torch.distributed import all_reduce, reduce_scatter_tensor -from fast_llm.core.distributed import ProcessGroup +from fast_llm.core.distributed import ProcessGroup, all_reduce, reduce_scatter_tensor from fast_llm.core.ops import gather_op from fast_llm.engine.config_utils.data_type import DataType from fast_llm.engine.config_utils.tensor_dim import TensorDim From 0ad59050cd56a79285ecbef70196c4720ba67352 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 9 Jul 2026 14:18:24 -0400 Subject: [PATCH 2/2] Fix flaky GSPO/GRPO metrics test tolerance _check_policy_metrics compared each metric with a relative-only threshold (min_threshold=0), so a near-zero-scale metric (e.g. kl_new_old ~4e-3) got an effective bound of 5e-5 * scale ~= 2e-7, below the float32 accumulation noise floor. Reference (naive loop) and implementation (fused kernel + tensor-parallel all-reduce) sum in different orders and differ by a few e-7 absolute, so the check tripped once float32 rounding shifted (surfaced on torch 2.13). Add a 1e-6 absolute floor, matching every other rms_close_relative call in the file (loss and 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. Co-Authored-By: Claude Opus 4.8 --- tests/layers/test_lm_losses.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 78d588cb1..07d42f02a 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -562,7 +562,7 @@ def _check_policy_metrics(ref: PolicyMetrics, got: PolicyMetrics, threshold: flo if ref_value is None: assert got_value is None, name else: - Assert.rms_close_relative(got_value, ref_value, threshold) + Assert.rms_close_relative(got_value, ref_value, threshold, 1e-6) def _test_grpo_metrics(