fix(nixl): select UCX copy transport by accelerator - #3215
Open
yafshar wants to merge 1 commit into
Open
Conversation
Choose cuda_copy for CUDA and ze_copy for XPU when setting UCX_TLS for NIXL. This avoids forcing the CUDA copy transport on Level Zero devices, where UCX cannot pin device memory through that path and registration fails with ibv_reg_dmabuf_mr: Invalid argument (md supports: host|ze-device). UCX_TLS is a whitelist, so an accelerator missing from the map omits the copy transport rather than naming one the local UCX build lacks, which would narrow transport selection instead of widening it. The CUDA path is unchanged. Use NIXL's generic VRAM mem_type instead of the CUDA-specific alias; both resolve to VRAM_SEG, so the transfer path is unaffected. Signed-off-by: Yaser Afshar <yaser.afshar@intel.com>
yafshar
marked this pull request as ready for review
August 8, 2026 11:43
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
set_ucx_env_defaults()hardcodedcuda_copyas the UCX device-memorycopy transport. UCX exposes one such transport per accelerator family —
cuda_copyfor CUDA,ze_copyfor Level Zero — so on XPU the NIC cannotpin device memory and every NIXL registration fails:
ibv_reg_dmabuf_mr(...) failed: Invalid argument failed to register address ... (ze-device) ... on md[0]=mlx5_0 registerMem: registration failed for the specified or all potential backendsChange
Derive the copy transport from
torch.accelerator.current_accelerator().UCX_TLSis a whitelist, so an accelerator missing from the map omits thetoken rather than naming a transport the local UCX build lacks — which
would narrow transport selection instead of widening it.
Also swaps the deprecated
mem_type="cuda"alias for"VRAM". Bothresolve to
VRAM_SEG(nixl/_api.py:255-261), so this is a no-op rename.Verified
Emitted
UCX_TLS, measured on two machines:UCX_TLScudarc_x,rc,dc_x,dc,cuda_copyxpurc_x,rc,dc_x,dc,ze_copyThe CUDA string is byte-identical to the previous hardcoded value. Both
match the transports the respective UCX builds advertise under
ucx_info -d(Transport: cuda_copy/Transport: ze_copy).Note for reviewers
DEFAULT_INFERENCE_ENV_VARSsetsUCX_TLS=all(utils/process.py:32),applied on every inference launch path. Since these are
setdefaultcalls, the derived value only takes effect on the trainer path today;
inference keeps
all, which already includesze_copy. Left alone here— happy to follow up if you'd rather inference use the derived string.
Separately,
scripts/install_nixl_from_source.shconfigures UCX with--with-ze=no, so a UCX built by that script cannot provideze_copy.Out of scope for this PR; needs a verified build.