From 9612df0e0f88381a43c2e90fcb321212dbf7632e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Tue, 14 Apr 2026 15:10:42 +0200 Subject: [PATCH] build: bump DeepEP to 34152ae (#4228) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig Co-authored-by: Claude Sonnet 4.6 (cherry picked from commit 123645bb7812525515d6a70b22da0f6e184f0f0a) --- docker/Dockerfile.ci.dev | 2 +- tests/unit_tests/conftest.py | 6 +++++- tests/unit_tests/test_utilities.py | 10 +++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 7a8b69c1297..5e3d7419b3a 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -71,7 +71,7 @@ RUN bash -ex <<"EOF" git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git pushd DeepEP - git checkout eb9cee7de5a24193bf09500668d3a619d3d3f3fb + git checkout 34152ae28f80bcc3ee38d7a12cb2ad87cfd4ea72 patch -p1 < /workspace/deepep.patch popd TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v DeepEP/. diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 62413313fe0..10aec38b15c 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os +from datetime import timedelta from pathlib import Path import pytest @@ -42,7 +43,10 @@ def pytest_sessionfinish(session, exitstatus): def cleanup(): yield if torch.distributed.is_initialized(): - torch.distributed.barrier() + try: + torch.distributed.barrier(timeout=timedelta(seconds=300)) + except Exception: + return torch.distributed.destroy_process_group() diff --git a/tests/unit_tests/test_utilities.py b/tests/unit_tests/test_utilities.py index 39c78efb2b9..f8fad3325f5 100644 --- a/tests/unit_tests/test_utilities.py +++ b/tests/unit_tests/test_utilities.py @@ -91,7 +91,15 @@ def destroy_model_parallel(): os.environ.pop('NVTE_UNFUSED_ATTN', None) if not Utils.inited: return - torch.distributed.barrier() + + try: + # Flush pending CUDA work before the barrier so slow ranks don't + # time out while fast ranks tear down process groups. + torch.cuda.synchronize() + torch.distributed.barrier(timeout=timedelta(seconds=300)) + except Exception: + Utils.inited = False + return ps.destroy_model_parallel() Utils.inited = False