diff --git a/examples/deepeyes_v2_agentic/run_deepeyes_v2_agentic_klx.sh b/examples/deepeyes_v2_agentic/run_deepeyes_v2_agentic_klx.sh new file mode 100644 index 000000000..530d5bd34 --- /dev/null +++ b/examples/deepeyes_v2_agentic/run_deepeyes_v2_agentic_klx.sh @@ -0,0 +1,318 @@ +#!/bin/bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -ex +set -o pipefail +now=$(date "+%Y-%m-%d-%H:%M:%S") +echo "当前时间: $now" + +############################################################################### +# klx args # +############################################################################### +export WANDB_API_KEY="${WANDB_API_KEY:=YOUR-KEY}" +export WORKDIR="${WORKDIR:-/workspace}" +export MODEL_DIR="${MODEL_DIR:-/workspace/deepeyes}" +export DATA_DIR="${DATA_DIR:-/workspace/deepeyes}" +export SAVE_DIR="${SAVE_DIR:-/workspace/deepeyes}" +export PROJECT_NAME=deepeyes_v2_agentic + +export MEGATRON=${WORKDIR}/Megatron-LM + +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export RL_MODEL_TYPE="Qwen36_35B_multimodal" +export XMLIR_USE_HYDRA_LINEAR=${XMLIR_USE_HYDRA_LINEAR:-1} +export XMLIR_ENABLE_FAST_FC=${XMLIR_ENABLE_FAST_FC:-1} +export XMLIR_MATMUL_FAST_MODE=${XMLIR_MATMUL_FAST_MODE:-1} +export XMLIR_MEMCPY_RETRY_SYNC=${XMLIR_MEMCPY_RETRY_SYNC:-true} + +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-eth0} +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-eth0} +export TP_SOCKET_IFNAME=${TP_SOCKET_IFNAME:-eth0} +export BKCL_RDMA_NICS=${BKCL_RDMA_NICS:-bond0,bond1,bond2,bond3,bond4,bond5,bond6,bond7} + +unset http_proxy +unset https_proxy + +############################################################################### +# ENVIRONMENT # +############################################################################### + +TIMESTAMP=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# Auto-source env.sh if present (gitignored, machine-specific overrides). +# shellcheck source=/dev/null +[ -f "${SCRIPT_DIR}/env.sh" ] && source "${SCRIPT_DIR}/env.sh" + +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../scripts/entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen36-35B-A3B.sh" + +############################################################################### +# DIRS # +############################################################################### + +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/deepeyes-v2}" +EXP_NAME="qwen36-35B-A3B-deepeyes-v2-agentic-${TIMESTAMP}" + +if [ -z "${MODEL_DIR:-}" ] || [ -z "${DATA_DIR:-}" ] || [ -z "${SAVE_DIR:-}" ]; then + echo "ERROR: MODEL_DIR, DATA_DIR, and SAVE_DIR must be set." + echo "Example: MODEL_DIR=/path/to/models DATA_DIR=/path/to/data SAVE_DIR=/path/to/save bash $0" + exit 1 +fi +mkdir -p ${SAVE_DIR} + +############################################################################### +# STARTUP CLEANUP # +############################################################################### +# Sessions that die by SIGKILL / OOM / uncaught crash never run +# ApptainerJupyterSession.close(), so their /tmp/relax-apptainer-* dirs +# (kernel conn files + bind-mounted imgs/inputs holding base64-decoded +# dataset images) leak. Over 2000 rollouts × 32 batch × 8 samples ≈ 512K +# sessions, even a 1% leak fills the container disk and triggers eviction. +# mtime > 240 min is well above any single session's max wall-clock (bounded +# by max_turns × code_timeout_s ≈ tens of minutes), so live sessions are +# never touched. +find /tmp -maxdepth 1 -name 'relax-apptainer-*' -mmin +240 -exec rm -rf {} + 2>/dev/null || true +# Per-run agent log dirs accumulate one file per session; keep 3 days. +find "${SCRIPT_DIR}/log/agent" -maxdepth 1 -type d -mtime +3 -exec rm -rf {} + 2>/dev/null || true + +# SIF path under DATA_DIR layout; user may override APPTAINER_IMAGE_PATH to +# point at a shared NFS copy elsewhere. +export APPTAINER_IMAGE_PATH="${APPTAINER_IMAGE_PATH:-${DATA_DIR}/sif/deepeyes_v2_kernel.sif}" +if [ ! -f "${APPTAINER_IMAGE_PATH}" ]; then + echo "ERROR: SIF not found at ${APPTAINER_IMAGE_PATH}." + echo "Run: DATA_DIR=${DATA_DIR} bash ${SCRIPT_DIR}/scripts/prepare.sh" + exit 1 +fi + +############################################################################### +# JUDGE MODEL API # +############################################################################### + +source "${SCRIPT_DIR}/sglang_judge_service_klx.sh" + +############################################################################### +# MODEL CONFIG # +############################################################################### + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.6-35B-A3B + --ref-load ${MODEL_DIR}/Qwen3.6-35B-A3B + --warm-hf-checkpoint-page-cache + --save ${SAVE_DIR}/Qwen3.6-35B-A3B-Checkpoint_v2 + --megatron-to-hf-mode bridge + --save-interval 200 + --max-actor-ckpt-to-keep 1 +) + +############################################################################### +# DATASETS # +############################################################################### + +# Layout produced by scripts/prepare.sh: +# ${DATA_DIR}/data/{perception_all_*,reason,search,vstar_test}.parquet +# ${DATA_DIR}/sif/deepeyes_v2_kernel.sif +TRAIN_FILES=( + "'${DATA_DIR}/data/perception_all_1.parquet@[0:5000]'" + "'${DATA_DIR}/data/reason.parquet@[0:5000]'" +) +TEST_FILES=("${DATA_DIR}/data/vstar_test.parquet@[0:256]") +PROMPT_SET="[$(IFS=,; echo "${TRAIN_FILES[*]}")]" + +############################################################################### +# ROLLOUT CONFIG # +############################################################################### + +NUM_ROLLOUT="${NUM_ROLLOUT:=200}" + +# Sandbox env vars propagated into every Ray worker so the per-session +# agent process can find apptainer / search cache. +# SANDBOX_CONFIG_PATH is required — the agent reads it in _build_executor +# to find the apptainer backend YAML config (image path, bind paths, etc). +EXTRA_ENV_VARS_JSON="\"SANDBOX_BACKEND\": \"apptainer_jupyter\", + \"SANDBOX_CONFIG_PATH\": \"${SCRIPT_DIR}/apptainer_env/apptainer_config.yaml\", + \"APPTAINER_IMAGE_PATH\": \"${APPTAINER_IMAGE_PATH}\", + \"DEEPEYES_V2_SEARCH_CACHE_PATHS\": \"${DEEPEYES_V2_SEARCH_CACHE_PATHS:-}\", + \"DEEPEYES_JUDGE_BASE_URL\": \"${DEEPEYES_JUDGE_BASE_URL:-}\", + \"DEEPEYES_JUDGE_MODELS\": \"${DEEPEYES_JUDGE_MODELS:-}\", + \"DEEPEYES_JUDGE_API_KEY\": \"${DEEPEYES_JUDGE_API_KEY:-}\", + \"XMLIR_ENABLE_H2D_SSE_COPY\": \"${XMLIR_ENABLE_H2D_SSE_COPY:-1}\", + \"USE_CAST_FC_FUSION\": \"${USE_CAST_FC_FUSION:-1}\"" +source "${SCRIPT_DIR}/../../scripts/entrypoint/runtime-env-klx.sh" + +ROLLOUT_ARGS=( + --prompt-data "${PROMPT_SET}" + --input-key prompt + --label-key reward_model + --multimodal-keys '{"image":"images"}' + --reward-key score + --metadata-key extra_info + --custom-rm-path examples.deepeyes_v2_agentic.reward_deepeyes_v2.reward_func + --use-agentic-rollout + --agent-command ". ${SCRIPT_DIR}/run_agent_app.sh" + --agent-cwd "${SCRIPT_DIR}" + # Per-run agent log dir: every session's stdout/stderr is tee'd to + # ${dir}/${session_id}.log (run_agent_app.sh). Successful AND failed + # sessions are both kept, which is what makes hang/timeout diagnosis + # possible — Relax's own tmpdir-based capture drops both. + --agent-env "AGENT_DEBUG_LOG_DIR=${SCRIPT_DIR}/log/agent/${TIMESTAMP}" + # 30-min default is too generous: normal sessions take 1-3 min, a + # zombie session still burning chat completions after 10 min is + # ~always doomed. Faster session-level SIGKILL clears prepare-gate + # IR backlog quicker so new groups actually get served. + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size ${ROLLOUT_BATCH_SIZE:-32} + --micro-batch-size 1 + --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT:-8} + --rollout-max-context-len 16384 + --rollout-max-response-len 4096 + --rollout-max-prompt-len 4096 + --rollout-temperature 1 + --global-batch-size 256 + --rollout-shuffle + --use-streaming-dataset + --agentic-prepare-pool-size 0 +) + +############################################################################### +# EVAL CONFIG # +############################################################################### + +EVAL_ARGS=( + --skip-eval-before-train + --eval-interval 500 + --eval-prompt-data vstar ${TEST_FILES} + --n-samples-per-eval-prompt 8 + --eval-max-response-len 4096 + --eval-top-p 0.7 + --agentic-eval-prepare-pool-size 32 +) + +############################################################################### +# ALGORITHM CONFIG # +############################################################################### + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --eps-clip-c 3 + --use-tis +) + +############################################################################### +# OPTIMIZER CONFIG # +############################################################################### + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + --no-rope-fusion +) + +############################################################################### +# SGLANG CONFIG # +############################################################################### + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 2 + --sglang-mem-fraction-static 0.7 + --sglang-disable-custom-all-reduce + --sglang-page-size 64 + --sglang-attention-backend kunlun + --sglang-disable-radix-cache + --sglang-max-running-requests 256 + # --sglang-disable-cuda-graph + --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) + --sglang-router-policy round_robin + + --sglang-mm-attention-backend fa3 + --sglang-mm-enable-dp-encoder +) + +############################################################################### +# LOGGING CONFIG # +############################################################################### + +LOG_ARGS=( + --tb-experiment-name deepeyes_v2_agentic-klx-${now} + --use-wandb + --wandb-project ${PROJECT_NAME} + --wandb-group deepeyes_v2_agentic-klx-${now} + --wandb-key ${WANDB_API_KEY} + --disable-wandb-random-suffix + --no-use-metrics-service +) + +############################################################################### +# MEGATRON CONFIG # +############################################################################### + +MEGATRON_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --max-tokens-per-gpu 16384 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-dynamic-batch-size + + --moe-flex-dispatcher-backend deepep + --moe-token-dispatcher-type flex + --moe-grouped-gemm true +) + +############################################################################### +# RESOURCE CONFIG # +############################################################################### + +RAY_RESOURCE_ARGS=( + --resource '{"actor": [1, 8], "rollout": [1, 8]}' + --max-staleness 0 + --num-data-storage-units 1 + # --use-health-check + --colocate +) + +############################################################################### +# LAUNCH JOB # +############################################################################### + +mkdir -p logs + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + --runtime-env-json "${RUNTIME_ENV_JSON}" \ + -- python3 relax/entrypoints/train.py \ + --selective-offload \ + "${RAY_RESOURCE_ARGS[@]}" \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${LOG_ARGS[@]}" \ + "${MEGATRON_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + 2>&1 | tee logs/${EXP_NAME}.log diff --git a/examples/deepeyes_v2_agentic/sglang_judge_service_klx.sh b/examples/deepeyes_v2_agentic/sglang_judge_service_klx.sh new file mode 100644 index 000000000..d01f263c1 --- /dev/null +++ b/examples/deepeyes_v2_agentic/sglang_judge_service_klx.sh @@ -0,0 +1,141 @@ +#!/bin/bash +# SGLang Judge 服务:脚本启动时自动部署,结束时自动终止 +# 用法: source "$(dirname "${BASH_SOURCE[0]}")/sglang_judge_service.sh" +# 依赖: MODEL_DIR, TIMESTAMP 需在 source 前已定义 + +# 设置默认端口,确保变量有初始值 +SGLANG_JUDGE_PORT=${SGLANG_JUDGE_PORT:-30000} +SGLANG_JUDGE_MODEL_PATH="${MODEL_DIR}/Qwen2.5-1.5B-Instruct" +SGLANG_JUDGE_PID="" +LOG_FILE="logs/sglang_judge_${TIMESTAMP}.log" + +cleanup_sglang_judge() { + if [ -n "$SGLANG_JUDGE_PID" ] && kill -0 "$SGLANG_JUDGE_PID" 2>/dev/null; then + echo "Stopping sglang judge service (PID: $SGLANG_JUDGE_PID)..." + kill -9 "$SGLANG_JUDGE_PID" 2>/dev/null || true + echo "Sglang judge service stopped successfully." + fi +} + +trap cleanup_sglang_judge EXIT + +# 检查必要依赖是否存在 +if ! command -v curl &> /dev/null; then + echo "Error: curl is required but not installed." + exit 1 +fi + +# 检查模型路径是否存在 +if [ ! -d "$SGLANG_JUDGE_MODEL_PATH" ]; then + echo "Error: Model path not found - $SGLANG_JUDGE_MODEL_PATH" + exit 1 +fi + +mkdir -p logs || { echo "Error: Failed to create logs directory"; exit 1; } + +# 启动 SGLang 服务 +echo "Starting sglang judge service on port $SGLANG_JUDGE_PORT..." +echo "Model path: $SGLANG_JUDGE_MODEL_PATH" +echo "Log file: $LOG_FILE" +python -m sglang.launch_server \ + --model-path "$SGLANG_JUDGE_MODEL_PATH" \ + --port "$SGLANG_JUDGE_PORT" \ + --api-key "EMPTY" \ + --mem-fraction-static 0.05 \ + --disable-cuda-graph \ + --attention-backend kunlun \ + --max-running-requests 32 \ + > "$LOG_FILE" 2>&1 & + +SGLANG_JUDGE_PID=$! +if [ -z "$SGLANG_JUDGE_PID" ]; then + echo "Error: Failed to get PID for sglang judge service" + exit 1 +fi +echo "Sglang judge service started with PID: $SGLANG_JUDGE_PID" + +# 等待服务就绪 +wait_for_sglang_ready() { + local max_attempts=60 # 最大等待次数(总计5分钟) + local attempt=1 + local retry_interval=5 # 重试间隔(秒) + local url="http://127.0.0.1:${SGLANG_JUDGE_PORT}/health" + + echo "Waiting for sglang judge service to be ready (max wait: $((max_attempts * retry_interval)) seconds)..." + + # 健康检查 + while [ $attempt -le $max_attempts ]; do + local http_status + http_status=$(curl -s -o /dev/null -w "%{http_code}" --max-time 2 "$url" 2>/dev/null || echo "000") + + if [ -z "$http_status" ]; then + http_status="000" + fi + + if [ "$http_status" = "200" ] || [ "$http_status" = "204" ]; then + echo "Sglang judge service is ready (health check HTTP 200)." + return 0 + fi + + # 仅在非最后一次尝试时输出重试信息 + if [ $attempt -lt $max_attempts ]; then + echo " Attempt $attempt/$max_attempts: Waiting for judge model service (HTTP status: $http_status), retrying in ${retry_interval}s..." + fi + + sleep $retry_interval + attempt=$((attempt + 1)) + done + + # 超时处理 + echo "Error: Sglang judge service failed to start within timeout (${max_attempts} attempts)." + cleanup_sglang_judge + return 1 +} + +# 执行等待逻辑,失败则退出 +if ! wait_for_sglang_ready; then + echo "Aborting due to sglang judge service startup failure." + exit 1 +fi + + +export DEEPEYES_JUDGE_API_KEY="EMPTY" +export DEEPEYES_JUDGE_BASE_URL="http://127.0.0.1:${SGLANG_JUDGE_PORT}/v1" +export DEEPEYES_JUDGE_MODELS="Qwen2.5-1.5B-Instruct" + +if [ -n "${RUNTIME_ENV_JSON:-}" ]; then + json_escape() { + local value="${1:-}" + value=${value//\\/\\\\} + value=${value//\"/\\\"} + value=${value//$'\n'/\\n} + value=${value//$'\r'/\\r} + value=${value//$'\t'/\\t} + printf '%s' "$value" + } + + runtime_env_prefix="${RUNTIME_ENV_JSON%$'\n}\n}'}" + export RUNTIME_ENV_JSON="${runtime_env_prefix}, + \"DEEPEYES_JUDGE_API_KEY\": \"$(json_escape "${DEEPEYES_JUDGE_API_KEY}")\", + \"DEEPEYES_JUDGE_BASE_URL\": \"$(json_escape "${DEEPEYES_JUDGE_BASE_URL}")\", + \"DEEPEYES_JUDGE_MODELS\": \"$(json_escape "${DEEPEYES_JUDGE_MODELS}")\" +} +}" +fi + +ray() { + if [ "$1" = "job" ] && [ "$2" = "submit" ] && [ -n "${RUNTIME_ENV_JSON:-}" ]; then + local arg="" + for arg in "$@"; do + if [ "$arg" = "--runtime-env-json" ] || [[ "$arg" == --runtime-env-json=* ]]; then + command ray "$@" + return + fi + done + command ray job submit ${RAY_NO_WAIT:+--no-wait} --runtime-env-json="${RUNTIME_ENV_JSON}" "${@:3}" + return + fi + command ray "$@" +} + +echo "Sglang judge service is fully ready for use."