Fix all_gather/reduce_scatter for raw process groups on torch 2.13#563
Open
jlamypoirier wants to merge 1 commit into
Open
Fix all_gather/reduce_scatter for raw process groups on torch 2.13#563jlamypoirier wants to merge 1 commit into
jlamypoirier wants to merge 1 commit into
Conversation
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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Claude Sonnet 5:
Summary
CI on
mainwas failing (alltests/models/test_model.pydistributed tests, cascading through pytest-depends) with:Root cause:
setup.cfgpinstorch>=2.9.0with no upper bound, so CI picked up torch 2.13.0 (released between the last two scheduled CI runs). In 2.13, the deprecatedtorch.distributed.all_gather_into_tensor/reduce_scatter_tensornow dispatch through newall_gather_single/reduce_scatter_singlemethods. Those are only bound on the genericProcessGroupwrapper returned bytorch.distributed.new_group()— not on the rawBackend-derived objects (ProcessGroupGloo,ProcessGroupNCCL) that Fast-LLM constructs directly (fast_llm/engine/distributed/distributed.pydeliberately bypassesinit_process_group). Confirmed via PyTorch's C++ pybind bindings at thev2.13.0tag.Fix
fast_llm/core/distributed.py: reimplementall_gather_into_tensor/reduce_scatter_tensorto callgroup._allgather_base(...)/group._reduce_scatter_base(...)directly, bypassing torch's deprecated public wrapper entirely. These are the actual underlying virtual methods (verified to have an identical signature onBackendacross torch 2.9.0, 2.13.0, and currentmain), so this works uniformly across supported torch versions without a version check.fast_llm/engine/multi_stage/fsdp.pyimportedreduce_scatter_tensorstraight fromtorch.distributed(same bug, on the FSDP gradient reduce-scatter path) — routed throughfast_llm.core.distributedinstead, consistent with the rest of the codebase.Test plan
tests/models/test_model.py --models gpt_2(30 tests, incl.dp2/tp2/sdp2distributed variants on Gloo) passes locally against torch 2.11.0 — no regression on the previously-working version.