Repository navigation
Commit 3d596fe
fix: answer gloo requests with the flagos backend and stage collectives over host memory (flagos-ai#435)
On a host with no FlagCX, no vendor torch and no vendor communication library
there was no working collective at all. `init_process_group(backend="flagos")`
raised "no suitable inner backend" from `ProcessGroupFlagOS._build_inner`, and
`init_process_group(backend="gloo")` -- what torchrun and the transformers FSDP2
tests ask for -- built a `ProcessGroupGloo` whose C++ collectives reject flagos
tensors with `ProcessGroupGloo::allreduce: unsupported device type flagos`
(issue flagos-ai#263). Six `Qwen3ModelTest::test_fsdp2_*` nodeids failed on that.
Two additions close the gap. Both live in `torch_fl/comm/process_group.py` and
both are gated by a runtime switch.
* `_try_build_staged_gloo` adds a last tier to `_build_inner`. gloo ships inside
c10d, so it is always present; the tier builds a real `ProcessGroupGloo` and
wraps it in `_HostStagedGloo`, a `ProcessGroup` facade that moves every flagos
operand through host memory (device -> host -> collective -> device) and
synthesises the two virtuals gloo does not implement
(`allgather_into_tensor_coalesced`, `reduce_scatter_tensor_coalesced`) from
gloo's `_allgather_base` / `_reduce_scatter_base`. It costs a copy per
collective, so it runs only after FlagCX and the vendor-native backend have
declined, and it warns once when it is used. `FLAGOS_DIST_STAGED_GLOO=0`
declines the tier so the final diagnostic names the reason instead.
* `redirect_gloo_requests` answers the request on the requested side. gloo
cannot be repaired from the gloo side: gloo's own `BackendConfig` expansion is
what registers a `ProcessGroupGloo` under the flagos device type, its creator
branch asserts `isinstance(backend_class, ProcessGroupGloo)`, and
`c10d.Backend` cannot be subclassed from Python at all. So the request is
rewritten to `"flagos"` before the backend config is built: `import torch_fl`
interposes on `torch.distributed.init_process_group` and `new_group` (and
mirrors the wrappers onto `torch.distributed.distributed_c10d`, or a
`from ... import init_process_group` call site would miss the redirect), but
only when `torch._C._get_accelerator().type` is the flagos device, because
gloo is a legitimate choice on a cpu/cuda host. `FLAGOS_DIST_REDIRECT_GLOO=0`
keeps the requested backend verbatim while debugging.
Measured on the MUSA MTT S5000 box with no FlagCX, no `torch_musa.distributed`
and no NCCL -- the worst case. Two ranks, `torch_fl/comm/process_group.py` as
committed:
init_process_group(backend="flagos") -> ProcessGroupFlagOS/_HostStagedGloo
all_reduce -> [3.0, 3.0] on both
init_process_group(backend="gloo") -> redirected; all_reduce,
all_gather_into_tensor,
reduce_scatter_tensor, cpu
all_reduce, barrier and
_c10d_functional.all_gather_into_tensor
all correct on both ranks
new_group([0, 1], backend="gloo") -> redirected; subgroup all_reduce ->
[3.0, 3.0] on both
FLAGOS_DIST_REDIRECT_GLOO=0 -> plain ProcessGroupGloo again, fails
on the flagos tensor: the switch is
load-bearing
Transformers FSDP2 re-measurement, same harness, before/after, six nodeids:
{"FAIL": 6} -> {"PASS": 3, "FAIL": 3}. The three that now pass are
`test_fsdp2_save_load`, `test_fsdp2_sharding_structure_0_untied` and
`test_fsdp2_sharding_structure_1_tied`. The other three get all the way through
the collectives (`[Gloo] Rank 1 is connected to 1 peer ranks`) and then fail on a
separate, newly root-caused defect at `_fsdp_param_group.py::FSDPCommContext.lazy_init`:
`torch_fl.flagos.Stream.__init__` routes every non-GCU platform without a CUDA
runtime to the Ascend ACL stream, so on MUSA it raises `RuntimeError:
libascendcl.so not found. ACL runtime is required`. That is a MUSA stream gap,
not the collective path this change fixes; it is recorded as such in
`docs/reference/hf-coverage.md` and `docs/architecture/distributed-flagcx.md`
section 0.4 rather than claimed as fixed. The new
`tests/unit/test_gloo_redirect.py` drives the decision table with fakes (no GPU,
no rendezvous): both switch values, both accelerator spellings, positional and
keyword backend arguments, that arguments are copied rather than mutated, that
the interposition is idempotent and reaches the defining module, and that a
staged-tier gloo failure is reported as a reason rather than raised.
Tested:
ruff check . -> All checks passed!
ruff format --check . -> 313 files already formatted
pytest tests/unit/test_env_registry.py tests/unit/test_vendor_routing.py
tests/unit/test_nccl_extension_fallback.py tests/unit/test_gloo_redirect.py
-> 87 passed
No operator routing changes: this change reroutes no operator through a
vendor-native backend or FlagGems, so the `docs/reference/operator-support.md`
per-platform record does not apply. For the same reason
`tests/manual/flaggems_overload_survey.py` was not re-run.
Fixes flagos-ai#263
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>1 parent 05be850 commit 3d596fe
6 files changed
Lines changed: 871 additions & 20 deletions
File tree
- docs
- architecture
- reference
- tests/unit
- torch_fl
- comm
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
22 | 22 | | |
23 | 23 | | |
24 | 24 | | |
25 | | - | |
26 | | - | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
27 | 28 | | |
28 | 29 | | |
29 | 30 | | |
| |||
80 | 81 | | |
81 | 82 | | |
82 | 83 | | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
| 99 | + | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
| 127 | + | |
| 128 | + | |
| 129 | + | |
| 130 | + | |
| 131 | + | |
| 132 | + | |
| 133 | + | |
| 134 | + | |
| 135 | + | |
83 | 136 | | |
84 | 137 | | |
85 | 138 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
112 | 112 | | |
113 | 113 | | |
114 | 114 | | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
115 | 127 | | |
116 | 128 | | |
117 | 129 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
83 | 83 | | |
84 | 84 | | |
85 | 85 | | |
86 | | - | |
87 | | - | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
88 | 90 | | |
89 | 91 | | |
90 | 92 | | |
| |||
108 | 110 | | |
109 | 111 | | |
110 | 112 | | |
111 | | - | |
112 | | - | |
113 | | - | |
114 | | - | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
| 127 | + | |
| 128 | + | |
115 | 129 | | |
116 | 130 | | |
117 | 131 | | |
118 | 132 | | |
119 | 133 | | |
120 | | - | |
121 | | - | |
122 | | - | |
| 134 | + | |
| 135 | + | |
| 136 | + | |
123 | 137 | | |
124 | 138 | | |
125 | 139 | | |
| |||
0 commit comments