Use dlopen to load NCCL EP - #3434
Conversation
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
1873e7b to
fdded06
Compare
Greptile SummaryThe PR moves NCCL EP behind runtime loading and packages its shared library and JIT headers with Transformer Engine.
Confidence Score: 4/5The PR is not yet safe to merge because availability detection can approve an NCCL EP installation that subsequently fails during native initialization. The Python check accepts any Files Needing Attention: transformer_engine/init.py and transformer_engine/common/ep/nccl_ep_provider.cpp Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[EP initialization] --> B[Validate NCCL runtime version]
B --> C[Search NCCL_EP_HOME and packaged paths]
C --> D[dlopen libnccl_ep]
D --> E[Resolve NCCL EP symbols]
E --> F[Configure JIT header environment]
F --> G[Create EP group and handles]
Reviews (2): Last reviewed commit: "license fix" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if any( | ||
| library.is_file() | ||
| for lib_dir in ("lib", "lib64") | ||
| for library in (home / lib_dir).glob("libnccl_ep.so*") |
There was a problem hiding this comment.
When NCCL_EP_HOME contains only a fully versioned file such as libnccl_ep.so.0.1, this glob reports NCCL EP as available, but the native loader searches only libnccl_ep.so and libnccl_ep.so.<major>, causing the first EP initialization to fail.
Knowledge Base Used: Build, extensions, and packaging
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
| shutil.rmtree(target_nccl_ep_dir) | ||
| shutil.copytree( | ||
| nccl_ep_dir, | ||
| target_nccl_ep_dir, | ||
| ) |
There was a problem hiding this comment.
Do we copy the whole nccl_ep dir instead of header files only?
|
|
||
|
|
||
| def test_nccl_ep_library_not_found(monkeypatch): | ||
| monkeypatch.delenv("NCCL_EP_HOME", raising=False) |
There was a problem hiding this comment.
Does it mean NCCL_EP_HOME needs to be set at runtime?
| ncclResult_t create_group(ncclEpGroup_t* ep_group, ncclComm_t comm, | ||
| const ncclEpGroupConfig_t* config) { | ||
| using FuncT = decltype(&ncclEpCreateGroup); | ||
| static FuncT func = reinterpret_cast<FuncT>(get_symbol("ncclEpCreateGroup")); | ||
| return func(ep_group, comm, config); | ||
| } | ||
|
|
||
| ncclResult_t group_destroy(ncclEpGroup_t ep_group) { | ||
| using FuncT = decltype(&ncclEpGroupDestroy); | ||
| static FuncT func = reinterpret_cast<FuncT>(get_symbol("ncclEpGroupDestroy")); | ||
| return func(ep_group); | ||
| } | ||
|
|
||
| ncclResult_t handle_destroy(ncclEpHandle_t handle) { | ||
| using FuncT = decltype(&ncclEpHandleDestroy); | ||
| static FuncT func = reinterpret_cast<FuncT>(get_symbol("ncclEpHandleDestroy")); | ||
| return func(handle); | ||
| } |
There was a problem hiding this comment.
I think we could use std::forward here to consolidate these:
template <auto FuncPtr, typename... Args>
auto call_symbol(const char* name, Args&&... args) {
using FuncT = decltype(FuncPtr);
static FuncT func = reinterpret_cast<FuncT>(get_symbol(name));
return func(std::forward<Args>(args)...);
}
| Path to shared object file for a Transformer Engine library. | ||
|
|
||
| TE libraries are 'core', 'torch', or 'jax'. This function first | ||
| TE libraries are 'core', 'torch', 'jax', or 'nccl_ep'. This function first |
There was a problem hiding this comment.
Not sure if nccl_ep listed here is appropriate as it's a 3rd-level dependency rather than a TE's extension.
2 participants
Description
NCCL EP is currently linked statically into the TE core shared library. This PR moves from statically linking to
dlopening the file at runtime, but only when NCCL EP features are needed.This solves two problems:
-z nowflag, which resolves all symbols at dynamic load time. This means that even if NCCL EP isn't used, it's dependencies symbols are still resolved. Because NCCL EP features require a newer version of NCCL than the rest of TE, this means that enabling BOLT optimization raises the minimum NCCL version. Bydlopening NCCL EP instead, resolution of these symbols controlled by TE, and can be limited to only be loaded when NCCL EP features are used.libandincludeare now shipped with TE as a part of this work.Type of change
Changes
Please list the changes introduced in this PR:
Checklist: