feat: add uccl-ep to docker built and add USING_UEP flag to run_pretrain and runner hook by zhenhuang12 · Pull Request #540 · AMD-AGI/Primus

zhenhuang12 · 2026-02-05T09:10:26Z

No description provided.

Copilot

Pull request overview

This PR adds support for UCCL-EP (Unified Collective Communication Library - Expert Parallelism) to the project, enabling an alternative backend for MoE dispatch/combine operations. The changes include Docker build integration, runtime configuration through environment variables, and conditional logic to select between TURBO and DEEP_EP backends.

Changes:

Adds UCCL-EP installation to Docker build process with configurable commit hash
Implements USING_UEP flag to enable UCCL-EP functionality in runner hooks and pretrain scripts
Updates validation logic to conditionally apply turbo_deepep_num_cu constraint based on backend selection

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
.github/workflows/ci.yaml	Adds UCCL_COMMIT environment variable and passes it to Docker builds
.github/workflows/docker/Dockerfile	Adds UCCL-EP installation steps during Docker image build
runner/helpers/hooks/05_using_uep.sh	New hook script to validate and configure UCCL-EP environment
examples/run_pretrain.sh	Adds UCCL rebuild option and UEP configuration logic
examples/run_local_pretrain.sh	Passes REBUILD_UCCL and USING_UEP environment variables to Docker container
examples/moe_package/run_deepseek_v2_pretrain_mi355x.sh	Enables UCCL-EP feature (option 8) in MoE features array
examples/moe_package/run_deepseek_v2_lite_pretrain_mi355x.sh	Enables UCCL-EP feature (option 8) in MoE features array
primus/modules/trainer/megatron/utils.py	Updates validation to only enforce CU constraint when using TURBO backend

runner/helpers/hooks/05_using_uep.sh

examples/run_pretrain.sh

.github/workflows/docker/Dockerfile

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

runner/helpers/hooks/05_using_uep.sh

Copilot · 2026-02-05T09:13:42Z

examples/run_pretrain.sh

+
+    git clone https://github.com/uccl-project/uccl.git
+    cd uccl || exit
+    cd ep && PYTORCH_ROCM_ARCH="gfx942;gfx950" python3 setup.py build && cd ..


Hardcoded GPU architectures 'gfx942;gfx950' should be configurable via an environment variable to support different hardware configurations.

Suggested change

cd ep && PYTORCH_ROCM_ARCH="gfx942;gfx950" python3 setup.py build && cd ..

cd ep && PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx942;gfx950}" python3 setup.py build && cd ..

Copilot · 2026-02-05T09:13:42Z

examples/run_pretrain.sh

+# ----------------- Using UCCL-EP -----------------
+if [ "$USING_UEP" == "1" ]; then
+    LOG_INFO "USING_UEP is enabled, checking required packages..."
+
+    if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then
+        LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."
+        exit 1
+    fi
+    LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"
+    LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"
+
+    if [ "$ENABLE_NUMA_BINDING" != "1" ]; then
+        LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."
+    fi
+
+    export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP
+    LOG_INFO "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to DEEP_EP"
+else
+    export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=TURBO
+    LOG_INFO "USING_UEP is disabled. PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to TURBO"
+fi


This logic block is duplicated between run_pretrain.sh and runner/helpers/hooks/05_using_uep.sh. Consider consolidating this logic into a single location to avoid maintenance issues.

Suggested change

# ----------------- Using UCCL-EP -----------------

if [ "$USING_UEP" == "1" ]; then

LOG_INFO "USING_UEP is enabled, checking required packages..."

if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"

if [ "$ENABLE_NUMA_BINDING" != "1" ]; then

LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."

fi

export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP

LOG_INFO "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to DEEP_EP"

else

export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=TURBO

LOG_INFO "USING_UEP is disabled. PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to TURBO"

fi

configure_using_uep() {

# ----------------- Using UCCL-EP -----------------

if [ "$USING_UEP" == "1" ]; then

LOG_INFO "USING_UEP is enabled, checking required packages..."

if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"

if [ "$ENABLE_NUMA_BINDING" != "1" ]; then

LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader worker exited unexpectedly."

fi

export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP

LOG_INFO "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to DEEP_EP"

else

export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=TURBO

LOG_INFO "USING_UEP is disabled. PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND set to TURBO"

fi

}

# Configure UCCL-EP / TURBO backend selection

configure_using_uep

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Copilot · 2026-02-05T09:14:35Z

examples/run_pretrain.sh

+    LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"
+
+    if [ "$ENABLE_NUMA_BINDING" != "1" ]; then
+        LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."


This duplicates the warning logic from the hook file (runner/helpers/hooks/05_using_uep.sh line 27). Consider using LOG_WARN for consistency with the hook implementation, as both locations serve the same warning purpose.

Suggested change

LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."

LOG_WARN "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Copilot · 2026-02-06T02:39:53Z

runner/helpers/hooks/05_using_uep.sh

+    LOG_INFO "USING_UEP is enabled, checking required packages..."
+
+    if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then
+        LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."


Error message only mentions 'uccl' but the check on line 19 validates both uccl and deep_ep packages. The error message should mention both packages to accurately reflect what is being validated.

Suggested change

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

LOG_ERROR "uccl and/or deep_ep are not installed! Please use pre-installed primus image or set REBUILD_UCCL=1 and ensure deep_ep is installed."

Copilot · 2026-02-06T02:39:53Z

primus/modules/trainer/megatron/utils.py

+            args.expert_model_parallel_size >= 16
+            and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND") == "TURBO"
+        ):
+            # Turbo DeepEP is not supported CUs > 32 when using internode dispatch/combine.


Grammar error in comment: 'is not supported CUs' should be 'does not support CUs' or 'is not supported for CUs'.

Suggested change

# Turbo DeepEP is not supported CUs > 32 when using internode dispatch/combine.

# Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Copilot · 2026-02-06T02:49:02Z

examples/run_pretrain.sh

+    LOG_INFO "USING_UEP is enabled, checking required packages..."
+
+    if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then
+        LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."


The error message on line 409 only mentions 'uccl is not installed' but the condition on line 408 checks for both uccl and deep_ep packages. The error message should clarify that either or both packages may be missing.

Suggested change

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

LOG_ERROR "uccl and/or deep_ep are not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

Copilot · 2026-02-06T02:49:02Z

examples/run_pretrain.sh

+# ----------------- Rebuild UCCL -----------------
+export REBUILD_UCCL=${REBUILD_UCCL:-0}
+if [ "$REBUILD_UCCL" == "1" ]; then
+    LOG_INFO "Rebuilding UCCL from source..."
+    apt update && apt install -y rdma-core libibverbs-dev libnuma-dev libgoogle-glog-dev
+    mkdir -p "/workspace/"
+    cd "/workspace" || exit
+
+    # Clean up old directory if exists to avoid git clone conflicts
+    if [ -d "uccl" ]; then
+        LOG_INFO "Removing existing uccl directory..."
+        rm -rf uccl
+    fi
+
+    git clone https://github.com/uccl-project/uccl.git
+    cd uccl || exit
+    cd ep && PYTORCH_ROCM_ARCH="gfx942;gfx950" python3 setup.py build && cd ..
+    cp ep/build/**/*.so uccl
+    pip3 install --no-build-isolation .
+    cd ep/deep_ep_wrapper && pip3 install --no-build-isolation . -v
+    cd "${PRIMUS_PATH}" || exit
+    LOG_INFO "Rebuilding UCCL from source done."
+else
+    LOG_INFO "Skip UCCL rebuild. REBUILD_UCCL=$REBUILD_UCCL"
+fi


This UCCL rebuild logic (lines 378-402) is duplicated between examples/run_pretrain.sh and the hook runner/helpers/hooks/05_using_uep.sh validation logic (lines 405-424). The validation logic in lines 405-424 is also duplicated in the hook file. Consider extracting this into a shared function or sourcing a common script to reduce duplication and ensure consistency.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Copilot · 2026-02-07T08:20:27Z

.github/workflows/docker/Dockerfile

+    cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd .. && \
+    cp ep/build/**/*.so uccl && \


TORCH_CUDA_ARCH_LIST is a CUDA-oriented env var and is unlikely to affect ROCm builds; additionally, cp ep/build/**/*.so relies on ** glob expansion, which typically won’t work in Docker RUN (default /bin/sh) and may fail to copy any .so files. Use the ROCm-appropriate arch env var (consistent with the runtime script) and replace the ** copy with a find ... -name '*.so' -exec ... (or similar) approach that works under /bin/sh.

Suggested change

cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd .. && \

cp ep/build/**/*.so uccl && \

cd ep && PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" python3 setup.py build && cd .. && \

find ep/build -name '*.so' -exec cp {} uccl \; && \

Copilot · 2026-02-07T08:20:28Z

examples/run_pretrain.sh

+    git clone https://github.com/uccl-project/uccl.git
+    cd uccl || exit
+    cd ep && PYTORCH_ROCM_ARCH="gfx942;gfx950" python3 setup.py build && cd ..
+    cp ep/build/**/*.so uccl


In bash, ** recursive globbing requires shopt -s globstar; without it, this copy may fail silently or not match any files (leaving uccl without the needed .so). Replace this with a find ep/build -name '*.so' -exec cp ... style copy, or enable globstar explicitly before using **.

Suggested change

cp ep/build/**/*.so uccl

find ep/build -name '*.so' -exec cp {} uccl \;

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 8 comments.

Copilot · 2026-02-07T11:29:03Z

runner/helpers/hooks/04_rebuild_uccl.sh

 UCCL_BUILD_DIR="${UCCL_BUILD_DIR:-/tmp/uccl_${HOSTNAME:-$(hostname)}}"
 UCCL_REF="${UCCL_REF:-}"
-GPU_ARCHS="${GPU_ARCHS:-gfx942;gfx950}"
+GPU_ARCHS="${GPU_ARCHS:-gfx942,gfx950}"


This changes both the arch-list separator and the environment variable used for ROCm builds. For PyTorch ROCm extensions, PYTORCH_ROCM_ARCH is typically expected (and list separators are commonly ;), while TORCH_CUDA_ARCH_LIST is CUDA-oriented and may be ignored on ROCm—potentially producing a build without the intended gfx targets. Consider reverting to PYTORCH_ROCM_ARCH and using the separator format expected by the build toolchain.

Copilot · 2026-02-07T11:29:03Z

runner/helpers/hooks/04_rebuild_uccl.sh


 LOG_INFO_RANK0 "[hook system] Building uccl ep"
-cd ep && PYTORCH_ROCM_ARCH="${GPU_ARCHS}" python3 setup.py build && cd ..
+cd ep && TORCH_CUDA_ARCH_LIST="${GPU_ARCHS}" python3 setup.py build && cd ..


This changes both the arch-list separator and the environment variable used for ROCm builds. For PyTorch ROCm extensions, PYTORCH_ROCM_ARCH is typically expected (and list separators are commonly ;), while TORCH_CUDA_ARCH_LIST is CUDA-oriented and may be ignored on ROCm—potentially producing a build without the intended gfx targets. Consider reverting to PYTORCH_ROCM_ARCH and using the separator format expected by the build toolchain.

Copilot · 2026-02-07T11:29:03Z

primus/modules/trainer/megatron/utils.py

-        if args.expert_model_parallel_size >= 16:
+        if (
+            args.expert_model_parallel_size >= 16
+            and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND") == "TURBO"


If PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND is unset, this condition becomes false and the turbo_deepep_num_cu <= 32 constraint is skipped even though the effective default backend elsewhere appears to be TURBO. To avoid silently bypassing validation, default the getenv to \"TURBO\" (or otherwise ensure the env var is always set before validation runs).

Suggested change

and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND") == "TURBO"

and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO"

Copilot · 2026-02-07T11:29:03Z

.github/workflows/docker/Dockerfile

+RUN if [ "$PRIMUS_TURBO_FRAMEWORK" != "JAX" ]; then \
+    cd /opt && \
+    git clone https://github.com/uccl-project/uccl.git && \
+    cd uccl && \
+    git checkout ${UCCL_COMMIT} && \
+    cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd .. && \
+    cp ep/build/**/*.so uccl && \
+    pip3 install --no-build-isolation . -v && \
+    cd ep/deep_ep_wrapper && pip3 install --no-build-isolation . -v && \
+    rm -rf /opt/uccl; \
+    fi


Two concrete issues here: (1) cp ep/build/**/*.so uccl relies on ** globstar expansion, which is not enabled by default in bash and may fail to match/copy anything in Docker builds. Prefer an explicit find ... -name '*.so' -exec cp ... pattern or a known build output path. (2) git checkout ${UCCL_COMMIT} will fail if UCCL_COMMIT is unset/empty in non-CI builds; guard with a non-empty check or provide a default to keep local builds working.

Copilot · 2026-02-07T11:29:04Z

examples/run_pretrain.sh

+    cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd ..
+    cp ep/build/**/*.so uccl


cp ep/build/**/*.so uccl depends on globstar (**) expansion, which is typically disabled by default and can lead to no .so files being copied (and subsequent install/runtime failures). Use a glob pattern that doesn't require globstar or switch to a find-based copy. Also consider aligning the arch env var/list format with whatever the ROCm build expects (see UCCL rebuild hook).

Suggested change

cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd ..

cp ep/build/**/*.so uccl

cd ep && TORCH_CUDA_ARCH_LIST="gfx942;gfx950" python3 setup.py build && cd ..

find ep/build -name '*.so' -exec cp -t uccl {} +

Copilot · 2026-02-07T11:29:04Z

examples/run_pretrain.sh

+    fi
+
+    git clone https://github.com/uccl-project/uccl.git
+    cd uccl || exit


The rebuild path clones UCCL from main without pinning to a commit/tag, which makes runs non-reproducible and can break unexpectedly over time. Consider supporting a UCCL_COMMIT/UCCL_REF env var (similar to CI) and checking out that ref when provided.

Suggested change

cd uccl || exit

cd uccl || exit

# Optionally pin UCCL to a specific ref/commit for reproducible rebuilds.

# If UCCL_REF is not set, UCCL_COMMIT can be used as a fallback.

UCCL_CHECKOUT_REF="${UCCL_REF:-$UCCL_COMMIT}"

if [ -n "${UCCL_CHECKOUT_REF}" ]; then

LOG_INFO "Checking out UCCL ref: ${UCCL_CHECKOUT_REF}"

git checkout "${UCCL_CHECKOUT_REF}" || exit

fi

Copilot · 2026-02-07T11:29:04Z

runner/helpers/hooks/05_using_uep.sh

+    if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then
+        LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."
+        exit 1
+    fi
+    LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"
+    LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"


Using pip directly can point at a different Python environment than python3/pip3 (especially in containers/venvs), leading to false negatives/positives when checking installation. Prefer python3 -m pip show ... (or consistently pip3) to ensure the check targets the same interpreter used elsewhere in the scripts.

Suggested change

if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"

if ! python3 -m pip show uccl &>/dev/null || ! python3 -m pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(python3 -m pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(python3 -m pip show deep_ep | grep Version)"

Copilot · 2026-02-07T11:29:04Z

runner/helpers/hooks/05_using_uep.sh

+    if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then
+        LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."
+        exit 1
+    fi
+    LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"
+    LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"


Using pip directly can point at a different Python environment than python3/pip3 (especially in containers/venvs), leading to false negatives/positives when checking installation. Prefer python3 -m pip show ... (or consistently pip3) to ensure the check targets the same interpreter used elsewhere in the scripts.

Suggested change

if ! pip show uccl &>/dev/null || ! pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(pip show deep_ep | grep Version)"

if ! python3 -m pip show uccl &>/dev/null || ! python3 -m pip show deep_ep &>/dev/null; then

LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."

exit 1

fi

LOG_INFO "uccl package is installed: $(python3 -m pip show uccl | grep Version)"

LOG_INFO "deep_ep package is installed: $(python3 -m pip show deep_ep | grep Version)"

zhenhuang12 requested a review from Xiaoming-AMD as a code owner February 5, 2026 09:10

Copilot AI review requested due to automatic review settings February 5, 2026 09:10

zhenhuang12 requested review from limou102 and wenxie-amd as code owners February 5, 2026 09:10

Copilot AI reviewed Feb 5, 2026

View reviewed changes

runner/helpers/hooks/05_using_uep.sh Show resolved Hide resolved

examples/run_pretrain.sh Show resolved Hide resolved

.github/workflows/docker/Dockerfile Outdated Show resolved Hide resolved

.github/workflows/docker/Dockerfile Outdated Show resolved Hide resolved

Copilot AI review requested due to automatic review settings February 5, 2026 09:12

AMD-AGI deleted a comment from Copilot AI Feb 5, 2026

Copilot AI reviewed Feb 5, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings February 5, 2026 09:13

Copilot AI reviewed Feb 5, 2026

View reviewed changes

zhenhuang12 added 3 commits February 6, 2026 02:38

feat: add uccl-ep

52a4a7c

update

8002521

fix

9f03657

wenxie-amd force-pushed the dev/zhuang12/uep branch from 0156bcc to 9f03657 Compare February 6, 2026 02:38

Update runner/helpers/hooks/05_using_uep.sh

e1ec1c4

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings February 6, 2026 02:39

Copilot AI reviewed Feb 6, 2026

View reviewed changes

Apply suggestion from @Copilot

6739335

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings February 6, 2026 02:48

Copilot AI reviewed Feb 6, 2026

View reviewed changes

zhenhuang12 added 2 commits February 6, 2026 07:30

fix build uccl-ep

9c574e1

fix

c7ca6a3

Copilot AI review requested due to automatic review settings February 7, 2026 08:19

Copilot AI reviewed Feb 7, 2026

View reviewed changes

zhenhuang12 added 3 commits February 7, 2026 11:21

skip build uccl for jax

8e8bf22

add test case

1aa5e7d

fix

b0ec0f5

Copilot AI review requested due to automatic review settings February 7, 2026 11:27

Copilot AI reviewed Feb 7, 2026

View reviewed changes

	cd ep && PYTORCH_ROCM_ARCH="gfx942;gfx950" python3 setup.py build && cd ..
	cd ep && PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx942;gfx950}" python3 setup.py build && cd ..

	LOG_INFO "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."
	LOG_WARN "ENABLE_NUMA_BINDING is not enabled! Please set ENABLE_NUMA_BINDING=1 to avoid dataloader workers exiting unexpectedly."

	LOG_ERROR "uccl is not installed! Please use pre-installed primus image or set REBUILD_UCCL=1."
	LOG_ERROR "uccl and/or deep_ep are not installed! Please use pre-installed primus image or set REBUILD_UCCL=1 and ensure deep_ep is installed."

	# Turbo DeepEP is not supported CUs > 32 when using internode dispatch/combine.
	# Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine.

		cd ep && TORCH_CUDA_ARCH_LIST="gfx942,gfx950" python3 setup.py build && cd .. && \
		cp ep/build/*/.so uccl && \

	cp ep/build/*/.so uccl
	find ep/build -name '*.so' -exec cp {} uccl \;

	and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND") == "TURBO"
	and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO"

-    cd uccl || exit
+    cd uccl || exit
+    # Optionally pin UCCL to a specific ref/commit for reproducible rebuilds.
+    # If UCCL_REF is not set, UCCL_COMMIT can be used as a fallback.
+    UCCL_CHECKOUT_REF="${UCCL_REF:-$UCCL_COMMIT}"
+    if [ -n "${UCCL_CHECKOUT_REF}" ]; then
+        LOG_INFO "Checking out UCCL ref: ${UCCL_CHECKOUT_REF}"
+        git checkout "${UCCL_CHECKOUT_REF}" || exit
+    fi

Conversation

zhenhuang12 commented Feb 5, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026