Skip to content

uv sync --extra moe rebuilds transformer-engine-torch and fails on bare metal #4020

Description

@bigben-dev777

Summary

Installing the moe extra on bare metal with uv sync --extra moe fails in the
transformer-engine-torch build path, even when the environment already has the
required CUDA Python wheels and a locally built transformer-engine-torch wheel.

Two separate issues appear in this path:

  1. transformer-engine-torch==2.19.0 probes a GitHub release wheel URL that returns HTTP 404.
  2. uv sync --extra moe still retries the TE source build instead of reusing a working prebuilt local wheel.

On this machine, a non-sync install path succeeds:

uv pip install -e '.[moe]'

after exporting the CUDA include/library paths described below.

Environment

  • Repo: NVIDIA-NeMo/Automodel
  • OS: Linux
  • Python: 3.12
  • GPU: 8x NVIDIA H200
  • Driver: 580.178.04
  • CUDA runtime/toolkit: 13.0
  • nvcc --version: V13.0.88
  • Install command that fails: uv sync --extra moe

Reproduction

From a fresh Automodel checkout:

uv venv
uv lock
uv sync --extra moe

On this system, deep_ep also needs RDMA headers and CUDA wheel include paths:

sudo apt-get update
sudo apt-get install -y libibverbs-dev rdma-core

export CUDNN_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cudnn"
export NCCL_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/nccl"
export C_INCLUDE_PATH="$CUDNN_HOME/include:$NCCL_HOME/include${C_INCLUDE_PATH:+:$C_INCLUDE_PATH}"
export CPLUS_INCLUDE_PATH="/usr/local/cuda/include/cccl:$CUDNN_HOME/include:$NCCL_HOME/include${CPLUS_INCLUDE_PATH:+:$CPLUS_INCLUDE_PATH}"
export CPATH="$CUDNN_HOME/include:$NCCL_HOME/include${CPATH:+:$CPATH}"
export LIBRARY_PATH="$CUDNN_HOME/lib:$NCCL_HOME/lib${LIBRARY_PATH:+:$LIBRARY_PATH}"
export LD_LIBRARY_PATH="$CUDNN_HOME/lib:$NCCL_HOME/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"

uv sync --extra moe

Actual behavior

uv sync --extra moe fails while building transformer-engine-torch.

The first failure mode is a missing prebuilt wheel asset probe:

https://github.com/NVIDIA/TransformerEngine/releases/download/v2.19.0/transformer_engine_torch-2.19.0+cu13torch2.10.0+cu130cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
HTTP/2 404

The second failure mode is the TE fallback source build, which may fail in the
NCCL EP path with:

fatal error: torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp: No such file or directory

even though the rest of the TE source build can succeed when that path is bypassed.

Expected behavior

One of the following should work reliably:

  1. uv sync --extra moe should consume a valid published transformer-engine-torch wheel.
  2. If no TE wheel exists for the current combo, the install path should fall back cleanly to source build.
  3. If a compatible local TE wheel is already installed, uv sync should not immediately force a rebuild that fails.

Additional investigation

The following was verified locally:

  • The v2.19.0 TransformerEngine release page exposes source archives, but the
    guessed wheel asset URL above returns 404.
  • A direct local TE build can succeed outside uv sync.
  • A direct non-sync Automodel install can also succeed.

Successful local path:

uv pip install /root/.cache/pip/wheels/.../transformer_engine_torch-2.19.0-cp312-cp312-linux_x86_64.whl
uv pip install -e '.[moe]'

Installed successfully after that:

  • transformer-engine 2.19.0
  • transformer-engine-torch 2.19.0
  • deep-ep 1.2.1+10d4dd7
  • mamba-ssm 2.3.1
  • causal-conv1d 1.7.0

Working imports after the non-sync install:

import transformer_engine
import deep_ep
import mamba_ssm
import causal_conv1d

Suggested fixes

Possible upstream fixes:

  1. Publish the expected transformer-engine-torch release wheel assets for the supported CUDA/PyTorch/Python combinations.
  2. Make the moe install docs prefer uv pip install -e '.[moe]' on bare metal when the TE wheel probe is known to fail.
  3. Adjust the install flow so uv sync can reuse an already installed compatible local TE wheel instead of forcing a source rebuild.
  4. Consider disabling the TE NCCL EP build path automatically when the required internal torch header is absent.

Workaround

This workaround succeeded on this machine:

sudo apt-get update
sudo apt-get install -y libibverbs-dev rdma-core

export CUDNN_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cudnn"
export NCCL_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/nccl"
export C_INCLUDE_PATH="$CUDNN_HOME/include:$NCCL_HOME/include${C_INCLUDE_PATH:+:$C_INCLUDE_PATH}"
export CPLUS_INCLUDE_PATH="/usr/local/cuda/include/cccl:$CUDNN_HOME/include:$NCCL_HOME/include${CPLUS_INCLUDE_PATH:+:$CPLUS_INCLUDE_PATH}"
export CPATH="$CUDNN_HOME/include:$NCCL_HOME/include${CPATH:+:$CPATH}"
export LIBRARY_PATH="$CUDNN_HOME/lib:$NCCL_HOME/lib${LIBRARY_PATH:+:$LIBRARY_PATH}"
export LD_LIBRARY_PATH="$CUDNN_HOME/lib:$NCCL_HOME/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export NVTE_PYTORCH_FORCE_BUILD=TRUE

uv pip install -e '.[moe]'

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions