-
Notifications
You must be signed in to change notification settings - Fork 173
chore: Bump vllm to 0.11.2, torch to 2.9, transformers to 4.57.1 #1563
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
82b6f95
39a9b03
fa2ccf4
baf37d6
fae9fc4
f61238c
eab6019
b671719
c5c3eca
ee4a84c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,7 +4,7 @@ | |
| # Self-contained build (remote NeMo RL source; no need for a local clone of NeMo RL): docker buildx build -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push https://github.com/NVIDIA-NeMo/RL.git | ||
| # Local NeMo RL source override: docker buildx build --build-context nemo-rl=. -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push . | ||
|
|
||
| ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.05-cuda12.9-devel-ubuntu24.04 | ||
| ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.11-cuda13.0-devel-ubuntu24.04 | ||
| FROM scratch AS nemo-rl | ||
| ARG NRL_GIT_REF=main | ||
| ADD --keep-git-dir=true https://github.com/NVIDIA-NeMo/RL.git#${NRL_GIT_REF} / | ||
|
|
@@ -129,4 +129,4 @@ RUN git rev-parse --is-shallow-repository | grep -q true && git fetch --unshallo | |
| RUN UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py | ||
| # NOTICES.txt file points to where the OSS source code is archived | ||
| RUN echo "This distribution includes open source which is archived at the following URL: https://opensource.nvidia.com/oss/teams/nvidia/nemo-rl/${RC_DATE}:linux-${TARGETARCH}/index.html" > NOTICES.txt && \ | ||
| echo "For further inquiries or assistance, contact us at [email protected]" >> NOTICES.txt | ||
| echo "For further inquiries or assistance, contact us at [email protected]" >> NOTICES.txt | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,6 +16,7 @@ | |
| import gc | ||
| import os | ||
| import sys | ||
| from importlib.util import find_spec | ||
| from typing import Any, Optional, cast | ||
|
|
||
| import ray | ||
|
|
@@ -156,63 +157,134 @@ def __init__( | |
| self.rank = 0 | ||
| self.world_size = 1 | ||
|
|
||
| # Monkey patch for vLLM to ensure RAY_ADDRESS is set in Ray actors. | ||
| try: | ||
| from vllm.logger import init_logger | ||
| # Monkey patches for vLLM behavior. We avoid importing vllm modules | ||
| # here to prevent side effects during initialization and instead | ||
| # locate the files via importlib metadata. | ||
|
|
||
| logger = init_logger("vllm_patch") | ||
| from vllm.logger import init_logger | ||
|
|
||
| def _patch_vllm_init_workers_ray(): | ||
| """Patch the vLLM ray_distributed_executor.py file. | ||
| logger = init_logger("vllm_patch") | ||
|
|
||
| 1. Pass custom runtime_env in _init_workers_ray call. | ||
| - This allows passing custom py_executable to worker initialization. | ||
| 2. Add NCCL_CUMEM_ENABLE and NCCL_NVLS_ENABLE to vLLM ADDITIONAL_ENV_VARS. | ||
| - This is a workaround to fix async vllm in some scenarios. | ||
| - See https://github.com/NVIDIA-NeMo/RL/pull/898 for more details. | ||
| """ | ||
| try: | ||
| import vllm.executor.ray_distributed_executor as ray_executor_module | ||
| def _get_vllm_file(relative_path: str) -> str: | ||
| """Return absolute path to a vLLM file or raise if it cannot be found. | ||
|
|
||
| The relative_path should be a POSIX-style path under the vllm | ||
| package root, e.g. "v1/executor/ray_executor.py" or | ||
| "attention/layer.py". | ||
| """ | ||
| spec = find_spec("vllm") | ||
| if spec is None or not spec.submodule_search_locations: | ||
| raise RuntimeError( | ||
| "vLLM package not found while attempting to patch " | ||
| f"'{relative_path}'. Ensure vLLM is installed and " | ||
| "available in this environment." | ||
| ) | ||
|
|
||
| file_to_patch = ray_executor_module.__file__ | ||
| base_dir = next(iter(spec.submodule_search_locations)) | ||
| file_path = os.path.join(base_dir, *relative_path.split("/")) | ||
|
|
||
| with open(file_to_patch, "r") as f: | ||
| content = f.read() | ||
| if not os.path.exists(file_path): | ||
| raise RuntimeError( | ||
| "Failed to locate expected vLLM file to patch. " | ||
| f"Looked for '{relative_path}' at '{file_path}'. " | ||
| "This likely indicates an unexpected vLLM installation " | ||
| "layout or version mismatch." | ||
| ) | ||
|
|
||
| return file_path | ||
|
|
||
| def _patch_vllm_init_workers_ray(): | ||
| """Patch the vLLM ray_distributed_executor.py file. | ||
|
|
||
| 1. Pass custom runtime_env in _init_workers_ray call. | ||
| - This allows passing custom py_executable to worker initialization. | ||
| 2. Add NCCL_CUMEM_ENABLE and NCCL_NVLS_ENABLE to vLLM ADDITIONAL_ENV_VARS. | ||
| - This is a workaround to fix async vllm in some scenarios. | ||
| - See https://github.com/NVIDIA-NeMo/RL/pull/898 for more details. | ||
| """ | ||
| file_to_patch = _get_vllm_file("v1/executor/ray_executor.py") | ||
|
|
||
| old_lines = [ | ||
| "self._init_workers_ray(placement_group)", | ||
| 'ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}', | ||
| ] | ||
| with open(file_to_patch, "r") as f: | ||
| content = f.read() | ||
|
|
||
| new_lines = [ | ||
| f'self._init_workers_ray(placement_group, runtime_env={{"py_executable": "{self.py_executable}"}})', | ||
| 'ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE", "RAY_ENABLE_UV_RUN_RUNTIME_ENV"}', | ||
| ] | ||
| old_lines = [ | ||
| "self._init_workers_ray(placement_group)", | ||
| 'ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}', | ||
| ] | ||
|
|
||
| need_replace = False | ||
| for old_line, new_line in zip(old_lines, new_lines): | ||
| if new_line in content or old_line not in content: | ||
| continue | ||
| content = content.replace(old_line, new_line) | ||
| need_replace = True | ||
| new_lines = [ | ||
| f'self._init_workers_ray(placement_group, runtime_env={{"py_executable": "{self.py_executable}"}})', | ||
| 'ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE", "RAY_ENABLE_UV_RUN_RUNTIME_ENV"}', | ||
| ] | ||
|
|
||
| need_replace = False | ||
| for old_line, new_line in zip(old_lines, new_lines): | ||
| if new_line in content or old_line not in content: | ||
| continue | ||
| content = content.replace(old_line, new_line) | ||
| need_replace = True | ||
|
|
||
| if not need_replace: | ||
| return | ||
|
|
||
| # Write back the patched content | ||
| with open(file_to_patch, "w") as f: | ||
| f.write(content) | ||
|
|
||
| def _patch_vllm_vit_flash_attn_backend(): | ||
| """Patch vLLM vision attention backend selection logic. | ||
|
|
||
| Modify the CUDA branch of maybe_get_vit_flash_attn_backend in | ||
| vllm.attention.layer to avoid overriding the backend when it | ||
| is already set to XFORMERS. This avoids flash attention related | ||
| errors when the ViT head dimension is not a multiple of 32. | ||
|
|
||
| Related issues: | ||
| - https://github.com/vllm-project/vllm/issues/27562 | ||
| - https://github.com/vllm-project/vllm/issues/26989 | ||
|
|
||
| This is properly fixed in https://github.com/vllm-project/vllm/pull/28763. We can remove this patch once we upgrade to a version of vllm that contains this fix. | ||
| """ | ||
| file_to_patch = _get_vllm_file("attention/layer.py") | ||
| with open(file_to_patch, "r") as f: | ||
| content = f.read() | ||
|
|
||
| old_snippet = ( | ||
| " elif current_platform.is_cuda():\n" | ||
| " if (\n" | ||
| " attn_backend != AttentionBackendEnum.FLASH_ATTN\n" | ||
| " and check_upstream_fa_availability(torch.get_default_dtype())\n" | ||
| " ):\n" | ||
| " attn_backend = AttentionBackendEnum.FLASH_ATTN\n" | ||
| " use_upstream_fa = True\n" | ||
| ) | ||
|
|
||
| new_snippet = ( | ||
| " elif current_platform.is_cuda():\n" | ||
| " if (\n" | ||
| " attn_backend != AttentionBackendEnum.FLASH_ATTN\n" | ||
| " and attn_backend != AttentionBackendEnum.XFORMERS\n" | ||
| " and check_upstream_fa_availability(torch.get_default_dtype())\n" | ||
| " ):\n" | ||
| " attn_backend = AttentionBackendEnum.FLASH_ATTN\n" | ||
| " use_upstream_fa = True\n" | ||
| ) | ||
|
|
||
| if not need_replace: | ||
| return | ||
| # Only patch if the file still has the old snippet and | ||
| # hasn't been patched already. | ||
| if new_snippet in content or old_snippet not in content: | ||
| return | ||
|
|
||
| # Write back the patched content | ||
| with open(file_to_patch, "w") as f: | ||
| f.write(content) | ||
| content = content.replace(old_snippet, new_snippet) | ||
|
|
||
| except (ImportError, FileNotFoundError, PermissionError): | ||
| # Allow failures gracefully | ||
| pass | ||
| with open(file_to_patch, "w") as f: | ||
| f.write(content) | ||
|
Comment on lines
+234
to
+281
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same file modification concerns apply. This patch function has the same operational risks as Since the comment indicates this is fixed in vllm PR #28763, consider adding a version check to skip this patch when using vLLM versions that include the fix. def _patch_vllm_vit_flash_attn_backend():
"""Patch vLLM vision attention backend selection logic.
Modify the CUDA branch of maybe_get_vit_flash_attn_backend in
vllm.attention.layer to avoid overriding the backend when it
is already set to XFORMERS. This avoids flash attention related
errors when the ViT head dimension is not a multiple of 32.
Related issues:
- https://github.com/vllm-project/vllm/issues/27562
- https://github.com/vllm-project/vllm/issues/26989
This is properly fixed in https://github.com/vllm-project/vllm/pull/28763. We can remove this patch once we upgrade to a version of vllm that contains this fix.
"""
+ # TODO: Add version check to skip patching when using vLLM versions >= X.Y.Z that include PR #28763
file_to_patch = _get_vllm_file("attention/layer.py")
🤖 Prompt for AI Agents |
||
|
|
||
| _patch_vllm_init_workers_ray() | ||
| logger.info("Successfully patched vllm _init_workers_ray.") | ||
| _patch_vllm_init_workers_ray() | ||
| logger.info("Successfully patched vllm _init_workers_ray.") | ||
|
|
||
| except (ImportError, AttributeError): | ||
| # vllm not installed or has a different structure, skipping patch. | ||
| pass | ||
| _patch_vllm_vit_flash_attn_backend() | ||
| logger.info("Successfully patched vllm vit flash attention backend.") | ||
|
Comment on lines
+283
to
+287
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Add error handling for patch invocations. The patch functions can raise Based on learnings from previous reviews, Ray actor methods should handle errors gracefully rather than propagating exceptions. Apply this diff: - _patch_vllm_init_workers_ray()
- logger.info("Successfully patched vllm _init_workers_ray.")
-
- _patch_vllm_vit_flash_attn_backend()
- logger.info("Successfully patched vllm vit flash attention backend.")
+ try:
+ _patch_vllm_init_workers_ray()
+ logger.info("Successfully patched vllm _init_workers_ray.")
+ except Exception:
+ import traceback
+ logger.error("Failed to patch vllm _init_workers_ray:")
+ traceback.print_exc()
+ # Consider whether to continue or return based on criticality
+
+ try:
+ _patch_vllm_vit_flash_attn_backend()
+ logger.info("Successfully patched vllm vit flash attention backend.")
+ except Exception:
+ import traceback
+ logger.error("Failed to patch vllm vit flash attention backend:")
+ traceback.print_exc()
+ # This patch may not be needed for all vLLM versions, so continueBased on learnings, Ray actor methods benefit from explicit error handling with print/traceback patterns. 🤖 Prompt for AI Agents |
||
|
|
||
| try: | ||
| import vllm | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,11 +18,11 @@ license = { text = "Apache 2.0" } | |
| dependencies = [ | ||
| "setuptools", | ||
| "ninja", # for flash-attn parallel build | ||
| "torch==2.8.0", | ||
| "torch==2.9.0", | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧩 Analysis chain🌐 Web query:
💡 Result: Short answer: Yes — vLLM v0.11.2 is compatible with PyTorch 2.9.0, but you must use a matching torch CUDA build (e.g. torch==2.9.0+cu129) or build vLLM from source if your CUDA/PyTorch build differs. [1][2] Sources:
References: 🌐 Web query:
💡 Result: Short answer: usually yes — transformers 4.57.1 can run with PyTorch 2.9.0, but it’s not guaranteed to be fully supported and several incompatibility bugs have been reported. Use a supported PyTorch release recommended by the Transformers docs (Transformers expects PyTorch 2.x; current docs note PyTorch 2.2+/2.0+ depending on doc version) for best results. [1][2] Details / evidence:
Recommendation:
Sources:
References: Verify torch 2.9.0 compatibility—known issues with transformers 4.57.1. vLLM 0.11.2 is compatible with torch 2.9.0, but ensure the CUDA build matches (e.g., torch==2.9.0+cu129). Transformers 4.57.1 with torch 2.9.0 has reported incompatibility issues (torchcodec loading, tensor-parallel errors) and is not in the officially tested matrix. Test this combination thoroughly before merging, or consider pinning to a tested PyTorch version (2.2+) if issues arise during integration testing. 🤖 Prompt for AI Agents |
||
| "triton; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", | ||
| "colored==2.2.3", | ||
| "ray[default]==2.49.2", | ||
| "transformers>=4.55.4", | ||
| "transformers>=4.57.1", | ||
| "wandb", | ||
| "numpy", | ||
| "datasets>=4.0.0", | ||
|
|
@@ -57,7 +57,7 @@ automodel = [ | |
| # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) | ||
| # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 | ||
| # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 | ||
| "vllm==0.11.0", # Remove this once https://github.com/NVIDIA-NeMo/RL/issues/811 resolved | ||
| "vllm==0.11.2", # Remove this once https://github.com/NVIDIA-NeMo/RL/issues/811 resolved | ||
| "flash-attn==2.8.1", | ||
| "mamba-ssm", | ||
| "causal-conv1d", | ||
|
|
@@ -69,7 +69,7 @@ vllm = [ | |
| # sudo apt-get update | ||
| # sudo apt-get install libibverbs-dev | ||
| "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@e3908bf5bd0cc6265bcb225d15cd8c996d4759ef", | ||
| "vllm==0.11.0", | ||
| "vllm==0.11.2", | ||
| "num2words>=0.5.14", | ||
| # Remove this once https://github.com/NVIDIA-NeMo/RL/issues/501 resolved | ||
| "flash-attn==2.8.1", | ||
|
|
@@ -92,7 +92,7 @@ mcore = [ | |
| "megatron-core", | ||
| "megatron-bridge", | ||
| # Remove this once https://github.com/NVIDIA-NeMo/RL/issues/501 resolved | ||
| "vllm==0.11.0", | ||
| "vllm==0.11.2", | ||
| # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) | ||
| # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 | ||
| # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 | ||
|
|
@@ -105,7 +105,7 @@ penguin = ["penguin"] | |
| # This is a default group so that we install these even with bare `uv sync` | ||
| build = [ | ||
| # Build requirement for TE | ||
| "torch==2.8.0", | ||
| "torch==2.9.0", | ||
| # Build requirement for TE | ||
| "setuptools", | ||
| "packaging", | ||
|
|
@@ -153,15 +153,15 @@ penguin = { workspace = true } | |
| nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" } | ||
| # torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels | ||
| torch = [ | ||
| { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pypi", marker = "sys_platform == 'darwin'" }, | ||
| ] | ||
| torchvision = [ | ||
| { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pypi", marker = "sys_platform == 'darwin'" }, | ||
| ] | ||
| triton = [ | ||
| { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, | ||
| { index = "pypi", marker = "sys_platform == 'darwin'" }, | ||
| ] | ||
| causal-conv1d = { git = "https://github.com/Dao-AILab/causal-conv1d", tag = "v1.5.0.post8" } | ||
|
|
@@ -187,8 +187,8 @@ url = "https://pypi.org/simple" | |
| explicit = true | ||
|
|
||
| [[tool.uv.index]] | ||
| name = "pytorch-cu129" | ||
| url = "https://download.pytorch.org/whl/cu129" | ||
| name = "pytorch-cu130" | ||
| url = "https://download.pytorch.org/whl/cu130" | ||
| explicit = true | ||
|
|
||
| [tool.uv] | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Critical: Code injection risk and concurrency hazard.
This function has several critical issues:
Code injection vulnerability (Line 216):
self.py_executableis directly embedded into an f-string that becomes Python source code. If the path contains quotes or special characters, it could break syntax or enable code injection.File modification race condition: Multiple workers or processes may attempt to patch the same installed vLLM files concurrently, risking corruption or incomplete patches.
Persistent modifications: Changes to installed package files persist across runs and affect all projects using the same vLLM installation, potentially causing unexpected behavior.
Missing strict parameter (Line 221): The
zip()call should includestrict=Truefor Python 3.10+.Apply these fixes:
Additional recommendations:
🧰 Tools
🪛 Ruff (0.14.6)
221-221:
zip()without an explicitstrict=parameterAdd explicit value for parameter
strict=(B905)