From 9d4f9df813393180e1c26fa11d4ce596f9802a83 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 6 Jul 2026 14:44:20 +0000 Subject: [PATCH] Add 8-GPU RL-Kernel linear_logp benchmark runner --- .../run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh | 333 ++++++++++++++++++ scripts/run-qwen3-30B-A3B.sh | 17 +- tests/utils/test_vllm_engine.py | 21 ++ vime-RLK-8gpu-single-op-runbook.md | 312 ++++++++++++++++ vime-RLK-final-metrics-config-matrix.md | 134 +++++++ vime/backends/vllm_utils/vllm_engine.py | 4 +- vime/ray/rollout.py | 121 ++++++- 7 files changed, 938 insertions(+), 4 deletions(-) create mode 100755 scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh create mode 100644 vime-RLK-8gpu-single-op-runbook.md create mode 100644 vime-RLK-final-metrics-config-matrix.md diff --git a/scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh b/scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh new file mode 100755 index 00000000..973612f3 --- /dev/null +++ b/scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh @@ -0,0 +1,333 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat >&2 <<'EOF' +usage: scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh CONFIG MODE + +CONFIG: T01 | T02 | T03 | T04 | T05 | T06 | T07 | T08 +MODE: baseline | cuda + +Important environment overrides: + TRAIN_SCOPE=full|output_layer default: full + TRACE_MODE=none|train|rollout|all default: all + TRACE_ROLLOUTS=3 nsys capture rollout ids + RUN_ROOT=/workspace log/trace root + NSYS=/path/to/nsys nsys binary +EOF +} + +CONFIG_ID="${1:-}" +MODE="${2:-}" +if [[ -z "${CONFIG_ID}" || -z "${MODE}" ]]; then + usage + exit 2 +fi +if [[ "${MODE}" != "baseline" && "${MODE}" != "cuda" ]]; then + usage + exit 2 +fi + +WORKSPACE_ROOT="${WORKSPACE_ROOT:-/workspace}" +RUN_ROOT="${RUN_ROOT:-${WORKSPACE_ROOT}}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" +VENV="${VIME_PYTHON_ENV:-${WORKSPACE_ROOT}/vime-rlk-env}" +NSYS="${NSYS:-/opt/nvidia/nsight-compute/2024.1.1/host/target-linux-x64/nsys}" + +case "${CONFIG_ID}" in + T01) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-2048}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-512}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.40}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-2048}" + ;; + T02) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-3072}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-1024}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.42}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-3072}" + ;; + T03) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-4096}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-1536}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + T04) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-4096}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-2048}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + T05) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-6144}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3072}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + T06) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3072}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + T07) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3968}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.46}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + T08) + ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}" + N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}" + GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}" + MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" + ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3584}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}" + VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}" + ;; + *) + usage + exit 2 + ;; +esac + +TRAIN_SCOPE="${TRAIN_SCOPE:-full}" +if [[ "${TRAIN_SCOPE}" != "full" && "${TRAIN_SCOPE}" != "output_layer" ]]; then + echo "TRAIN_SCOPE must be full or output_layer, got ${TRAIN_SCOPE}" >&2 + exit 2 +fi + +TRACE_MODE="${TRACE_MODE:-all}" +case "${TRACE_MODE}" in + none|train|rollout|all) ;; + *) + echo "TRACE_MODE must be none, train, rollout, or all; got ${TRACE_MODE}" >&2 + exit 2 + ;; +esac + +RUN_NAME="${RUN_NAME:-8gpu_${CONFIG_ID}_${TRAIN_SCOPE}_${MODE}_$(date +%Y%m%d_%H%M%S)}" +LOG_DIR="${LOG_DIR:-${RUN_ROOT}/logs/${RUN_NAME}}" +TRACE_DIR="${TRACE_DIR:-${RUN_ROOT}/nsys_traces/${RUN_NAME}}" +mkdir -p "${LOG_DIR}" "${TRACE_DIR}" + +export NUM_GPUS="${NUM_GPUS:-8}" +export MEGATRON_TP="${MEGATRON_TP:-2}" +export MEGATRON_EP="${MEGATRON_EP:-8}" +export MEGATRON_CP="${MEGATRON_CP:-1}" +export ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-8}" +export NUM_ROLLOUT="${NUM_ROLLOUT:-12}" +export VIME_DISABLE_SAVE="${VIME_DISABLE_SAVE:-1}" +export VIME_SKIP_EVAL_BEFORE_TRAIN="${VIME_SKIP_EVAL_BEFORE_TRAIN:-1}" +export VIME_USE_KL_LOSS="${VIME_USE_KL_LOSS:-0}" +export VIME_SKIP_ZERO_ENTROPY_METRIC="${VIME_SKIP_ZERO_ENTROPY_METRIC:-1}" +export VIME_VLLM_ENFORCE_EAGER="${VIME_VLLM_ENFORCE_EAGER:-1}" +export VIME_RL_KERNEL_OPS="${VIME_RL_KERNEL_OPS:-linear_logp}" +export VIME_RL_KERNEL_LINEAR_LOGP_BACKEND="${VIME_RL_KERNEL_LINEAR_LOGP_BACKEND:-cuda}" +export VIME_RL_KERNEL_CUDA_EVENT_TIMER="${VIME_RL_KERNEL_CUDA_EVENT_TIMER:-1}" +export VIME_LINEAR_LOGP_MEMORY_PROBE="${VIME_LINEAR_LOGP_MEMORY_PROBE:-1}" +export VIME_TRAIN_MEMORY_MARGIN_BYTES="${VIME_TRAIN_MEMORY_MARGIN_BYTES:-1073741824}" +export MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ="${MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ:-1}" +export VIME_USE_DISTRIBUTED_OPTIMIZER="${VIME_USE_DISTRIBUTED_OPTIMIZER:-1}" +export VIME_USE_PRECISION_AWARE_OPTIMIZER="${VIME_USE_PRECISION_AWARE_OPTIMIZER:-1}" +export VIME_OPTIMIZER_CPU_OFFLOAD="${VIME_OPTIMIZER_CPU_OFFLOAD:-1}" +export VIME_USE_FP32_GRAD_BUFFER="${VIME_USE_FP32_GRAD_BUFFER:-0}" +export VIME_GRAD_REDUCE_IN_BF16="${VIME_GRAD_REDUCE_IN_BF16:-1}" +export VIME_SEQUENCE_PARALLEL="${VIME_SEQUENCE_PARALLEL:-0}" +export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-0}" +export NCCL_CUMEM_ENABLE="${NCCL_CUMEM_ENABLE:-0}" +export CUDA_MODULE_LOADING="${CUDA_MODULE_LOADING:-EAGER}" +export VIME_SKIP_PROCESS_CLEANUP="${VIME_SKIP_PROCESS_CLEANUP:-0}" + +export ROLLOUT_BATCH_SIZE +export N_SAMPLES_PER_PROMPT +export GLOBAL_BATCH_SIZE +export MAX_TOKENS_PER_GPU +export ROLLOUT_MAX_RESPONSE_LEN +export VLLM_GPU_MEMORY_UTILIZATION +export VLLM_MAX_MODEL_LEN + +if [[ "${MODE}" == "cuda" ]]; then + export VIME_RL_KERNEL=1 + export VIME_BASELINE_LINEAR_LOGP_TIMER=0 + export VIME_BASELINE_CUDA_EVENT_TIMER=0 + if [[ "${TRAIN_SCOPE}" == "output_layer" ]]; then + export VIME_ONLY_TRAIN_PARAMS_NAME_LIST="${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-output_layer}" + export VIME_SYNC_TRAINABLE_WEIGHTS_ONLY="${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY:-1}" + export VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN="${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-1}" + export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16="${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-1}" + export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL="${RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL:-0}" + else + export VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN="${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-0}" + export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16="${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-0}" + export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL="${RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL:-1}" + fi +else + export VIME_RL_KERNEL=0 + export VIME_BASELINE_LINEAR_LOGP_TIMER=1 + export VIME_BASELINE_CUDA_EVENT_TIMER=1 + export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0 + export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=0 + if [[ "${TRAIN_SCOPE}" == "output_layer" ]]; then + export VIME_ONLY_TRAIN_PARAMS_NAME_LIST="${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-output_layer}" + export VIME_SYNC_TRAINABLE_WEIGHTS_ONLY="${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY:-1}" + export VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN="${VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN:-1}" + else + export VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN="${VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN:-0}" + fi +fi + +case "${TRACE_MODE}" in + none) + export VIME_NSYS_CAPTURE_ROLLOUTS="" + export VIME_NSYS_CAPTURE_ROLE="${VIME_NSYS_CAPTURE_ROLE:-actor}" + ;; + train) + export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}" + export VIME_NSYS_CAPTURE_ROLE="actor" + ;; + rollout) + export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}" + export VIME_NSYS_CAPTURE_ROLE="rollout" + ;; + all) + export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}" + export VIME_NSYS_CAPTURE_ROLE="all" + ;; +esac + +if [[ "${TRACE_MODE}" == "rollout" || "${TRACE_MODE}" == "all" ]]; then + if [[ -z "${VLLM_PROFILER_CONFIG:-}" ]]; then + export VLLM_PROFILER_CONFIG='{"profiler":"cuda"}' + else + export VLLM_PROFILER_CONFIG + fi +fi + +TOTAL_SAMPLES_PER_ROLLOUT=$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT)) +TRAIN_DP_SIZE=$((NUM_GPUS / MEGATRON_TP / MEGATRON_CP)) +if (( GLOBAL_BATCH_SIZE > TOTAL_SAMPLES_PER_ROLLOUT )); then + echo "GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE} exceeds ROLLOUT_BATCH_SIZE*N_SAMPLES_PER_PROMPT=${TOTAL_SAMPLES_PER_ROLLOUT}" >&2 + exit 2 +fi + +if (( TOTAL_SAMPLES_PER_ROLLOUT < TRAIN_DP_SIZE )); then + echo "RBS*NSP=${TOTAL_SAMPLES_PER_ROLLOUT} is smaller than training DP size ${TRAIN_DP_SIZE}; each DP rank needs at least one sample" >&2 + exit 2 +fi + +export CONFIG_ID MODE TRAIN_SCOPE TRACE_MODE RUN_NAME LOG_DIR TRACE_DIR + +echo "run_name=${RUN_NAME}" | tee "${LOG_DIR}/run_env.log" +env | sort | grep -E '^(CONFIG_ID|MODE|TRAIN_SCOPE|TRACE_MODE|NUM_GPUS|MEGATRON_|ROLLOUT_|N_SAMPLES|GLOBAL_BATCH|MAX_TOKENS|VLLM_|VIME_|RL_KERNEL_|NCCL_|CUDA_MODULE_LOADING|LOG_DIR|TRACE_DIR|RUN_NAME)=' \ + | tee -a "${LOG_DIR}/run_env.log" + +run_train() { + bash "${VIME_ROOT}/scripts/run-qwen3-30B-A3B.sh" +} + +NSYS_RAY_PID="" +cleanup() { + set +e + if [[ -n "${NSYS_RAY_PID}" ]] && kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then + "${VENV}/bin/ray" stop --force >/dev/null 2>&1 || true + for _ in $(seq 1 120); do + if ! kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then + break + fi + sleep 1 + done + if kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then + kill -TERM "${NSYS_RAY_PID}" >/dev/null 2>&1 || true + fi + wait "${NSYS_RAY_PID}" >/dev/null 2>&1 || true + fi + pkill -TERM -f "[v]llm serve" >/dev/null 2>&1 || true + pkill -TERM -f "[V]LLM::" >/dev/null 2>&1 || true +} +trap cleanup EXIT + +cd "${VIME_ROOT}" +if [[ "${TRACE_MODE}" == "none" ]]; then + run_train 2>&1 | tee "${LOG_DIR}/ray_job_${MODE}.log" +else + if [[ ! -x "${NSYS}" ]]; then + echo "nsys not found or not executable: ${NSYS}" >&2 + exit 1 + fi + + "${VENV}/bin/ray" stop --force >/dev/null 2>&1 || true + pkill -9 -f "[v]llm serve" >/dev/null 2>&1 || true + pkill -9 -f "[V]LLM::" >/dev/null 2>&1 || true + sleep 2 + + "${NSYS}" profile \ + --trace=cuda,nvtx,osrt,cublas,cudnn \ + --sample=none \ + --cpuctxsw=none \ + --capture-range=cudaProfilerApi \ + --capture-range-end=repeat:64 \ + --flush-on-cudaprofilerstop=true \ + --wait=all \ + --force-overwrite=true \ + --export=sqlite \ + --output="${TRACE_DIR}/${RUN_NAME}" \ + "${VENV}/bin/ray" start \ + --head \ + --node-ip-address 127.0.0.1 \ + --num-gpus "${NUM_GPUS}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 \ + --block \ + >"${LOG_DIR}/nsys_ray_start.log" 2>&1 & + NSYS_RAY_PID=$! + + for _ in $(seq 1 120); do + if curl -fsS "http://127.0.0.1:8265/api/version" >/dev/null 2>&1; then + break + fi + if ! kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then + echo "nsys/ray process exited early; see ${LOG_DIR}/nsys_ray_start.log" >&2 + exit 1 + fi + sleep 1 + done + if ! curl -fsS "http://127.0.0.1:8265/api/version" >/dev/null 2>&1; then + echo "ray dashboard did not become ready; see ${LOG_DIR}/nsys_ray_start.log" >&2 + exit 1 + fi + + VIME_SKIP_PROCESS_CLEANUP=1 VIME_SKIP_RAY_START=1 run_train 2>&1 | tee "${LOG_DIR}/ray_job_${MODE}.log" +fi + +cleanup +trap - EXIT + +echo "trace files:" +find "${TRACE_DIR}" -maxdepth 1 -type f -print | sort || true +echo "logs:" +find "${LOG_DIR}" -maxdepth 1 -type f -print | sort || true diff --git a/scripts/run-qwen3-30B-A3B.sh b/scripts/run-qwen3-30B-A3B.sh index 03a0f5ce..bc5186c5 100644 --- a/scripts/run-qwen3-30B-A3B.sh +++ b/scripts/run-qwen3-30B-A3B.sh @@ -128,6 +128,7 @@ ROLLOUT_NUM_GPUS_PER_ENGINE=${ROLLOUT_NUM_GPUS_PER_ENGINE:-${NUM_GPUS}} VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION:-0.5} VLLM_MAX_MODEL_LEN=${VLLM_MAX_MODEL_LEN:-} VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS:-} +VLLM_PROFILER_CONFIG=${VLLM_PROFILER_CONFIG:-} VIME_CKPT_DIR=${VIME_CKPT_DIR:-${WORKSPACE_ROOT}/Qwen3-30B-A3B_vime_tp2_dev} VIME_DISABLE_SAVE=${VIME_DISABLE_SAVE:-1} VIME_SKIP_EVAL_BEFORE_TRAIN=${VIME_SKIP_EVAL_BEFORE_TRAIN:-1} @@ -197,6 +198,7 @@ echo "ROLLOUT_MAX_RESPONSE_LEN: $ROLLOUT_MAX_RESPONSE_LEN" echo "VLLM_GPU_MEMORY_UTILIZATION: $VLLM_GPU_MEMORY_UTILIZATION" echo "VLLM_MAX_MODEL_LEN: ${VLLM_MAX_MODEL_LEN:-}" echo "VLLM_MAX_NUM_SEQS: ${VLLM_MAX_NUM_SEQS:-}" +echo "VLLM_PROFILER_CONFIG: ${VLLM_PROFILER_CONFIG:-}" echo "WORKSPACE_ROOT: $WORKSPACE_ROOT" echo "MEGATRON_ROOT: $MEGATRON_ROOT" echo "QWEN3_30B_A3B_HF_DIR: $QWEN3_30B_A3B_HF_DIR" @@ -347,6 +349,9 @@ fi if [[ -n "${VLLM_MAX_NUM_SEQS}" ]]; then VLLM_ARGS+=(--vllm-max-num-seqs "${VLLM_MAX_NUM_SEQS}") fi +if [[ -n "${VLLM_PROFILER_CONFIG}" ]]; then + VLLM_ARGS+=(--vllm-profiler-config "${VLLM_PROFILER_CONFIG}") +fi if [[ "${VIME_VLLM_ENFORCE_EAGER:-0}" == "1" ]]; then VLLM_ARGS+=(--vllm-enforce-eager) else @@ -401,7 +406,9 @@ fi # launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} cd "${VIME_ROOT}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +if [[ "${VIME_SKIP_RAY_START:-0}" != "1" ]]; then + ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +fi # Build the runtime environment JSON with proper variable substitution RUNTIME_ENV_JSON="{ @@ -423,13 +430,21 @@ RUNTIME_ENV_JSON="{ \"VIME_RL_KERNEL_LINEAR_LOGP_BACKEND\": \"${VIME_RL_KERNEL_LINEAR_LOGP_BACKEND:-}\", \"VIME_RL_KERNEL_CUDA_EVENT_TIMER\": \"${VIME_RL_KERNEL_CUDA_EVENT_TIMER:-0}\", \"VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN\": \"${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-}\", + \"VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN\": \"${VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN:-}\", \"VIME_RL_KERNEL_VALIDATE_TP_TARGETS\": \"${VIME_RL_KERNEL_VALIDATE_TP_TARGETS:-0}\", \"RL_KERNEL_LINEAR_LOGP_FUSED_BACKWARD\": \"${RL_KERNEL_LINEAR_LOGP_FUSED_BACKWARD:-1}\", \"RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16\": \"${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-0}\", + \"RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_ROWWISE_DLOGITS\": \"${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_ROWWISE_DLOGITS:-0}\", + \"RL_KERNEL_LINEAR_LOGP_FUSED_BWD_ROWWISE_DLOGITS\": \"${RL_KERNEL_LINEAR_LOGP_FUSED_BWD_ROWWISE_DLOGITS:-0}\", + \"RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL\": \"${RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL:-0}\", + \"RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_VOCAB_TILE\": \"${RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_VOCAB_TILE:-}\", \"RL_KERNEL_LINEAR_LOGP_VALIDATE_TP_TARGETS\": \"${RL_KERNEL_LINEAR_LOGP_VALIDATE_TP_TARGETS:-0}\", \"VIME_LINEAR_LOGP_MEMORY_PROBE\": \"${VIME_LINEAR_LOGP_MEMORY_PROBE:-0}\", \"VIME_BASELINE_LINEAR_LOGP_TIMER\": \"${VIME_BASELINE_LINEAR_LOGP_TIMER:-0}\", \"VIME_BASELINE_CUDA_EVENT_TIMER\": \"${VIME_BASELINE_CUDA_EVENT_TIMER:-0}\", + \"VIME_SKIP_ZERO_ENTROPY_METRIC\": \"${VIME_SKIP_ZERO_ENTROPY_METRIC:-0}\", + \"VIME_NSYS_CAPTURE_ROLLOUTS\": \"${VIME_NSYS_CAPTURE_ROLLOUTS:-}\", + \"VIME_NSYS_CAPTURE_ROLE\": \"${VIME_NSYS_CAPTURE_ROLE:-actor}\", \"MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ\": \"${MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ:-0}\", \"MEGATRON_ALLOW_MOE_TP_WITHOUT_SP\": \"${MEGATRON_ALLOW_MOE_TP_WITHOUT_SP}\", \"TENSORBOARD_DIR\": \"${TENSORBOARD_DIR:-}\" diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py index fccb5b4f..b283de28 100644 --- a/tests/utils/test_vllm_engine.py +++ b/tests/utils/test_vllm_engine.py @@ -347,6 +347,27 @@ def fake_post(endpoint: str, payload: dict): assert calls == [("finish_weight_update", {})] +@pytest.mark.unit +def test_vllm_profile_endpoints_return_plain_dict(vllm_engine, monkeypatch): + seen: list[str] = [] + + def fake_post(url, *, json=None, timeout=30, params=None): + seen.append(url) + assert json == {} + assert timeout == 30 + assert params is None + return _MockResponse(text="") + + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.start_profile() == {"ok": True} + assert vllm_engine.stop_profile() == {"ok": True} + assert seen == [ + "http://127.0.0.1:8765/start_profile", + "http://127.0.0.1:8765/stop_profile", + ] + + @pytest.mark.unit def test_update_weights_from_tensor_posts_ipc_payload_and_records_version(vllm_engine, monkeypatch): posted: list[tuple[str, dict]] = [] diff --git a/vime-RLK-8gpu-single-op-runbook.md b/vime-RLK-8gpu-single-op-runbook.md new file mode 100644 index 00000000..cf5b78af --- /dev/null +++ b/vime-RLK-8gpu-single-op-runbook.md @@ -0,0 +1,312 @@ +# vime + RL-Kernel linear_logp 8 卡训推 12 轮运行步骤 + +日期：2026-07-06 UTC + +## 0. 运行原则 + +- 8 卡正式结果必须完整训推 12 轮：`NUM_ROLLOUT=12`。 +- 模型固定 Qwen3-30B-A3B。 +- 不设置 `LOAD_DEBUG_ROLLOUT_DATA`，不走 debug train-only。 +- 先跑 candidate full-gradient 找最大不 OOM，再补 baseline。 +- 每个成功配置必须完成 rollout 0-11；取 rollout 3-11 的均值填表。 +- trace 用 `TRACE_MODE=train|rollout|all`；默认 `all`，需要少 overhead 时才改成 `train` 或 `none`。 + +## 1. 准备环境 + +```bash +cd /workspace/vime-rlk-tp2 +git status --short +nvidia-smi -L +``` + +确认 RL-Kernel 扩展： + +```bash +/workspace/vime-rlk-env/bin/python - <<'PY' +import torch +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +print("torch", torch.__version__, "cuda", torch.version.cuda) +print("rl_kernel_ext", _EXT_AVAILABLE) +print("fused_bwd", hasattr(_C, "fused_linear_logp_sm90_backward")) +PY +``` + +确认模型和数据： + +```bash +test -d /workspace/Qwen3-30B-A3B && echo "HF ok" +test -d /workspace/Qwen3-30B-A3B_torch_dist && echo "torch_dist ok" +test -f /workspace/dapo-math-17k/dapo-math-17k.jsonl && echo "prompt ok" +``` + +## 2. 运行脚本 + +统一使用： + +```bash +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh CONFIG MODE +``` + +参数： + +| 参数 | 可选值 | +| --- | --- | +| `CONFIG` | `T01`, `T02`, `T03`, `T04`, `T05`, `T06`, `T07`, `T08` | +| `MODE` | `cuda`, `baseline` | +| `TRAIN_SCOPE` | `full`, `output_layer` | +| `TRACE_MODE` | `none`, `train`, `rollout`, `all` | +| `TRACE_ROLLOUTS` | `3`, `3-5`, `all` | + +默认： + +- `TRAIN_SCOPE=full` +- `TRACE_MODE=all` +- `TRACE_ROLLOUTS=3` +- `NUM_GPUS=8` +- `MEGATRON_TP=2` +- `MEGATRON_EP=8` +- `NUM_ROLLOUT=12` + +trace 文件写到： + +```text +/workspace/nsys_traces// +``` + +日志写到： + +```text +/workspace/logs// +``` + +配置： + +| 配置 | 参数 | +| --- | --- | +| T01 | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=2048; RESP=512; VLLM_MEM=0.40; VLLM_MAX_MODEL_LEN=2048` | +| T02 | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=3072; RESP=1024; VLLM_MEM=0.42; VLLM_MAX_MODEL_LEN=3072` | +| T03 | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=4096; RESP=1536; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | +| T04 | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=4096; RESP=2048; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | +| T05 | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=6144; RESP=3072; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | +| T06 | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3072; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | +| T07 | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3968; VLLM_MEM=0.46; VLLM_MAX_MODEL_LEN=4096` | +| T08 | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3584; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | + +约束： + +- `GBS <= RBS * NSP`。 +- 8 卡 TP=2 默认 DP=4，`RBS * NSP` 至少为 4。 +- T06/T07/T08 是主搜索档；T02/T04/T05 只做回退定位。 + +## 3. 通用清理命令 + +每次 OOM、hang、手动中断后执行： + +```bash +/workspace/vime-rlk-env/bin/ray stop --force || true +pkill -9 -f "[v]llm serve" || true +pkill -9 -f "[V]LLM::" || true +pkill -9 -f "[r]ay" || true +pkill -9 -f "[p]ython.*train.py" || true +pkill -9 -f "[r]edis" || true +sleep 3 +nvidia-smi +``` + +## 4. Smoke + +先跑 T01 candidate full-gradient： + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T01 cuda +``` + +检查： + +```bash +rg -n "Using RL-Kernel linear_logp op|fused-tile|rl_kernel_fallback_count|rl_kernel_linear_logp|perf 11|step 11" \ + /workspace/logs/8gpu_T01_full_cuda_*/ray_job_cuda.log +``` + +通过标准： + +- 完成 rollout 0-11。 +- 出现 `Using RL-Kernel linear_logp op`。 +- 出现 full-gradient fast path。 +- `train/rl_kernel_fallback_count` 为 0。 + +## 5. 最大不 OOM 搜索 + +搜索顺序： + +```text +T01 -> T03 -> T06 -> T07 +``` + +分支规则： + +| 结果 | 下一步 | +| --- | --- | +| T03 OOM | 跑 T02；T02 成功后补 T02 baseline | +| T03 成功，T06 OOM | 跑 T05；T05 成功则补 T05 baseline，否则补 T03 baseline | +| T06 成功，T07 成功 | T07 是 candidate 最大候选，补 T07 baseline | +| T06 成功，T07 OOM | 跑 T08；T08 成功则补 T08 baseline，否则补 T06 baseline | + +命令： + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T03 cuda + +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda + +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T07 cuda +``` + +T06/T07 是最能体现 kernel 性能的配置： + +- `RESP` 长，接近模型 `seq_length=4096` 上限。 +- `MAX_TOKENS_PER_GPU=8192`，允许动态 batch 把更多 packed tokens 放进一次 `linear_logp`。 +- `RBS=4, NSP=2, GBS=8`，每轮 8 条样本，满足 8 卡 TP=2 的 DP=4 调度要求，也避免 `GBS > RBS*NSP`。 +- `VLLM_MEM` 没有盲目拉高，避免 rollout KV cache 抢训练显存。 + +## 6. 训推 trace + +只抓 actor train： + +```bash +TRACE_MODE=train TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +只抓 rollout generate： + +```bash +TRACE_MODE=rollout TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +训推都抓： + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +多轮 trace： + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3-5 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +说明： + +- `TRACE_MODE=train` 捕获 actor train 的 `cudaProfilerStart/Stop` range。 +- `TRACE_MODE=rollout` 会自动设置 `VLLM_PROFILER_CONFIG='{"profiler":"cuda"}'`，并在目标 rollout 前后调用 vLLM `/start_profile` 和 `/stop_profile`。 +- `TRACE_MODE=all` 两者都捕获。 +- Nsight 使用 `--capture-range-end=repeat:64`，能覆盖多个 rollout 的多个 capture range。 +- rollout trace 里看 vLLM worker 的 NVTX annotation，例如 `execute_context_*_generation_*`。 + +## 7. 补 baseline + +只给最大 candidate 成功档和一个中等档补 baseline。 + +示例：最大成功档是 T06。 + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 baseline + +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +如果 baseline OOM： + +1. 记录当前档为 `baseline OOM / candidate success`。 +2. 回退到上一档补 baseline。 +3. 在矩阵里同时写 candidate 最大成功档和 baseline 最大成功档。 + +## 8. 可选 output-layer-only + +只在 full-gradient 主结果完成后跑最大成功档。 + +```bash +TRAIN_SCOPE=output_layer TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 baseline + +TRAIN_SCOPE=output_layer TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T06 cuda +``` + +## 9. 结果提取 + +确认 12 轮完成： + +```bash +rg -n "step 11:|perf 11:" /workspace/logs//ray_job_*.log +``` + +baseline 核心字段： + +```bash +rg -n "step [3-9]:|step 1[01]:|perf [3-9]:|perf 1[01]:|baseline_linear_logp|train_rollout_logprob_abs_diff|rollout_time|raw_reward" \ + /workspace/logs//ray_job_baseline.log +``` + +candidate 核心字段： + +```bash +rg -n "step [3-9]:|step 1[01]:|perf [3-9]:|perf 1[01]:|Using RL-Kernel linear_logp op|fused-tile|save-probs|rl_kernel_fallback_count|rl_kernel_linear_logp|train_rollout_logprob_abs_diff|rollout_time|raw_reward" \ + /workspace/logs//ray_job_cuda.log +``` + +填表规则： + +- `run_status`：有 `step 11` 和 `perf 11` 才填 success。 +- 时间类指标：取 rollout 3-11 均值。 +- 显存峰值：取 rollout 0-11 最大值。 +- CUDA event 秒转毫秒：`value * 1000`。 +- `tokens/call`：填 `*_tokens_per_call_delta`。 +- `token_count_delta`：填 `*_token_count_delta`。 +- fallback：candidate 必须是 0。 + +full-gradient 主指标： + +```text +train/baseline_linear_logp_forward_backward_cuda_event_elapsed_s_delta +train/rl_kernel_linear_logp_forward_backward_cuda_event_elapsed_s_delta +``` + +output-layer-only 主指标： + +```text +train/baseline_linear_logp_forward_cuda_event_elapsed_s_delta +train/rl_kernel_linear_logp_forward_cuda_event_elapsed_s_delta +``` + +## 10. OOM 记录 + +OOM 时记录： + +| 字段 | 内容 | +| --- | --- | +| 配置ID | T01/T03/T05/T06/T07/T08 | +| mode | baseline 或 candidate | +| 触发点 | rollout init / rollout generate / weight sync / train forward / train backward / optimizer | +| 完成轮数 | 最后一个 `step N` 或 `perf N` | +| 最后一条 log | OOM 前最后 20 行 | +| GPU 峰值 | 日志 peak reserved 或 `nvidia-smi` | +| 下一步 | 回退到哪个配置 | + +常见 OOM grep： + +```bash +rg -n "out of memory|CUDA error|RayOutOfMemoryError|WorkerCrashedError|NCCL|killed|ActorDiedError" /workspace/logs/ +``` + +每次 OOM 后执行第 3 节清理命令。 diff --git a/vime-RLK-final-metrics-config-matrix.md b/vime-RLK-final-metrics-config-matrix.md new file mode 100644 index 00000000..475cca96 --- /dev/null +++ b/vime-RLK-final-metrics-config-matrix.md @@ -0,0 +1,134 @@ +# vime + RL-Kernel linear_logp 8 卡训推 12 轮配置矩阵 + +日期：2026-07-06 UTC + +## 最终方案 + +本轮使用 Qwen3-30B-A3B 跑 8 卡完整训推 12 轮，不使用 debug rollout。关注点是完整 rollout + train 流程里的 `linear_logp` 单算子指标。 + +| 对比项 | baseline | candidate | 训练 scope | 目标 fast path | 结论 | +| --- | --- | --- | --- | --- | --- | +| full-gradient | native output layer + native logprob | `save-logits` / fused-tile full backward | `TRAIN_SCOPE=full` | `Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path.` | 主结果；用来搜索最大不 OOM | +| output-layer-only | native output layer + native logprob | `save-prob` | `TRAIN_SCOPE=output_layer`，`--only-train-params-name-list output_layer` | `Using save-probs bf16 output-only tensor-parallel linear_logp fast path.` | 可选补充；不作为主 OOM 边界 | + +原因： + +- 2 卡结果是 debug train-only 单算子，只能证明 isolated kernel 能快；8 卡需要验证完整训推 12 轮下是否仍稳定。 +- 为最大体现 kernel 性能，配置优先放大训练侧 `tokens_per_call`：先提高 `RESP/MAX_TOKENS`，再提高 `GBS/NSP`，同时把 `VLLM_MEM` 控制在不抢训练显存的范围。 +- 最大不 OOM 搜索先跑 candidate full-gradient，再对最大成功档补 baseline。 + +运行脚本：`scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh` + +详细步骤：`vime-RLK-8gpu-single-op-runbook.md` + +## 固定环境 + +| 项 | 值 | +| --- | --- | +| GPU | 8x H100 80GB | +| 模型 | Qwen3-30B-A3B | +| checkpoint | HF: `/workspace/Qwen3-30B-A3B`; Megatron load: `/workspace/Qwen3-30B-A3B_vime_tp2_dev/`; ref: `/workspace/Qwen3-30B-A3B_torch_dist` | +| TP/PP/CP/EP | TP=2, PP=1, CP=1, EP=8 | +| rollout | full vLLM rollout，不使用 debug rollout | +| num rollout | `NUM_ROLLOUT=12` | +| prompt data | `/workspace/dapo-math-17k/dapo-math-17k.jsonl` | +| baseline timer | `VIME_BASELINE_LINEAR_LOGP_TIMER=1`, `VIME_BASELINE_CUDA_EVENT_TIMER=1` | +| candidate timer | `VIME_RL_KERNEL=1`, `VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=cuda`, `VIME_RL_KERNEL_CUDA_EVENT_TIMER=1` | +| trace | `TRACE_MODE=train|rollout|all`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint | +| correctness | `kl_loss_coef=0`, `entropy_coef=0`, `VIME_SKIP_ZERO_ENTROPY_METRIC=1` | + +## 配置矩阵 + +所有配置固定 `NUM_ROLLOUT=12`。8 卡 TP=2 时 DP 通常是 4，每轮至少要 `RBS*NSP >= 4`，且 `GBS <= RBS*NSP`。T05/T06/T07 是主要 kernel-heavy 档位：`RESP` 接近 4096 上限，`MAX_TOKENS` 提高到 6144/8192，让一个 `linear_logp` call 尽量覆盖更多 packed tokens。 + +| 配置ID | 规模 | 关键配置 | 运行策略 | 状态 | +| --- | --- | --- | --- | --- | +| T01 | smoke | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=2048; RESP=512; VLLM_MEM=0.40; VLLM_MAX_MODEL_LEN=2048` | candidate full 必须先跑通 | 待跑 | +| T02 | small | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=3072; RESP=1024; VLLM_MEM=0.42; VLLM_MAX_MODEL_LEN=3072` | T01 失败后定位用 | 待跑 | +| T03 | target-low | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=4096; RESP=1536; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | candidate full -> baseline full | 待跑 | +| T04 | target-mid | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=4096; RESP=2048; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | T05 失败后回退定位 | 待跑 | +| T05 | kernel-heavy | `RBS=2; NSP=2; GBS=4; MAX_TOKENS=6144; RESP=3072; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | candidate full -> baseline full if max | 待跑 | +| T06 | max-token-call probe | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3072; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | candidate full，成功后补 baseline | 待跑 | +| T07 | OOM probe | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3968; VLLM_MEM=0.46; VLLM_MAX_MODEL_LEN=4096` | T06 成功且显存余量足够时跑 | 待跑 | +| T08 | bisect | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3584; VLLM_MEM=0.45; VLLM_MAX_MODEL_LEN=4096` | T06 成功但 T07 OOM 时跑 | 待跑 | + +推荐搜索顺序： + +```text +T01 -> T03 -> T06 -> T07 +``` + +若 `T07` OOM 且 `T06` 成功，再跑 `T08`。若 `T06` OOM，回退 `T05 -> T04`。 + +## 训推 12 轮总览 + +统计口径： + +- `run_status`：12 个 rollout 全部完成才算 success。 +- `step_time_s`、`rollout_time_s`、`train_time_s`、`actor_train_time_s`：取 rollout 3-11 的均值，避开前 3 轮 warmup。 +- `peak_vram_gb`：取全 12 轮最大值。 +- `raw_reward`、`train_rollout_logprob_abs_diff`：取 rollout 3-11 均值。 + +| 配置ID | run_status baseline | run_status candidate | peak_vram_gb baseline | peak_vram_gb candidate | step_time_s baseline | step_time_s candidate | train_time_s baseline | train_time_s candidate | actor_train_time_s baseline | actor_train_time_s candidate | rollout_time_s baseline | rollout_time_s candidate | tokens_per_gpu_per_sec baseline | tokens_per_gpu_per_sec candidate | raw_reward baseline | raw_reward candidate | abs_diff baseline | abs_diff candidate | loss_finite_pass baseline | loss_finite_pass candidate | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | --- | +| T01 | | | | | | | | | | | | | | | | | | | | | +| T03 | | | | | | | | | | | | | | | | | | | | | +| T05 | | | | | | | | | | | | | | | | | | | | | +| T06 | | | | | | | | | | | | | | | | | | | | | +| T07 | | | | | | | | | | | | | | | | | | | | | +| T08 | | | | | | | | | | | | | | | | | | | | | + +## full-gradient: baseline vs save-logits + +统计口径： + +- 主表填 rollout 3-11 均值。 +- 需要看单轮 trace 时，额外列 rollout 3 或 `TRACE_ROLLOUTS` 指定的轮次。 + +| 配置ID | run_status baseline | run_status save-logits | tokens/call baseline | tokens/call save-logits | token_count_delta baseline | token_count_delta save-logits | baseline fwd CUDA ms | save-logits fwd CUDA ms | fwd speedup | baseline fwd+bwd CUDA ms | save-logits fwd+bwd CUDA ms | fwd+bwd speedup | baseline dispatch ms | save-logits dispatch ms | peak_alloc_delta MB baseline | peak_alloc_delta MB save-logits | peak_reserved_delta MB baseline | peak_reserved_delta MB save-logits | fallback | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| T01 | | | | | | | | | | | | | | | | | | | | +| T03 | | | | | | | | | | | | | | | | | | | | +| T05 | | | | | | | | | | | | | | | | | | | | +| T06 | | | | | | | | | | | | | | | | | | | | +| T07 | | | | | | | | | | | | | | | | | | | | +| T08 | | | | | | | | | | | | | | | | | | | | + +解读： + +- `save-logits` 必须命中 fused-tile full-gradient fast path，fallback=0。 +- 主指标是 `fwd+bwd CUDA ms`；`dispatch ms` 只作辅助。 +- 如果 candidate 能跑通而 baseline OOM，单独记录为 candidate 最大不 OOM 优势。 + +## output-layer-only: baseline vs save-prob + +只在 full-gradient 主结果完成后按需补测。 + +| 配置ID | run_status baseline | run_status save-prob | tokens/call | token_count_delta | baseline fwd CUDA ms | save-prob fwd CUDA ms | fwd speedup | baseline dispatch ms | save-prob dispatch ms | actor_train_s baseline | actor_train_s save-prob | peak_alloc_delta MB baseline | peak_alloc_delta MB save-prob | peak_reserved_delta MB baseline | peak_reserved_delta MB save-prob | abs_diff baseline | abs_diff save-prob | fallback | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| TMAX | | | | | | | | | | | | | | | | | | | + +## 运行日志 + +| 配置 | log | +| --- | --- | +| T01 full candidate | | +| T03 full baseline | | +| T03 full save-logits | | +| T06 full baseline | | +| T06 full save-logits | | +| T07 full save-logits | | +| T08 full baseline | | +| T08 full save-logits | | +| TMAX output baseline | | +| TMAX output save-prob | | + +## 后续建议 + +| 目标 | 建议 | +| --- | --- | +| 快速找最大不 OOM | candidate full-gradient 跑 `T01 -> T03 -> T06 -> T07` | +| 最终 baseline 对比 | 只在最大成功档和一个中等档补 baseline | +| 训推 trace | `TRACE_MODE=all TRACE_ROLLOUTS=3` 同时抓 rollout generate 和 actor train | +| 12 轮稳定性 | 以 rollout 3-11 均值为主，同时确认 12 轮无 OOM/no hang | +| output-layer-only | 主结果完成后再补 `TMAX` 一组 | diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 4d0b1b3f..5c82e4c4 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -445,12 +445,12 @@ def start_profile( ): response = requests.post(f"http://{self.server_host}:{self.server_port}/start_profile", json={}) response.raise_for_status() - return response + return {"ok": True} def stop_profile(self): response = requests.post(f"http://{self.server_host}:{self.server_port}/stop_profile", json={}) response.raise_for_status() - return response + return {"ok": True} def simulate_crash(self): if self.args.rollout_external or not getattr(self, "process", None): diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index cb206511..86b5fef0 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -2,6 +2,7 @@ import itertools import logging import multiprocessing +import os import random import time from pathlib import Path @@ -39,6 +40,123 @@ logger = logging.getLogger(__name__) +def _parse_capture_rollouts(value: str) -> set[int] | str: + text = value.strip().lower() + if text in {"all", "*"}: + return "all" + result: set[int] = set() + for part in text.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + start_s, end_s = part.split("-", 1) + start, end = int(start_s), int(end_s) + result.update(range(start, end + 1)) + else: + result.add(int(part)) + return result + + +def _should_nsys_capture(role: str, rollout_id: int) -> bool: + value = os.getenv("VIME_NSYS_CAPTURE_ROLLOUTS", "").strip() + if not value: + return False + capture_role = os.getenv("VIME_NSYS_CAPTURE_ROLE", "actor").strip().lower() + if capture_role not in {"all", role.lower()}: + return False + try: + rollouts = _parse_capture_rollouts(value) + except ValueError: + logger.warning("Ignoring invalid VIME_NSYS_CAPTURE_ROLLOUTS=%r", value) + return False + return rollouts == "all" or rollout_id in rollouts + + +class _NsysCudaProfilerCapture: + def __init__(self, role: str, rollout_id: int, name: str): + self.role = role + self.rollout_id = rollout_id + self.name = name + self.enabled = _should_nsys_capture(role, rollout_id) + + def __enter__(self): + if not self.enabled: + return self + try: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push(f"{self.role}_{self.name}_rollout_{self.rollout_id}") + logger.info( + "Nsight Systems cudaProfilerStart: role=%s rollout_id=%s range=%s", + self.role, + self.rollout_id, + self.name, + ) + else: + self.enabled = False + except Exception: + logger.warning("Failed to start Nsight Systems profiler capture.", exc_info=True) + self.enabled = False + return self + + def __exit__(self, exc_type, exc, tb): + if not self.enabled: + return False + try: + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + torch.cuda.cudart().cudaProfilerStop() + logger.info( + "Nsight Systems cudaProfilerStop: role=%s rollout_id=%s range=%s", + self.role, + self.rollout_id, + self.name, + ) + except Exception: + logger.warning("Failed to stop Nsight Systems profiler capture.", exc_info=True) + return False + + +class _VLLMCudaProfilerCapture: + def __init__(self, manager: "RolloutManager", rollout_id: int): + self.manager = manager + self.rollout_id = rollout_id + self.enabled = _should_nsys_capture("rollout", rollout_id) + + def __enter__(self): + if not self.enabled: + return self + try: + engines = self.manager.rollout_engines + if not engines: + self.enabled = False + return self + logger.info("Starting vLLM CUDA profiler for rollout_id=%s", self.rollout_id) + ray.get([engine.start_profile.remote() for engine in engines]) + except Exception: + logger.warning("Failed to start vLLM CUDA profiler capture.", exc_info=True) + try: + if "engines" in locals(): + ray.get([engine.stop_profile.remote() for engine in engines]) + except Exception: + logger.warning("Failed to clean up partially started vLLM profiler capture.", exc_info=True) + self.enabled = False + return self + + def __exit__(self, exc_type, exc, tb): + if not self.enabled: + return False + try: + engines = self.manager.rollout_engines + logger.info("Stopping vLLM CUDA profiler for rollout_id=%s", self.rollout_id) + ray.get([engine.stop_profile.remote() for engine in engines]) + except Exception: + logger.warning("Failed to stop vLLM CUDA profiler capture.", exc_info=True) + return False + + @dataclasses.dataclass class ServerGroup: """A group of homogeneous vLLM engines with the same configuration. @@ -485,7 +603,8 @@ def generate(self, rollout_id): self.health_monitoring_resume() if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: self._try_ci_fault_injection() - data, metrics = self._get_rollout_data(rollout_id=rollout_id) + with _VLLMCudaProfilerCapture(self, rollout_id), _NsysCudaProfilerCapture("rollout", rollout_id, "generate"): + data, metrics = self._get_rollout_data(rollout_id=rollout_id) self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False) _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) if self.args.debug_rollout_only: