nemo_rl.models.generation.vllm.patches#

Module Contents#

Functions#

_get_vllm_file

Return absolute path to a vLLM file or raise if it cannot be found.

_locked_file_patch

Yield (content, writer) under an exclusive file lock.

_patch_vllm_init_workers_ray

Patch the vLLM ray_distributed_executor.py file.

_patch_vllm_llama_eagle3_own_lm_head

Patch LlamaEagle3 to keep truncated draft lm_head ownership.

_patch_vllm_hermes_tool_parser_thread_safety

Patch Hermes2ProToolParser.init to cache tokenizer calls.

_apply_vllm_patches

API#

nemo_rl.models.generation.vllm.patches._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”.

nemo_rl.models.generation.vllm.patches._locked_file_patch(file_path: str)#

Yield (content, writer) under an exclusive file lock.

nemo_rl.models.generation.vllm.patches._patch_vllm_init_workers_ray(
py_executable: str,
extra_env_vars: list[str] | None,
) → None#

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.

nemo_rl.models.generation.vllm.patches._patch_vllm_llama_eagle3_own_lm_head(logger) → None#

Patch LlamaEagle3 to keep truncated draft lm_head ownership.

nemo_rl.models.generation.vllm.patches._patch_vllm_hermes_tool_parser_thread_safety(logger) → None#

Patch Hermes2ProToolParser.init to cache tokenizer calls.

The HuggingFace tokenizer’s Rust backend does not support concurrent access. When multiple async requests call _preprocess_chat concurrently, each one constructs a new Hermes2ProToolParser which calls tokenizer.encode() and tokenizer.decode() in init, causing “RuntimeError: Already borrowed”.

A lock alone is insufficient because the tool parser’s encode() can race with render_chat_async() in another concurrent request - two different codepaths sharing the same tokenizer instance.

This patch caches the encode/decode results so only the first instantiation (protected by a lock) touches the tokenizer. All subsequent instantiations read from cache without any tokenizer access.

Related:

  • https://github.com/vllm-project/vllm/pull/30264

  • https://github.com/huggingface/tokenizers/issues/537

  • https://github.com/PrimeIntellect-ai/prime-rl/pull/1837

nemo_rl.models.generation.vllm.patches._apply_vllm_patches(
py_executable: str,
*,
extra_env_vars: list[str] | None = None,
) → None#