chore: major version bump (torch 2.8, vllm 0.11, ray 2.49) & SP fixes (#1334)

Signed-off-by: Terry Kong <terryk@nvidia.com>

Author: terrykong

.github/workflows/cicd-main.yml

@@ -162,7 +162,7 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@v5
         with:
-          version: "0.7.2"
+          version: "0.9.1"
           enable-cache: true
           prune-cache: false
       # Faster than uv python install since it caches python alongside runner

3rdparty/Automodel-workspace/Automodel

@@ -25,7 +25,9 @@
 
 
 CACHED_DEPENDENCIES = [
-    "openai<=1.97.1",
+    # TODO: openai was un-pinned in https://github.com/NVIDIA-NeMo/RL/pull/1334 to comply with the version vllm mandates
+    # When penguin is actually added, this install will fail and we'll need penguin to ensure compatibility.
+    "openai",
     "tqdm",
     "pydantic",
     "pydantic_core",

3rdparty/Penguin-workspace/setup.py

@@ -131,9 +131,8 @@ For more examples and setup details, continue to the [Prerequisites](#prerequisi
     <tr>
       <td colspan="2" style="border:1px solid #d0d7de; padding:8px; vertical-align:top; word-break:break-word; overflow-wrap:anywhere; white-space:normal;">
         <strong>Clone and create the environment</strong>
-        <pre style="white-space:pre-wrap; word-break:break-word; overflow-wrap:anywhere;"><code class="language-sh">git clone git@github.com:NVIDIA-NeMo/RL.git nemo-rl
+        <pre style="white-space:pre-wrap; word-break:break-word; overflow-wrap:anywhere;"><code class="language-sh">git clone git@github.com:NVIDIA-NeMo/RL.git nemo-rl --recursive
 cd nemo-rl
-git submodule update --init --recursive
 uv venv</code></pre>
         <em>Note:</em> If you previously ran without checking out the submodules, you may need to rebuild virtual environments by setting <code>NRL_FORCE_REBUILD_VENVS=true</code>. See <a href="#tips-and-tricks">Tips and Tricks</a>.
       </td>
@@ -203,19 +202,6 @@ uv venv
 > Please do not use `-p/--python` and instead allow `uv venv` to read it from `.python-version`.
 > This ensures that the version of python used is always what we prescribe.
 
-If working outside a container, it can help to build [flash-attn](https://github.com/Dao-AILab/flash-attention) and warm the uv cache before your first run.
-```sh
-bash tools/build-flash-attn-in-uv-cache.sh
-```
-> [!NOTE]
-> On the first install, `flash-attn` can take a while to install (~45min with 48 CPU hyperthreads). After it is built once, it is cached in your uv's cache directory, making subsequent installs much quicker.
-
-> [!TIP]
-> The NeMo RL Dockerfile will warm the uv cache with flash-attn.
-> See https://docs.nvidia.com/nemo/rl/latest/docker.html for instructions if you are looking for the NeMo RL container.
-
-If successful, you should see `✅ flash-attn successfully added to uv cache`.
-
 Use `uv run` to launch all commands. It handles pip installing implicitly and ensures your environment is up to date with our lock file.
 > [!NOTE]
 > - It is not recommended to activate the `venv`, and you should use `uv run <command>` instead to execute scripts within the managed environment.

README.md

@@ -41,7 +41,7 @@ rm -rf /var/lib/apt/lists/*
 EOF
 
 # Install uv and python
-ARG UV_VERSION=0.7.2
+ARG UV_VERSION=0.9.1
 ARG PYTHON_VERSION=3.12
 ENV PATH="/root/.local/bin:$PATH"
 RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
@@ -50,6 +50,10 @@ RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
 # Disable usage stats by default for users who are sensitive to sharing usage.
 # Users are encouraged to enable if the wish.
 ENV RAY_USAGE_STATS_ENABLED=0
+# After ray>=2.47, this feature is enabled by default which creates uv venvs for any py_executable starting with `uv run`.
+# There is severe contention and performance issues with this enabled considering our dependencies are so large and occasionally
+# need to be compiled, so NeMo RL has an implementation in nemo_rl/utils/venv.py that does it once per node as opposed to once per task.
+ENV RAY_ENABLE_UV_RUN_RUNTIME_ENV=0
 ENV NEMO_RL_VENV_DIR=/opt/ray_venvs
 
 
@@ -67,15 +71,6 @@ ARG BUILD_CUSTOM_VLLM
 ENV UV_PROJECT_ENVIRONMENT=/opt/nemo_rl_venv
 ENV UV_LINK_MODE=copy
 
-# This step is to warm the uv cache with flash-attn without invalidating it due to COPY layers
-# This layer has to be manually updated
-RUN <<"EOF" bash -exu
-uv venv ${UV_PROJECT_ENVIRONMENT}
-
-VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT uv pip install --link-mode symlink setuptools torch==2.7.0 psutil ninja --torch-backend=cu128
-VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT uv pip install --link-mode symlink flash-attn==2.7.4.post1 --no-build-isolation
-EOF
-
 # First copy only the dependency files
 COPY --from=nemo-rl pyproject.toml uv.lock ./
 COPY --from=nemo-rl tools/build-custom-vllm.sh ./tools/build-custom-vllm.sh

docker/Dockerfile

@@ -38,12 +38,16 @@ rm -rf /var/lib/apt/lists/*
 EOF
 
 # Install uv at /usr/local/bin in case the root home directory is bind mounted
-ARG UV_VERSION=0.7.2
+ARG UV_VERSION=0.9.1
 RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | XDG_BIN_HOME=/usr/local/bin sh
 
 # Disable usage stats by default for users who are sensitive to sharing usage.
 # Users are encouraged to enable if they wish.
 ENV RAY_USAGE_STATS_ENABLED=0
+# After ray>=2.47, this feature is enabled by default which creates uv venvs for any py_executable starting with `uv run`.
+# There is severe contention and performance issues with this enabled considering our dependencies are so large and occasionally
+# need to be compiled, so NeMo RL has an implementation in nemo_rl/utils/venv.py that does it once per node as opposed to once per task.
+ENV RAY_ENABLE_UV_RUN_RUNTIME_ENV=0
 ENV NEMO_RL_VENV_DIR=/opt/ray_venvs
 
 # Build vLLM from source to use with the NVIDIA PyTorch base image

docker/Dockerfile.ngc_pytorch

@@ -113,6 +113,9 @@ policy:
     moe_permute_fusion: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
+    # gives ~25% training perf speedup with sequence packing and apply_rope_fusion
+    bias_activation_fusion: True
+    defer_fp32_logits: null
 
     optimizer:
       optimizer: "adam"

examples/configs/dpo.yaml

@@ -84,7 +84,7 @@ policy:
 
   megatron_cfg:
     enabled: false
-    empty_unused_memory_level: 0
+    empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1

examples/configs/grpo_math_1B.yaml

@@ -77,7 +77,7 @@ policy:
 
   megatron_cfg:
     enabled: true
-    empty_unused_memory_level: 0
+    empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1

examples/configs/grpo_math_1B_megatron.yaml

@@ -30,7 +30,7 @@ policy:
 
   megatron_cfg:
     enabled: true
-    empty_unused_memory_level: 0
+    empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     converter_type: "LlamaForCausalLM"
     tensor_model_parallel_size: 1
     # On H100, can run with pp=1 for better performance with expandable segments (which is enabled by default)