diff --git a/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_disagg_sc.sh b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_disagg_sc.sh new file mode 100755 index 00000000000..105d18cec16 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_disagg_sc.sh @@ -0,0 +1,188 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# launch_3k_disagg_sc.sh +# +# The BACKUP 3K test arm: SingleController with disaggregated sandboxes, but the +# judges served on the allocation instead of from NVCF. Use this if NVCF proves +# unusable; the primary arm is launch_3k_nvcf_disagg_sc.sh. +# +# Shape (GB200 NVL72, 4 GPUs/node), deliberately identical to the v1 baseline +# launch_3k_pipeclean.sh: +# Hetgroup 0: 256 training + 494 generation + 2 Gym = 752 nodes (segment 16) +# Hetgroup 1: 8 GenRM + 2 NL2Bash = 10 nodes (segment 2) +# Total: 762 nodes / 3048 GPUs +# +# Judges are deployed exactly as the baseline deploys them — GenRM and NL2Bash on +# the external hetgroup, the safety judge in Gym — so this arm differs from the +# baseline only in the execution path (v1 -> SC) and the sandbox backend +# (colocated -> disaggregated), and differs from the NVCF arm only in judge +# placement. That makes NVCF a single isolated variable. +# +# This departs from the 6K disagg recipe, which serves all three judges inside +# Gym on 64 nodes. Reproducing that here would both change the node total and +# carry the 6K judge parallelism (GenRM TP4 x DP16 = 64 GPUs) into a half-size +# job, giving this arm roughly double the baseline's judge capacity. Set +# EXTERNAL_JUDGES=0 with NUM_GYM_NODES=32 to get the old shape back. +# +# NO CHECKPOINTING. The SC path raises if checkpointing.enabled is true, so size +# NRL_MAX_STEPS to fit one allocation. +# +# Required, in addition to the usual site variables: +# OPENSANDBOX_BASE_URL sandbox service endpoint +# OPENSANDBOX_API_KEY keep in the environment or a sourced creds file — +# never on the command line or in configs +# NS_SANDBOX_IMAGE nemo-skills sandbox image the service can pull +# +# Recommended at this scale (see disagg_overrides.yaml for sizing): +# NS_SANDBOX_POOL_REF= claim PREWARMED pods from a server-side +# Pool created before the job +# NS_SANDBOX_POOL_SIZE=128 keep sessions/pod ≈ CPUs/pod +# +# The Gym overlay must provide the sandbox_pool backend (NVIDIA-NeMo/Gym branch +# hemild/rlvr-osb-473f446f (the 6K Gym pin + sandbox commit) until merged). +# +# Optional: NRL_MAX_STEPS, WALLTIME, CONTEXT_PARALLEL_SIZE, DRY_RUN, +# NUM_TRAIN_NODES / NUM_GEN_NODES / NUM_GYM_NODES, STREAM_MIN_GROUPS, +# NUM_STORAGE_UNITS, REFIT_TRANSPORT — as launch_3k_pipeclean_sc.sh. +# +# Extra positional args are forwarded as Hydra overrides to ultra_launch.sh. +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) + +# Default config — callers may still override CONFIG_PATH explicitly. +export CONFIG_PATH="${CONFIG_PATH:-${SCRIPT_DIR}/pipeclean_3k_disagg_sc.yaml}" + +# The SC driver. data_plane.enabled=true (set in the config) is mandatory for it. +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-./examples/run_grpo_single_controller.py}" + +# ============================================================================= +# Disaggregated sandbox service +# ============================================================================= +export NO_COLOCATED_SANDBOX=1 +: "${OPENSANDBOX_BASE_URL:?OPENSANDBOX_BASE_URL is required (sandbox service endpoint)}" +: "${OPENSANDBOX_API_KEY:?OPENSANDBOX_API_KEY is required (sandbox service credential)}" +: "${NS_SANDBOX_IMAGE:?NS_SANDBOX_IMAGE is required (sandbox image the service can pull)}" +export OPENSANDBOX_BASE_URL OPENSANDBOX_API_KEY NS_SANDBOX_IMAGE +# Unset-only default (`-` not `:-`): export NS_SANDBOX_POOL_REF="" forces +# direct creates at Gym startup, matching the header contract. +export NS_SANDBOX_POOL_REF="${NS_SANDBOX_POOL_REF-ns-tools-warm}" +export LEAN_SANDBOX_POOL_REF="${LEAN_SANDBOX_POOL_REF-math-lean-warm}" +# Half of 6K's 256, tracking the halved 8192-sample cohort. +export NS_SANDBOX_POOL_SIZE="${NS_SANDBOX_POOL_SIZE:-128}" + +# ============================================================================= +# Site defaults — the checkpoint and blend of the other 3K arms +# ============================================================================= +export MODEL_PATH="${MODEL_PATH:-/lustre/fsw/portfolios/llmservice/users/jiaqiz/models/ultra_stage2sft_step300}" +export TRAIN_PATH="${TRAIN_PATH:-/lustre/fsw/portfolios/llmservice/users/jiaqiz/data/gym/rl-data-tools/blends/curriculum_v35_inescapable-sawfly.train.efforts0p15_qamathcode.jsonl}" +# The reference runs validate on the training blend; there is no separate split. +export VAL_PATH="${VAL_PATH:-${TRAIN_PATH}}" + +# Image with the Gym venvs prebaked (python interpreter included — the stock +# nightly ships a venv tree whose interpreter is absent) and the sandbox deps +# installed. Pinned to a resolved image, not the moving nightly symlink. +# +# Byte-identical restripe of hemild/scale-6k/...plus_sandbox_venvs.sqsh. The +# original is laid out over 8 OSTs at 1 MB, and 752 nodes pulling 63 GB through +# 8 OSTs is the I/O concentration that took down an earlier run. This copy +# inherits the 350-OST / 16 MB layout the other arms' image already uses. +export CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/llmservice/users/sauramishra/images-striped/main_ultra_recipes_prebaked_venvs_20260730.plus_sandbox_venvs.sqsh}" +# Same stale-image delta as the v1 disagg wrapper, plus the SC driver itself: +# ultra_launch does not mount examples/, so run_grpo_single_controller.py has to +# be bind-mounted in or the image's older copy runs against the newer library. +GYM_EXAMPLES_DIR=$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd) +export EXTRA_MOUNTS="${EXTRA_MOUNTS:-/lustre:/lustre,${GYM_EXAMPLES_DIR}:/opt/nemo-rl/examples/nemo_gym,${GYM_EXAMPLES_DIR}/../run_grpo_single_controller.py:/opt/nemo-rl/examples/run_grpo_single_controller.py}" +export PERSISTENT_CACHE="${PERSISTENT_CACHE:-/lustre/fsw/portfolios/llmservice/users/${USER}/.cache/nemotron_ultra}" +export HF_HOME="${HF_HOME:-/lustre/fsw/portfolios/llmservice/users/${USER}/hf_home}" + +export GENRM_MODEL="${GENRM_MODEL:-/lustre/fsw/portfolios/llmservice/users/ansubramania/models/qwen235b_principle_comparison_genrm_step1230}" +export NL2BASH_JUDGE_MODEL="${NL2BASH_JUDGE_MODEL:-/lustre/fsw/portfolios/llmservice/users/ansubramania/models/Qwen3-235B-A22B-Instruct-2507-FP8}" +export SAFETY_JUDGE_MODEL="${SAFETY_JUDGE_MODEL:-/lustre/fsw/portfolios/llmservice/users/ansubramania/super_v3/model_checkpoints/Nemotron-Content-Safety-Reasoning-4B}" + +export SLURM_ACCOUNT="${SLURM_ACCOUNT:-nemotron_sw_pre}" +export SLURM_PARTITION="${SLURM_PARTITION:-batch}" + +# Serve GenRM and NL2Bash on the external hetgroup and the safety judge in Gym, +# at the baseline's pool sizes, so judge capacity matches the arm this is +# compared against. Gym needs only 2 nodes because the safety judge at TP=4 x +# DP=2 is its sole GPU consumer. +export EXTERNAL_JUDGES="${EXTERNAL_JUDGES:-1}" +export GENRM_REPLICAS="${GENRM_REPLICAS:-8}" +export GENRM_TENSOR_PARALLEL_SIZE="${GENRM_TENSOR_PARALLEL_SIZE:-4}" +export GENRM_REASONING_PARSER_NAME="${GENRM_REASONING_PARSER_NAME:-deepseek_r1}" +export GENRM_ENABLE_EXPERT_PARALLEL="${GENRM_ENABLE_EXPERT_PARALLEL:-0}" +export NL2BASH_REPLICAS="${NL2BASH_REPLICAS:-2}" +export NL2BASH_TENSOR_PARALLEL_SIZE="${NL2BASH_TENSOR_PARALLEL_SIZE:-4}" +export EXTERNAL_VLLM_SEGMENT_SIZE="${EXTERNAL_VLLM_SEGMENT_SIZE:-2}" +export EXTERNAL_VLLM_SKIP_PREFLIGHT="${EXTERNAL_VLLM_SKIP_PREFLIGHT:-1}" + +# Node split, identical to launch_3k_pipeclean.sh. Keep the Ray total a multiple +# of 16 and NUM_TRAIN_NODES a multiple of 32 (TP * CP * PP = 128 GPUs). +export NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-256}" +export NUM_GEN_NODES="${NUM_GEN_NODES:-494}" +export NUM_GYM_NODES="${NUM_GYM_NODES:-2}" + +# CP=16 is inherited from pipeclean_6k.yaml; allow an override for memory +# experiments. +CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE:-16}" + +# The 6K values halved with the cohort: 128 is a quarter of the 512-group +# cohort, and storage units track per-step volume. +STREAM_MIN_GROUPS="${STREAM_MIN_GROUPS:-128}" +NUM_STORAGE_UNITS="${NUM_STORAGE_UNITS:-32}" + +# Data-plane backend. mooncake_cpu carries rollouts over CPU RDMA and has no TCP +# fallback — it raises at data-plane setup if no RDMA device is usable, rather +# than degrading silently. DP_BACKEND=simple reverts to the SimpleStorage Ray +# actors, the only backend NUM_STORAGE_UNITS applies to. +DP_BACKEND="${DP_BACKEND:-mooncake_cpu}" + +# Shard-to-shard weight refit, on by default in this variant. REFIT_TRANSPORT=null +# restores the full-tensor broadcast. +REFIT_TRANSPORT="${REFIT_TRANSPORT:-nccl_reshard}" + +export EXP_NAME="${EXP_NAME:-ultra-3k-disagg-sandboxes-sc}" +export WALLTIME="${WALLTIME:-4:00:00}" + +# Reuse the container's prebaked venvs instead of building in-tree on first run +# (GYM_VENV_DIR= empty restores the stock in-tree behavior). +GYM_VENV_DIR="${GYM_VENV_DIR:-/opt/gym_venvs}" +VENV_OVERRIDE=() +if [ -n "${GYM_VENV_DIR}" ]; then + VENV_OVERRIDE=("++env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR}") +fi + +# Preflight the sandbox dependency this variant introduces: ultra_launch overlays +# $PWD/3rdparty/Gym-workspace/Gym whenever that directory exists (even empty — +# an uninitialized submodule silently masks the container's baked Gym). +GYM_OVERLAY="$PWD/3rdparty/Gym-workspace/Gym" +if [ ! -s "${GYM_OVERLAY}/resources_servers/ns_tools/sandbox_pool.py" ]; then + echo "ERROR: ${GYM_OVERLAY} is missing the sandbox_pool backend (empty or wrong Gym overlay?)." >&2 + echo " Populate it with NVIDIA-NeMo/Gym branch hemild/rlvr-osb-473f446f and re-run." >&2 + exit 1 +fi + +exec bash "${SCRIPT_DIR}/../ultra_launch.sh" \ + "policy.megatron_cfg.context_parallel_size=${CONTEXT_PARALLEL_SIZE}" \ + "async_rl.min_groups_for_streaming_train=${STREAM_MIN_GROUPS}" \ + "data_plane.backend=${DP_BACKEND}" \ + "data_plane.simple.num_storage_units=${NUM_STORAGE_UNITS}" \ + "policy.generation.refit_transport=${REFIT_TRANSPORT}" \ + "${VENV_OVERRIDE[@]}" \ + "$@" diff --git a/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_nvcf_disagg_sc.sh b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_nvcf_disagg_sc.sh new file mode 100755 index 00000000000..4e7cfb01c1f --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/launch_3k_nvcf_disagg_sc.sh @@ -0,0 +1,203 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# launch_3k_nvcf_disagg_sc.sh +# +# The primary 3K test arm: SingleController, NVCF-hosted judges, disaggregated +# sandboxes. Judges and tool execution both leave the allocation, so every GPU +# in the job belongs to policy training or generation. +# +# Shape (GB200 NVL72, 4 GPUs/node): +# 256 training + 488 generation + 8 Gym = 752 nodes / 3008 GPUs, one hetgroup +# +# Compared with the v1 baseline (launch_3k_pipeclean.sh, 762 nodes): the 10-node +# external judge hetgroup is gone because GenRM and NL2Bash are hosted, and Gym's +# GPUs go unused because the safety judge is hosted too. Training and the batch +# are identical, so loss and reward curves remain comparable; judge latency and +# tool-execution latency do not, since both are now network round trips. +# +# The backup arm if NVCF proves unusable is launch_3k_disagg_sc.sh, which keeps +# the same SC + disagg stack but serves the judges on the allocation. +# +# NO CHECKPOINTING (inherited from the SC overlay). Size NRL_MAX_STEPS to fit +# one allocation. +# +# Required, in addition to the usual site variables: +# NVIDIA_API_KEY NVCF credential, read by the config as +# ${oc.env:NVIDIA_API_KEY}. Keep it in the environment +# or a sourced creds file — never on a command line. +# OPENSANDBOX_BASE_URL sandbox service endpoint +# OPENSANDBOX_API_KEY sandbox service credential, same handling +# NS_SANDBOX_IMAGE nemo-skills sandbox image the service can pull +# +# Recommended (see disagg_overrides.yaml for sizing): +# NS_SANDBOX_POOL_REF= claim PREWARMED pods from a server-side +# Pool created before the job; set empty +# ("") for direct creates at Gym startup +# NS_SANDBOX_POOL_SIZE=128 keep sessions/pod ≈ CPUs/pod +# +# Judge endpoints and model names live in nvcf_judges.yaml; sandbox backends live +# in disagg_overrides.yaml. Neither is duplicated here, so they cannot drift. +# +# The Gym overlay must provide the sandbox_pool backend (NVIDIA-NeMo/Gym branch +# hemild/rlvr-osb-473f446f (the 6K Gym pin + sandbox commit) until merged). +# +# Optional: NRL_MAX_STEPS, WALLTIME, CONTEXT_PARALLEL_SIZE, DRY_RUN, +# NUM_TRAIN_NODES / NUM_GEN_NODES / NUM_GYM_NODES, STREAM_MIN_GROUPS, +# NUM_STORAGE_UNITS, REFIT_TRANSPORT — as launch_3k_pipeclean_sc.sh. +# +# Extra positional args are forwarded as Hydra overrides to ultra_launch.sh. +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) + +# Default config — callers may still override CONFIG_PATH explicitly. +export CONFIG_PATH="${CONFIG_PATH:-${SCRIPT_DIR}/pipeclean_3k_nvcf_disagg_sc.yaml}" + +# The SC driver. data_plane.enabled=true (set in the config) is mandatory for it. +export TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-./examples/run_grpo_single_controller.py}" + +# ============================================================================= +# NVCF judges +# ============================================================================= +: "${NVIDIA_API_KEY:?NVIDIA_API_KEY is required (NVCF credential read by the config as \${oc.env:NVIDIA_API_KEY})}" +# Deployment-specific NVCF function routes consumed by nvcf_judges.yaml. +: "${GENRM_NVCF_MODEL:?GENRM_NVCF_MODEL is required (NVCF route for the GenRM judge)}" +: "${NL2BASH_NVCF_MODEL:?NL2BASH_NVCF_MODEL is required (NVCF route for the NL2Bash judge)}" +: "${SAFETY_NVCF_MODEL:?SAFETY_NVCF_MODEL is required (NVCF route for the safety judge)}" +export GENRM_NVCF_MODEL NL2BASH_NVCF_MODEL SAFETY_NVCF_MODEL +export NVIDIA_API_KEY + +# Hand the judge wiring entirely to nvcf_judges.yaml. ultra_launch.sh emits a +# `.model=` or `.base_url=` Hydra override for any of these that is non-empty, +# and a CLI override beats the config — it would point a judge back at a local +# checkpoint while the YAML still looked correct. Empty rather than unset also +# defends against values inherited from the caller's shell. +export EXTERNAL_JUDGES=0 +export GENRM_MODEL="" +export NL2BASH_JUDGE_MODEL="" +export SAFETY_JUDGE_MODEL="" +export GENRM_BASE_URL="" +export NL2BASH_BASE_URL="" + +# ============================================================================= +# Disaggregated sandbox service +# ============================================================================= +export NO_COLOCATED_SANDBOX=1 +: "${OPENSANDBOX_BASE_URL:?OPENSANDBOX_BASE_URL is required (sandbox service endpoint)}" +: "${OPENSANDBOX_API_KEY:?OPENSANDBOX_API_KEY is required (sandbox service credential)}" +: "${NS_SANDBOX_IMAGE:?NS_SANDBOX_IMAGE is required (sandbox image the service can pull)}" +export OPENSANDBOX_BASE_URL OPENSANDBOX_API_KEY NS_SANDBOX_IMAGE +# Unset-only default (`-` not `:-`): export NS_SANDBOX_POOL_REF="" forces direct +# creates at Gym startup, matching the header contract. +export NS_SANDBOX_POOL_REF="${NS_SANDBOX_POOL_REF-ns-tools-warm}" +export LEAN_SANDBOX_POOL_REF="${LEAN_SANDBOX_POOL_REF-math-lean-warm}" +# Half of 6K's 256, tracking the halved 8192-sample cohort. +export NS_SANDBOX_POOL_SIZE="${NS_SANDBOX_POOL_SIZE:-128}" + +# ============================================================================= +# Site defaults — the checkpoint and blend of the other 3K arms +# ============================================================================= +export MODEL_PATH="${MODEL_PATH:-/lustre/fsw/portfolios/llmservice/users/jiaqiz/models/ultra_stage2sft_step300}" +export TRAIN_PATH="${TRAIN_PATH:-/lustre/fsw/portfolios/llmservice/users/jiaqiz/data/gym/rl-data-tools/blends/curriculum_v35_inescapable-sawfly.train.efforts0p15_qamathcode.jsonl}" +# The reference runs validate on the training blend; there is no separate split. +export VAL_PATH="${VAL_PATH:-${TRAIN_PATH}}" + +# Image with the Gym venvs prebaked (python interpreter included — the stock +# nightly ships a venv tree whose interpreter is absent) and the sandbox deps +# installed. Pinned to a resolved image, not the moving nightly symlink. +# +# This is the disagg image the 6K disagg recipe uses. It is older than the +# nightly the pipeclean arms pin, but verified equivalent where it matters: its +# /opt/ray_venvs generation workers carry vLLM 0.25.1, which is what this +# branch's code needs (a pre-bump venv fails at import with "cannot import name +# ServingTokenization"). Its /opt/gym_venvs judge venvs are on 0.20.0 — the same +# as the nightly's, and unused on this arm anyway, since NVCF judges mean no +# local vLLM is ever launched. +# +# Byte-identical restripe of hemild/scale-6k/...plus_sandbox_venvs.sqsh. The +# original is laid out over 8 OSTs at 1 MB, and 752 nodes pulling 63 GB through +# 8 OSTs is the I/O concentration that took down an earlier run. This copy +# inherits the 350-OST / 16 MB layout the other arms' image already uses. +export CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/llmservice/users/sauramishra/images-striped/main_ultra_recipes_prebaked_venvs_20260730.plus_sandbox_venvs.sqsh}" +# Part of the sandbox-image delta: the image predates the tree, and ultra_launch +# mounts nemo_rl/ and examples/configs/ but NOT examples/, so both the Gym +# examples and the SC driver have to be bind-mounted in or the image's stale +# copies run against the mounted newer library. +GYM_EXAMPLES_DIR=$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd) +export EXTRA_MOUNTS="${EXTRA_MOUNTS:-/lustre:/lustre,${GYM_EXAMPLES_DIR}:/opt/nemo-rl/examples/nemo_gym,${GYM_EXAMPLES_DIR}/../run_grpo_single_controller.py:/opt/nemo-rl/examples/run_grpo_single_controller.py}" +export PERSISTENT_CACHE="${PERSISTENT_CACHE:-/lustre/fsw/portfolios/llmservice/users/${USER}/.cache/nemotron_ultra}" +export HF_HOME="${HF_HOME:-/lustre/fsw/portfolios/llmservice/users/${USER}/hf_home}" + +export SLURM_ACCOUNT="${SLURM_ACCOUNT:-nemotron_sw_pre}" +export SLURM_PARTITION="${SLURM_PARTITION:-batch}" + +# Training matches every other 3K arm. Gym keeps 8 nodes, as the NVCF reference +# run did: no judge is served here, so their GPUs sit idle, but Gym still needs +# CPU capacity of its own for the resource servers, and borrowing it from the +# training or generation nodes would contend with the thing being measured. +# Generation absorbs the remainder to keep the total on the 16-node segment +# boundary, which costs it 6 nodes against the other two arms. +export NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-256}" +export NUM_GEN_NODES="${NUM_GEN_NODES:-488}" +export NUM_GYM_NODES="${NUM_GYM_NODES:-8}" + +# CP=16 is inherited from pipeclean_6k.yaml; allow an override for memory +# experiments. +CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE:-16}" + +# Same SC knobs as the other SC arms, so they differ only in judge and sandbox +# placement. +STREAM_MIN_GROUPS="${STREAM_MIN_GROUPS:-128}" +NUM_STORAGE_UNITS="${NUM_STORAGE_UNITS:-32}" +REFIT_TRANSPORT="${REFIT_TRANSPORT:-nccl_reshard}" + +# Data-plane backend. mooncake_cpu carries rollouts over CPU RDMA and has no TCP +# fallback — it raises at data-plane setup if no RDMA device is usable, rather +# than degrading silently. DP_BACKEND=simple reverts to the SimpleStorage Ray +# actors, the only backend NUM_STORAGE_UNITS applies to. +DP_BACKEND="${DP_BACKEND:-mooncake_cpu}" + +export EXP_NAME="${EXP_NAME:-ultra-3k-sc-nvcf-disagg}" +export WALLTIME="${WALLTIME:-4:00:00}" + +# Reuse the container's prebaked venvs instead of building in-tree on first run +# (GYM_VENV_DIR= empty restores the stock in-tree behavior). +GYM_VENV_DIR="${GYM_VENV_DIR:-/opt/gym_venvs}" +VENV_OVERRIDE=() +if [ -n "${GYM_VENV_DIR}" ]; then + VENV_OVERRIDE=("++env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR}") +fi + +# Preflight the sandbox dependency this variant introduces: ultra_launch overlays +# $PWD/3rdparty/Gym-workspace/Gym whenever that directory exists (even empty — +# an uninitialized submodule silently masks the container's baked Gym). +GYM_OVERLAY="$PWD/3rdparty/Gym-workspace/Gym" +if [ ! -s "${GYM_OVERLAY}/resources_servers/ns_tools/sandbox_pool.py" ]; then + echo "ERROR: ${GYM_OVERLAY} is missing the sandbox_pool backend (empty or wrong Gym overlay?)." >&2 + echo " Populate it with NVIDIA-NeMo/Gym branch hemild/rlvr-osb-473f446f and re-run." >&2 + exit 1 +fi + +exec bash "${SCRIPT_DIR}/../ultra_launch.sh" \ + "policy.megatron_cfg.context_parallel_size=${CONTEXT_PARALLEL_SIZE}" \ + "async_rl.min_groups_for_streaming_train=${STREAM_MIN_GROUPS}" \ + "data_plane.backend=${DP_BACKEND}" \ + "data_plane.simple.num_storage_units=${NUM_STORAGE_UNITS}" \ + "policy.generation.refit_transport=${REFIT_TRANSPORT}" \ + "${VENV_OVERRIDE[@]}" \ + "$@" diff --git a/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_disagg_sc.yaml b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_disagg_sc.yaml new file mode 100755 index 00000000000..e596f83e3de --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_disagg_sc.yaml @@ -0,0 +1,15 @@ +# ============================================================================= +# Nemotron 3 Ultra — 3K-GPU pipeclean, SingleController + DISAGGREGATED sandboxes +# ============================================================================= +# pipeclean_3k_sc.yaml (streaming forward/backward, TransferQueue data plane, +# checkpointing off) with tool execution on an external sandbox service. The +# disagg fragment is merged AFTER the SC overlay so both sets of deltas apply; +# they touch disjoint keys, and disagg_overrides.yaml has no parent of its own +# to drag the non-SC defaults back in. +# +# Launch with disagg_sandboxes/launch_3k_disagg_sc.sh. See disagg_overrides.yaml +# for requirements, prewarming and sizing notes. +# ============================================================================= +defaults: + - ../pipeclean_3k_sc.yaml + - disagg_overrides.yaml diff --git a/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_nvcf_disagg_sc.yaml b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_nvcf_disagg_sc.yaml new file mode 100755 index 00000000000..3c8fae5fda0 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/disagg_sandboxes/pipeclean_3k_nvcf_disagg_sc.yaml @@ -0,0 +1,25 @@ +# ============================================================================= +# Nemotron 3 Ultra — 3K pipeclean: SingleController + NVCF judges + DISAGG sandboxes +# ============================================================================= +# The fully off-loaded arm: SC execution (streaming forward/backward, +# TransferQueue data plane, checkpointing off), all three judge/reward models +# served from NVCF, and tool execution on an external sandbox service. Nothing +# but policy training and generation runs on the allocation. +# +# Model, parallelism, batch shape, sequence length and the Gym blend are +# inherited unchanged from pipeclean_3k.yaml through the SC overlay, so reward +# and loss curves stay comparable against the v1 baseline. +# +# Both fragments merge AFTER the SC overlay and touch disjoint keys: the NVCF +# fragment rewrites judge base_url/api_key/model, the disagg fragment rewrites +# the ns_tools and math_formal_lean backends. That ordering is safe only because +# neither fragment declares a parent of its own — a second parent inheriting +# pipeclean_6k.yaml would re-enable checkpointing and undo the SC deltas. +# +# Launch with launch_3k_nvcf_disagg_sc.sh, which requires NVIDIA_API_KEY plus +# the sandbox service variables and carries the 752-node split. +# ============================================================================= +defaults: + - ../pipeclean_3k_sc.yaml + - ../nvcf_judges.yaml + - disagg_overrides.yaml diff --git a/examples/nemo_gym/nemotron-3-ultra/nvcf_judges.yaml b/examples/nemo_gym/nemotron-3-ultra/nvcf_judges.yaml new file mode 100755 index 00000000000..37afdfd0c87 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/nvcf_judges.yaml @@ -0,0 +1,62 @@ +# ============================================================================= +# NVCF-hosted judges — mergeable fragment +# ============================================================================= +# Routes all three judge/reward models at NVCF instead of serving them on the +# allocation. Translated from verified-working overrides (13 clean steps on the +# reference run). +# +# The NVCF function routes are deployment-specific: export GENRM_NVCF_MODEL, +# NL2BASH_NVCF_MODEL and SAFETY_NVCF_MODEL alongside NVIDIA_API_KEY. +# +# base_url is the ONLY switch. When it is truthy Gym skips the local vLLM launch +# entirely (responses_api_models/local_vllm_model/app.py:194), and every +# vllm_serve_kwargs / vllm_serve_env_vars setting on the block below it becomes +# inert — inherited, but never read. +# +# This file has NO `defaults` key on purpose. It is a fragment, so a recipe can +# list it as a second parent without dragging a base config in behind it and +# clobbering the first parent's deltas. +# +# Four things that bite: +# +# 1. The nesting is asymmetric. GenRM's inner key repeats the outer name +# (genrm_model.responses_api_models.genrm_model); the other two use +# local_vllm_model. A wrong key silently creates an unused block instead of +# erroring — the only symptom is that the run serves a local vLLM anyway. +# 2. No `_override_: true`. That directive replaces a block rather than merging +# into it, which would discard the judge's entrypoint, +# uses_reasoning_parser, ray_worker_py_executable and the rest. Default deep +# merge is what is wanted here. (Overlaying onto a base whose genrm_model +# block already carries _override_: true, such as rlvr_dolphin.yaml, merges +# into that already-replaced block instead of the Ultra base.) +# 3. ${oc.env:NVIDIA_API_KEY} is written plainly, with no backslash. YAML is +# not shell-expanded. The escaping that a launcher needs — an earlier attempt died +# with "bad substitution" when bash tried to expand it inside +# driver_command.sh — would land in the literal string here and break it. +# The caller must export NVIDIA_API_KEY; it reaches the in-container driver +# the same way WANDB_API_KEY does. +# 4. Shrink the Gym GPU allocation to match, or the freed nodes are wasted. +# That is the companion change this file cannot express: with the judges +# hosted, nothing on those GPUs is serving them. The 3K NVCF wrapper sets +# NUM_GYM_NODES=0 for exactly this reason. +# ============================================================================= +env: + nemo_gym: + genrm_model: + responses_api_models: + genrm_model: # note: nested key repeats the block name + base_url: ["https://integrate.api.nvidia.com/v1"] + api_key: ${oc.env:NVIDIA_API_KEY} + model: ${oc.env:GENRM_NVCF_MODEL} + nl2bash_judge_model: + responses_api_models: + local_vllm_model: # note: NOT nl2bash_judge_model + base_url: ["https://integrate.api.nvidia.com/v1"] + api_key: ${oc.env:NVIDIA_API_KEY} + model: ${oc.env:NL2BASH_NVCF_MODEL} + safety_judge_model: + responses_api_models: + local_vllm_model: + base_url: ["https://integrate.api.nvidia.com/v1"] + api_key: ${oc.env:NVIDIA_API_KEY} + model: ${oc.env:SAFETY_NVCF_MODEL} diff --git a/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k.yaml b/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k.yaml new file mode 100755 index 00000000000..583d451d079 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k.yaml @@ -0,0 +1,36 @@ +# ============================================================================= +# Nemotron 3 Ultra — 3K-GPU pipeclean +# ============================================================================= +# Half-scale sibling of pipeclean_6k.yaml. Model, parallelism, sequence length, +# loss and the Gym blend are inherited unchanged; only the cohort size changes, +# because the node split halves with it (see launch_3k_pipeclean.sh). +# +# Why the batch halves: training drops from 512 nodes to 256, so DP falls from +# 16 to 8. Halving the cohort keeps per-rank work identical to the 6K run — +# 8192 / 8 = 1024 microbatches per optimizer step, the same 1024 that 6K +# accumulates at 16384 / 16 — so step time and activation memory behave as they +# do at 6K. The optimization is NOT 6K's: an 8192-sample update is a different +# update, so read convergence curves as their own experiment rather than as a +# 6K reproduction. +# +# Parallelism is deliberately untouched (TP=8, EP=64, ETP=1, PP=1, CP=16), and +# 1024 training GPUs satisfies both divisibility rules with room to spare: DP +# needs world_size % (TP * CP * PP) == 0, i.e. 128-GPU (32-node) granularity, +# giving DP=8; the expert grid needs world_size % (ETP * EP * PP) == 0, i.e. 64 +# GPUs, giving expert DP = 1024 / 64 = 16. CP does not enter the second rule — +# Megatron sizes expert DP straight off world_size (parallel_state.py) — so it is +# the 32-node step, not EP=64, that binds if training shrinks further. Lowering +# CONTEXT_PARALLEL_SIZE relaxes that step. +# +# checkpointing.checkpoint_dir is inherited but inert: ultra_launch.sh always +# overrides it from EXP_NAME, so 3K runs never land in the 6K results tree. +# ============================================================================= +defaults: pipeclean_6k.yaml + +grpo: + num_prompts_per_step: 512 + +policy: + # num_prompts_per_step * num_generations_per_prompt (512 * 16). Nothing + # interpolates this, so both keys have to move together. + train_global_batch_size: 8192 diff --git a/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k_sc.yaml b/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k_sc.yaml new file mode 100755 index 00000000000..db05f8955aa --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/pipeclean_3k_sc.yaml @@ -0,0 +1,198 @@ +# ============================================================================= +# Nemotron 3 Ultra — 3K-GPU pipeclean on the SingleController path +# ============================================================================= +# Overlay on pipeclean_3k.yaml that swaps the legacy async-GRPO driver for the +# SingleController (SC) architecture: streaming forward/backward, the +# TransferQueue data plane, gated in-order rollout sampling and shard-to-shard +# weight refit. Model, parallelism, batch shape, sequence length and the Gym +# blend come from pipeclean_3k.yaml, so a run of this config is comparable +# against a pipeclean_3k.yaml run on the same node split. +# +# Entrypoint: examples/run_grpo_single_controller.py (NOT run_grpo_nemo_gym.py). +# launch_3k_pipeclean_sc.sh sets that via TRAIN_ENTRYPOINT. +# +# Why the SC deltas are restated here instead of inheriting pipeclean_6k_sc.yaml: +# `defaults` merges each parent's fully resolved tree in order, and every path to +# a 3K batch runs back through pipeclean_6k.yaml. Listing pipeclean_6k_sc.yaml +# alongside pipeclean_3k.yaml would therefore re-apply pipeclean_6k.yaml's +# checkpointing.enabled: true and gpu_memory_utilization: 0.85 on top of the SC +# deltas, silently re-enabling checkpointing (which SC rejects outright). The +# disagg fragments can be merged that way only because disagg_overrides.yaml has +# no parent of its own. +# +# Group vs sample units: async_rl counts PROMPT GROUPS, not generations. One +# optimizer cohort here is 8192 / 16 = 512 groups, half of 6K's 1024. +# ============================================================================= +defaults: pipeclean_3k.yaml + +# ============================================================================= +# Checkpointing — must be off +# ============================================================================= +# setup_single_controller raises NotImplementedError when checkpointing.enabled +# is true (single_controller_utils/setup.py). Saving is still a TODO on the SC +# path, so the run cannot resume across the wall clock: size max_num_steps +# (NRL_MAX_STEPS) to fit one allocation. +# ============================================================================= +checkpointing: + enabled: false + +grpo: + # Required because loss_fn.reference_policy_kl_penalty is 0.0. No reference + # model is built, but SC still computes reference logprobs unless this is set, + # and then fails with AttributeError: 'MegatronPolicyWorker' has no + # 'reference_state_dict'. + skip_reference_policy_logprobs_calculation: true + + # val_period=-1 / val_at_start=false / val_at_end=false are inherited and + # already satisfy the SC guard (it rejects val_period > 0). + # + # grpo.async_grpo is inherited and dead on this path — SC never reads it — but + # env.nemo_gym interpolates ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # for num_groups_nemo_rl, so removing the block would break config load. + +policy: + # --------------------------------------------------------------------------- + # Speculative-decoding draft head. run_grpo_single_controller.py reads + # config.policy["draft"]["enabled"] unconditionally, so omitting this block is + # a KeyError before Ray starts. The nemo_gym runner guards it with an `in` + # check, which is why pipeclean_3k.yaml does not need it. + # --------------------------------------------------------------------------- + draft: + enabled: false + model_name: null + loss_weight: 0.1 + num_layers: null + aux_layer_indices: null + + generation: + # ------------------------------------------------------------------------- + # Weight refit transport. The inherited default (null) broadcasts every full + # parameter tensor from training to all 1976 generation ranks. nccl_reshard + # instead moves each parameter shard-to-shard between the two parallelism + # layouts over NCCL's M2N reshard, so no rank ever materializes a full + # tensor. The bulk path covers the MoE FFN projections, which are 97-98% of + # the weights at EP=64. + # + # Declared here rather than passed with `+` because refit_transport is + # NotRequired in VllmConfig; with the key present, REFIT_TRANSPORT=null in + # launch_3k_pipeclean_sc.sh falls back to the broadcast path. + # + # It degrades quietly: if the container lacks the nccl4py M2N integration, + # xferdtensor.py logs "nccl.m2n.reshard not found" and uses a Python + # implementation instead. Check the first run's logs for that line. + # + # The config satisfies every constraint in + # docs/design-docs/nccl-reshard-refit.md: non-colocated, Megatron train, + # vLLM gen, ETP=1, PP=1 with no custom layout, BF16 end to end, and vLLM + # expert_parallel_size == tensor_parallel_size (8). + # ------------------------------------------------------------------------- + refit_transport: nccl_reshard + + vllm_cfg: + # 0.85 -> 0.80, matching pipeclean_6k_sc.yaml. The NCCL staging buffer for + # update_weights_from_collective grows with generation rank count; a + # smaller-scale run OOM'd there at 0.85 with 128 generation ranks after + # running fine at 32. This recipe has 1976. Costs KV cache, so watch + # generation throughput. + # + # This is the pressure refit_transport: nccl_reshard removes, so 0.85 may + # be recoverable once a reshard run is confirmed healthy. Left at 0.80 so + # the v1 and SC arms differ only in the transport. + gpu_memory_utilization: 0.80 + +# ============================================================================= +# SingleController async RL — required by MasterConfig +# ============================================================================= +# in_order stamps each dispatched cohort with its exact target optimizer step. +# The lookahead is the ceiling on off-policyness: generation may start a cohort +# at most this many optimizer steps before the step that trains on it, so the +# reported `lag` settles there once generation is the bottleneck. +# +# 1, not the 4 the 6K recipe uses, and every 3K arm shares it so that lag is a +# constant of this campaign rather than a per-arm variable: the NVCF and disagg +# arms inherit this block untouched, and the ready_first arm restates the same +# depth under its own key. It is also the depth validated at small scale, and the +# shallowest that still overlaps generation with training. +# +# The loss does not correct what lag remains: force_on_policy_ratio is true here, +# which pins the PPO ratio at 1.0 and leaves use_importance_sampling_correction +# and truncated_importance_sampling_ratio inert. That is tolerable at a lookahead +# of 1 and is the reason to prefer it, but it means token_mult_prob_error and +# is_oob_ratio are not measuring a live correction on this arm. +# ============================================================================= +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + + # Begin forward/backward once this many groups are ready rather than waiting + # for all 512; the remainder streams in behind it and accumulates into the + # single GBS=8192 optimizer step. 128 is 25% of the cohort, the same ratio the + # 6K recipe uses (256 of 1024), so a step dispatches in ~4 chunks. + min_groups_for_streaming_train: 128 + + # The current cohort plus its single lookahead cohort. The gate stops + # generation at trainer_version + 1, so anything above 512 * 2 is unreachable + # rather than merely generous. This is also the NemoGym actor's rollout fan-in, + # which its Ray max_concurrency is sized from. + max_inflight_prompts: 1024 + + # Hard floor, not a preference: validate_sampler_buffer_capacity raises if this + # is below num_prompts_per_step * (max_lookahead_versions + 1) = 512 * 2 = + # 1024, because the rollout pump would deadlock waiting for buffer slots. It + # scales with the lookahead, so raising one raises both. + max_buffered_rollouts: 1024 + + diagnostics: false + +# ============================================================================= +# TransferQueue data plane — mandatory for the SC path, over CPU RDMA +# ============================================================================= +# run_grpo_single_controller.py refuses to start unless enabled is true. +# +# backend=mooncake_cpu moves rollout tensors over RDMA instead of through the +# SimpleStorage Ray actors. There is no TCP mode and no fallback: if no usable +# RDMA device is found the run raises at data-plane setup rather than quietly +# degrading, so a misconfigured fabric fails in seconds instead of showing up +# as unexplained step time. Set DP_BACKEND=simple on the launcher to go back. +# +# Sizing is PER CLIENT PROCESS (one per GPU), and under RDMA it is pinned and +# resident from setup — not lazily faulted in. At 4 GPUs/node this recipe pins +# 4 x (64 + 4) = 272 GiB of the 952 GiB these nodes carry. Under-sizing surfaces +# as "batch_get_tensor returned None", so raise it if that appears rather than +# guessing upward now. +# +# Storage-unit sharding does not apply to this backend: MooncakeStore has no +# unit count and no sample cap, sizing itself entirely from the two values +# below. num_storage_units and storage_capacity therefore live in the simple: +# block, which is read only when that backend is selected. +# +# On this fleet each node exposes 4 InfiniBand rails and 2 Ethernet rails; the +# adapter prefers InfiniBand and offers all 4, never mixing the two fabrics. +# ============================================================================= +data_plane: + enabled: true + impl: transfer_queue + backend: "mooncake_cpu" + claim_meta_poll_interval_s: 0.5 + # Read only under backend=simple, kept so falling back is a one-word change + # and keeps the sharding this recipe was tuned with. + simple: + storage_capacity: 1000000 + # Half of the 6K recipe's 64, tracking the halved per-step volume (8192 rows + # instead of 16384, at the same 131072 tokens). Like that value it is + # proportional rather than measured — raise it first if the TQ storage actors + # bottleneck or OOM. + num_storage_units: 32 + mooncake_cpu: + global_segment_size: 68719476736 # 64 GiB per client process + local_buffer_size: 4294967296 # 4 GiB per client process + # Keep RDMA-registered buffers alive across transfers instead of + # registering one per call. Set false to fall back to upstream's behavior. + reuse_registered_buffers: true + # Pooling threshold, not a cap: a larger payload still transfers, with a + # transient registration. Slots ratchet up to the largest payload admitted + # and never shrink, so raise this only if a single sample of a single field + # genuinely exceeds it. + staging_buffer_size: 268435456 # 256 MiB per pool slot diff --git a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh index dc2dd439b99..7992db6bb03 100755 --- a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh +++ b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh @@ -136,7 +136,21 @@ fi # ============================================================================= # Project root and code root # ============================================================================= -PROJECT_ROOT=$(realpath "$PWD") +if [[ -n "${PROJECT_ROOT_OVERRIDE:-}" ]]; then + if [[ "${PROJECT_ROOT_OVERRIDE}" != /* || "${PROJECT_ROOT_OVERRIDE}" == *[,[:space:]:]* ]]; then + echo "ERROR: PROJECT_ROOT_OVERRIDE must be a plain absolute path" >&2 + exit 1 + fi + PROJECT_ROOT="${PROJECT_ROOT_OVERRIDE}" +else + PROJECT_ROOT=$(realpath "$PWD") +fi +SLURM_WORKDIR="${SLURM_WORKDIR:-${PROJECT_ROOT}}" +if [[ "${SLURM_WORKDIR}" != /* || "${SLURM_WORKDIR}" == *[,[:space:]:]* ]]; then + echo "ERROR: SLURM_WORKDIR must be a plain absolute path" >&2 + exit 1 +fi +export SLURM_WORKDIR cd "${PROJECT_ROOT}" # ============================================================================= @@ -1003,6 +1017,7 @@ if [[ "${INTERACTIVE}" == "1" ]]; then SBATCH_OUTPUT=$(sbatch \ --nodes="${NUM_TOTAL_NODES}" \ + --chdir="${SLURM_WORKDIR}" \ --account="${SLURM_ACCOUNT}" \ --job-name="interactive-${JOB_NAME}" \ --partition="${SLURM_PARTITION}" \ @@ -1066,6 +1081,7 @@ DEPENDENCY="singleton" if (( NUM_EXTERNAL_SERVICE_NODES > 0 )); then SBATCH_OUTPUT=$(sbatch \ --nodes="${NUM_RAY_NODES}" \ + --chdir="${SLURM_WORKDIR}" \ --account="${SLURM_ACCOUNT}" \ --job-name="${JOB_NAME}" \ --partition="${SLURM_PARTITION}" \ @@ -1083,6 +1099,7 @@ if (( NUM_EXTERNAL_SERVICE_NODES > 0 )); then --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"other","description":"batch training run"}}' \ : \ --nodes="${NUM_EXTERNAL_SERVICE_NODES}" \ + --chdir="${SLURM_WORKDIR}" \ --account="${SLURM_ACCOUNT}" \ --job-name="${JOB_NAME}-services" \ --partition="${SLURM_PARTITION}" \ @@ -1098,6 +1115,7 @@ if (( NUM_EXTERNAL_SERVICE_NODES > 0 )); then else SBATCH_OUTPUT=$(sbatch \ --nodes="${NUM_TOTAL_NODES}" \ + --chdir="${SLURM_WORKDIR}" \ --account="${SLURM_ACCOUNT}" \ --job-name="${JOB_NAME}" \ --partition="${SLURM_PARTITION}" \ diff --git a/ray.sub b/ray.sub index 91d20a2cf9b..e5711531cd4 100644 --- a/ray.sub +++ b/ray.sub @@ -24,6 +24,9 @@ set -eoux pipefail +# Slurm records the submit host's physical cwd, which may not exist on workers. +SLURM_WORKDIR="${SLURM_WORKDIR:-${SLURM_SUBMIT_DIR}}" + # Record job-start epoch so Python can measure pre-Python overhead export NRL_JOB_START_EPOCH=$(date +%s.%N) @@ -98,7 +101,7 @@ detect_cpus_per_node() { export OMP_NUM_THREADS=${OMP_NUM_THREADS:-16} CONTAINER=$CONTAINER MOUNTS=$MOUNTS -COMMAND=${COMMAND:-} # This is a script relative to the SLURM_SUBMIT_DIR. If left empty, it will leave the cluster idle after it's brought up. +COMMAND=${COMMAND:-} # This is a script relative to SLURM_WORKDIR. If left empty, it will leave the cluster idle after it's brought up. SETUP_COMMAND=${SETUP_COMMAND:-} # Setup commands to run on all nodes before starting Ray. ######################################################## # Ports for all nodes (should be odd numbers since we place head/worker[0] on the same node) so all workers get the odd ports, but the head will get +1 the ports @@ -190,7 +193,7 @@ fi # NOTE: LOG_DIR must be on a shared filesystem visible to both the submit host and # all compute nodes. File-based signaling between the submit host and the containers # (STARTED_RAY_HEAD, ray_worker_units, ENDED, sandbox readiness) relies on this. -BASE_LOG_DIR=${BASE_LOG_DIR:-$SLURM_SUBMIT_DIR} +BASE_LOG_DIR=${BASE_LOG_DIR:-$SLURM_WORKDIR} # On a Slurm requeue/restart, suffix the restart count so we don't clobber the # previous attempt's logs (and signal files) under the same job ID. if [[ -n "${SLURM_RESTART_COUNT:-}" ]]; then @@ -249,7 +252,7 @@ COMMON_SRUN_ARGS+=" --no-container-mount-home" COMMON_SRUN_ARGS+=" --mpi=pmix" COMMON_SRUN_ARGS+=" --container-mounts=$MOUNTS" COMMON_SRUN_ARGS+=" --container-image=$CONTAINER" -COMMON_SRUN_ARGS+=" --container-workdir=$SLURM_SUBMIT_DIR" +COMMON_SRUN_ARGS+=" --container-workdir=$SLURM_WORKDIR" # Pass partition/account explicitly for overlapping srun calls. COMMON_SRUN_ARGS+=" -p $SLURM_JOB_PARTITION" COMMON_SRUN_ARGS+=" -A $SLURM_JOB_ACCOUNT" @@ -265,6 +268,14 @@ else CPUS_PER_WORKER=$(detect_cpus_per_node) echo "[INFO] Auto-detected CPUS_PER_WORKER=$CPUS_PER_WORKER from SLURM (CPUTot of first allocated node)." fi +# Clusters with reserved system cores (CoreSpecCount) allocate fewer CPUs than +# the node physically has; a step asking for more than the allocation blocks +# in "step creation temporarily disabled" retries forever. +if [[ "${SLURM_CPUS_ON_NODE:-}" =~ ^[0-9]+$ ]] && (( SLURM_CPUS_ON_NODE > 0 )) \ + && (( CPUS_PER_WORKER > SLURM_CPUS_ON_NODE )); then + echo "[INFO] Clamping CPUS_PER_WORKER=$CPUS_PER_WORKER to the allocation's $SLURM_CPUS_ON_NODE CPUs per node." + CPUS_PER_WORKER=$SLURM_CPUS_ON_NODE +fi # --------------------------------------------------------------------------- # GCS burst-tolerance tuning for larger jobs @@ -1057,7 +1068,7 @@ fi # Start the Ray head. log_phase "launching Ray head + worker sruns" -srun $COMMON_SRUN_ARGS --container-name=ray-head --nodes=1 --ntasks=1 --cpus-per-task=$CPUS_PER_WORKER -w "$head_node" -o $LOG_DIR/ray-head.log bash -x -c "$head_cmd" & +srun $COMMON_SRUN_ARGS --container-name=ray-head --overlap --nodes=1 --ntasks=1 --cpus-per-task=$CPUS_PER_WORKER -w "$head_node" -o $LOG_DIR/ray-head.log bash -x -c "$head_cmd" & SRUN_PIDS["ray-head"]=$! # Start the Ray workers (all nodes except the head) with a single batched srun, so @@ -1140,7 +1151,7 @@ else log_phase "generating attach helper (scontrol show job)" CONTAINER_CWD=$(scontrol show job $SLURM_JOB_ID | grep -oP 'WorkDir=\K[^ ]+' | head -1) echo "[INFO]: Ray Cluster is idled, run this on the slurm head node to get a shell to the head node:" - cat <$SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh + cat <$SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh # No args launches on the head node (node 0) # Args 1-N launch on worker nodes (nodes 1 through N-1) # Optional: set COMMAND='...' to run non-interactively instead of opening an interactive shell @@ -1170,12 +1181,12 @@ else fi fi EOF - chmod +x $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh - echo " COMMAND='echo hello' bash $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh # run a non-interactive command on head node" - echo " bash $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh # to attach to head node (i.e., 'worker 0')" + chmod +x $SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh + echo " COMMAND='echo hello' bash $SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh # run a non-interactive command on head node" + echo " bash $SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh # to attach to head node (i.e., 'worker 0')" if [[ $SLURM_JOB_NUM_NODES -gt 1 ]]; then - echo " bash $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh 1 # to attach to worker 1" - echo " bash $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh 2 # to attach to worker 2, etc." + echo " bash $SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh 1 # to attach to worker 1" + echo " bash $SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh 2 # to attach to worker 2, etc." fi sleep infinity fi diff --git a/tests/unit/tools/test_external_gym_vllm.py b/tests/unit/tools/test_external_gym_vllm.py index 677fc6e1211..d2b917e529d 100644 --- a/tests/unit/tools/test_external_gym_vllm.py +++ b/tests/unit/tools/test_external_gym_vllm.py @@ -867,3 +867,41 @@ def test_submission_validation_checks_placeholders_paths_and_node_total(): assert ( "skipping external hetgroup node-count validation" in missing_node_count.stderr ) + + +def test_slurm_workdir_controls_batch_host_paths(): + ray_sub = (REPO_ROOT / "ray.sub").read_text() + wrapper = (REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh").read_text() + fallback = 'SLURM_WORKDIR="${SLURM_WORKDIR:-${SLURM_SUBMIT_DIR}}"' + + assert fallback in ray_sub + assert "BASE_LOG_DIR=${BASE_LOG_DIR:-$SLURM_WORKDIR}" in ray_sub + assert "--container-workdir=$SLURM_WORKDIR" in ray_sub + assert "$SLURM_WORKDIR/${SLURM_JOB_ID}-attach.sh" in ray_sub + assert fallback in wrapper + assert 'RAY_SUB="${RAY_SUB:-${SLURM_WORKDIR}/ray.sub}"' in wrapper + assert '--container-workdir="${SLURM_WORKDIR}"' in wrapper + + +def test_service_steps_do_not_hold_ray_head_resources(): + ray_sub = (REPO_ROOT / "ray.sub").read_text() + wrapper = (REPO_ROOT / "tools/external_gym_vllm/run_in_allocation.sh").read_text() + + head_line = next( + line + for line in ray_sub.splitlines() + if "--container-name=ray-head" in line and "$COMMON_SRUN_ARGS" in line + ) + assert "--overlap" in head_line + + lb_flags = wrapper.split('--container-name="external-vllm-lb-', 1)[1] + lb_flags = lb_flags.split("bash -lc", 1)[0] + assert "--overlap" in lb_flags + assert "--gres=gpu:0" in lb_flags + assert "--mem=16G" in lb_flags + + +def test_cpus_per_worker_is_clamped_to_the_allocation(): + ray_sub = (REPO_ROOT / "ray.sub").read_text() + assert "(( CPUS_PER_WORKER > SLURM_CPUS_ON_NODE ))" in ray_sub + assert "CPUS_PER_WORKER=$SLURM_CPUS_ON_NODE" in ray_sub diff --git a/tools/external_gym_vllm/run_in_allocation.sh b/tools/external_gym_vllm/run_in_allocation.sh index 19020b74bbf..3d4c14454c3 100755 --- a/tools/external_gym_vllm/run_in_allocation.sh +++ b/tools/external_gym_vllm/run_in_allocation.sh @@ -26,6 +26,9 @@ set -euo pipefail : "${SLURM_JOB_ACCOUNT:?SLURM_JOB_ACCOUNT is required}" : "${SLURM_JOB_PARTITION:?SLURM_JOB_PARTITION is required}" : "${SLURM_SUBMIT_DIR:?SLURM_SUBMIT_DIR is required}" + +# Slurm records the submit host's physical cwd, which may not exist on workers. +SLURM_WORKDIR="${SLURM_WORKDIR:-${SLURM_SUBMIT_DIR}}" : "${BASE_LOG_DIR:?BASE_LOG_DIR is required}" : "${CONTAINER:?CONTAINER is required}" : "${MOUNTS:?MOUNTS is required}" @@ -38,7 +41,7 @@ if [[ "${SLURM_HET_SIZE}" != "2" ]]; then exit 1 fi -RAY_SUB="${RAY_SUB:-${SLURM_SUBMIT_DIR}/ray.sub}" +RAY_SUB="${RAY_SUB:-${SLURM_WORKDIR}/ray.sub}" export GPUS_PER_NODE="${GPUS_PER_NODE:-4}" EXTERNAL_VLLM_LB_PYTHON="${EXTERNAL_VLLM_LB_PYTHON:-/opt/nemo_rl_venv/bin/python}" EXTERNAL_VLLM_SHARED_ROOT="${EXTERNAL_VLLM_SHARED_ROOT:-/lustre}" @@ -553,11 +556,13 @@ else --no-container-mount-home \ --container-image="${CONTAINER}" \ --container-mounts="${lb_mounts}" \ - --container-workdir="${SLURM_SUBMIT_DIR}" \ + --container-workdir="${SLURM_WORKDIR}" \ --mpi=pmix \ -A "${SLURM_JOB_ACCOUNT}" \ -p "${SLURM_JOB_PARTITION}" \ --overlap \ + --gres=gpu:0 \ + --mem=16G \ --kill-on-bad-exit=1 \ --nodelist="${ray_head_node}" \ --nodes=1 \ @@ -621,17 +626,21 @@ ray_head_ip=$(resolve_node_ip "${ray_head_node}") for pool in "${pool_names[@]}"; do pool_urls["${pool}"]="http://${ray_head_ip}:${lb_ports[${pool}]}/v1" echo "[INFO] Starting ${display_names[${pool}]} load balancer at ${pool_urls[${pool}]}" + # Without explicit --gres/--mem a step inherits the node's full GPU and + # memory allocation, which blocks the later ray-head step on this node. srun \ --het-group=0 \ --no-container-mount-home \ --container-name="external-vllm-lb-${pool,,}-${SLURM_JOB_ID}" \ --container-image="${CONTAINER}" \ --container-mounts="${lb_mounts}" \ - --container-workdir="${SLURM_SUBMIT_DIR}" \ + --container-workdir="${SLURM_WORKDIR}" \ --mpi=pmix \ -A "${SLURM_JOB_ACCOUNT}" \ -p "${SLURM_JOB_PARTITION}" \ --overlap \ + --gres=gpu:0 \ + --mem=16G \ --nodelist="${ray_head_node}" \ --nodes=1 \ --ntasks=1 \