feat(wheel): self-contained wheels that run on a stock torch+cpu - #68
Open
lvyufeng wants to merge 2 commits into
Open
feat(wheel): self-contained wheels that run on a stock torch+cpu#68lvyufeng wants to merge 2 commits into
lvyufeng wants to merge 2 commits into
Conversation
A plugin wheel previously required the *vendor's* torch to be installed:
libtorch_fl.so resolves ~2100 undefined symbols out of the vendor's forked
libtorch_cpu.so, and the build baked ${PYTORCH_INSTALL_DIR}/lib into RUNPATH,
so the .so only ever worked on the machine that built it.
This bundles the forked libtorch into the wheel and relinks the active torch at
import, so one wheel installs into a clean env holding only stock
torch==2.10.0+cpu, with no vendor torch package present.
Mechanism, in three parts:
* RUNPATH generalisation (cmake/FlagosRpath.cmake). CMakeLists.txt
unconditionally appended ${PYTORCH_INSTALL_DIR}/lib -- the single source of
build-machine conda paths in shipped .so. Replaced with a per-target override,
$ORIGIN:$ORIGIN/../<bundle>:<driver paths>, applied to torch_fl,
torch_bindings and flagos with INSTALL_RPATH_USE_LINK_PATH OFF. torch_fl._C
gets the bundle dir too, else its auditwheel-mangled deps go missing. Driver
runtime is deliberately left to the target machine.
* Vendor libtorch relink (torch_fl/accelerator/_vendor_libtorch.py), generalised
from the MetaX-only version, with per-backend .so manifests for dcu/ppu/metax.
Replaces physical files in torch/lib with symlinks, originals preserved under
torch/lib/_orig_backup/, then dlopens the set RTLD_GLOBAL in dependency order.
A ctypes preload alone cannot do this: the stock wheel's _C.so carries $ORIGIN
RUNPATH and pulls upstream libc10 back in by full path, and two libc10 in one
process means duplicate static-init registration. CUDA is exempt -- its core
libs already are the upstream ones, so the existing ctypes preload suffices.
* Bundle scripts (scripts/lib/bundle_common.sh + per-backend). Extracts the fork
plus its non-driver third-party deps: metax 9 files, ppu 12, dcu 28 (including
the 17 auditwheel-mangled torch.libs/, which the stock +cpu wheel ships none of
and whose system counterparts have different sonames).
Four defects this surfaced, all specific to front-ending a stock +cpu torch:
* ensure_maca_libtorch_links() self-gated on FLAGOS_METAX_BOXING, making the
self-contained path a silent no-op -- stock libtorch_cpu.so stayed in place and
libtorch_cuda.so then could not resolve at::maca symbols that only the forked
CPU runtime defines. Gating now belongs to the caller.
* GetFlagosDefaultCudaGenerator indexed torch.cuda.default_generators, which CUDA
lazy-init populates and nothing had triggered, so it was empty and indexing
raised IndexError. Every device-side RNG op (randn/rand/normal_) failed while
empty/zeros/ones worked. Forces the idempotent init first.
* libcaffe2_nvrtc.so is what that init dlopens, and DTK ships it while stock +cpu
does not -- so on DCU the fix above just turned IndexError into "Error in
dlopen: libcaffe2_nvrtc.so". Now bundled.
* torch.version is pure Python generated at torch build time, so swapping .so
cannot change it: a self-contained DCU wheel reported hip=None, and triton's
hcu backend gates is_active() on exactly that, so every flag_gems op died with
"0 active drivers ([])". The bundle carries the vendor version.py and
torch_fl reads hip/rocm back at import.
Also: the unconditional torch.cuda.init() at import now checks the *build*
backend, not torch.cuda.is_available() -- relinking a hipified libtorch into a
stock +cpu torch makes is_available() true while the CUDA runtime libs are
absent.
Verified end to end, each in an env holding only stock torch 2.10.0+cpu, with
RUNPATH free of build-machine paths, ldd clean, and /proc/<pid>/maps showing
every libtorch mapped from the bundle:
MetaX 0.1.0+metax cp312 642M 8 devices 836 passed
CUDA 0.1.0 cp312 728M 8 devices 848 passed
DCU 0.1.0+dtk cp310 986M 8 devices 868 passed
PPU 0.1.0+ppu cp312 303M 18 devices 63 passed (vendor bmm(1,1,1) abort)
randn/rand/normal_ pass on all four. add/mul maxdiff 0.0; silu <= 2.4e-07; mm
<= 9.5e-06 (1.1e-02 on MetaX, fp32 accumulation order).
Two known limits of the stock-+cpu-front-end approach, neither a packaging bug:
* torch.compile does not work -- Inductor asks triton for a GPU driver and a
stock +cpu torch has none registered ("Could not find an active GPU backend" /
"libcuda.so cannot found"). Eager mode is unaffected. Accounts for ~8
test_compile.py failures per backend. On DCU one case survives even with the
hip-version fix, because Inductor's compile workers are bare sys.executable
subprocesses that import torch without torch_fl;
TORCHINDUCTOR_COMPILE_THREADS=1 works around it.
* import torch_fl must precede import torch. PyTorch caches CUDAHooks on first
import, so the preload and relink have to run first.
lvyufeng
force-pushed
the
metax-2.10
branch
2 times, most recently
from
August 8, 2026 18:20
99852e5 to
bf6030c
Compare
libtorch_python.so carries DT_NEEDED for libshm.so (torch.multiprocessing's shared-memory manager). When the _LOAD_ORDER list omits libshm.so, dlopen resolves that dependency through the *loader's* search path (not through the bundle dir), and in a fresh install nothing has placed libshm.so anywhere the loader looks yet -- so libtorch_python.so fails to load with "libshm.so: cannot open shared object file" even though the file sits in the bundle. Affects metax/ppu/dcu equally; the CI caught it on metax first because CUDA happened to be tested before it. Also translate all Chinese comments and user-facing strings to English (repo convention).
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
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.
What
Makes a plugin wheel installable into an environment holding only stock
torch==2.10.0+cpu— no vendor torch package required. The vendor's torch is nowonly a source to extract
.sofrom at build time.Previously a wheel needed the vendor's torch installed (
libtorch_fl.soresolves~2100 undefined symbols out of the vendor's forked
libtorch_cpu.so), andCMakeLists.txtbaked${PYTORCH_INSTALL_DIR}/libinto RUNPATH, so the built.soonly ever worked on the machine that produced it.How
RUNPATH generalisation (
cmake/FlagosRpath.cmake) — replaces theunconditional
${PYTORCH_INSTALL_DIR}/libappend with a per-target override,$ORIGIN:$ORIGIN/../<bundle>:<driver paths>, ontorch_fl,torch_bindingsandflagos, withINSTALL_RPATH_USE_LINK_PATH OFF.torch_fl._Cgets the bundledir too, otherwise its auditwheel-mangled deps go missing. Driver runtime is
deliberately left to the target machine.
Vendor libtorch relink (
torch_fl/accelerator/_vendor_libtorch.py) —generalised from the MetaX-only version, with per-backend
.somanifests fordcu/ppu/metax. Replaces physical files in
torch/libwith symlinks (originalspreserved under
torch/lib/_orig_backup/, so it is reversible), then dlopens theset
RTLD_GLOBALin dependency order.A ctypes preload alone cannot do this: the stock wheel's
_C.socarries$ORIGINRUNPATH and pulls upstreamlibc10.soback in by full path, and twolibc10in one process means duplicate static-init registration. CUDA is exempt— its core libs already are the upstream ones, so the existing ctypes preload
suffices.
Bundle scripts (
scripts/lib/bundle_common.sh+ per-backend) — extracts thefork plus its non-driver third-party deps. metax 9 files, ppu 12, dcu 28
(including the 17 auditwheel-mangled
torch.libs/, which the stock+cpuwheelships none of and whose system counterparts have different sonames).
Verified
Each in an env holding only stock
torch 2.10.0+cpu, run from/tmp(outsidethe repo), with RUNPATH free of build-machine paths,
lddclean, and/proc/<pid>/mapsshowing every libtorch mapped from the bundle:tests/integration0.1.0+metaxcp312, 642M0.1.0cp312, 728M0.1.0+dtkcp310, 986M0.1.0+ppucp312, 303M¹ aborts on
bmm(1,1,1)— a vendor kernel bug on degenerate shapes; the same opsequence also aborts under the vendor torch natively.
randn/rand/normal_pass on all four.add/mulmaxdiff 0.0;silu≤2.4e-07;
mm≤ 9.5e-06 (1.1e-02 on MetaX, fp32 accumulation order). Realdispatch line, e.g.
[flagos dispatch] add.Tensor -> flagos_python.Four defects this surfaced
All specific to front-ending a stock
+cputorch, and each only reachable oncethe previous one was fixed:
ensure_maca_libtorch_links()self-gated onFLAGOS_METAX_BOXING, making theself-contained path a silent no-op — stock
libtorch_cpu.sostayed in placeand
libtorch_cuda.sothen could not resolveat::macasymbols that only theforked CPU runtime defines. Gating now belongs to the caller.
GetFlagosDefaultCudaGeneratorindexedtorch.cuda.default_generators, whichCUDA lazy-init populates and nothing had triggered, so it was empty and
indexing raised
IndexError. Every device-side RNG op failed whileempty/zeros/onesworked. Forces the idempotent init first.libcaffe2_nvrtc.sois what that init dlopens, and DTK ships it while stock+cpudoes not — so on DCU fix 2 just turnedIndexErrorintoError in dlopen: libcaffe2_nvrtc.so. Now bundled.torch.versionis pure Python generated at torch build time, so swapping.socannot change it: a self-contained DCU wheel reportedhip=None, andtriton's hcu backend gates
is_active()on exactly that, so every flag_gemsop died with
0 active drivers ([]). The bundle now carries the vendorversion.pyandtorch_flreadship/rocmback at import.Also: the unconditional
torch.cuda.init()at import now checks the buildbackend rather than
torch.cuda.is_available()— relinking a hipified libtorchinto a stock
+cputorch makesis_available()true while the CUDA runtime libsare absent.
Two known limits of this approach
Neither is a packaging bug; both are inherent to a stock-
+cpufront-end.torch.compiledoes not work. Inductor asks triton for a GPU driver and astock
+cputorch has none registered (Could not find an active GPU backend/
libcuda.so cannot found!). Eager mode is unaffected. Accounts for ~8test_compile.pyfailures per backend. On DCU one case survives even with fix4, because Inductor's compile workers are bare
sys.executablesubprocessesthat import
torchwithouttorch_fl;TORCHINDUCTOR_COMPILE_THREADS=1works around it.
import torch_flmust precedeimport torch. PyTorch caches CUDAHooks onfirst import, so the preload and relink have to run first.
Build
Both
bdist_wheelcalls needACCELERATOR, else setup.py falls back tocudaand CMake fails hunting for nvcc. Stage 2 is easy to skip and yields a wheel that
looks fine but has an empty bundle dir.
ACCELERATOR=metax FLAGOS_METAX_BOXING=1-DMETAX_KERNEL=OFF; a bareMETAX_KERNEL=OFFis not read by setup.pyACCELERATOR=cuda FLAGGEMS_KERNEL=OFFPPU_SDK/CUDA_SDK; distinguished at runtime bylib_ppu/ACCELERATOR=dcu FLAGGEMS_KERNEL=ONACCELERATOR=cuda FLAGGEMS_KERNEL=OFFBundle scripts require
patchelf.Not in scope
FlagGems' C++ path (
liboperators.so/libtriton_jit.so) is not madeself-contained, which is why the CUDA wheel is built
FLAGGEMS_KERNEL=OFF; theFlagGems Python path does not need those libs.
.confop routing is unchanged.