From bf6030cc42affb832c41af068bb8ca35f6fb131e Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Fri, 7 Aug 2026 13:27:39 +0800 Subject: [PATCH 1/2] feat(wheel): self-contained wheels that run on a stock torch+cpu 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/../:, 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//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. --- CMakeLists.txt | 53 +++- cmake/FlagosRpath.cmake | 76 +++++ csrc/CMakeLists.txt | 31 +- csrc/aten/backends/flagos/python_op_caller.cc | 9 + csrc/aten/copy_ops.cc | 17 +- csrc/runtime/accelerator/CMakeLists.txt | 5 + scripts/bundle_dcu_libtorch.sh | 152 ++++++++++ scripts/bundle_maca_libtorch.sh | 49 ++- scripts/bundle_ppu_libtorch.sh | 85 ++++++ scripts/lib/bundle_common.sh | 141 +++++++++ setup.py | 86 ++++-- torch_fl/__init__.py | 134 +++++++- torch_fl/accelerator/_vendor_libtorch.py | 285 ++++++++++++++++++ torch_fl/accelerator/dcu/__init__.py | 15 + .../accelerator/dcu/_dcu_libtorch_link.py | 124 ++++++++ .../accelerator/metax/_metax_libtorch_link.py | 244 +++++---------- torch_fl/accelerator/ppu/__init__.py | 15 + .../accelerator/ppu/_ppu_libtorch_link.py | 119 ++++++++ torch_fl/csrc/CMakeLists.txt | 6 + 19 files changed, 1386 insertions(+), 260 deletions(-) create mode 100644 cmake/FlagosRpath.cmake create mode 100644 scripts/bundle_dcu_libtorch.sh create mode 100644 scripts/bundle_ppu_libtorch.sh create mode 100644 scripts/lib/bundle_common.sh create mode 100644 torch_fl/accelerator/_vendor_libtorch.py create mode 100644 torch_fl/accelerator/dcu/__init__.py create mode 100644 torch_fl/accelerator/dcu/_dcu_libtorch_link.py create mode 100644 torch_fl/accelerator/ppu/__init__.py create mode 100644 torch_fl/accelerator/ppu/_ppu_libtorch_link.py diff --git a/CMakeLists.txt b/CMakeLists.txt index a1e5a1e7..805feceb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,6 +16,15 @@ cmake_minimum_required(VERSION 3.18 FATAL_ERROR) set(ACCELERATOR "cuda" CACHE STRING "Accelerator platform: cuda, metax, ascend, tsingmicro, dcu, gcu, or musa") +# Directory inside the wheel holding a bundled forked libtorch, when the backend +# ships one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate bundle +# dir": the CUDA backend drops its extra .so straight into torch_fl/lib/. +# Set per backend below; consumed by cmake/FlagosRpath.cmake. +set(FLAGOS_BUNDLE_LIBDIR "lib") +# Vendor driver/runtime dirs to embed in RUNPATH. Never bundled into the wheel +# (a machine with the card has the driver), mirroring MetaX's /opt/maca. +set(FLAGOS_VENDOR_RPATH_DIRS "") + # Kernel build options (can be overridden via -D flags from setup.py) option(CUDA_KERNEL "Build CUDA kernel implementations" ON) option(FLAGGEMS_KERNEL "Build FlagGems C++ kernel wrappers (requires liboperators.so)" ON) @@ -54,6 +63,9 @@ elseif(ACCELERATOR STREQUAL "metax") elseif(NOT FLAGGEMS_KERNEL) set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) endif() + # Self-contained wheel: forked libtorch bundled by + # scripts/bundle_maca_libtorch.sh (the maca runtime stays on the target). + set(FLAGOS_BUNDLE_LIBDIR "lib_maca") project(TORCH_FLAGOS CXX C) elseif(ACCELERATOR STREQUAL "tsingmicro") project(TORCH_FLAGOS CXX C) @@ -67,6 +79,10 @@ elseif(ACCELERATOR STREQUAL "dcu") # No nvcc and no CUDA language: nothing outside backends/metax/ is a .cu. set(CUDA_KERNEL OFF CACHE BOOL "Build CUDA kernel implementations" FORCE) set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) + # Self-contained wheel: DTK's forked libtorch (libtorch_cpu.so carries hip + # symbols and needs libgalaxyhip.so.5) is bundled here by + # scripts/bundle_dcu_libtorch.sh. The DTK driver itself stays on the target. + set(FLAGOS_BUNDLE_LIBDIR "lib_dcu") project(TORCH_FLAGOS CXX C) elseif(ACCELERATOR STREQUAL "gcu") project(TORCH_FLAGOS CXX C) @@ -80,6 +96,15 @@ elseif(ACCELERATOR STREQUAL "musa") set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE) project(TORCH_FLAGOS CXX C) else() + # PPU rides this branch: it has no accelerator value of its own because it is + # built against PPU_SDK/CUDA_SDK, so the CUDA boxing kernels apply unchanged + # (see the README's PPU section). It still needs its own bundle dir, since that + # libtorch is a local PPU build and not the upstream one. Detected by the SDK + # env vars, matching setup.py::_vendor_supplies_triton(). + if(DEFINED ENV{PPU_SDK} OR DEFINED ENV{PPU_HOME}) + set(FLAGOS_BUNDLE_LIBDIR "lib_ppu") + message(STATUS "PPU SDK detected: bundling libtorch into torch_fl/lib_ppu") + endif() # project(... CUDA) runs nvcc compiler-id before find_package(CUDAToolkit); conda # layouts need CUDAToolkit_ROOT (headers under targets/x86_64-linux/include). if(NOT CUDAToolkit_ROOT) @@ -222,7 +247,23 @@ elseif(ACCELERATOR STREQUAL "dcu") set(CUDA_RUNTIME_LIB cuda_runtime_compat) # Embed the compat + DTK library paths so LD_LIBRARY_PATH is not needed. - list(APPEND CMAKE_INSTALL_RPATH "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib") + # Measured layout of the deps the forked libtorch actually pulls. They are all + # on the DTK container's LD_LIBRARY_PATH, which a clean env does NOT inherit -- + # so the RUNPATH has to name them or a self-contained wheel gets "not found": + # cuda/cuda-*/lib64 libcudart.so.12 (shim over libgalaxyhip) + # lib libgalaxyhip.so.5 libhipnn librocfft.so.0 + # librocrand.so.1 librocsparse.so.1 libunwind.so.8 + # libMIOpen-recommend.so + # hip/lib, lib64 hip runtime + # aillvm/lib libomp.so + # .hyhal/rocm_smi/lib librocm_smi64.so.2 + # Keep in sync with VENDOR_RPATH in scripts/bundle_dcu_libtorch.sh. + set(_dcu_rpath_dirs + "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib" "${DTK_ROOT}/hip/lib" + "${DTK_ROOT}/lib64" "${DTK_ROOT}/aillvm/lib" "${DTK_ROOT}/llvm/lib" + "${DTK_ROOT}/.hyhal/rocm_smi/lib" "/opt/hyhal/lib") + list(APPEND CMAKE_INSTALL_RPATH ${_dcu_rpath_dirs}) + set(FLAGOS_VENDOR_RPATH_DIRS ${_dcu_rpath_dirs}) # Pre-create torch::cudart so PyTorch's cuda.cmake is skipped entirely # (guard: if(TARGET torch::cudart) return()). @@ -358,6 +399,12 @@ elseif(ACCELERATOR STREQUAL "metax") endif() message(STATUS "MetaX SDK path: ${METAX_PATH}") + # The maca runtime (mcblas/mcdnn/...) comes from the target's /opt/maca driver + # install; it is never bundled. /opt/maca is kept as an explicit fallback in + # case the build machine's METAX_PATH is a non-default location. + set(FLAGOS_VENDOR_RPATH_DIRS + "${METAX_PATH}/lib" "${METAX_PATH}/lib64" "/opt/maca/lib" "/opt/maca/lib64") + add_library(cuda_runtime_compat INTERFACE IMPORTED) set_target_properties(cuda_runtime_compat PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${METAX_PATH}/tools/cu-bridge/include;${METAX_PATH}/include;${METAX_PATH}/include/mcr" @@ -513,11 +560,15 @@ else() endif() if(PYTORCH_INSTALL_DIR) + # Needed for the *build* link step (and for an in-place/editable install, where + # the built .so stay in the source tree). flagos_set_portable_rpath() strips it + # from the installed targets' RUNPATH so the wheel is relocatable. list(APPEND CMAKE_INSTALL_RPATH "${PYTORCH_INSTALL_DIR}/lib") endif() include_directories(${CMAKE_CURRENT_SOURCE_DIR}) include(${PROJECT_SOURCE_DIR}/cmake/TorchPythonTargets.cmake) +include(${PROJECT_SOURCE_DIR}/cmake/FlagosRpath.cmake) add_subdirectory(${PROJECT_SOURCE_DIR}/csrc/runtime/accelerator) add_subdirectory(${PROJECT_SOURCE_DIR}/csrc) diff --git a/cmake/FlagosRpath.cmake b/cmake/FlagosRpath.cmake new file mode 100644 index 00000000..d408e519 --- /dev/null +++ b/cmake/FlagosRpath.cmake @@ -0,0 +1,76 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Portable RUNPATH for the installed native libs. +# +# The problem: CMakeLists.txt appends "${PYTORCH_INSTALL_DIR}/lib" to +# CMAKE_INSTALL_RPATH so an in-place build finds the active torch wheel. That is +# the *build machine's* interpreter path (e.g. +# /nfs/.../envs/torch-fl-210/lib/python3.12/site-packages/torch/lib), which does +# not exist on a target machine -- and worse, on a machine that does have that +# path it silently wins over the interpreter actually running, so a py3.12 env +# can load another env's libtorch. +# +# Dropping it is safe: libtorch_fl.so is only ever loaded through +# `import torch_fl` -> torch_fl._C -> libtorch_bindings.so, which happens after +# torch_fl/__init__.py has already run `import torch`. libc10 / libtorch_cpu are +# mapped by then, so the loader satisfies those DT_NEEDED entries by soname from +# the already-loaded set. This is why the MetaX self-contained wheel works today. +# +# Anything else the backend put on CMAKE_INSTALL_RPATH (vendor driver dirs, the +# FlagGems liboperators dir) is preserved, so in-place builds keep working. +# +# flagos_set_portable_rpath( [EXTRA_DIRS ...]) +function(flagos_set_portable_rpath _target) + cmake_parse_arguments(_ARG "" "" "EXTRA_DIRS" ${ARGN}) + + if(NOT UNIX OR APPLE) + return() + endif() + if(NOT TARGET ${_target}) + return() + endif() + + # $ORIGIN is torch_fl/lib/ for every installed target. A sibling bundle dir + # (torch_fl/lib_maca, lib_dcu, lib_ppu) is reached via $ORIGIN/../. + set(_rpath "$ORIGIN" "$ORIGIN/lib") + if(FLAGOS_BUNDLE_LIBDIR AND NOT FLAGOS_BUNDLE_LIBDIR STREQUAL "lib") + list(APPEND _rpath "$ORIGIN/../${FLAGOS_BUNDLE_LIBDIR}") + endif() + + foreach(_dir IN LISTS _ARG_EXTRA_DIRS FLAGOS_VENDOR_RPATH_DIRS) + if(_dir) + list(APPEND _rpath "${_dir}") + endif() + endforeach() + + # Inherit whatever the backend branches appended, minus the build machine's + # torch/lib and minus the $ORIGIN entries already placed above. + foreach(_dir IN LISTS CMAKE_INSTALL_RPATH) + if(NOT _dir MATCHES "^\\$ORIGIN" + AND NOT (PYTORCH_INSTALL_DIR AND _dir STREQUAL "${PYTORCH_INSTALL_DIR}/lib")) + list(APPEND _rpath "${_dir}") + endif() + endforeach() + + list(REMOVE_DUPLICATES _rpath) + string(REPLACE ";" ":" _rpath_str "${_rpath}") + + set_target_properties(${_target} PROPERTIES + INSTALL_RPATH "${_rpath_str}" + # Stop CMake appending the imported torch target's build-machine link dir. + INSTALL_RPATH_USE_LINK_PATH OFF + BUILD_WITH_INSTALL_RPATH ON + ) +endfunction() diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 520e2cd6..af5c4f01 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -325,29 +325,18 @@ if(ACCELERATOR STREQUAL "metax") message(WARNING "mcblas not found, MetaX mm will fallback to naive kernel") endif() - # Self-contained wheel: the installed libtorch_fl.so must NOT bake the build - # machine's absolute torch/lib path into its RUNPATH (that path does not exist - # on target machines). Resolve the forked libtorch bundled under - # torch_fl/lib_maca/ via $ORIGIN, and the maca runtime (mcblas/mcdnn/...) from - # the target's /opt/maca driver install. Disabling INSTALL_RPATH_USE_LINK_PATH - # stops CMake from appending the imported torch target's build-machine link - # dir. Mirrors scripts/bundle_maca_libtorch.sh's patchelf pass. - set(_metax_rpath - "$ORIGIN:$ORIGIN/../lib_maca:${METAX_PATH}/lib:${METAX_PATH}/lib64:/opt/maca/lib:/opt/maca/lib64") - # Setting INSTALL_RPATH explicitly discards the entries the top-level - # CMakeLists appended to CMAKE_INSTALL_RPATH, including FlagGems' - # liboperators.so dir. Re-add it so the C++ FlagGems path (FLAGGEMS_KERNEL, - # kFlagOs) resolves liboperators.so without LD_LIBRARY_PATH. - if(FLAGGEMS_KERNEL AND _flaggems_libdir) - set(_metax_rpath "${_metax_rpath}:${_flaggems_libdir}") - endif() - set_target_properties(${LIBRARY_NAME} PROPERTIES - INSTALL_RPATH "${_metax_rpath}" - INSTALL_RPATH_USE_LINK_PATH OFF - BUILD_WITH_INSTALL_RPATH ON - ) endif() +# Self-contained wheel: the installed libtorch_fl.so must NOT bake the build +# machine's absolute torch/lib path into its RUNPATH -- that path does not exist +# on a target machine, and on one where it does exist it silently overrides the +# interpreter actually running. Resolve a bundled forked libtorch via $ORIGIN +# instead. Everything else the backend put on CMAKE_INSTALL_RPATH is kept, +# including FlagGems' liboperators.so dir, so the C++ FlagGems path +# (FLAGGEMS_KERNEL, kFlagOs) still resolves without LD_LIBRARY_PATH. Mirrors the +# patchelf pass in scripts/bundle_*_libtorch.sh. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} diff --git a/csrc/aten/backends/flagos/python_op_caller.cc b/csrc/aten/backends/flagos/python_op_caller.cc index 9c196342..b8eb7aea 100644 --- a/csrc/aten/backends/flagos/python_op_caller.cc +++ b/csrc/aten/backends/flagos/python_op_caller.cc @@ -610,6 +610,15 @@ at::Generator GetFlagosDefaultCudaGenerator(int64_t device_index) { } py::gil_scoped_acquire gil; py::module_ torch_cuda = py::module_::import("torch.cuda"); + // torch.cuda.default_generators is populated by CUDA lazy-init, not at import. + // A self-contained wheel front-ends a stock torch+cpu whose CUDA state nothing + // else touches, so the tuple is still empty here and indexing it would raise + // IndexError -- which surfaces as a failure of every device-side RNG op + // (randn/rand/normal_) while plain factories (empty/zeros/ones) work fine. + // Force lazy-init first; it is idempotent and cheap once initialized. + if (py::len(torch_cuda.attr("default_generators")) == 0) { + torch_cuda.attr("init")(); + } py::object gens = torch_cuda.attr("default_generators"); py::object py_gen = gens[py::cast(device_index)]; // torch.Generator -> at::Generator via THPGenerator unpack. diff --git a/csrc/aten/copy_ops.cc b/csrc/aten/copy_ops.cc index eb56e719..1f8f7778 100644 --- a/csrc/aten/copy_ops.cc +++ b/csrc/aten/copy_ops.cc @@ -26,8 +26,16 @@ // On the CUDA-family backends (including MetaX boxing) the flagos device shares // the vendor's CUDA streams, so the current stream is readable from c10::cuda. +// +// USE_DCU is excluded for the same reason runtime/guard.h:22 excludes it: the +// DCU wheel is hipified, so resolves through DTK's +// CUDA-compat shim against USE_ROCM-hipified torch headers and expands to hip* +// symbols the shim never declares (`'hipStreamCaptureStatus' was not declared`), +// and c10::cuda::getCurrentCUDAStream would not link there anyway -- DTK exports +// c10::hip with zero c10::cuda symbols. DCU still shares the vendor's streams, +// so it needs *some* barrier; see SyncCurrentStreamBeforeBlockingCopy below. #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \ - !defined(USE_MUSA) + !defined(USE_MUSA) && !defined(USE_DCU) #define FLAGOS_COPY_HAS_CUDA_STREAM 1 #include #endif @@ -53,6 +61,13 @@ inline void SyncCurrentStreamBeforeBlockingCopy() { if (stream.stream() != nullptr) { stream.synchronize(); } +#elif defined(USE_DCU) + // DCU shares the vendor's streams, so it needs this barrier just as much as + // the other CUDA-family backends -- but it cannot ask which stream is current + // (no c10::cuda symbols in a hipified wheel; see the guard above). Fall back to + // a device-wide sync: a superset of the per-stream wait, so still correct, + // just coarser. Only reached on the blocking-copy paths. + ::DeviceSynchronize(); #endif } diff --git a/csrc/runtime/accelerator/CMakeLists.txt b/csrc/runtime/accelerator/CMakeLists.txt index ab285eae..998bea55 100644 --- a/csrc/runtime/accelerator/CMakeLists.txt +++ b/csrc/runtime/accelerator/CMakeLists.txt @@ -123,6 +123,11 @@ else() target_link_libraries(${LIBRARY_NAME} PRIVATE CUDA::cudart) endif() +# libflagos.so needs no torch symbols at all (only the vendor runtime), so the +# inherited ${PYTORCH_INSTALL_DIR}/lib RUNPATH entry is pure build-machine +# leakage. Vendor driver dirs are preserved. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} diff --git a/scripts/bundle_dcu_libtorch.sh b/scripts/bundle_dcu_libtorch.sh new file mode 100644 index 00000000..6dce7230 --- /dev/null +++ b/scripts/bundle_dcu_libtorch.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# 把 DTK fork 的 libtorch C++ .so 打进 torch_fl/lib_dcu/，做成自包含单 wheel。 +# +# 为什么要打 core 库（不是只打 libtorch_hip.so）： +# 实测 libtorch_fl.so 的未定义符号里，DTK 的 libtorch_cpu.so 提供 2101 个， +# libtorch_hip.so 提供 0 个 —— fork 在 core 库这一侧（DTK 的 libtorch_cpu.so +# 带 128 个 hip 符号、DT_NEEDED 写着 libgalaxyhip.so.5）。libtorch_hip.so 仍然 +# 必须在场：boxing kernel 要求 CUDA dispatch key 上已注册厂商 kernel。 +# +# 为什么要打 torch.libs/：那 17 个文件是 auditwheel 改名的通用库 +# （glog/gflags/MKL/OpenMPI/hwloc/libxml2...），其中 5 个是 DTK libtorch_cpu.so 的 +# 直接 DT_NEEDED。实测它们对 hip/dtk/rocm 的引用数为 0，不引入 SDK 绑定；但官方 +# torch+cpu wheel 的 torch.libs 里一个都没有（名字带 hash，系统库也对不上）， +# 不打就跑不起来。它们和 core 库放同一目录，靠 $ORIGIN 解析。 +# +# 注意 $ORIGIN 的实测语义：glibc 按"打开这个对象时用的路径"展开 $ORIGIN，不是按 +# 真实路径。所以从 torch/lib 的软链打开 libc10.so 时 $ORIGIN 是 torch/lib，找不到 +# 同在 lib_dcu 的 libgflags-8aee0f6c.so.2.1.2。因此运行期预加载走 bundle 原路径 +# （_vendor_libtorch._preload_global），不走软链 —— 那里有同样的说明。 +# +# 不打包：DTK 驱动栈（libgalaxyhip.so.5 libMIOpen.so.1 librocblas.so.4 +# libhipblas.so.2 librccl.so.1 等 12 个 soname）留在目标机 /opt/dtk。 +# SDK 版本绑定本来就存在且更强（libtorch_hip.so 的 DT_NEEDED 写死 +# librocblas.so.4），多打那 129 MB 不会让绑定更紧。 +# +# 用法: +# FLAGOS_DCU_TORCH_LIB= bash scripts/bundle_dcu_libtorch.sh +# DTK_ROOT=/opt/dtk bash scripts/bundle_dcu_libtorch.sh +# +# 应在 `python setup.py bdist_wheel`（ACCELERATOR=dcu）之后、打 wheel 之前跑。幂等。 + +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + +LIB_DCU="${REPO_DIR}/torch_fl/lib_dcu" +TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" +DTK_ROOT="${DTK_ROOT:-${ROCM_PATH:-/opt/dtk}}" + +SRC="${FLAGOS_DCU_TORCH_LIB:-}" +if [ -z "${SRC}" ]; then + SRC="$(bundle_find_vendor_torch_lib libtorch_hip.so dtk hip das || true)" +fi + +if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then + echo "error: 找不到 DTK torch/lib。设 FLAGOS_DCU_TORCH_LIB=" >&2 + exit 1 +fi +if [ ! -f "${SRC}/libtorch_hip.so" ]; then + echo "error: ${SRC} 里没有 libtorch_hip.so，不是 DTK torch/lib" >&2 + exit 1 +fi + +bundle_require_patchelf + +# 与 _dcu_libtorch_link._CORE_SO + _HIP_SO 对齐的自洽集合。 +# libshm.so：libtorch_python.so 的直接 DT_NEEDED（torch.multiprocessing 的共享内存 +# 管理器），DTK 的 torch/lib 里有实体文件，但它不在软链清单里，所以必须打进 bundle +# 让 $ORIGIN 找得到。 +# libcaffe2_nvrtc.so：torch.cuda.init() 会 dlopen 它。DTK torch 带这个文件， +# stock +cpu wheel 没有，所以不打进来的话 GetFlagosDefaultCudaGenerator 里那次 +# 按需 init 就死在 "Error in dlopen: libcaffe2_nvrtc.so"（RNG 修复反而暴露了它）。 +CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so libcaffe2_nvrtc.so) +HIP_SO=(libc10_hip.so libtorch_hip.so libmagma.so libshm.so) + +# DTK 驱动栈实测分布（容器 LD_LIBRARY_PATH 与 find 结果一致）： +# lib/ libhipnn librocfft.so.0 librocrand.so.1 librocsparse.so.1 +# libMIOpen-recommend.so libunwind.so.8 libgalaxyhip.so.5 +# hip/lib/ hip runtime +# aillvm/lib/ libomp.so +# .hyhal/rocm_smi/lib/ librocm_smi64.so.2 +# 全部留在目标机（装了 DCU 卡就有 /opt/dtk），但 RPATH 必须覆盖到，否则 +# LD_LIBRARY_PATH 没设时就 not found。 +VENDOR_RPATH="${DTK_ROOT}/lib:${DTK_ROOT}/hip/lib:${DTK_ROOT}/lib64" +VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/aillvm/lib:${DTK_ROOT}/.hyhal/rocm_smi/lib" +VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/llvm/lib:/opt/hyhal/lib" + +echo "源 DTK torch/lib : ${SRC}" +echo "目标 lib_dcu : ${LIB_DCU}" +echo "DTK 驱动路径 : ${DTK_ROOT}" + +# bundle 内的库要能从两条路径被打开： +# 1. 直接从 lib_dcu/（运行期预加载走这条，见 _vendor_libtorch._preload_global） +# 2. 通过 torch/lib/ 的软链（`import torch` 自己加载 libtorch_global_deps.so 时） +# glibc 按"打开这个对象时用的路径"展开 $ORIGIN，所以第 2 条路上 $ORIGIN 是 +# torch/lib，找不到同在 lib_dcu 的 libmpi-3fcb240d.so.40.40.3 等 auditwheel 改名库。 +# 两个目录都在 site-packages 下同级（torch/lib -> ../../torch_fl/lib_dcu）， +# 所以再加一条相对路径就同时覆盖两种情况。实测：不加这条，先 `import torch` +# 再 import torch_fl 会死在 libmpi not found。 +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_dcu" + +bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" +bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${HIP_SO[@]}" + +# torch.libs/：与 torch/ 同级的 auditwheel 目录，文件名带 hash 后缀，只能 glob。 +TORCH_LIBS="$(cd "${SRC}/../.." && pwd)/torch.libs" +if [ -d "${TORCH_LIBS}" ]; then + echo "源 torch.libs : ${TORCH_LIBS}" + _names=() + while IFS= read -r f; do + _names+=("$(basename "${f}")") + done < <(find "${TORCH_LIBS}" -maxdepth 1 -type f -name '*.so*' | sort) + if [ ${#_names[@]} -gt 0 ]; then + bundle_copy_so "${TORCH_LIBS}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${_names[@]}" + else + echo "warning: ${TORCH_LIBS} 里没有 .so，跳过" >&2 + fi +else + echo "warning: 找不到 ${TORCH_LIBS}；若 libtorch_cpu.so 的 DT_NEEDED 里有带 hash" >&2 + echo " 的通用库（libglog-*.so.0 等），目标机上会缺库。" >&2 +fi + +# libtorch_fl.so / libflagos.so 还需要 DTK 的 CUDA 兼容层 libcudart.so.12 +# （cuda_runtime_compat，见 CMakeLists.txt 的 DCU_CUDA_ROOT）。cmake 已经把它写进 +# RUNPATH 了，这里重写时不能丢 —— 否则干净环境里 libcudart.so.12 就 not found。 +_DCU_CUDA_LIB64="" +for _c in "${DTK_ROOT}"/cuda/cuda-*/lib64; do + if [ -f "${_c}/libcudart.so.12" ]; then + _DCU_CUDA_LIB64="${_c}" + break + fi +done +if [ -z "${_DCU_CUDA_LIB64}" ]; then + echo "warning: ${DTK_ROOT}/cuda/cuda-*/lib64 里没有 libcudart.so.12" >&2 +fi +PLUGIN_RPATH="\$ORIGIN:\$ORIGIN/../lib_dcu" +[ -n "${_DCU_CUDA_LIB64}" ] && PLUGIN_RPATH="${PLUGIN_RPATH}:${_DCU_CUDA_LIB64}" +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" "${PLUGIN_RPATH}:${VENDOR_RPATH}" + +# torch/version.py 是构建期生成的纯 Python，换掉 .so 改不了它。自包含 DCU wheel +# 前面是 stock torch+cpu，torch.version.hip 报 None，而 triton 的 hcu backend +# 恰好按这个值判活（backends/hcu/driver.py is_active(): torch.cuda.is_available() +# and torch.version.hip is not None）。None 就永不激活，任何 flag_gems 算子都死在 +# triton 的 driver factory："0 active drivers ([])"。把厂商 torch 自己的 version.py +# 带上，import 时由 _restore_dcu_hip_version() 读回 hip/rocm 字符串。 +_VENDOR_VERSION_PY="$(cd "${SRC}/.." && pwd)/version.py" +if [ -f "${_VENDOR_VERSION_PY}" ]; then + cp -fL "${_VENDOR_VERSION_PY}" "${LIB_DCU}/vendor_version.py" + echo "已复制 vendor version.py -> lib_dcu/vendor_version.py" + grep -E "^\s*(hip|rocm)\s*(:|=)" "${LIB_DCU}/vendor_version.py" || true +else + echo "warning: 找不到 ${_VENDOR_VERSION_PY}，triton hcu backend 可能不激活" >&2 +fi + +bundle_summary "${LIB_DCU}" +bundle_check_needed "${LIB_DCU}" \ + "${DTK_ROOT}/lib" "${DTK_ROOT}/hip/lib" "${DTK_ROOT}/lib64" \ + "${DTK_ROOT}/aillvm/lib" "${DTK_ROOT}/llvm/lib" \ + "${DTK_ROOT}/.hyhal/rocm_smi/lib" "${_DCU_CUDA_LIB64:-/nonexistent}" \ + "/opt/hyhal/lib" "/usr/lib64" "/usr/lib" "/usr/lib/x86_64-linux-gnu" diff --git a/scripts/bundle_maca_libtorch.sh b/scripts/bundle_maca_libtorch.sh index 9d1e8b5b..25c624c9 100644 --- a/scripts/bundle_maca_libtorch.sh +++ b/scripts/bundle_maca_libtorch.sh @@ -21,6 +21,9 @@ set -euo pipefail REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + LIB_MACA="${REPO_DIR}/torch_fl/lib_maca" TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" MACA_PATH="${MACA_PATH:-${METAX_HOME:-${MACA_HOME:-/opt/maca}}}" @@ -28,17 +31,7 @@ MACA_PATH="${MACA_PATH:-${METAX_HOME:-${MACA_HOME:-/opt/maca}}}" # 沐曦 torch/lib 来源：显式 env，或从 conda 里找 +metax torch。 SRC="${FLAGOS_MACA_TORCH_LIB:-}" if [ -z "${SRC}" ]; then - SRC=$(python - <<'PY' 2>/dev/null || true -import importlib.util, os -spec = importlib.util.find_spec("torch") -if spec and spec.submodule_search_locations: - lib = os.path.join(spec.submodule_search_locations[0], "lib") - ver = os.path.join(spec.submodule_search_locations[0], "version.py") - txt = open(ver).read() if os.path.isfile(ver) else "" - if ("metax" in txt or "maca" in txt) and os.path.exists(os.path.join(lib, "libtorch_cuda.so")): - print(lib) -PY -) + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so metax maca || true)" fi if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then @@ -50,35 +43,27 @@ if [ ! -f "${SRC}/libtorch_cuda.so" ]; then exit 1 fi -command -v patchelf >/dev/null 2>&1 || { echo "error: 需要 patchelf（pip install patchelf）" >&2; exit 1; } +bundle_require_patchelf # 与 _metax_libtorch_link._CORE_SO + _CUDA_SO 对齐的自洽集合。 CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) -CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so) +CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) + +VENDOR_RPATH="${MACA_PATH}/lib:${MACA_PATH}/lib64" +# 同 DCU：bundle 内的库要能从 lib_maca/ 和 torch/lib/ 软链两条路径被打开。 +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_maca" echo "源沐曦 torch/lib : ${SRC}" echo "目标 lib_maca : ${LIB_MACA}" echo "maca runtime 路径: ${MACA_PATH}/lib" -mkdir -p "${LIB_MACA}" -for so in "${CORE_SO[@]}" "${CUDA_SO[@]}"; do - src="${SRC}/${so}" - if [ ! -f "${src}" ]; then - echo " 跳过 ${so}（源不存在）" - continue - fi - # 解引用软链，拷实体文件。 - cp -fL "${src}" "${LIB_MACA}/${so}" - patchelf --set-rpath "\$ORIGIN:${MACA_PATH}/lib:${MACA_PATH}/lib64" "${LIB_MACA}/${so}" - echo " 打包 ${so} ($(du -h "${LIB_MACA}/${so}" | cut -f1))" -done +bundle_copy_so "${SRC}" "${LIB_MACA}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 \ + "${CORE_SO[@]}" "${CUDA_SO[@]}" # torch_fl.so 去掉构建机写死的沐曦 torch/lib 绝对路径，改为从包内 lib_maca 找 fork libtorch。 -for so in libtorch_fl.so libtorch_bindings.so; do - target="${TORCH_FL_LIB}/${so}" - [ -f "${target}" ] || continue - patchelf --set-rpath "\$ORIGIN:\$ORIGIN/../lib_maca:${MACA_PATH}/lib:${MACA_PATH}/lib64" "${target}" - echo " 重写 RPATH ${so}" -done +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ + "\$ORIGIN:\$ORIGIN/../lib_maca:${VENDOR_RPATH}" -echo "完成。lib_maca 总大小: $(du -sh "${LIB_MACA}" | cut -f1)" +bundle_summary "${LIB_MACA}" +bundle_check_needed "${LIB_MACA}" \ + "${MACA_PATH}/lib" "${MACA_PATH}/lib64" "/usr/lib64" "/usr/lib" diff --git a/scripts/bundle_ppu_libtorch.sh b/scripts/bundle_ppu_libtorch.sh new file mode 100644 index 00000000..a874ce52 --- /dev/null +++ b/scripts/bundle_ppu_libtorch.sh @@ -0,0 +1,85 @@ +#!/usr/bin/env bash +# 把 PPU 本地构建的 libtorch C++ .so 打进 torch_fl/lib_ppu/，做成自包含单 wheel。 +# +# PPU 是对着 PPU_SDK/CUDA_SDK 编的，所以 ACCELERATOR=cuda、CUDA boxing kernel +# 原样可用。跟真 NVIDIA 机器的区别在 libtorch：它是本地 USE_CUDA=1 源码构建，不是 +# 上游 wheel。实测 libtorch_fl.so 的未定义符号里它的 libtorch_cpu.so 提供 2092 个， +# libtorch_cuda.so 提供 0 个、libc10_cuda.so 10 个 —— 所以 core 库必须换掉， +# libtorch_cuda.so 只需在场（CUDA dispatch key 上要有已注册的厂商 kernel）。 +# +# 那个本地构建还链了 /usr/local/lib 下的系统 MKL +# （libmkl_core / libmkl_gnu_thread / libmkl_intel_lp64，~171 MB），官方 +# torch+cpu wheel 里没有对应文件，所以一并打进 lib_ppu/。 +# +# 不打包：PPU SDK runtime 留在目标机 /usr/local/PPU_SDK。 +# +# 用法: +# FLAGOS_PPU_TORCH_LIB= bash scripts/bundle_ppu_libtorch.sh +# PPU_SDK=/usr/local/PPU_SDK bash scripts/bundle_ppu_libtorch.sh +# FLAGOS_PPU_MKL_DIR=/usr/local/lib bash scripts/bundle_ppu_libtorch.sh +# +# 应在 `python setup.py bdist_wheel` 之后、打 wheel 之前跑。幂等。 + +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=scripts/lib/bundle_common.sh +source "${REPO_DIR}/scripts/lib/bundle_common.sh" + +LIB_PPU="${REPO_DIR}/torch_fl/lib_ppu" +TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" +PPU_SDK="${PPU_SDK:-${PPU_HOME:-/usr/local/PPU_SDK}}" +MKL_DIR="${FLAGOS_PPU_MKL_DIR:-/usr/local/lib}" + +SRC="${FLAGOS_PPU_TORCH_LIB:-}" +if [ -z "${SRC}" ]; then + # PPU 的 torch 是本地构建，version.py 里不一定有 "ppu" 字样，所以先按标记找， + # 找不到就退化成"当前解释器的 torch 只要有 libtorch_cuda.so 就算"。 + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so ppu || true)" + if [ -z "${SRC}" ]; then + SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so || true)" + fi +fi + +if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then + echo "error: 找不到 PPU torch/lib。设 FLAGOS_PPU_TORCH_LIB=" >&2 + exit 1 +fi +if [ ! -f "${SRC}/libtorch_cuda.so" ]; then + echo "error: ${SRC} 里没有 libtorch_cuda.so，不是 CUDA 构建的 torch/lib" >&2 + exit 1 +fi + +bundle_require_patchelf + +# 与 _ppu_libtorch_link._CORE_SO + _CUDA_SO 对齐的自洽集合。 +CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) +CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) +MKL_SO=(libmkl_core.so.1 libmkl_gnu_thread.so.1 libmkl_intel_lp64.so.1) + +VENDOR_RPATH="${PPU_SDK}/CUDA_SDK/lib64:${PPU_SDK}/lib:${PPU_SDK}/lib64" +# 同 DCU：bundle 内的库要能从 lib_ppu/ 和 torch/lib/ 软链两条路径被打开。 +BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_ppu" + +echo "源 PPU torch/lib : ${SRC}" +echo "目标 lib_ppu : ${LIB_PPU}" +echo "PPU SDK 路径 : ${PPU_SDK}" + +bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" +bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${CUDA_SO[@]}" + +# 系统 MKL：libtorch_cpu.so 的直接 DT_NEEDED，官方 wheel 里没有同名文件。 +if [ -d "${MKL_DIR}" ]; then + echo "源 MKL : ${MKL_DIR}" + bundle_copy_so "${MKL_DIR}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${MKL_SO[@]}" +else + echo "warning: ${MKL_DIR} 不存在，跳过 MKL；目标机上若无 MKL 会缺库" >&2 +fi + +bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ + "\$ORIGIN:\$ORIGIN/../lib_ppu:${VENDOR_RPATH}" + +bundle_summary "${LIB_PPU}" +bundle_check_needed "${LIB_PPU}" \ + "${PPU_SDK}/CUDA_SDK/lib64" "${PPU_SDK}/lib" "${PPU_SDK}/lib64" \ + "${MKL_DIR}" "/usr/lib64" "/usr/lib" "/usr/local/lib" diff --git a/scripts/lib/bundle_common.sh b/scripts/lib/bundle_common.sh new file mode 100644 index 00000000..5e1bdde9 --- /dev/null +++ b/scripts/lib/bundle_common.sh @@ -0,0 +1,141 @@ +#!/usr/bin/env bash +# 自包含 wheel 打包的公共部分，被 scripts/bundle_{maca,dcu,ppu}_libtorch.sh source。 +# +# 背景：metax / dcu / ppu 三个后端跑的都是厂商 fork 的 libtorch —— 实测 +# libtorch_cpu.so 提供 libtorch_fl.so 两千多个未定义符号，而厂商库 +# （libtorch_hip.so / libtorch_cuda.so）提供 0 个。所以自包含 wheel 必须把 core +# 库一起打进来，不能只打厂商库。做法统一为： +# - bundle 目录内的 .so -> RPATH = $ORIGIN:<厂商驱动目录...> +# （驱动运行时留在目标机，不打包：装了卡的机器必有驱动） +# - torch_fl/lib/libtorch_fl.so -> RPATH = $ORIGIN:$ORIGIN/../:<驱动...> +# （去掉构建机写死的绝对路径。cmake/FlagosRpath.cmake 已经这么设了， +# 这里的 patchelf 是对已有 .so / 非 cmake 构建产物的兜底，保持幂等。） +# +# 运行期把 bundle 里的 core 库软链到 stock torch/lib 的逻辑在 +# torch_fl/accelerator/_vendor_libtorch.py，两边的 .so 清单必须对齐。 +# +# 提供: +# bundle_require_patchelf +# bundle_find_vendor_torch_lib +# bundle_copy_so +# bundle_rewrite_plugin_rpath +# bundle_summary + +set -euo pipefail + +# patchelf 四个节点默认都没有。bw1000/810e 上 `pip install patchelf` 直接可用； +# mc550 的默认 index 没有这个包，要指定源。 +bundle_require_patchelf() { + if command -v patchelf >/dev/null 2>&1; then + return 0 + fi + cat >&2 <<'EOF' +error: 找不到 patchelf，无法重写 RPATH。装法（任选）: + pip install patchelf + pip install -i https://pypi.tuna.tsinghua.edu.cn/simple patchelf # 默认源没有时 + conda install -c conda-forge patchelf +EOF + return 1 +} + +# 定位厂商 torch/lib：优先当前解释器的 torch（校验 version.py 含厂商标记且 +# probe_so 存在），失败返回非 0。调用方应先看自己的 env 覆盖变量。 +# 用法: bundle_find_vendor_torch_lib libtorch_hip.so dtk hip +bundle_find_vendor_torch_lib() { + local probe="$1" + shift + local markers="$*" + local py + py="$(command -v python || command -v python3 || true)" + [ -n "${py}" ] || return 1 + PROBE_SO="${probe}" VENDOR_MARKERS="${markers}" "${py}" - <<'PY' 2>/dev/null || return 1 +import importlib.util, os + +probe = os.environ["PROBE_SO"] +markers = os.environ["VENDOR_MARKERS"].split() +spec = importlib.util.find_spec("torch") +if spec and spec.submodule_search_locations: + root = spec.submodule_search_locations[0] + lib = os.path.join(root, "lib") + ver = os.path.join(root, "version.py") + txt = open(ver).read() if os.path.isfile(ver) else "" + if (not markers or any(m in txt for m in markers)) and os.path.exists( + os.path.join(lib, probe) + ): + print(lib) +PY +} + +# 拷 .so 并设 RPATH。cp -fL 解引用软链，拷实体文件。 +# required=1 时源缺失即失败，=0 时跳过（stock +cpu wheel 本来就没有厂商库）。 +bundle_copy_so() { + local src_dir="$1" dst_dir="$2" rpath="$3" required="$4" + shift 4 + local so src + mkdir -p "${dst_dir}" + for so in "$@"; do + src="${src_dir}/${so}" + if [ ! -f "${src}" ]; then + if [ "${required}" = "1" ]; then + echo "error: 缺少必需的 ${src}" >&2 + return 1 + fi + echo " 跳过 ${so}（源不存在）" + continue + fi + cp -fL "${src}" "${dst_dir}/${so}" + # 非 ELF（极少数 .so 其实是 linker script）patchelf 会失败，不致命。 + if ! patchelf --set-rpath "${rpath}" "${dst_dir}/${so}" 2>/dev/null; then + echo " ${so}: patchelf 跳过（非 ELF?）" + fi + echo " 打包 ${so} ($(du -h "${dst_dir}/${so}" | cut -f1))" + done +} + +# 重写 libtorch_fl.so / libtorch_bindings.so / libflagos.so 的 RPATH。 +bundle_rewrite_plugin_rpath() { + local lib_dir="$1" rpath="$2" so target + for so in libtorch_fl.so libtorch_bindings.so libflagos.so; do + target="${lib_dir}/${so}" + [ -f "${target}" ] || continue + patchelf --set-rpath "${rpath}" "${target}" + echo " 重写 RPATH ${so} -> ${rpath}" + done +} + +bundle_summary() { + local dst_dir="$1" + echo "完成。$(basename "${dst_dir}") 总大小: $(du -sh "${dst_dir}" | cut -f1)（$(find "${dst_dir}" -type f | wc -l | tr -d ' ') 个文件）" +} + +# 目标机上一定有的基线库，自检时不报。 +_BUNDLE_BASELINE_RE='^(libc|libm|libdl|librt|libpthread|libstdc\+\+|libgcc_s|ld-linux.*|libutil|libresolv|libnsl|libcrypt|libatomic|libgomp|libz|libnuma)\.so' + +# 自检：bundle 里每个 .so 的 DT_NEEDED，凡是既不在 bundle 内、也不在给定的目标机 +# 目录里、又不是基线系统库的，都报出来 —— 那就是目标机上会 "not found" 的候选。 +# 用法: bundle_check_needed +bundle_check_needed() { + local dst_dir="$1" + shift + local so dep d found + echo "---- DT_NEEDED 自检（只报可能在目标机缺失的项）----" + local missing=0 + for so in "${dst_dir}"/*.so*; do + [ -f "${so}" ] || continue + while read -r dep; do + [ -n "${dep}" ] || continue + [[ "${dep}" =~ ${_BUNDLE_BASELINE_RE} ]] && continue + [ -e "${dst_dir}/${dep}" ] && continue + found=0 + for d in "$@"; do + if [ -e "${d}/${dep}" ]; then found=1; break; fi + done + [ "${found}" = "1" ] && continue + echo " $(basename "${so}") -> ${dep}" + missing=1 + done < <(patchelf --print-needed "${so}" 2>/dev/null || true) + done + if [ "${missing}" = "0" ]; then + echo " 无（bundle + 驱动目录已覆盖全部非基线依赖）" + fi +} diff --git a/setup.py b/setup.py index de46c287..e26de1c7 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,17 @@ # "dcu", "gcu", or "musa" ACCELERATOR = os.environ.get("ACCELERATOR", "cuda").lower() +# Directory inside the wheel holding a bundled forked libtorch, for the backends +# that ship one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate +# bundle dir": the CUDA backend drops its extra .so straight into torch_fl/lib/. +# Must match FLAGOS_BUNDLE_LIBDIR in CMakeLists.txt -- _C.so's RUNPATH has to +# reach the bundle or its auditwheel-mangled deps (libglog-*.so.0) go missing. +_BUNDLE_LIBDIR = {"metax": "lib_maca", "dcu": "lib_dcu"}.get(ACCELERATOR, "lib") +if _BUNDLE_LIBDIR == "lib" and ( + os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME") +): + _BUNDLE_LIBDIR = "lib_ppu" + BASE_DIR = os.path.dirname(os.path.realpath(__file__)) # Only run cmake build for actual build commands, not metadata collection @@ -363,17 +374,23 @@ def build_deps(): ] ) elif ACCELERATOR == "dcu": - # Pure boxing build. The DCU torch wheel is a hipified build whose HIP - # kernels are registered under the CUDA dispatch key, so the generated + # Boxing build. The DCU torch wheel is a hipified build whose HIP kernels + # are registered under the CUDA dispatch key, so the generated # PrivateUse1 -> CUDA boxing kernels reach them with no hand-written - # kernels of our own. FLAGGEMS_KERNEL needs liboperators.so (not built - # for DTK); FLAGGEMS_PYTHON needs a DTK triton, which is a separate - # install -- both stay off unless explicitly requested below. + # kernels of our own. FLAGGEMS_KERNEL needs liboperators.so, which is not + # built for DTK, and stays off. + # + # FLAGGEMS_PYTHON defaults ON, same as metax/cuda: DTK ships a working + # triton (hcu backend) that flag_gems runs on, so the wheel compiles the + # FlagGems Python-path kernels too and the choice becomes a runtime one + # (FLAGOS_USE_FLAGGEMS -> backends_dcu_flaggems.conf). python_op_caller + # links torch_python_library, already in the link set, so this adds + # nothing to the wheel size. Set FLAGGEMS_PYTHON=0 for a slim pure-boxing + # build; the generic pass-through below honors that. cmake_args.extend( [ "-DCUDA_KERNEL=OFF", "-DFLAGGEMS_KERNEL=OFF", - "-DFLAGGEMS_PYTHON=OFF", "-DMETAX_KERNEL=OFF", "-DASCEND_KERNEL=OFF", ] @@ -604,19 +621,33 @@ def run(self): os.remove(os.path.join(dirpath, filename)) +def _extension_rpath_args(): + """RUNPATH for torch_fl._C: torch_fl/lib plus the bundle dir when separate. + + _C.so links libtorch_bindings.so out of torch_fl/lib, which in turn pulls the + bundled vendor libtorch and its auditwheel-mangled deps out of the bundle dir. + Without the second entry a self-contained wheel fails at import with e.g. + "libglog.so.0: cannot open shared object file". + """ + args = make_relative_rpath_args("lib") + if _BUNDLE_LIBDIR != "lib": + args += make_relative_rpath_args(_BUNDLE_LIBDIR) + return args + + def _extension_compile_args(): if IS_WINDOWS: # /NODEFAULTLIB makes sure we only link to DLL runtime # and matches the flags set for protobuf and ONNX - extra_link_args: list[str] = ["/NODEFAULTLIB:LIBCMT.LIB"] + [ - *make_relative_rpath_args("lib") - ] + extra_link_args: list[str] = [ + "/NODEFAULTLIB:LIBCMT.LIB" + ] + _extension_rpath_args() # /MD links against DLL runtime # and matches the flags set for protobuf and ONNX # /EHsc is about standard C++ exception handling extra_compile_args = ["/MD", "/FS", "/EHsc"] else: - extra_link_args = [*make_relative_rpath_args("lib")] + extra_link_args = _extension_rpath_args() extra_compile_args = [ "-Wall", "-Wextra", @@ -649,10 +680,20 @@ def _get_setup_kwargs(): "lib/*.dylib*", "lib/*.dll", "lib/*.lib", - # MetaX self-contained wheel: forked libtorch C++ .so bundled here so - # the process loads the MetaX C++ runtime without a separate metax - # torch wheel (see scripts/bundle_maca_libtorch.sh). + # Self-contained wheels: the vendor's forked libtorch C++ .so bundled + # here so the process loads that C++ runtime without a separate + # vendor torch wheel (see scripts/bundle_*_libtorch.sh, and + # torch_fl/accelerator/_vendor_libtorch.py for the relink at import). + # The trailing * matters for lib_dcu: DTK's auditwheel-mangled + # torch.libs deps end in a version suffix (libglog-6ed04f2c.so.0.0.0). "lib_maca/*.so*", + "lib_dcu/*.so*", + # DTK torch's own version.py, carried so _restore_dcu_hip_version() + # can hand triton's hcu backend the hip/rocm strings the stock +cpu + # torch in front does not have. Needed explicitly: the globs above + # only match *.so*. + "lib_dcu/vendor_version.py", + "lib_ppu/*.so*", # All backend configs, not just the default: runtime op-routing # configs selected via FLAGOS_USE_FLAGGEMS (backends_flaggems.conf) # and boxing modes via FLAGOS_BACKEND_CONFIG (backends_cuda.conf / @@ -663,11 +704,20 @@ def _get_setup_kwargs(): } version = "0.1.0" - if ACCELERATOR == "metax": - # Local version segment tags the wheel as a MetaX build (self-contained - # forked libtorch). Overridable via FLAGOS_WHEEL_LOCAL for a concrete - # MACA/driver version, e.g. FLAGOS_WHEEL_LOCAL=metax3.8.1. - local = os.environ.get("FLAGOS_WHEEL_LOCAL", "metax") + # A local version segment tags which vendor a self-contained wheel bundles a + # forked libtorch for. That bundle is SDK-version-bound whether we say so or + # not -- DTK's libtorch_hip.so has librocblas.so.4 written into its + # DT_NEEDED -- so making the binding visible in the filename is strictly + # better than leaving two incompatible wheels both called 0.1.0. Override + # with FLAGOS_WHEEL_LOCAL to pin the exact SDK, e.g. + # FLAGOS_WHEEL_LOCAL=metax3.8.1 / FLAGOS_WHEEL_LOCAL=dtk2604. + _default_local = {"metax": "metax", "dcu": "dtk"}.get(ACCELERATOR) + if _default_local is None and ( + os.environ.get("PPU_SDK") or os.environ.get("PPU_HOME") + ): + _default_local = "ppu" + local = os.environ.get("FLAGOS_WHEEL_LOCAL", _default_local) + if local: version = f"{version}+{local}" return dict( diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 169f3ba3..64ff660e 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. import os +import re import sys @@ -165,17 +166,73 @@ def _select_backend_config() -> None: ensure_cudart_shim() -# When reusing PyTorch's CUDA boxing kernels on MetaX with a stock +cpu torch -# wheel, the active wheel's torch/lib must point at the MetaX C++ runtime .so. -# This MUST run before `import torch` (afterwards libc10 is already mapped and -# relinking is too late). Gated on FLAGOS_METAX_BOXING=1; idempotent; no-op when -# torch already IS the MetaX wheel. -if os.environ.get("FLAGOS_METAX_BOXING", "0") == "1": - from torch_fl.accelerator.metax._metax_libtorch_link import ( - ensure_maca_libtorch_links, - ) - ensure_maca_libtorch_links() +def _relink_vendor_libtorch() -> None: + """Point the active torch wheel's torch/lib at this wheel's bundled libtorch. + + MetaX, DCU and PPU all run on a *forked* libtorch whose core .so + (libc10/libtorch_cpu/libtorch_python/...) differ from the upstream ones a + stock ``torch==X.Y.Z+cpu`` wheel ships. A self-contained wheel bundles them + under torch_fl/lib_{maca,dcu,ppu}/ and symlinks them over the stock files; + see torch_fl.accelerator._vendor_libtorch for why a ctypes preload alone is + not enough there. + + This MUST run before `import torch` -- afterwards libc10 is already mapped + and relinking is too late. Every backend's entry point is idempotent and a + no-op when its bundle dir is absent (a plain in-place build, where torch + already IS the vendor wheel), so this is safe to call unconditionally. + + MetaX: FLAGOS_METAX_BOXING=1 triggers relink unconditionally (for in-place + MetaX builds that want to test boxing). When accel=="metax" and the bundle + dir exists, relink regardless of the env var (self-contained wheel). DCU + and PPU have no native-kernel mode, so bundle-dir presence alone decides. + The CUDA backend is not here: the official +cpu wheel's core .so ARE the + upstream ones, so only the extra CUDA libs are missing and + _preload_cuda_assets() below handles those with ctypes. + """ + accel = _build_accelerator() + + if os.environ.get("FLAGOS_METAX_BOXING", "0") == "1": + from torch_fl.accelerator.metax._metax_libtorch_link import ( + ensure_maca_libtorch_links, + ) + + ensure_maca_libtorch_links() + return + + if accel == "metax": + from torch_fl.accelerator._vendor_libtorch import bundled_lib_dir + + if bundled_lib_dir("lib_maca", "libtorch_cuda.so"): + from torch_fl.accelerator.metax._metax_libtorch_link import ( + ensure_maca_libtorch_links, + ) + + ensure_maca_libtorch_links() + return + + if accel == "dcu": + from torch_fl.accelerator.dcu._dcu_libtorch_link import ( + ensure_dcu_libtorch_links, + ) + + ensure_dcu_libtorch_links() + return + + # PPU builds as ACCELERATOR=cuda (it targets PPU_SDK/CUDA_SDK), so the only + # distinguishing signal at import time is its own bundle dir. + if accel in ("cuda", ""): + from torch_fl.accelerator._vendor_libtorch import bundled_lib_dir + + if bundled_lib_dir("lib_ppu", "libtorch_cuda.so"): + from torch_fl.accelerator.ppu._ppu_libtorch_link import ( + ensure_ppu_libtorch_links, + ) + + ensure_ppu_libtorch_links() + + +_relink_vendor_libtorch() def _preload_cuda_assets() -> None: @@ -322,7 +379,6 @@ def _check_privateuse1_unclaimed() -> None: import torch # noqa: E402 - if sys.platform == "win32": from ._utils import _load_dll_libraries @@ -371,7 +427,6 @@ def _check_privateuse1_unclaimed() -> None: from . import flagos # noqa: E402 - torch.utils.rename_privateuse1_backend("flagos") torch._register_device_module("flagos", flagos) torch.utils.generate_methods_for_privateuse1_backend(for_storage=True) @@ -458,6 +513,43 @@ def _patched(increment, generator=None): pass +def _restore_dcu_hip_version() -> None: + """Set torch.version.hip/rocm for a self-contained DCU wheel. + + See the DCU branch of _patch_flaggems_codegen_config() for why this matters: + the bundled libtorch is DTK's HIP build, but torch/version.py comes from the + stock torch+cpu wheel in front and reports hip=None, which switches triton's + hcu backend off. scripts/bundle_dcu_libtorch.sh copies the vendor torch's own + version.py next to the bundled .so as vendor_version.py; read the strings + back from there. No-op when a real vendor torch is in front. + """ + import torch + + if getattr(torch.version, "hip", None): + return # a real DTK torch is in front; leave its values alone. + + hip_ver = os.environ.get("FLAGOS_DCU_HIP_VERSION", "").strip() + rocm_ver = "" + if not hip_ver: + ver_py = os.path.join(os.path.dirname(__file__), "lib_dcu", "vendor_version.py") + try: + with open(ver_py, encoding="utf-8") as f: + for line in f: + m = re.match(r"\s*hip\s*(?::[^=]*)?=\s*'([^']+)'", line) + if m: + hip_ver = m.group(1) + continue + m = re.match(r"\s*rocm\s*(?::[^=]*)?=\s*'([^']+)'", line) + if m: + rocm_ver = m.group(1) + except OSError: + return # not a bundled build (source checkout); nothing to restore. + if hip_ver: + torch.version.hip = hip_ver + if rocm_ver: + torch.version.rocm = rocm_ver + + def _patch_flaggems_codegen_config(): """ Configure FlagGems' vendor + torch.cuda shim for the flagos device. @@ -522,6 +614,18 @@ def _patch_flaggems_codegen_config(): # fallback. setdefault so an explicit GEMS_VENDOR still wins. if _build_accelerator() == "dcu" and os.environ.get("GEMS_VENDOR") != "ascend": os.environ.setdefault("GEMS_VENDOR", "hygon") + # torch.version is pure Python (torch/version.py), generated when the + # wheel is built -- swapping the bundled DTK .so files cannot change it. + # A self-contained DCU wheel therefore front-ends a stock torch+cpu whose + # torch.version.hip is None, while the DTK torch it replaces reports + # e.g. "6.3.26113". triton's hcu backend gates on exactly that value + # (backends/hcu/driver.py is_active(): torch.cuda.is_available() and + # torch.version.hip is not None), so with None the driver never activates + # and any flag_gems op dies in triton's driver factory with + # "0 active drivers ([]). There should only be one." Restore the attribute + # from the bundled libtorch's own version so triton sees a HIP torch, + # matching what the vendor wheel reported. + _restore_dcu_hip_version() return # --- Enflame GCU branch --- @@ -640,8 +744,14 @@ def _patched_cuda_device_init(self, device): _patch_cuda_device_context() # Initialize CUDA runtime only when FlagGems Python path needs it (CUDA backend ops). +# The check must be against the *build* backend, not torch.cuda.is_available(): +# a DCU/PPU self-contained wheel relinks a hipified/cuda libtorch into a stock +# +cpu torch, which makes is_available() return True even though the CUDA runtime +# libs are absent, and torch.cuda.init() would fail with "libcaffe2_nvrtc.so: not +# found". Only actual CUDA-backend builds need this init. if ( os.environ.get("FLAGOS_DISABLE_FLAGGEMS_PY", "0") != "1" + and _build_accelerator() in ("cuda", "") and torch.cuda.is_available() ): torch.cuda.init() diff --git a/torch_fl/accelerator/_vendor_libtorch.py b/torch_fl/accelerator/_vendor_libtorch.py new file mode 100644 index 00000000..3a3d471d --- /dev/null +++ b/torch_fl/accelerator/_vendor_libtorch.py @@ -0,0 +1,285 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Point the active (stock) torch wheel's ``torch/lib`` at a bundled vendor libtorch. + +Why this exists +--------------- +Several backends run on a *forked* libtorch: MetaX (``at::maca::*``), Hygon DCU +(DTK's hipified build -- its ``libtorch_cpu.so`` carries hip symbols and needs +``libgalaxyhip.so.5``), and PPU (a local ``USE_CUDA=1`` build). Measured symbol +attribution says the fork lives in the *core* libs: for both DCU and PPU, +``libtorch_cpu.so`` resolves ~2100 of ``libtorch_fl.so``'s undefined symbols while +the vendor lib (``libtorch_hip.so`` / ``libtorch_cuda.so``) resolves 0. So a +self-contained wheel must ship the core libs, not just the vendor one. + +Pure ``ctypes`` preloading does NOT work for core libs. The stock wheel's +``_C.so`` / ``libtorch_python.so`` carry an ``$ORIGIN`` RUNPATH that pulls the +upstream ``libc10.so`` back in *by full path*, so the process ends up with two +libc10 and dies in duplicate static init (``Key already registered ... +caffe2_report_cpu_memory_usage``). The robust fix is to make the physical files +that RUNPATH resolves to *be* the vendor ones: replace the stock wheel's +``torch/lib/`` with symlinks into the bundle dir. Originals move to +``torch/lib/_orig_backup/``, so the operation is fully reversible. + +The CUDA backend is the one exception and does not use this module: the official +``+cpu`` wheel's core libs *are* the upstream ones, so only the extra CUDA libs +are missing and a ctypes preload (``torch_fl.__init__._preload_cuda_assets``) +suffices. + +Callers must invoke this from ``torch_fl/__init__.py`` BEFORE ``import torch`` +(afterwards libc10 is already mapped and relinking is too late). Every entry +point here is idempotent. + +Note on ``$ORIGIN`` and symlinks: glibc expands ``$ORIGIN`` from the path the +object was *loaded by*, NOT from its resolved target. Opening +``torch/lib/libtorch_cpu.so`` (a symlink) therefore gives ``$ORIGIN`` = +``torch/lib``, where a bundle-internal dependency does not exist -- measured on +DCU, whose ``libc10.so`` needs the auditwheel-mangled ``libgflags-8aee0f6c.so`` +that ships in the bundle dir: + + ctypes.CDLL(".../torch/lib/libc10.so") -> libgflags-...so: not found + ctypes.CDLL(".../torch_fl/lib_dcu/libc10.so") -> OK + +So ``_preload_global`` dlopens the *bundle* paths, not the symlinks. That also +covers the symlinks: glibc keys loaded objects by (device, inode), and a symlink +shares both with its target, so a later lookup that resolves through +``torch/lib`` finds the object already mapped instead of re-opening it. +""" + +import ctypes +import importlib.util +import os +import sys + +# One flag per bundle dir: a process only ever relinks for its own backend, but +# keying by name keeps the module reentrant and makes the no-op cheap. +_done = set() +# dlopen handles kept alive for the process lifetime (see _preload_global). +_runtime_handles = [] + + +def active_torch_lib(): + """``torch/lib`` of the importable torch, WITHOUT importing torch.""" + spec = importlib.util.find_spec("torch") + if spec is None or not spec.submodule_search_locations: + return None + lib = os.path.join(spec.submodule_search_locations[0], "lib") + return lib if os.path.isdir(lib) else None + + +def bundled_lib_dir(bundle_dirname, probe_so): + """The bundle dir inside this wheel, if the bundling step actually ran. + + ``scripts/bundle__libtorch.sh`` copies the vendor libtorch .so into + ``torch_fl//``. When present this is the preferred source: + the target machine then needs only the official ``torch+cpu`` wheel plus the + vendor driver runtime, no vendor torch wheel at all. Absent (a plain + in-place/dev build) every entry point here becomes a no-op. + """ + # this file: torch_fl/accelerator/_vendor_libtorch.py -> torch_fl/ + pkg_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + libdir = os.path.join(pkg_root, bundle_dirname) + if os.path.isdir(libdir) and os.path.exists(os.path.join(libdir, probe_so)): + return libdir + return None + + +def _scan_sibling_envs(probe_so, vendor_markers): + """Fallback for multi-env dev setups: a sibling conda env's vendor torch. + + Matches on ``torch/version.py`` containing one of ``vendor_markers`` (e.g. + "metax"/"maca", "dtk"/"hip", "ppu") so we never pick up a stock wheel. + """ + if not vendor_markers: + return None + prefix = os.environ.get("CONDA_PREFIX") or os.path.dirname( + os.path.dirname(os.__file__) + ) + envs_root = os.path.dirname(prefix) # .../envs + if not os.path.isdir(envs_root): + return None + py = "python{}.{}".format(*sys.version_info[:2]) + try: + names = sorted(os.listdir(envs_root)) + except OSError: + return None + for name in names: + cand = os.path.join(envs_root, name, "lib", py, "site-packages", "torch") + libdir = os.path.join(cand, "lib") + ver_file = os.path.join(cand, "version.py") + if not os.path.isfile(ver_file) or not os.path.isdir(libdir): + continue + try: + with open(ver_file) as f: + txt = f.read() + except OSError: + continue + if any(m in txt for m in vendor_markers) and os.path.exists( + os.path.join(libdir, probe_so) + ): + return libdir + return None + + +def discover_vendor_torch_lib( + bundle_dirname, probe_so, env_override=None, vendor_markers=() +): + """Locate the vendor libtorch .so dir. + + Priority: bundled in this wheel, then ``env_override``, then sibling conda + envs whose torch is a vendor build. + """ + bundled = bundled_lib_dir(bundle_dirname, probe_so) + if bundled: + return bundled + if env_override: + env = os.environ.get(env_override) + if env and os.path.isdir(env): + return env + return _scan_sibling_envs(probe_so, vendor_markers) + + +def _link_one(dst_dir, backup_dir, name, target, required, vendor): + """Idempotently point ``dst_dir/name`` at ``target`` (a vendor .so).""" + dst = os.path.join(dst_dir, name) + if not os.path.exists(target): + if required: + raise FileNotFoundError(f"{vendor} so missing: {target}") + return + # Already correctly linked? + if os.path.islink(dst) and os.path.realpath(dst) == os.path.realpath(target): + return + # Back up a real (non-symlink) original once. + if os.path.exists(dst) and not os.path.islink(dst): + os.makedirs(backup_dir, exist_ok=True) + bak = os.path.join(backup_dir, name) + if not os.path.exists(bak): + os.replace(dst, bak) + else: + os.remove(dst) + elif os.path.islink(dst): + os.remove(dst) # stale/incorrect link + os.symlink(target, dst) + + +def _preload_global(lib_dir, load_order, core_so, vendor): + """dlopen the vendor set RTLD_GLOBAL, in dependency order. + + A CPU-only torch wheel never loads the forked runtime itself. Symlinking the + files is not sufficient on its own: symbols the plugin needs may live in the + forked *CPU* runtime (``GetFlagosDefaultCudaGenerator`` is the measured case + on MetaX) and loading only the vendor library can leave its CPU dependency + RTLD_LOCAL, after which ``libtorch_fl.so`` cannot resolve that symbol. + Loading the whole set globally, core first, avoids that. + + ``lib_dir`` must be the *source* dir (the bundle), never the ``torch/lib`` + symlink dir -- see the ``$ORIGIN`` note in the module docstring. + """ + handles = [] + for name in load_order: + path = os.path.join(lib_dir, name) + if not os.path.exists(path): + if name in core_so: + raise FileNotFoundError(f"{vendor} libtorch runtime missing: {path}") + continue + try: + handles.append(ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)) + except OSError as exc: + raise RuntimeError( + f"Failed to load {vendor} libtorch runtime: {path}" + ) from exc + return handles + + +def ensure_vendor_libtorch_links( + bundle_dirname, + core_so, + extra_so=(), + env_override=None, + vendor_markers=(), + probe_so=None, + vendor=None, + load_order=None, +): + """Symlink the active torch wheel's core .so to the vendor wheel's copies. + + Args: + bundle_dirname: dir inside the wheel holding the bundle ("lib_maca", ...). + core_so: .so that MUST be present; a missing one raises. + extra_so: .so the stock ``+cpu`` wheel may not ship at all (the vendor + libs); linked in fresh when available, skipped silently otherwise. + env_override: env var naming an explicit source dir. + vendor_markers: substrings identifying a vendor torch in ``version.py``. + probe_so: file whose presence proves a dir is a real vendor libtorch dir. + Defaults to the first entry of ``extra_so``, else of ``core_so``. + vendor: label used in error messages. + load_order: when given, dlopen these RTLD_GLOBAL after linking, in this + order (see ``_preload_global``). Names absent from ``core_so`` may be + missing; a missing core .so raises. + + Returns True if links are in place (or already were), False if there was + nothing to do (no bundle, no vendor torch found, or already running on it). + """ + if bundle_dirname in _done: + return True + probe = probe_so or (extra_so[0] if extra_so else core_so[0]) + label = vendor or bundle_dirname + + active = active_torch_lib() + src = discover_vendor_torch_lib( + bundle_dirname, probe, env_override=env_override, vendor_markers=vendor_markers + ) + if active is None or src is None: + return False + # Already running on the vendor wheel itself -> nothing to do. + if os.path.realpath(active) == os.path.realpath(src): + _done.add(bundle_dirname) + return True + + backup = os.path.join(active, "_orig_backup") + for name in core_so: + _link_one(active, backup, name, os.path.join(src, name), True, label) + for name in extra_so: + _link_one(active, backup, name, os.path.join(src, name), False, label) + + if load_order: + # dlopen from `src` (the bundle), not from `active`: loading through the + # torch/lib symlinks would expand $ORIGIN to torch/lib and lose the + # bundle-internal deps. Keep the handles alive for the process lifetime. + _runtime_handles.extend(_preload_global(src, load_order, core_so, label)) + + _done.add(bundle_dirname) + return True + + +def restore_original_libtorch(core_so, extra_so=(), bundle_dirname=None): + """Undo ensure_vendor_libtorch_links(): drop links, restore the backups.""" + active = active_torch_lib() + if active is None: + return + backup = os.path.join(active, "_orig_backup") + for name in tuple(core_so) + tuple(extra_so): + dst = os.path.join(active, name) + if os.path.islink(dst): + os.remove(dst) + if os.path.isdir(backup): + for name in os.listdir(backup): + os.replace(os.path.join(backup, name), os.path.join(active, name)) + try: + os.rmdir(backup) + except OSError: + pass + if bundle_dirname: + _done.discard(bundle_dirname) diff --git a/torch_fl/accelerator/dcu/__init__.py b/torch_fl/accelerator/dcu/__init__.py new file mode 100644 index 00000000..cc840712 --- /dev/null +++ b/torch_fl/accelerator/dcu/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hygon DCU (DTK) support for the flagos backend.""" diff --git a/torch_fl/accelerator/dcu/_dcu_libtorch_link.py b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py new file mode 100644 index 00000000..7a98ade9 --- /dev/null +++ b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py @@ -0,0 +1,124 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Symlink DTK's libtorch .so into the active (official) torch wheel's lib dir. + +DCU runs the CUDA boxing kernels on top of DTK's hipified libtorch: HIP kernels +are registered under the CUDA dispatch key, so PrivateUse1 -> CUDA re-dispatch +works unchanged. That libtorch is a fork -- measured, DTK's ``libtorch_cpu.so`` +exports 128 hip symbols and carries ``DT_NEEDED: libgalaxyhip.so.5`` -- and it +resolves 2101 of ``libtorch_fl.so``'s undefined symbols, while ``libtorch_hip.so`` +resolves 0. So the core libs are what has to be swapped in; ``libtorch_hip.so`` +still has to be *present* for the dispatch-key registration to exist. + +See ``torch_fl.accelerator._vendor_libtorch`` for the symlink mechanism and why a +ctypes preload cannot replace core libs. + +Unlike MetaX there is no env gate: a DCU wheel is only ever installed on a DCU +box, and the relink is skipped automatically when ``torch_fl/lib_dcu/`` was not +bundled (a plain in-place build) or when torch already IS the DTK wheel. + +The DTK driver stack itself (``libgalaxyhip.so.5``, ``libMIOpen.so.1``, +``librocblas.so.4``, ``librccl.so.1``, ...) stays on the target under +``/opt/dtk`` and is reached via the RUNPATH baked in by +``cmake/FlagosRpath.cmake`` / ``scripts/bundle_dcu_libtorch.sh``. +""" + +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_dcu" + +# Core C++ .so that must come from the DTK wheel as a self-consistent set. +_CORE_SO = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libtorch_python.so", +) +# HIP-side .so the stock +cpu wheel does not ship at all. libmagma.so is a +# DT_NEEDED of DTK's libtorch_hip.so; libshm.so is one of libtorch_python.so's +# (torch.multiprocessing's shared-memory manager) and is a *different* build in +# the DTK wheel, so it has to come from the same set. +_HIP_SO = ( + "libc10_hip.so", + "libtorch_hip.so", + "libmagma.so", + "libshm.so", +) + +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the HIP +# side, then libtorch_python. Same reasoning as MetaX -- symbols the plugin needs +# live in the forked CPU runtime, and loading only the HIP lib can leave its CPU +# dependency RTLD_LOCAL. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_hip.so", + "libtorch_hip.so", + "libtorch_python.so", +) + +_MARKERS = ("dtk", "hip", "das") + + +def _bundled_dcu_lib(): + """Forked libtorch bundled inside this wheel (self-contained DCU build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_hip.so") + + +def _discover_dcu_torch_lib(): + """Locate the DTK libtorch .so dir. + + Priority: bundled lib_dcu/, then FLAGOS_DCU_TORCH_LIB, then sibling conda + envs whose torch is a DTK build. + """ + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_hip.so", + env_override="FLAGOS_DCU_TORCH_LIB", + vendor_markers=_MARKERS, + ) + + +def ensure_dcu_libtorch_links(): + """Symlink the active torch wheel's core .so to DTK's copies. + + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do. + """ + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_HIP_SO, + env_override="FLAGOS_DCU_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_hip.so", + vendor="DCU/DTK", + load_order=_LOAD_ORDER, + ) + + +def restore_original_libtorch(): + """Undo ensure_dcu_libtorch_links(): remove links, restore backups.""" + _restore(_CORE_SO, _HIP_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/accelerator/metax/_metax_libtorch_link.py b/torch_fl/accelerator/metax/_metax_libtorch_link.py index d872e27d..9705d5ba 100644 --- a/torch_fl/accelerator/metax/_metax_libtorch_link.py +++ b/torch_fl/accelerator/metax/_metax_libtorch_link.py @@ -14,34 +14,30 @@ """Symlink MetaX libtorch .so into the active (official) torch wheel's lib dir. -Rationale ---------- On MetaX we reuse PyTorch's CUDA boxing kernels (FLAGOS_METAX_BOXING) by running the *MetaX* C++ runtime (libtorch_cpu.so / libtorch_cuda.so / libc10.so ...), -which is a hard fork exporting ``at::maca::*`` symbols. When the Python front-end -is a *stock* ``torch==X.Y.Z+cpu`` wheel (no CUDA, clean pip env), its own -``torch/lib`` ships the upstream C++ .so. We must make the process load the -MetaX C++ runtime instead. - -Pure ``ctypes`` preloading does NOT reliably work: the official ``_C.so`` / -``libtorch_python.so`` carry an ``$ORIGIN`` RUNPATH that pulls the upstream -``libc10.so`` back in by full path, giving a *second* libc10 in the process and a -duplicate static-init crash (``Key already registered ... caffe2_report_cpu_memory_usage``). - -The robust fix is to make the physical files the RUNPATH resolves to *be* the -MetaX ones -- i.e. replace the stock wheel's ``torch/lib/`` with symlinks to -the MetaX wheel's copies. Originals are backed up to ``torch/lib/_orig_backup/`` -so the operation is fully reversible. - -This runs from ``torch_fl/__init__.py`` BEFORE ``import torch`` (once torch is -imported its libc10 is already mapped and relinking is too late). It is -idempotent, gated on ``FLAGOS_METAX_BOXING=1``, and a no-op when the active torch -already IS the MetaX wheel. +which is a hard fork exporting ``at::maca::*`` symbols. With a stock +``torch==X.Y.Z+cpu`` front-end, the process must load that fork instead of the +upstream .so shipped in ``torch/lib``. + +The mechanism -- symlink replacement, ``_orig_backup/``, the RTLD_GLOBAL preload, +and why a pure ctypes preload cannot do this on its own -- lives in +``torch_fl.accelerator._vendor_libtorch``. This module is just the MetaX .so +lists; whether to relink at all is decided by +``torch_fl.__init__._relink_vendor_libtorch`` (FLAGOS_METAX_BOXING=1 for an +in-place MetaX build, or lib_maca/ being present for a self-contained wheel). """ -import ctypes -import importlib.util -import os +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_maca" # Core C++ .so that must come from the MetaX wheel as a self-consistent set. # libtorch_python.so is included because the stock one references symbols @@ -54,182 +50,80 @@ "libtorch_python.so", ) # CUDA .so the stock +cpu wheel does not ship at all; symlinked in fresh. +# libshm.so is a DT_NEEDED of MetaX's libtorch_python.so (torch.multiprocessing's +# shared-memory manager) and is a *different* build in the MetaX wheel, so it has +# to come from the same set. _CUDA_SO = ( "libc10_cuda.so", "libtorch_cuda.so", "libtorch_cuda_linalg.so", + "libshm.so", ) -_done = False -_runtime_handles = [] - +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the CUDA +# side, then libtorch_python. GetFlagosDefaultCudaGenerator lives in the forked +# ATen CPU runtime, not in libtorch_cuda.so, so loading only the CUDA lib would +# leave that symbol unresolvable from libtorch_fl.so. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_cuda.so", + "libtorch_cuda_linalg.so", + "libtorch_cuda.so", + "libtorch_python.so", +) -def _active_torch_lib(): - """torch/lib of the importable torch, WITHOUT importing torch.""" - spec = importlib.util.find_spec("torch") - if spec is None or not spec.submodule_search_locations: - return None - lib = os.path.join(spec.submodule_search_locations[0], "lib") - return lib if os.path.isdir(lib) else None +_MARKERS = ("metax", "maca") def _bundled_maca_lib(): - """Forked libtorch bundled inside this wheel (self-contained MetaX build). - - ``scripts/bundle_maca_libtorch.sh`` copies the MetaX libtorch .so into - ``torch_fl/lib_maca/``. When present this is the preferred source: the - target machine then needs only the official ``torch+cpu`` wheel plus the - ``/opt/maca`` driver runtime, no separate MetaX torch wheel. - """ - here = os.path.dirname(os.path.abspath(__file__)) - # this file: torch_fl/accelerator/metax/_metax_libtorch_link.py - pkg_root = os.path.dirname(os.path.dirname(here)) # -> torch_fl/ - libdir = os.path.join(pkg_root, "lib_maca") - if os.path.isdir(libdir) and os.path.exists( - os.path.join(libdir, "libtorch_cuda.so") - ): - return libdir - return None + """Forked libtorch bundled inside this wheel (self-contained MetaX build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_cuda.so") def _discover_maca_torch_lib(): """Locate the MetaX libtorch .so dir. - Priority: forked libtorch bundled in this wheel (lib_maca/), then an - explicit env var, then sibling conda envs whose torch is a - ``+metax``/``+maca`` build (fallback for multi-env dev setups). + Priority: bundled lib_maca/, then FLAGOS_MACA_TORCH_LIB, then sibling conda + envs whose torch is a ``+metax``/``+maca`` build. """ - bundled = _bundled_maca_lib() - if bundled: - return bundled - - env = os.environ.get("FLAGOS_MACA_TORCH_LIB") - if env and os.path.isdir(env): - return env - - # Scan conda envs next to the current prefix for a MetaX torch build. - prefix = os.environ.get("CONDA_PREFIX") or os.path.dirname( - os.path.dirname(os.__file__) + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_cuda.so", + env_override="FLAGOS_MACA_TORCH_LIB", + vendor_markers=_MARKERS, ) - envs_root = os.path.dirname(prefix) # .../envs - if not os.path.isdir(envs_root): - return None - py = "python{}.{}".format(*__import__("sys").version_info[:2]) - for name in sorted(os.listdir(envs_root)): - cand = os.path.join(envs_root, name, "lib", py, "site-packages", "torch") - libdir = os.path.join(cand, "lib") - ver_file = os.path.join(cand, "version.py") - if not os.path.isfile(ver_file) or not os.path.isdir(libdir): - continue - try: - with open(ver_file) as f: - txt = f.read() - except OSError: - continue - if ("metax" in txt or "maca" in txt) and os.path.exists( - os.path.join(libdir, "libtorch_cuda.so") - ): - return libdir - return None - - -def _link_one(dst_dir, backup_dir, name, target, required): - """Idempotently point dst_dir/name at target (a MetaX .so).""" - dst = os.path.join(dst_dir, name) - if not os.path.exists(target): - if required: - raise FileNotFoundError(f"MetaX so missing: {target}") - return - # Already correctly linked? - if os.path.islink(dst) and os.path.realpath(dst) == os.path.realpath(target): - return - # Back up a real (non-symlink) original once. - if os.path.exists(dst) and not os.path.islink(dst): - os.makedirs(backup_dir, exist_ok=True) - bak = os.path.join(backup_dir, name) - if not os.path.exists(bak): - os.replace(dst, bak) - else: - os.remove(dst) - elif os.path.islink(dst): - os.remove(dst) # stale/incorrect link - os.symlink(target, dst) def ensure_maca_libtorch_links(): """Symlink the active torch wheel's core .so to the MetaX wheel's copies. - No-op unless FLAGOS_METAX_BOXING=1. Idempotent; reversible via _orig_backup. - Returns True if links are in place (or already were), False if skipped. + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do (no + bundle, no MetaX torch found, or torch already IS the MetaX wheel). + + Deciding *whether* to relink is the caller's job -- see + ``torch_fl.__init__._relink_vendor_libtorch``, which gates on + FLAGOS_METAX_BOXING=1 for an in-place MetaX build and on lib_maca/ being + present for a self-contained wheel. This used to self-gate on + FLAGOS_METAX_BOXING, which made the self-contained path a silent no-op: + the stock libtorch_cpu.so stayed in place and libtorch_cuda.so then failed + to resolve at::maca symbols that only the forked CPU runtime defines. """ - global _done, _runtime_handles - if _done: - return True - if os.environ.get("FLAGOS_METAX_BOXING", "0") != "1": - return False - - active = _active_torch_lib() - maca = _discover_maca_torch_lib() - if active is None or maca is None: - return False - # Already running on the MetaX wheel itself -> nothing to do. - if os.path.realpath(active) == os.path.realpath(maca): - _done = True - return True - - backup = os.path.join(active, "_orig_backup") - for name in _CORE_SO: - _link_one(active, backup, name, os.path.join(maca, name), required=True) - for name in _CUDA_SO: - _link_one(active, backup, name, os.path.join(maca, name), required=False) - - # A CPU-only torch wheel does not load the forked runtime itself. Load the - # complete bundled dependency set globally before torch_fl._C loads. In - # particular, GetFlagosDefaultCudaGenerator is provided by the forked - # ATen runtime and is not guaranteed to be in libtorch_cuda.so itself. - # Loading only the CUDA library can leave its CPU dependency RTLD_LOCAL, - # so libtorch_fl.so cannot resolve that symbol later. - load_order = ( - "libc10.so", - "libtorch_cpu.so", - "libtorch.so", - "libtorch_global_deps.so", - "libc10_cuda.so", - "libtorch_cuda_linalg.so", - "libtorch_cuda.so", - "libtorch_python.so", + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_CUDA_SO, + env_override="FLAGOS_MACA_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_cuda.so", + vendor="MetaX", + load_order=_LOAD_ORDER, ) - for name in load_order: - path = os.path.join(active, name) - if not os.path.exists(path): - if name in _CORE_SO: - raise FileNotFoundError(f"MetaX libtorch runtime missing: {path}") - continue - try: - _runtime_handles.append(ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)) - except OSError as exc: - raise RuntimeError( - f"Failed to load MetaX libtorch runtime: {path}" - ) from exc - - _done = True - return True def restore_original_libtorch(): """Undo ensure_maca_libtorch_links(): remove links, restore backups.""" - active = _active_torch_lib() - if active is None: - return - backup = os.path.join(active, "_orig_backup") - for name in _CORE_SO + _CUDA_SO: - dst = os.path.join(active, name) - if os.path.islink(dst): - os.remove(dst) - if os.path.isdir(backup): - for name in os.listdir(backup): - os.replace(os.path.join(backup, name), os.path.join(active, name)) - try: - os.rmdir(backup) - except OSError: - pass + _restore(_CORE_SO, _CUDA_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/accelerator/ppu/__init__.py b/torch_fl/accelerator/ppu/__init__.py new file mode 100644 index 00000000..da8ad8f7 --- /dev/null +++ b/torch_fl/accelerator/ppu/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""T-Head PPU support for the flagos backend.""" diff --git a/torch_fl/accelerator/ppu/_ppu_libtorch_link.py b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py new file mode 100644 index 00000000..dc8d6b86 --- /dev/null +++ b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py @@ -0,0 +1,119 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Symlink the PPU libtorch .so into the active (official) torch wheel's lib dir. + +PPU builds against ``PPU_SDK/CUDA_SDK``, so ``ACCELERATOR=cuda`` and the CUDA +boxing kernels apply unchanged. What differs from a real NVIDIA box is the +libtorch: it is a local ``USE_CUDA=1`` source build, not the upstream wheel, and +it resolves 2092 of ``libtorch_fl.so``'s undefined symbols (``libtorch_cuda.so`` +resolves 0, ``libc10_cuda.so`` 10) -- so the core libs must be swapped in, and +``libtorch_cuda.so`` must merely be present for the CUDA dispatch keys. + +That local build also links the system MKL from ``/usr/local/lib`` +(``libmkl_core``/``libmkl_gnu_thread``/``libmkl_intel_lp64``), which the bundling +script copies into ``lib_ppu/`` alongside the core libs. + +See ``torch_fl.accelerator._vendor_libtorch`` for the symlink mechanism and why a +ctypes preload cannot replace core libs. No env gate: the relink is skipped when +``torch_fl/lib_ppu/`` was not bundled or torch already IS the PPU build. The PPU +SDK runtime stays on the target under ``/usr/local/PPU_SDK``. +""" + +from torch_fl.accelerator._vendor_libtorch import ( + bundled_lib_dir, + discover_vendor_torch_lib, + ensure_vendor_libtorch_links, +) +from torch_fl.accelerator._vendor_libtorch import ( + restore_original_libtorch as _restore, +) + +_BUNDLE_DIR = "lib_ppu" + +# Core C++ .so that must come from the PPU build as a self-consistent set. +_CORE_SO = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libtorch_python.so", +) +# CUDA-side .so the stock +cpu wheel does not ship at all. +# libshm.so is a DT_NEEDED of PPU's libtorch_python.so (torch.multiprocessing's +# shared-memory manager) and is a *different* build in the PPU wheel, so it has +# to come from the same set. +_CUDA_SO = ( + "libc10_cuda.so", + "libtorch_cuda.so", + "libtorch_cuda_linalg.so", + "libshm.so", +) + +# Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the CUDA +# side, then libtorch_python. Same reasoning as MetaX. +_LOAD_ORDER = ( + "libc10.so", + "libtorch_cpu.so", + "libtorch.so", + "libtorch_global_deps.so", + "libc10_cuda.so", + "libtorch_cuda_linalg.so", + "libtorch_cuda.so", + "libtorch_python.so", +) + +_MARKERS = ("ppu",) + + +def _bundled_ppu_lib(): + """PPU libtorch bundled inside this wheel (self-contained PPU build).""" + return bundled_lib_dir(_BUNDLE_DIR, "libtorch_cuda.so") + + +def _discover_ppu_torch_lib(): + """Locate the PPU libtorch .so dir. + + Priority: bundled lib_ppu/, then FLAGOS_PPU_TORCH_LIB, then sibling conda + envs whose torch is a PPU build. + """ + return discover_vendor_torch_lib( + _BUNDLE_DIR, + "libtorch_cuda.so", + env_override="FLAGOS_PPU_TORCH_LIB", + vendor_markers=_MARKERS, + ) + + +def ensure_ppu_libtorch_links(): + """Symlink the active torch wheel's core .so to the PPU build's copies. + + Idempotent; reversible via ``torch/lib/_orig_backup/``. Returns True if + links are in place (or already were), False if there was nothing to do. + """ + return ensure_vendor_libtorch_links( + _BUNDLE_DIR, + _CORE_SO, + extra_so=_CUDA_SO, + env_override="FLAGOS_PPU_TORCH_LIB", + vendor_markers=_MARKERS, + probe_so="libtorch_cuda.so", + vendor="PPU", + load_order=_LOAD_ORDER, + ) + + +def restore_original_libtorch(): + """Undo ensure_ppu_libtorch_links(): remove links, restore backups.""" + _restore(_CORE_SO, _CUDA_SO, bundle_dirname=_BUNDLE_DIR) diff --git a/torch_fl/csrc/CMakeLists.txt b/torch_fl/csrc/CMakeLists.txt index 6e7fa3e5..46993560 100644 --- a/torch_fl/csrc/CMakeLists.txt +++ b/torch_fl/csrc/CMakeLists.txt @@ -36,6 +36,12 @@ endif() target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib) +# Same reasoning as csrc/CMakeLists.txt: keep the build machine's torch/lib out +# of the installed RUNPATH and resolve a bundled forked libtorch via $ORIGIN. +# libtorch_bindings.so is loaded as torch_fl._C, i.e. always after +# torch_fl/__init__.py has imported torch. +flagos_set_portable_rpath(${LIBRARY_NAME}) + install(TARGETS ${LIBRARY_NAME} ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} From dd38cd34435d096903fbf1250ac308050f64ffa2 Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Sun, 9 Aug 2026 02:05:11 +0800 Subject: [PATCH 2/2] fix: fall back to the stock libshm.so when the vendor image lacks one The MetaX CI integration test died at import with OSError: libshm.so: cannot open shared object file RuntimeError: Failed to load MetaX libtorch runtime: .../lib_maca/libtorch_python.so libtorch_python.so carries a hard DT_NEEDED on libshm.so, and _preload_global dlopens it from the bundle dir, where the loader searches only the bundle's RUNPATH. On a dev box the vendor torch ships libshm.so so the bundle has one and nothing breaks; the CI image's /opt/vendor-libtorch/lib holds exactly the eight libs set_env_metax.sh asserts on, libshm.so not among them, so bundle_copy_so skips it (required=0) and the dlopen fails. Fix: for a non-core .so absent from the bundle, fall back to the stock wheel's torch/lib. dlopening it RTLD_GLOBAL before libtorch_python.so satisfies the DT_NEEDED by soname against the already-loaded object. Measured on mc550: of the 79 symbols a stock libshm.so leaves undefined, every torch/c10 one is provided by the vendor core libs (the other 67 are libstdc++/libc), so the stock build is ABI compatible with the fork. Its own deps resolve through torch/lib, where they are symlinks into the bundle, so they share an inode with what is already mapped and no second copy appears. Verified on mc550 by reshaping flagos-wheel-py312 into the CI layout (bundle without libshm.so, a real stock libshm.so in torch/lib): before the change import fails exactly as CI does; after it, torch 2.10.0+cpu / avail True / 8 devices, add maxdiff 0.0, mm maxdiff 1.0e-02, randn ok, with /proc/self/maps showing libshm from torch/lib and libtorch_cpu from the bundle. The env was restored and re-verified afterwards. Also make bundle_copy_so's skip a stderr warning instead of a silent info line -- that silence is what let a bundle missing a DT_NEEDED lib ship as if it were fine. --- scripts/bundle_dcu_libtorch.sh | 140 ++++++++++-------- scripts/bundle_maca_libtorch.sh | 52 ++++--- scripts/bundle_ppu_libtorch.sh | 56 ++++--- scripts/lib/bundle_common.sh | 78 +++++----- setup.py | 3 +- torch_fl/__init__.py | 10 +- torch_fl/accelerator/_vendor_libtorch.py | 25 +++- .../accelerator/dcu/_dcu_libtorch_link.py | 8 +- .../accelerator/metax/_metax_libtorch_link.py | 7 + .../accelerator/ppu/_ppu_libtorch_link.py | 4 +- 10 files changed, 226 insertions(+), 157 deletions(-) diff --git a/scripts/bundle_dcu_libtorch.sh b/scripts/bundle_dcu_libtorch.sh index 6dce7230..4e52055c 100644 --- a/scripts/bundle_dcu_libtorch.sh +++ b/scripts/bundle_dcu_libtorch.sh @@ -1,33 +1,43 @@ #!/usr/bin/env bash -# 把 DTK fork 的 libtorch C++ .so 打进 torch_fl/lib_dcu/，做成自包含单 wheel。 +# Bundle the DTK-forked libtorch C++ .so into torch_fl/lib_dcu/ for a +# self-contained single wheel. # -# 为什么要打 core 库（不是只打 libtorch_hip.so）： -# 实测 libtorch_fl.so 的未定义符号里，DTK 的 libtorch_cpu.so 提供 2101 个， -# libtorch_hip.so 提供 0 个 —— fork 在 core 库这一侧（DTK 的 libtorch_cpu.so -# 带 128 个 hip 符号、DT_NEEDED 写着 libgalaxyhip.so.5）。libtorch_hip.so 仍然 -# 必须在场：boxing kernel 要求 CUDA dispatch key 上已注册厂商 kernel。 +# Why bundle the core libs (not just libtorch_hip.so): +# Measured symbol attribution: DTK's libtorch_cpu.so resolves 2101 undefined +# symbols in libtorch_fl.so, while libtorch_hip.so resolves 0 -- the fork is on +# the core-lib side (DTK's libtorch_cpu.so exports 128 hip symbols and carries +# DT_NEEDED: libgalaxyhip.so.5). libtorch_hip.so still has to be present: the +# boxing kernels require vendor kernels to be registered under the CUDA dispatch +# key. # -# 为什么要打 torch.libs/：那 17 个文件是 auditwheel 改名的通用库 -# （glog/gflags/MKL/OpenMPI/hwloc/libxml2...），其中 5 个是 DTK libtorch_cpu.so 的 -# 直接 DT_NEEDED。实测它们对 hip/dtk/rocm 的引用数为 0，不引入 SDK 绑定；但官方 -# torch+cpu wheel 的 torch.libs 里一个都没有（名字带 hash，系统库也对不上）， -# 不打就跑不起来。它们和 core 库放同一目录，靠 $ORIGIN 解析。 +# Why bundle torch.libs/: those 17 files are auditwheel-mangled common libs +# (glog/gflags/MKL/OpenMPI/hwloc/libxml2...), 5 of which are direct DT_NEEDED of +# DTK's libtorch_cpu.so. Measured, they carry 0 hip/dtk/rocm references and +# introduce no SDK binding; but the official torch+cpu wheel's torch.libs +# contains none of them (names carry a hash, system libs do not match either), +# so the wheel cannot run without bundling them. They sit in the same dir as the +# core libs and are resolved via $ORIGIN. # -# 注意 $ORIGIN 的实测语义：glibc 按"打开这个对象时用的路径"展开 $ORIGIN，不是按 -# 真实路径。所以从 torch/lib 的软链打开 libc10.so 时 $ORIGIN 是 torch/lib，找不到 -# 同在 lib_dcu 的 libgflags-8aee0f6c.so.2.1.2。因此运行期预加载走 bundle 原路径 -# （_vendor_libtorch._preload_global），不走软链 —— 那里有同样的说明。 +# Note on $ORIGIN semantics: glibc expands $ORIGIN from the path the object was +# *opened by*, not from the resolved real path. So opening libc10.so through the +# torch/lib symlink gives $ORIGIN = torch/lib, and it cannot find +# libgflags-8aee0f6c.so.2.1.2 that sits alongside it in lib_dcu. Runtime +# preloading therefore walks the bundle's original paths +# (_vendor_libtorch._preload_global), not the symlinks -- the same rationale is +# documented there. # -# 不打包：DTK 驱动栈（libgalaxyhip.so.5 libMIOpen.so.1 librocblas.so.4 -# libhipblas.so.2 librccl.so.1 等 12 个 soname）留在目标机 /opt/dtk。 -# SDK 版本绑定本来就存在且更强（libtorch_hip.so 的 DT_NEEDED 写死 -# librocblas.so.4），多打那 129 MB 不会让绑定更紧。 +# Not bundled: DTK driver stack (libgalaxyhip.so.5 libMIOpen.so.1 +# librocblas.so.4 libhipblas.so.2 librccl.so.1, etc., 12 sonames) stays on the +# target machine under /opt/dtk. SDK version binding already exists and is +# stronger (libtorch_hip.so's DT_NEEDED hard-codes librocblas.so.4); bundling +# that extra 129 MB would not tighten the binding. # -# 用法: +# Usage: # FLAGOS_DCU_TORCH_LIB= bash scripts/bundle_dcu_libtorch.sh # DTK_ROOT=/opt/dtk bash scripts/bundle_dcu_libtorch.sh # -# 应在 `python setup.py bdist_wheel`（ACCELERATOR=dcu）之后、打 wheel 之前跑。幂等。 +# Should run after `python setup.py bdist_wheel` (ACCELERATOR=dcu) and before +# packing the wheel. Idempotent. set -euo pipefail @@ -45,59 +55,64 @@ if [ -z "${SRC}" ]; then fi if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then - echo "error: 找不到 DTK torch/lib。设 FLAGOS_DCU_TORCH_LIB=" >&2 + echo "error: DTK torch/lib not found. Set FLAGOS_DCU_TORCH_LIB=" >&2 exit 1 fi if [ ! -f "${SRC}/libtorch_hip.so" ]; then - echo "error: ${SRC} 里没有 libtorch_hip.so，不是 DTK torch/lib" >&2 + echo "error: ${SRC} does not contain libtorch_hip.so, not a DTK torch/lib" >&2 exit 1 fi bundle_require_patchelf -# 与 _dcu_libtorch_link._CORE_SO + _HIP_SO 对齐的自洽集合。 -# libshm.so：libtorch_python.so 的直接 DT_NEEDED（torch.multiprocessing 的共享内存 -# 管理器），DTK 的 torch/lib 里有实体文件，但它不在软链清单里，所以必须打进 bundle -# 让 $ORIGIN 找得到。 -# libcaffe2_nvrtc.so：torch.cuda.init() 会 dlopen 它。DTK torch 带这个文件， -# stock +cpu wheel 没有，所以不打进来的话 GetFlagosDefaultCudaGenerator 里那次 -# 按需 init 就死在 "Error in dlopen: libcaffe2_nvrtc.so"（RNG 修复反而暴露了它）。 +# Self-consistent set aligned with _dcu_libtorch_link._CORE_SO + _HIP_SO. +# libshm.so: direct DT_NEEDED of libtorch_python.so (torch.multiprocessing's +# shared-memory manager); DTK's torch/lib has the real file, but it is not on +# the symlink manifest, so it must be bundled for $ORIGIN resolution. +# libcaffe2_nvrtc.so: torch.cuda.init() dlopens it. DTK torch ships this file, +# stock +cpu wheel does not, so without bundling it the on-demand init inside +# GetFlagosDefaultCudaGenerator dies on "Error in dlopen: libcaffe2_nvrtc.so" +# (the RNG fix exposed it). CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so libcaffe2_nvrtc.so) HIP_SO=(libc10_hip.so libtorch_hip.so libmagma.so libshm.so) -# DTK 驱动栈实测分布（容器 LD_LIBRARY_PATH 与 find 结果一致）： +# DTK driver stack measured layout (container LD_LIBRARY_PATH matches find results): # lib/ libhipnn librocfft.so.0 librocrand.so.1 librocsparse.so.1 # libMIOpen-recommend.so libunwind.so.8 libgalaxyhip.so.5 # hip/lib/ hip runtime # aillvm/lib/ libomp.so # .hyhal/rocm_smi/lib/ librocm_smi64.so.2 -# 全部留在目标机（装了 DCU 卡就有 /opt/dtk），但 RPATH 必须覆盖到，否则 -# LD_LIBRARY_PATH 没设时就 not found。 +# All stay on the target machine (a box with DCU cards has /opt/dtk), but RPATH +# must cover them, else not-found when LD_LIBRARY_PATH is unset. VENDOR_RPATH="${DTK_ROOT}/lib:${DTK_ROOT}/hip/lib:${DTK_ROOT}/lib64" VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/aillvm/lib:${DTK_ROOT}/.hyhal/rocm_smi/lib" VENDOR_RPATH="${VENDOR_RPATH}:${DTK_ROOT}/llvm/lib:/opt/hyhal/lib" -echo "源 DTK torch/lib : ${SRC}" -echo "目标 lib_dcu : ${LIB_DCU}" -echo "DTK 驱动路径 : ${DTK_ROOT}" +echo "Source DTK torch/lib : ${SRC}" +echo "Target lib_dcu : ${LIB_DCU}" +echo "DTK driver path : ${DTK_ROOT}" -# bundle 内的库要能从两条路径被打开： -# 1. 直接从 lib_dcu/（运行期预加载走这条，见 _vendor_libtorch._preload_global） -# 2. 通过 torch/lib/ 的软链（`import torch` 自己加载 libtorch_global_deps.so 时） -# glibc 按"打开这个对象时用的路径"展开 $ORIGIN，所以第 2 条路上 $ORIGIN 是 -# torch/lib，找不到同在 lib_dcu 的 libmpi-3fcb240d.so.40.40.3 等 auditwheel 改名库。 -# 两个目录都在 site-packages 下同级（torch/lib -> ../../torch_fl/lib_dcu）， -# 所以再加一条相对路径就同时覆盖两种情况。实测：不加这条，先 `import torch` -# 再 import torch_fl 会死在 libmpi not found。 +# Libs inside the bundle must be openable from two paths: +# 1. Directly from lib_dcu/ (runtime preload walks this one, see +# _vendor_libtorch._preload_global) +# 2. Through the torch/lib/ symlinks (when `import torch` loads +# libtorch_global_deps.so itself) +# glibc expands $ORIGIN from the path the object was opened by, so path #2 gives +# $ORIGIN = torch/lib and cannot find libmpi-3fcb240d.so.40.40.3 and other +# auditwheel-mangled libs sitting in lib_dcu. Both dirs are siblings under +# site-packages (torch/lib -> ../../torch_fl/lib_dcu), so adding one more +# relative path covers both cases. Measured: without this, `import torch` before +# `import torch_fl` dies on libmpi not found. BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_dcu" bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" bundle_copy_so "${SRC}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${HIP_SO[@]}" -# torch.libs/：与 torch/ 同级的 auditwheel 目录，文件名带 hash 后缀，只能 glob。 +# torch.libs/: auditwheel dir sibling to torch/, filenames carry hash suffixes, +# can only glob. TORCH_LIBS="$(cd "${SRC}/../.." && pwd)/torch.libs" if [ -d "${TORCH_LIBS}" ]; then - echo "源 torch.libs : ${TORCH_LIBS}" + echo "Source torch.libs : ${TORCH_LIBS}" _names=() while IFS= read -r f; do _names+=("$(basename "${f}")") @@ -105,16 +120,17 @@ if [ -d "${TORCH_LIBS}" ]; then if [ ${#_names[@]} -gt 0 ]; then bundle_copy_so "${TORCH_LIBS}" "${LIB_DCU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${_names[@]}" else - echo "warning: ${TORCH_LIBS} 里没有 .so，跳过" >&2 + echo "warning: ${TORCH_LIBS} has no .so, skipping" >&2 fi else - echo "warning: 找不到 ${TORCH_LIBS}；若 libtorch_cpu.so 的 DT_NEEDED 里有带 hash" >&2 - echo " 的通用库（libglog-*.so.0 等），目标机上会缺库。" >&2 + echo "warning: ${TORCH_LIBS} not found; the target will lack hash-suffixed" >&2 + echo " common libs (libglog-*.so.0 etc.) from libtorch_cpu.so's DT_NEEDED." >&2 fi -# libtorch_fl.so / libflagos.so 还需要 DTK 的 CUDA 兼容层 libcudart.so.12 -# （cuda_runtime_compat，见 CMakeLists.txt 的 DCU_CUDA_ROOT）。cmake 已经把它写进 -# RUNPATH 了，这里重写时不能丢 —— 否则干净环境里 libcudart.so.12 就 not found。 +# libtorch_fl.so / libflagos.so also need DTK's CUDA compatibility layer +# libcudart.so.12 (cuda_runtime_compat, see CMakeLists.txt's DCU_CUDA_ROOT). +# cmake already wrote it into RUNPATH; cannot drop it when rewriting here -- +# else libcudart.so.12 becomes not-found in a clean environment. _DCU_CUDA_LIB64="" for _c in "${DTK_ROOT}"/cuda/cuda-*/lib64; do if [ -f "${_c}/libcudart.so.12" ]; then @@ -123,25 +139,27 @@ for _c in "${DTK_ROOT}"/cuda/cuda-*/lib64; do fi done if [ -z "${_DCU_CUDA_LIB64}" ]; then - echo "warning: ${DTK_ROOT}/cuda/cuda-*/lib64 里没有 libcudart.so.12" >&2 + echo "warning: ${DTK_ROOT}/cuda/cuda-*/lib64 has no libcudart.so.12" >&2 fi PLUGIN_RPATH="\$ORIGIN:\$ORIGIN/../lib_dcu" [ -n "${_DCU_CUDA_LIB64}" ] && PLUGIN_RPATH="${PLUGIN_RPATH}:${_DCU_CUDA_LIB64}" bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" "${PLUGIN_RPATH}:${VENDOR_RPATH}" -# torch/version.py 是构建期生成的纯 Python，换掉 .so 改不了它。自包含 DCU wheel -# 前面是 stock torch+cpu，torch.version.hip 报 None，而 triton 的 hcu backend -# 恰好按这个值判活（backends/hcu/driver.py is_active(): torch.cuda.is_available() -# and torch.version.hip is not None）。None 就永不激活，任何 flag_gems 算子都死在 -# triton 的 driver factory："0 active drivers ([])"。把厂商 torch 自己的 version.py -# 带上，import 时由 _restore_dcu_hip_version() 读回 hip/rocm 字符串。 +# torch/version.py is pure Python generated at build time, so swapping .so +# cannot change it. In the self-contained DCU wheel the base is stock torch+cpu, +# whose torch.version.hip reports None, and triton's hcu backend gates +# is_active() on that value (backends/hcu/driver.py is_active(): +# torch.cuda.is_available() and torch.version.hip is not None). None means never +# activate, and any flag_gems op dies in triton's driver factory: "0 active +# drivers ([])". Carry the vendor torch's own version.py; at import time +# _restore_dcu_hip_version() reads back the hip/rocm strings. _VENDOR_VERSION_PY="$(cd "${SRC}/.." && pwd)/version.py" if [ -f "${_VENDOR_VERSION_PY}" ]; then cp -fL "${_VENDOR_VERSION_PY}" "${LIB_DCU}/vendor_version.py" - echo "已复制 vendor version.py -> lib_dcu/vendor_version.py" + echo "Copied vendor version.py -> lib_dcu/vendor_version.py" grep -E "^\s*(hip|rocm)\s*(:|=)" "${LIB_DCU}/vendor_version.py" || true else - echo "warning: 找不到 ${_VENDOR_VERSION_PY}，triton hcu backend 可能不激活" >&2 + echo "warning: ${_VENDOR_VERSION_PY} not found, triton hcu backend may not activate" >&2 fi bundle_summary "${LIB_DCU}" diff --git a/scripts/bundle_maca_libtorch.sh b/scripts/bundle_maca_libtorch.sh index 25c624c9..598a5ac7 100644 --- a/scripts/bundle_maca_libtorch.sh +++ b/scripts/bundle_maca_libtorch.sh @@ -1,22 +1,28 @@ #!/usr/bin/env bash -# 把沐曦 fork 的 libtorch C++ .so 打进 torch_fl/lib_maca/，做成自包含单 wheel。 +# Bundle the MetaX-forked libtorch C++ .so into torch_fl/lib_maca/ for a +# self-contained single wheel. # -# 官方 torch+cpu wheel 缺沐曦 fork 的 libtorch（带 at::maca::* / wcuda* 符号）， -# boxing 产物 libtorch_fl.so 又必须链接这些符号。本脚本从沐曦 torch wheel 拷贝那批 -# libtorch .so 到 torch_fl/lib_maca/，并用 patchelf 重写 RPATH： -# - lib_maca 内 libtorch .so -> $ORIGIN:/opt/maca/lib:/opt/maca/lib64 -# （运行期从目标机 /opt/maca 找 mcblas/mcdnn 等 maca runtime；不打包 runtime） -# - torch_fl/lib/libtorch_fl.so -> $ORIGIN:$ORIGIN/../lib_maca -# （从包内 lib_maca 找 fork libtorch，去掉构建机写死的绝对路径） +# The official torch+cpu wheel lacks the MetaX-forked libtorch (which exports +# at::maca::* / wcuda* symbols), and the boxing artifact libtorch_fl.so must +# link against those symbols. This script copies that batch of libtorch .so +# from the MetaX torch wheel into torch_fl/lib_maca/, and rewrites RPATH with +# patchelf: +# - libtorch .so inside lib_maca -> $ORIGIN:/opt/maca/lib:/opt/maca/lib64 +# (at runtime find mcblas/mcdnn and other maca runtime from /opt/maca on +# the target machine; the runtime is not bundled) +# - torch_fl/lib/libtorch_fl.so -> $ORIGIN:$ORIGIN/../lib_maca +# (find the forked libtorch from lib_maca inside the package, removing the +# absolute path hard-coded on the build machine) # -# maca runtime（libmcblas 等，~4.9G）不打包：装了沐曦卡的机器必有 /opt/maca 驱动。 +# maca runtime (libmcblas etc., ~4.9G) is not bundled: machines with MetaX +# cards installed must have the /opt/maca driver. # -# 用法: +# Usage: # FLAGOS_MACA_TORCH_LIB= bash scripts/bundle_maca_libtorch.sh -# MACA_PATH=/opt/maca bash scripts/bundle_maca_libtorch.sh # 覆盖 maca 路径 +# MACA_PATH=/opt/maca bash scripts/bundle_maca_libtorch.sh # override maca path # -# 应在 `python setup.py bdist_wheel`（ACCELERATOR=metax）之后、打 wheel 之前跑， -# 或跑完再重打 wheel。幂等。 +# Should run after `python setup.py bdist_wheel` (ACCELERATOR=metax) and before +# packing the wheel, or re-pack the wheel after running. Idempotent. set -euo pipefail @@ -28,39 +34,41 @@ LIB_MACA="${REPO_DIR}/torch_fl/lib_maca" TORCH_FL_LIB="${REPO_DIR}/torch_fl/lib" MACA_PATH="${MACA_PATH:-${METAX_HOME:-${MACA_HOME:-/opt/maca}}}" -# 沐曦 torch/lib 来源：显式 env，或从 conda 里找 +metax torch。 +# MetaX torch/lib source: explicit env, or find +metax torch from conda. SRC="${FLAGOS_MACA_TORCH_LIB:-}" if [ -z "${SRC}" ]; then SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so metax maca || true)" fi if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then - echo "error: 找不到沐曦 torch/lib。设 FLAGOS_MACA_TORCH_LIB=" >&2 + echo "error: MetaX torch/lib not found. Set FLAGOS_MACA_TORCH_LIB=" >&2 exit 1 fi if [ ! -f "${SRC}/libtorch_cuda.so" ]; then - echo "error: ${SRC} 里没有 libtorch_cuda.so，不是沐曦 torch/lib" >&2 + echo "error: ${SRC} does not contain libtorch_cuda.so, not a MetaX torch/lib" >&2 exit 1 fi bundle_require_patchelf -# 与 _metax_libtorch_link._CORE_SO + _CUDA_SO 对齐的自洽集合。 +# Self-consistent set aligned with _metax_libtorch_link._CORE_SO + _CUDA_SO. CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) VENDOR_RPATH="${MACA_PATH}/lib:${MACA_PATH}/lib64" -# 同 DCU：bundle 内的库要能从 lib_maca/ 和 torch/lib/ 软链两条路径被打开。 +# Same as DCU: libs inside the bundle must be openable from both lib_maca/ +# and torch/lib/ symlink paths. BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_maca" -echo "源沐曦 torch/lib : ${SRC}" -echo "目标 lib_maca : ${LIB_MACA}" -echo "maca runtime 路径: ${MACA_PATH}/lib" +echo "Source MetaX torch/lib : ${SRC}" +echo "Target lib_maca : ${LIB_MACA}" +echo "maca runtime path : ${MACA_PATH}/lib" bundle_copy_so "${SRC}" "${LIB_MACA}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 \ "${CORE_SO[@]}" "${CUDA_SO[@]}" -# torch_fl.so 去掉构建机写死的沐曦 torch/lib 绝对路径，改为从包内 lib_maca 找 fork libtorch。 +# Strip the build machine's hard-coded MetaX torch/lib absolute path from +# torch_fl.so, rewrite to find the forked libtorch from lib_maca inside the package. bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ "\$ORIGIN:\$ORIGIN/../lib_maca:${VENDOR_RPATH}" diff --git a/scripts/bundle_ppu_libtorch.sh b/scripts/bundle_ppu_libtorch.sh index a874ce52..12b15dd5 100644 --- a/scripts/bundle_ppu_libtorch.sh +++ b/scripts/bundle_ppu_libtorch.sh @@ -1,24 +1,29 @@ #!/usr/bin/env bash -# 把 PPU 本地构建的 libtorch C++ .so 打进 torch_fl/lib_ppu/，做成自包含单 wheel。 +# Bundle the locally built PPU libtorch C++ .so into torch_fl/lib_ppu/ for a +# self-contained single wheel. # -# PPU 是对着 PPU_SDK/CUDA_SDK 编的，所以 ACCELERATOR=cuda、CUDA boxing kernel -# 原样可用。跟真 NVIDIA 机器的区别在 libtorch：它是本地 USE_CUDA=1 源码构建，不是 -# 上游 wheel。实测 libtorch_fl.so 的未定义符号里它的 libtorch_cpu.so 提供 2092 个， -# libtorch_cuda.so 提供 0 个、libc10_cuda.so 10 个 —— 所以 core 库必须换掉， -# libtorch_cuda.so 只需在场（CUDA dispatch key 上要有已注册的厂商 kernel）。 +# PPU is compiled against PPU_SDK/CUDA_SDK, so ACCELERATOR=cuda and the CUDA +# boxing kernels work as-is. The difference from a real NVIDIA machine lies in +# libtorch: it is a local USE_CUDA=1 source build, not an upstream wheel. +# Measured undefined symbols in libtorch_fl.so show that its libtorch_cpu.so +# provides 2092 of them, libtorch_cuda.so provides 0, and libc10_cuda.so +# provides 10 — so the core libs must be replaced, while libtorch_cuda.so just +# needs to be present (the CUDA dispatch key must have registered vendor kernels). # -# 那个本地构建还链了 /usr/local/lib 下的系统 MKL -# （libmkl_core / libmkl_gnu_thread / libmkl_intel_lp64，~171 MB），官方 -# torch+cpu wheel 里没有对应文件，所以一并打进 lib_ppu/。 +# That local build also links against system MKL in /usr/local/lib +# (libmkl_core / libmkl_gnu_thread / libmkl_intel_lp64, ~171 MB); the official +# torch+cpu wheel does not ship these files, so they are bundled into lib_ppu/ +# together. # -# 不打包：PPU SDK runtime 留在目标机 /usr/local/PPU_SDK。 +# Not bundled: PPU SDK runtime stays on the target machine at /usr/local/PPU_SDK. # -# 用法: +# Usage: # FLAGOS_PPU_TORCH_LIB= bash scripts/bundle_ppu_libtorch.sh # PPU_SDK=/usr/local/PPU_SDK bash scripts/bundle_ppu_libtorch.sh # FLAGOS_PPU_MKL_DIR=/usr/local/lib bash scripts/bundle_ppu_libtorch.sh # -# 应在 `python setup.py bdist_wheel` 之后、打 wheel 之前跑。幂等。 +# Should run after `python setup.py bdist_wheel` and before packing the wheel. +# Idempotent. set -euo pipefail @@ -33,8 +38,9 @@ MKL_DIR="${FLAGOS_PPU_MKL_DIR:-/usr/local/lib}" SRC="${FLAGOS_PPU_TORCH_LIB:-}" if [ -z "${SRC}" ]; then - # PPU 的 torch 是本地构建，version.py 里不一定有 "ppu" 字样，所以先按标记找， - # 找不到就退化成"当前解释器的 torch 只要有 libtorch_cuda.so 就算"。 + # PPU torch is locally built; version.py may not contain the "ppu" marker, + # so first try finding it by marker, fall back to "any torch with + # libtorch_cuda.so from the current interpreter" if not found. SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so ppu || true)" if [ -z "${SRC}" ]; then SRC="$(bundle_find_vendor_torch_lib libtorch_cuda.so || true)" @@ -42,38 +48,40 @@ if [ -z "${SRC}" ]; then fi if [ -z "${SRC}" ] || [ ! -d "${SRC}" ]; then - echo "error: 找不到 PPU torch/lib。设 FLAGOS_PPU_TORCH_LIB=" >&2 + echo "error: PPU torch/lib not found. Set FLAGOS_PPU_TORCH_LIB=" >&2 exit 1 fi if [ ! -f "${SRC}/libtorch_cuda.so" ]; then - echo "error: ${SRC} 里没有 libtorch_cuda.so，不是 CUDA 构建的 torch/lib" >&2 + echo "error: ${SRC} does not contain libtorch_cuda.so, not a CUDA-built torch/lib" >&2 exit 1 fi bundle_require_patchelf -# 与 _ppu_libtorch_link._CORE_SO + _CUDA_SO 对齐的自洽集合。 +# Self-consistent set aligned with _ppu_libtorch_link._CORE_SO + _CUDA_SO. CORE_SO=(libc10.so libtorch_cpu.so libtorch.so libtorch_global_deps.so libtorch_python.so) CUDA_SO=(libc10_cuda.so libtorch_cuda.so libtorch_cuda_linalg.so libshm.so) MKL_SO=(libmkl_core.so.1 libmkl_gnu_thread.so.1 libmkl_intel_lp64.so.1) VENDOR_RPATH="${PPU_SDK}/CUDA_SDK/lib64:${PPU_SDK}/lib:${PPU_SDK}/lib64" -# 同 DCU：bundle 内的库要能从 lib_ppu/ 和 torch/lib/ 软链两条路径被打开。 +# Same as DCU: libs inside the bundle must be openable from both lib_ppu/ +# and torch/lib/ symlink paths. BUNDLE_ORIGIN="\$ORIGIN:\$ORIGIN/../../torch_fl/lib_ppu" -echo "源 PPU torch/lib : ${SRC}" -echo "目标 lib_ppu : ${LIB_PPU}" -echo "PPU SDK 路径 : ${PPU_SDK}" +echo "Source PPU torch/lib : ${SRC}" +echo "Target lib_ppu : ${LIB_PPU}" +echo "PPU SDK path : ${PPU_SDK}" bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 1 "${CORE_SO[@]}" bundle_copy_so "${SRC}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${CUDA_SO[@]}" -# 系统 MKL：libtorch_cpu.so 的直接 DT_NEEDED，官方 wheel 里没有同名文件。 +# System MKL: a direct DT_NEEDED of libtorch_cpu.so, with no same-named file +# in the official wheel. if [ -d "${MKL_DIR}" ]; then - echo "源 MKL : ${MKL_DIR}" + echo "Source MKL : ${MKL_DIR}" bundle_copy_so "${MKL_DIR}" "${LIB_PPU}" "${BUNDLE_ORIGIN}:${VENDOR_RPATH}" 0 "${MKL_SO[@]}" else - echo "warning: ${MKL_DIR} 不存在，跳过 MKL；目标机上若无 MKL 会缺库" >&2 + echo "warning: ${MKL_DIR} does not exist, skipping MKL; libs will be missing if the target machine has no MKL" >&2 fi bundle_rewrite_plugin_rpath "${TORCH_FL_LIB}" \ diff --git a/scripts/lib/bundle_common.sh b/scripts/lib/bundle_common.sh index 5e1bdde9..8cdf023a 100644 --- a/scripts/lib/bundle_common.sh +++ b/scripts/lib/bundle_common.sh @@ -1,20 +1,23 @@ #!/usr/bin/env bash -# 自包含 wheel 打包的公共部分，被 scripts/bundle_{maca,dcu,ppu}_libtorch.sh source。 +# common parts for self-contained wheel bundling, sourced by scripts/bundle_{maca,dcu,ppu}_libtorch.sh. # -# 背景：metax / dcu / ppu 三个后端跑的都是厂商 fork 的 libtorch —— 实测 -# libtorch_cpu.so 提供 libtorch_fl.so 两千多个未定义符号，而厂商库 -# （libtorch_hip.so / libtorch_cuda.so）提供 0 个。所以自包含 wheel 必须把 core -# 库一起打进来，不能只打厂商库。做法统一为： -# - bundle 目录内的 .so -> RPATH = $ORIGIN:<厂商驱动目录...> -# （驱动运行时留在目标机，不打包：装了卡的机器必有驱动） -# - torch_fl/lib/libtorch_fl.so -> RPATH = $ORIGIN:$ORIGIN/../:<驱动...> -# （去掉构建机写死的绝对路径。cmake/FlagosRpath.cmake 已经这么设了， -# 这里的 patchelf 是对已有 .so / 非 cmake 构建产物的兜底，保持幂等。） +# background: metax / dcu / ppu all run on a vendor fork of libtorch -- measured +# symbol attribution shows libtorch_cpu.so provides 2000+ undefined symbols for +# libtorch_fl.so, while the vendor lib (libtorch_hip.so / libtorch_cuda.so) +# provides 0. so a self-contained wheel must bundle the core libs together, not +# just the vendor lib. unified approach: +# - .so inside the bundle dir -> RPATH = $ORIGIN: +# (driver runtime stays on the target machine, not bundled: a box with the +# card already has the driver) +# - torch_fl/lib/libtorch_fl.so -> RPATH = $ORIGIN:$ORIGIN/../: +# (removes the build machine's hard-coded absolute paths. cmake/FlagosRpath.cmake +# already sets this; the patchelf here is fallback for existing .so / non-cmake +# build artifacts, kept idempotent.) # -# 运行期把 bundle 里的 core 库软链到 stock torch/lib 的逻辑在 -# torch_fl/accelerator/_vendor_libtorch.py，两边的 .so 清单必须对齐。 +# the runtime logic that symlinks the bundle's core libs into stock torch/lib is +# in torch_fl/accelerator/_vendor_libtorch.py; the .so list on both sides must align. # -# 提供: +# provides: # bundle_require_patchelf # bundle_find_vendor_torch_lib # bundle_copy_so @@ -23,24 +26,25 @@ set -euo pipefail -# patchelf 四个节点默认都没有。bw1000/810e 上 `pip install patchelf` 直接可用； -# mc550 的默认 index 没有这个包，要指定源。 +# patchelf is missing by default on all four nodes. on bw1000/810e `pip install patchelf` +# works directly; mc550's default index lacks the package and needs an explicit source. bundle_require_patchelf() { if command -v patchelf >/dev/null 2>&1; then return 0 fi cat >&2 <<'EOF' -error: 找不到 patchelf，无法重写 RPATH。装法（任选）: +error: patchelf not found, cannot rewrite RPATH. install with (any): pip install patchelf - pip install -i https://pypi.tuna.tsinghua.edu.cn/simple patchelf # 默认源没有时 + pip install -i https://pypi.tuna.tsinghua.edu.cn/simple patchelf # when default index lacks it conda install -c conda-forge patchelf EOF return 1 } -# 定位厂商 torch/lib：优先当前解释器的 torch（校验 version.py 含厂商标记且 -# probe_so 存在），失败返回非 0。调用方应先看自己的 env 覆盖变量。 -# 用法: bundle_find_vendor_torch_lib libtorch_hip.so dtk hip +# locate the vendor torch/lib: prefer the current interpreter's torch (verify +# version.py contains a vendor marker and probe_so exists), return non-zero on +# failure. caller should check their own env override variable first. +# usage: bundle_find_vendor_torch_lib libtorch_hip.so dtk hip bundle_find_vendor_torch_lib() { local probe="$1" shift @@ -66,8 +70,9 @@ if spec and spec.submodule_search_locations: PY } -# 拷 .so 并设 RPATH。cp -fL 解引用软链，拷实体文件。 -# required=1 时源缺失即失败，=0 时跳过（stock +cpu wheel 本来就没有厂商库）。 +# copy .so and set RPATH. cp -fL dereferences symlinks, copies the real file. +# when required=1 a missing source fails; =0 skips it (the stock +cpu wheel does +# not ship vendor libs anyway). bundle_copy_so() { local src_dir="$1" dst_dir="$2" rpath="$3" required="$4" shift 4 @@ -77,48 +82,49 @@ bundle_copy_so() { src="${src_dir}/${so}" if [ ! -f "${src}" ]; then if [ "${required}" = "1" ]; then - echo "error: 缺少必需的 ${src}" >&2 + echo "error: missing required ${src}" >&2 return 1 fi - echo " 跳过 ${so}（源不存在）" + echo "warning: ${src} does not exist, not bundled" >&2 continue fi cp -fL "${src}" "${dst_dir}/${so}" - # 非 ELF（极少数 .so 其实是 linker script）patchelf 会失败，不致命。 + # non-ELF (rare .so that are actually linker scripts) will fail patchelf, not fatal. if ! patchelf --set-rpath "${rpath}" "${dst_dir}/${so}" 2>/dev/null; then - echo " ${so}: patchelf 跳过（非 ELF?）" + echo " ${so}: patchelf skipped (non-ELF?)" fi - echo " 打包 ${so} ($(du -h "${dst_dir}/${so}" | cut -f1))" + echo " bundled ${so} ($(du -h "${dst_dir}/${so}" | cut -f1))" done } -# 重写 libtorch_fl.so / libtorch_bindings.so / libflagos.so 的 RPATH。 +# rewrite RPATH for libtorch_fl.so / libtorch_bindings.so / libflagos.so. bundle_rewrite_plugin_rpath() { local lib_dir="$1" rpath="$2" so target for so in libtorch_fl.so libtorch_bindings.so libflagos.so; do target="${lib_dir}/${so}" [ -f "${target}" ] || continue patchelf --set-rpath "${rpath}" "${target}" - echo " 重写 RPATH ${so} -> ${rpath}" + echo " rewrote RPATH ${so} -> ${rpath}" done } bundle_summary() { local dst_dir="$1" - echo "完成。$(basename "${dst_dir}") 总大小: $(du -sh "${dst_dir}" | cut -f1)（$(find "${dst_dir}" -type f | wc -l | tr -d ' ') 个文件）" + echo "done. $(basename "${dst_dir}") total size: $(du -sh "${dst_dir}" | cut -f1) ($(find "${dst_dir}" -type f | wc -l | tr -d ' ') files)" } -# 目标机上一定有的基线库，自检时不报。 +# baseline libs guaranteed present on the target machine, not reported during self-check. _BUNDLE_BASELINE_RE='^(libc|libm|libdl|librt|libpthread|libstdc\+\+|libgcc_s|ld-linux.*|libutil|libresolv|libnsl|libcrypt|libatomic|libgomp|libz|libnuma)\.so' -# 自检：bundle 里每个 .so 的 DT_NEEDED，凡是既不在 bundle 内、也不在给定的目标机 -# 目录里、又不是基线系统库的，都报出来 —— 那就是目标机上会 "not found" 的候选。 -# 用法: bundle_check_needed +# self-check: for each .so in the bundle, report every DT_NEEDED that is neither in +# the bundle, nor in the given target machine directories, nor a baseline system lib +# -- those are candidates for "not found" on the target machine. +# usage: bundle_check_needed bundle_check_needed() { local dst_dir="$1" shift local so dep d found - echo "---- DT_NEEDED 自检（只报可能在目标机缺失的项）----" + echo "---- DT_NEEDED self-check (only reports items potentially missing on target) ----" local missing=0 for so in "${dst_dir}"/*.so*; do [ -f "${so}" ] || continue @@ -136,6 +142,6 @@ bundle_check_needed() { done < <(patchelf --print-needed "${so}" 2>/dev/null || true) done if [ "${missing}" = "0" ]; then - echo " 无（bundle + 驱动目录已覆盖全部非基线依赖）" + echo " none (bundle + driver directories already cover all non-baseline deps)" fi } diff --git a/setup.py b/setup.py index e26de1c7..d5db7956 100644 --- a/setup.py +++ b/setup.py @@ -534,7 +534,8 @@ def _bundle_cuda_assets() -> None: externally-supplied libtorch_cuda.so (CPU-only pip torch does not ship it). Historically this was LD_PRELOAD-ed by scripts/with_cuda_libtorch.sh; for a single self-contained wheel we bundle the assets and preload them from - torch_fl/__init__.py before `import torch` (see docs §约束1). CUDA only. + torch_fl/__init__.py before `import torch` (see that doc, constraint 1). + CUDA only. Set FLAGOS_SKIP_CUDA_ASSETS=1 to skip (e.g. a slim build for a machine that supplies libtorch_cuda.so out-of-band). diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 64ff660e..4863c580 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -238,11 +238,11 @@ def _relink_vendor_libtorch() -> None: def _preload_cuda_assets() -> None: """Load the bundled CUDA .so into this process BEFORE `import torch`. - Hard constraint (docs/cpu_torch_external_libtorch_cuda.md §约束1): PyTorch - caches its CUDAHooks on first `import torch`. If libtorch_cuda.so is loaded - afterwards, device init fails with "Cannot initialize CUDA without ATen_cuda - library" even though the kernels register. So we ctypes-dlopen it here, at - the very top of torch_fl, before torch is imported. + Hard constraint (docs/cpu_torch_external_libtorch_cuda.md, constraint 1): + PyTorch caches its CUDAHooks on first `import torch`. If libtorch_cuda.so is + loaded afterwards, device init fails with "Cannot initialize CUDA without + ATen_cuda library" even though the kernels register. So we ctypes-dlopen it + here, at the very top of torch_fl, before torch is imported. libtorch_cuda.so has unresolved deps on the NVIDIA runtime libs (libcudart, libcublas, libcudnn, libnvshmem_host, ...) shipped by the pip nvidia-*-cu12 diff --git a/torch_fl/accelerator/_vendor_libtorch.py b/torch_fl/accelerator/_vendor_libtorch.py index 3a3d471d..8151481f 100644 --- a/torch_fl/accelerator/_vendor_libtorch.py +++ b/torch_fl/accelerator/_vendor_libtorch.py @@ -174,7 +174,7 @@ def _link_one(dst_dir, backup_dir, name, target, required, vendor): os.symlink(target, dst) -def _preload_global(lib_dir, load_order, core_so, vendor): +def _preload_global(lib_dir, load_order, core_so, vendor, fallback_dir=None): """dlopen the vendor set RTLD_GLOBAL, in dependency order. A CPU-only torch wheel never loads the forked runtime itself. Symlinking the @@ -186,6 +186,17 @@ def _preload_global(lib_dir, load_order, core_so, vendor): ``lib_dir`` must be the *source* dir (the bundle), never the ``torch/lib`` symlink dir -- see the ``$ORIGIN`` note in the module docstring. + + ``fallback_dir`` (the stock wheel's ``torch/lib``) covers a non-core .so the + vendor image simply does not ship. Measured: the MetaX CI's + ``/opt/vendor-libtorch/lib`` has no ``libshm.so``, so the bundle has none + either, yet ``libtorch_python.so`` carries a hard ``DT_NEEDED: libshm.so``. + dlopening the stock copy RTLD_GLOBAL *before* ``libtorch_python.so`` + satisfies that DT_NEEDED by soname against the already-loaded object; without + it the loader only searches the bundle's RUNPATH and dies with "libshm.so: + cannot open shared object file". Its own deps (libc10, libtorch_cpu) resolve + through ``torch/lib``, where they are symlinks into the bundle, so they share + an inode with what is already mapped and no second copy appears. """ handles = [] for name in load_order: @@ -193,7 +204,9 @@ def _preload_global(lib_dir, load_order, core_so, vendor): if not os.path.exists(path): if name in core_so: raise FileNotFoundError(f"{vendor} libtorch runtime missing: {path}") - continue + path = os.path.join(fallback_dir, name) if fallback_dir else "" + if not path or not os.path.exists(path): + continue try: handles.append(ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)) except OSError as exc: @@ -257,8 +270,12 @@ def ensure_vendor_libtorch_links( if load_order: # dlopen from `src` (the bundle), not from `active`: loading through the # torch/lib symlinks would expand $ORIGIN to torch/lib and lose the - # bundle-internal deps. Keep the handles alive for the process lifetime. - _runtime_handles.extend(_preload_global(src, load_order, core_so, label)) + # bundle-internal deps. `active` is only the fallback for a non-core .so + # the vendor image does not ship. Keep the handles alive for the process + # lifetime. + _runtime_handles.extend( + _preload_global(src, load_order, core_so, label, fallback_dir=active) + ) _done.add(bundle_dirname) return True diff --git a/torch_fl/accelerator/dcu/_dcu_libtorch_link.py b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py index 7a98ade9..cd64de84 100644 --- a/torch_fl/accelerator/dcu/_dcu_libtorch_link.py +++ b/torch_fl/accelerator/dcu/_dcu_libtorch_link.py @@ -66,9 +66,10 @@ ) # Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the HIP -# side, then libtorch_python. Same reasoning as MetaX -- symbols the plugin needs -# live in the forked CPU runtime, and loading only the HIP lib can leave its CPU -# dependency RTLD_LOCAL. +# side, then libshm and libtorch_python. Same reasoning as MetaX -- symbols the +# plugin needs live in the forked CPU runtime, and loading only the HIP lib can +# leave its CPU dependency RTLD_LOCAL -- including why libshm.so has to be listed +# explicitly ahead of libtorch_python.so. _LOAD_ORDER = ( "libc10.so", "libtorch_cpu.so", @@ -76,6 +77,7 @@ "libtorch_global_deps.so", "libc10_hip.so", "libtorch_hip.so", + "libshm.so", "libtorch_python.so", ) diff --git a/torch_fl/accelerator/metax/_metax_libtorch_link.py b/torch_fl/accelerator/metax/_metax_libtorch_link.py index 9705d5ba..78697c66 100644 --- a/torch_fl/accelerator/metax/_metax_libtorch_link.py +++ b/torch_fl/accelerator/metax/_metax_libtorch_link.py @@ -64,6 +64,12 @@ # side, then libtorch_python. GetFlagosDefaultCudaGenerator lives in the forked # ATen CPU runtime, not in libtorch_cuda.so, so loading only the CUDA lib would # leave that symbol unresolvable from libtorch_fl.so. +# +# libshm.so must precede libtorch_python.so: it is a direct DT_NEEDED of it, and +# dlopen resolves that by soname through the *loader's* search path, not through +# this list. In a fresh install nothing has put libshm.so anywhere the loader +# looks yet, so omitting it here fails with "libshm.so: cannot open shared object +# file" even though the file is sitting in the bundle dir. _LOAD_ORDER = ( "libc10.so", "libtorch_cpu.so", @@ -72,6 +78,7 @@ "libc10_cuda.so", "libtorch_cuda_linalg.so", "libtorch_cuda.so", + "libshm.so", "libtorch_python.so", ) diff --git a/torch_fl/accelerator/ppu/_ppu_libtorch_link.py b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py index dc8d6b86..09e4e103 100644 --- a/torch_fl/accelerator/ppu/_ppu_libtorch_link.py +++ b/torch_fl/accelerator/ppu/_ppu_libtorch_link.py @@ -62,7 +62,8 @@ ) # Dependency order for the RTLD_GLOBAL preload: core (CPU) first, then the CUDA -# side, then libtorch_python. Same reasoning as MetaX. +# side, then libshm and libtorch_python. Same reasoning as MetaX, including why +# libshm.so has to be listed explicitly ahead of libtorch_python.so. _LOAD_ORDER = ( "libc10.so", "libtorch_cpu.so", @@ -71,6 +72,7 @@ "libc10_cuda.so", "libtorch_cuda_linalg.so", "libtorch_cuda.so", + "libshm.so", "libtorch_python.so", )