Skip to content

Commit 1adbc97

Browse files
authored
fix: serialize PPU FlagTree compilation (#148)
FlagTree's PPU driver queries torch.cuda.current_device while selecting compiler hints. Inductor workers may be forked after the parent initializes the PPU CUDA context, causing CUDA reinitialization failures in backward and FP16 compilation.\n\nDefault PPU FlagTree to one compile thread through a scoped Inductor config patch while preserving explicit user settings and all other compiler backends. Document the workaround and cover its scope with regression tests.\n\nTested: ruff check .\nTested: ruff format --check .\nTested: pytest tests/unit/ -q\nTested: pytest tests/integration/test_compile.py -v --tb=short\nTested: PPU FlagTree 0.6.0 full compile suite - 18 passed\n\nRelated: flagos-ai/FlagTree#1031\n\nCo-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 6559bca commit 1adbc97

4 files changed

Lines changed: 107 additions & 0 deletions

File tree

‎docs/architecture/torch-compile-integration.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,15 @@ Detection uses the FlagTree-only module `triton._flagtree_spec`. Note that
136136
empty string on nvidia/amd builds, since upstream directs you not to set
137137
`FLAGTREE_BACKEND` for those.
138138

139+
PPU FlagTree currently queries `torch.cuda.current_device()` while selecting
140+
compiler hints. Inductor's asynchronous compile workers can make that query
141+
after a PPU CUDA context was initialized in the parent, which PyTorch rejects as
142+
CUDA reinitialization after `fork` ([FlagTree #1031](https://github.com/flagos-ai/FlagTree/issues/1031)).
143+
The flagos backend therefore defaults PPU FlagTree to one compile thread until
144+
the upstream driver no longer initializes CUDA in a worker. An explicit
145+
`TORCHINDUCTOR_COMPILE_THREADS` value or `compile_threads` compile option remains
146+
authoritative.
147+
139148
## Architecture
140149

141150
### Phase 1: Inductor Integration

‎docs/vendors/ppu/installation.md‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,34 @@ print(f'Sample result: {y.cpu()[0, 0].item():.4f}')
7575

7676
Expected output shows `torch.cuda available: True`, `torch.version.cuda: 13.0`, and a floating-point result.
7777

78+
### `torch.compile` with FlagTree
79+
80+
Build FlagTree from source with its PPU backend; the resulting `flagtree` wheel
81+
installs a `triton` module, so it must replace or shadow the vendor Triton:
82+
83+
```bash
84+
git clone https://github.com/flagos-ai/FlagTree.git
85+
cd FlagTree
86+
pip install -r python/requirements.txt
87+
FLAGTREE_BACKEND=ppu MAX_JOBS=32 \
88+
pip install . --no-build-isolation -v
89+
```
90+
91+
Require the active compiler to be FlagTree and run the compile contract:
92+
93+
```bash
94+
FLAGOS_DISABLE_CUDA_ASSETS=1 \
95+
FLAGOS_USE_FLAGTREE=1 \
96+
pytest tests/integration/test_compile.py -v --tb=short
97+
```
98+
99+
PPU FlagTree currently initializes CUDA while selecting compiler hints in an
100+
asynchronous Inductor worker, which can fail after the parent process has already
101+
initialized the PPU context ([FlagTree #1031](https://github.com/flagos-ai/FlagTree/issues/1031)).
102+
Torch-FL defaults PPU FlagTree to serial compilation as a compatibility measure.
103+
Set `TORCHINDUCTOR_COMPILE_THREADS` explicitly only when testing an upstream fix
104+
or deliberately selecting a different worker configuration.
105+
78106
### AMP and GradScaler
79107

80108
PPU uses the shared CUDA-boxing implementation, while the public AMP device name

‎tests/integration/test_compile.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,54 @@ def test_flagtree_requires_flagtree_install():
280280
require_flagtree()
281281

282282

283+
def test_ppu_flagtree_defaults_to_serial_compile(monkeypatch):
284+
"""PPU FlagTree must avoid CUDA initialization in forked compile workers."""
285+
from torch_fl.compile import inductor_backend
286+
287+
monkeypatch.delenv("TORCHINDUCTOR_COMPILE_THREADS", raising=False)
288+
monkeypatch.setattr(
289+
"torch_fl.compile.flagtree_shim.flagtree_backend", lambda: "ppu"
290+
)
291+
292+
patches = {}
293+
inductor_backend._patch_ppu_flagtree_compile_workers(patches)
294+
295+
assert patches["compile_threads"] == 1
296+
297+
298+
def test_ppu_flagtree_preserves_explicit_compile_threads(monkeypatch):
299+
"""An explicit per-compile or environment setting remains authoritative."""
300+
from torch_fl.compile import inductor_backend
301+
302+
monkeypatch.setattr(
303+
"torch_fl.compile.flagtree_shim.flagtree_backend", lambda: "ppu"
304+
)
305+
306+
patches = {"compile_threads": 4}
307+
inductor_backend._patch_ppu_flagtree_compile_workers(patches)
308+
assert patches["compile_threads"] == 4
309+
310+
monkeypatch.setenv("TORCHINDUCTOR_COMPILE_THREADS", "8")
311+
patches = {}
312+
inductor_backend._patch_ppu_flagtree_compile_workers(patches)
313+
assert "compile_threads" not in patches
314+
315+
316+
def test_non_ppu_flagtree_keeps_default_compile_threads(monkeypatch):
317+
"""Other FlagTree backends keep Inductor's asynchronous compilation."""
318+
from torch_fl.compile import inductor_backend
319+
320+
monkeypatch.delenv("TORCHINDUCTOR_COMPILE_THREADS", raising=False)
321+
monkeypatch.setattr(
322+
"torch_fl.compile.flagtree_shim.flagtree_backend", lambda: "hcu"
323+
)
324+
325+
patches = {}
326+
inductor_backend._patch_ppu_flagtree_compile_workers(patches)
327+
328+
assert "compile_threads" not in patches
329+
330+
283331
def test_flagtree_is_never_importable_as_flagtree():
284332
"""Guard the packaging trap: the wheel is 'flagtree', the module is 'triton'.
285333

‎torch_fl/compile/inductor_backend.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,26 @@ def _resolve_config_patches(
104104
return patches
105105

106106

107+
def _patch_ppu_flagtree_compile_workers(config_patches: Dict[str, Any]) -> None:
108+
"""Keep PPU FlagTree compilation in the parent process by default.
109+
110+
FlagTree's PPU driver queries ``torch.cuda.current_device()`` while choosing
111+
compiler hints. Inductor's asynchronous workers may be forked after the PPU
112+
CUDA context is initialized, where that query triggers PyTorch's fork-safety
113+
error. Serial compilation avoids the unsafe reinitialization until the
114+
upstream driver can discover its device without touching CUDA in the worker.
115+
"""
116+
if "compile_threads" in config_patches:
117+
return
118+
if os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") is not None:
119+
return
120+
121+
from torch_fl.compile.flagtree_shim import flagtree_backend
122+
123+
if flagtree_backend() == "ppu":
124+
config_patches["compile_threads"] = 1
125+
126+
107127
def flagos_compile_backend(
108128
gm: torch.fx.GraphModule,
109129
example_inputs: List[torch.Tensor],
@@ -151,6 +171,8 @@ def flagos_compile_backend(
151171

152172
require_flagtree()
153173

174+
_patch_ppu_flagtree_compile_workers(config_patches)
175+
154176
# Make inductor treat flagos as a GPU device. Order matters: is_gpu() must
155177
# answer True and the device interface must be resolvable before the
156178
# codegen backend registration reads them.

0 commit comments

Comments
 (0)