Conversation
Micky774
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
Micky774
force-pushed
the
zain/aiter-static
branch
from
5c994a9 to
8451fc9
Compare
Micky774
force-pushed
the
zain/aiter-static
branch
from
8451fc9 to
52163c5
Compare
Micky774
force-pushed
the
zain/aiter-static
branch
from
52163c5 to
2c4c11b
Compare
| exit 1 | ||
| fi | ||
|
|
||
| mapfile -d '' obj_files < <(find "${src_obj_dir}" -type f -name '*.o' -print0) |
There was a problem hiding this comment.
Why so complicated approach, just do what build systems do -> call ar cq @file_list
|
|
||
| # Create static archives for both forward and backward passes | ||
| batch_size=${AITER_ARCHIVE_BATCH_SIZE:-400} | ||
| for lib in fwd bwd; do |
There was a problem hiding this comment.
Since it is built and used by us, will it be more efficient to make single lib?
| file(MAKE_DIRECTORY "${EXTRACT_DIR}") | ||
| message(STATUS "[AITER-PREBUILT] Caching locally built libs to ${EXTRACT_DIR}") | ||
| file(COPY "${SOURCE_DIR}/libmha_fwd.so" "${SOURCE_DIR}/libmha_bwd.so" DESTINATION "${EXTRACT_DIR}") | ||
| if(EXISTS "${SOURCE_DIR}/libmha_fwd.so" AND EXISTS "${SOURCE_DIR}/libmha_bwd.so") |
There was a problem hiding this comment.
They should always exist because they are built by AITER build script before creating .a
|
|
||
|
|
||
| set(AITER_MHA_INSTALL_PREFIX "transformer_engine" CACHE STRING "aiter mha shared lib install prefix in TE") | ||
| option(NVTE_AITER_STATIC_LINK "Statically link AITER MHA libs into ck_fused_attn" OFF) |
There was a problem hiding this comment.
By contract NVTE prefix is for env vriables. CMake options shoud rather be just AITER_STATIC_LINK
|
|
||
| if [[ -z "${AITER_DIR}" || -z "${AITER_TEST_DIR}" || -z "${GPU_ARCHS_VAL}" ]]; then | ||
| echo "[AITER-PREBUILT] --aiter-dir, --aiter-test-dir, and --gpu-archs are required." >&2 | ||
| echo "[AITER-BUILD] --aiter-dir, --aiter-test-dir, and --gpu-archs are required." >&2 |
There was a problem hiding this comment.
For consistency it should probably also be changed in CMakeLists.txt
|
|
||
| find_package(hip) | ||
| list(APPEND ck_fused_attn_LINKER_LIBS hip::host hip::device roctx64 ${__AITER_MHA_PATH}/libmha_fwd.so ${__AITER_MHA_PATH}/libmha_bwd.so) | ||
| if(NVTE_AITER_STATIC_LINK) |
There was a problem hiding this comment.
Let's add keep hip and roc libs here and only aiter libs under conditions
| if(NVTE_AITER_STATIC_LINK) | ||
| set(__AITER_MHA_FWD_LIB "${__AITER_MHA_PATH}/libmha_fwd.a") | ||
| set(__AITER_MHA_BWD_LIB "${__AITER_MHA_PATH}/libmha_bwd.a") | ||
| if(NOT EXISTS "${__AITER_MHA_FWD_LIB}" OR NOT EXISTS "${__AITER_MHA_BWD_LIB}") |
There was a problem hiding this comment.
wouldn't link FAIL anyways if libraries are not there and aren't those files checked when detected AITER cache?
|
This current impl depends on detailed assumptions in aiter. We can add a rule archive into this section and pass an option to compile static lib: |
3 participants
Description
Provides a TE-side build solution to assemble a static archive from intermediate object files generated during the AITER backend building process. These files are batch-processed into an archive, which is then used directly (and privately) in the compilation. This is done as an effort to avoid namespace collision for AITER runtime-loaded symbols provided by both TE-built AITER and third-party-built AITER libraries (e.g. via PyTorch).
The use of static-linking at build time is mediated via
NVTE_AITER_STATIC_LINK, with a default set to1to proactively prevent the namespace collision issue. The archiving is batched, with batch-size controlled byAITER_ARCHIVE_BATCH_SIZEand set to400by default. Higher batch sizes are marginally more efficient, but provide less granular progress updates.Note that using the static libraries does require a marginal increase in compilation time (a couple of minutes for the archiving process) but it also results in a smaller overall size for the necessary runtime libraries (saves ~300MB).
Type of change
Changes
Please list the changes introduced in this PR:
NVTE_AITER_STATIC_LINKas an environment variable to mediate the use of static linking at build-time, defaulted to1.AITER_ARCHIVE_BATCH_SIZEwhich is defaulted to400.Checklist: