Skip to content

feat(wheel): self-contained wheels that run on a stock torch+cpu - #68

Open
lvyufeng wants to merge 2 commits into
mainfrom
metax-2.10
Open

feat(wheel): self-contained wheels that run on a stock torch+cpu#68
lvyufeng wants to merge 2 commits into
mainfrom
metax-2.10

Conversation

@lvyufeng

@lvyufeng lvyufeng commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

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 now
only a source to extract .so from at build time.

Previously a wheel needed the vendor's torch installed (libtorch_fl.so resolves
~2100 undefined symbols out of the vendor's forked libtorch_cpu.so), and
CMakeLists.txt baked ${PYTORCH_INSTALL_DIR}/lib into RUNPATH, so the built
.so only ever worked on the machine that produced it.

How

RUNPATH generalisation (cmake/FlagosRpath.cmake) — replaces the
unconditional ${PYTORCH_INSTALL_DIR}/lib append with a per-target override,
$ORIGIN:$ORIGIN/../<bundle>:<driver paths>, on torch_fl, torch_bindings and
flagos, with INSTALL_RPATH_USE_LINK_PATH OFF. torch_fl._C gets the bundle
dir 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 .so manifests for
dcu/ppu/metax. Replaces physical files in torch/lib with symlinks (originals
preserved under torch/lib/_orig_backup/, so it is reversible), 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.so 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).

Verified

Each in an env holding only stock torch 2.10.0+cpu, run from /tmp (outside
the repo), with RUNPATH free of build-machine paths, ldd clean, and
/proc/<pid>/maps showing every libtorch mapped from the bundle:

backend wheel bundle devices tests/integration
MetaX 0.1.0+metax cp312, 642M 9 8 836 passed
CUDA 0.1.0 cp312, 728M ctypes 8 848 passed
DCU 0.1.0+dtk cp310, 986M 28 8 868 passed
PPU 0.1.0+ppu cp312, 303M 12 18 63 passed¹

¹ aborts on bmm(1,1,1) — a vendor kernel bug on degenerate shapes; the same op
sequence also aborts under the vendor torch natively.

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). Real
dispatch line, e.g. [flagos dispatch] add.Tensor -> flagos_python.

Four defects this surfaced

All specific to front-ending a stock +cpu torch, and each only reachable once
the previous one was fixed:

  1. 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.
  2. 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 failed while
    empty/zeros/ones worked. Forces the idempotent init first.
  3. libcaffe2_nvrtc.so is what that init dlopens, and DTK ships it while stock
    +cpu does not — so on DCU fix 2 just turned IndexError into
    Error in dlopen: libcaffe2_nvrtc.so. Now bundled.
  4. 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 now 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 rather than torch.cuda.is_available() — relinking a hipified libtorch
into a stock +cpu torch makes is_available() true while the CUDA runtime libs
are absent.

Two known limits of this approach

Neither is a packaging bug; both are inherent to a stock-+cpu front-end.

  • 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 fix
    4, 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.

Build

export ACCELERATOR=metax FLAGOS_METAX_BOXING=1   # per-backend, see table below
python setup.py bdist_wheel                       # 1. compile
bash scripts/bundle_maca_libtorch.sh              # 2. bundle vendor libtorch
python setup.py bdist_wheel                       # 3. repack with bundle

Both bdist_wheel calls need ACCELERATOR, else setup.py falls back to cuda
and CMake fails hunting for nvcc. Stage 2 is easy to skip and yields a wheel that
looks fine but has an empty bundle dir.

backend variables notes
metax ACCELERATOR=metax FLAGOS_METAX_BOXING=1 maps to -DMETAX_KERNEL=OFF; a bare METAX_KERNEL=OFF is not read by setup.py
ppu ACCELERATOR=cuda FLAGGEMS_KERNEL=OFF builds against PPU_SDK/CUDA_SDK; distinguished at runtime by lib_ppu/
dcu ACCELERATOR=dcu FLAGGEMS_KERNEL=ON needs Python 3.10 (DTK torch is cp310-only)
cuda ACCELERATOR=cuda FLAGGEMS_KERNEL=OFF no bundle step

Bundle scripts require patchelf.

Not in scope

FlagGems' C++ path (liboperators.so / libtriton_jit.so) is not made
self-contained, which is why the CUDA wheel is built FLAGGEMS_KERNEL=OFF; the
FlagGems Python path does not need those libs. .conf op routing is unchanged.

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
lvyufeng force-pushed the metax-2.10 branch 2 times, most recently from 99852e5 to bf6030c Compare August 8, 2026 18:20
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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant