build: enable multi-node HybridEP in container images - #4038
Conversation
Signed-off-by: seonjinn <sna@nvidia.com>
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 116ac87 |
Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test ff04168 |
Signed-off-by: seonjinn <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test 67f7159 |
1 similar comment
|
/ok to test 67f7159 |
Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test 8dde640 |
Signed-off-by: seonjinn <sna@nvidia.com>
Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test 36275f4 |
Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test 87151f9 |
Signed-off-by: seonjinn <sna@nvidia.com> # Conflicts: # docker/Dockerfile
|
/ok to test 28f4927 |
|
/ok to test e21d91e |
seonjinn
left a comment
seonjinn
left a comment
There was a problem hiding this comment.
Four notes on the NVML apt bracket and the DeepEP cache clean, plus a test diagnostics nit. Nothing blocking - the purge is in-layer so the shipped image is clean.
| # Everything is hardlinked against the uv cache; this layer holds exactly one real copy of every wheel. | ||
| # The venv is symlinked to avoid bloating the layer size | ||
| apt-get update | ||
| apt-get install -y --no-install-recommends libnvidia-ml-dev |
There was a problem hiding this comment.
libnvidia-ml-dev is not a stub install. On this PR's CI build apt pulled 7 packages, 136 MB download / 451 MB on disk, because DeepEP links with -l:libnvidia-ml.so.1 (setup.py:143) and that exact .so.1 ships only in libnvidia-compute-580, which the -dev package depends on.
It stays installed from #131 3.7s to #131 1950.6s, roughly 25 min after deep-ep finishes linking at 452.4s, so it also spans the sglang/mcore/automodel/modelopt/all-groups syncs. The purge and autoremove are in the same RUN, so the shipped image is provably clean (450 MB freed at #131 1953.1) - this is build hygiene, not a ship risk.
Worth checking whether symlinking the CUDA NVML stub to libnvidia-ml.so.1 and setting LIBRARY_PATH satisfies the -l: link, which would avoid the 451 MB entirely. Also note the package resolves to 12.0.140~12.0.1-4build4 from noble/multiverse inside a cuda-dl-base:26.05-cuda13.2 image.
| # The venv is symlinked to avoid bloating the layer size | ||
| apt-get update | ||
| apt-get install -y --no-install-recommends libnvidia-ml-dev | ||
| uv cache clean deep-ep |
There was a problem hiding this comment.
uv cache clean deep-ep runs after the seed rsync at L220, so it evicts the seeded wheel on every build, not once.
It was a no-op on this run (#131 11.60 No cache entries found - the seed was skipped on a cache-key mismatch), but once the :uv-cache tag is refreshed on main this becomes a recurring ~395s rebuild (Building deep-ep at 57.7s to Built at 452.4s, inside a 1958s layer that gates every downstream CI job).
uv already tracks this for you: [tool.uv.extra-build-variables] (pyproject.toml:542) participates in the build cache key, so changes to it trigger a rebuild of just the affected package. Adding
deep_ep = { HYBRID_EP_MULTINODE = "1" }invalidates once, then stays cached, and makes the flag authoritative for any uv sync rather than only these two Dockerfiles. Keep ENV HYBRID_EP_MULTINODE=1 for runtime JIT.
Please don't fold this into CACHE_KEY instead - L215 skips the entire seed on mismatch, so that would rebuild every source package to invalidate one wheel.
Side note: the mirrored clean at Dockerfile.ngc_pytorch:123 is always a no-op. That stage sets UV_CACHE_DIR=/opt/uv_cache (L89) with no seed stage and no cache mount, so the cache is empty when it runs.
| uv venv --system-site-packages ${UV_PROJECT_ENVIRONMENT} | ||
| uv pip install --no-cache-dir --no-deps /tmp/build_vllm/vllm/vllm*.whl | ||
| # Ensure nvshmem is installed before building DeepEP | ||
| apt-get update |
There was a problem hiding this comment.
# Ensure nvshmem is installed before building DeepEP (L120) used to sit directly above the uv sync that pulls nvidia-nvshmem-cu13 (pyproject.toml:62). These apt lines were inserted between them, so the comment now reads as if libnvidia-ml-dev provides nvshmem - it doesn't, that's NVML.
Suggest moving the nvshmem comment back down to L124, and adding a separate comment above the apt block explaining why NVML headers are needed for the multi-node DeepEP build. That rationale isn't in the PR description either, which still only mentions the env var and the cache clean.
| cache_clean = "uv cache clean deep-ep" | ||
| nvml_stub_purge = "apt-get purge -y libnvidia-ml-dev" | ||
| nvml_stub_dependencies_purge = "apt-get autoremove -y" | ||
| nvml_stub_purge_index = lines.index(nvml_stub_purge) |
There was a problem hiding this comment.
nvml_stub_purge_index = lines.index(nvml_stub_purge) runs before the membership asserts, which makes the assert nvml_stub_purge in lines, ... below it dead code - it can never fail. Remove the purge line from a Dockerfile and you get ValueError: 'list.index(x): x not in list' here instead of the message you wrote (Python 3.11+ omits the value from that error, so the traceback is genuinely opaque).
Same for first_sync_index = next(...) at L45 with no default: a missing uv sync raises a bare StopIteration rather than the ordering assert's message.
Moving both lookups below the membership asserts, and using next(..., None) plus an explicit assert for first_sync_index - the pattern you already use for nvml_stub_dependencies_purge_index - fixes both.
Minor: nvml_stub_* is a misnomer. libnvidia-ml-dev pulls in the real driver library, not a stub.
Plain "deepep" hangs with "DeepEP error: timeout (dispatch CPU)" on this EP=32, 4-node/32-rank layout under the deep_ep pin bumped by NVIDIA-NeMo#3436 (deep_ep 29d31c0 -> 17cfb81, "enable HybridEP on H100/B200"). The container's deep_ep is already built with HYBRID_EP_MULTINODE=1 (NVIDIA-NeMo#4038, merged), and "hybridep" is already a supported nemo_automodel BackendConfig.dispatcher value -- just not yet the repo-wide default for AutoModel MoE recipes (that's NVIDIA-NeMo#4037, still open). Switching these two recipes now rather than waiting on NVIDIA-NeMo#4037: verified with a live 4n8g eagle3 run past the exact point (policy.get_logprobs()'s first post-refit call) that failed twice under plain deepep. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
5 participants
Summary
HYBRID_EP_MULTINODE=1in both production Dockerfiles.HYBRID_EP_MULTINODEin the uv seed compatibility key so stale single-node wheels are skipped without deleting reusable multi-node wheels.docker/install_apptainer.shunchanged because it installs Apptainer and does not build DeepEP.Why
DeepEP reads
HYBRID_EP_MULTINODEwhile building its extension. Setting the variable only at runtime cannot add inter-node transport support to an existing single-node wheel.This change also enables multi-node AutoModel HybridEP recipes proposed in #4037.
Testing
ruff checkruff format --checkdocker buildx build --checkfordocker/Dockerfiledocker buildx build --checkfordocker/Dockerfile.ngc_pytorchgit diff --checkA full image build and multi-node runtime validation will be recorded separately.