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..51670f77 --- /dev/null +++ b/scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh @@ -0,0 +1,349 @@ +#!/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}" + +export HF_HOME="${HF_HOME:-${WORKSPACE_ROOT}/.cache/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export PIP_CACHE_DIR="${PIP_CACHE_DIR:-${WORKSPACE_ROOT}/.cache/pip}" +export TORCH_HOME="${TORCH_HOME:-${WORKSPACE_ROOT}/.cache/torch}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${WORKSPACE_ROOT}/.cache}" +export TMPDIR="${TMPDIR:-${WORKSPACE_ROOT}/.cache/tmp}" +export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${WORKSPACE_ROOT}}" +export FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST:-9.0}" +mkdir -p "${HF_HOME}" "${HUGGINGFACE_HUB_CACHE}" "${TRANSFORMERS_CACHE}" \ + "${PIP_CACHE_DIR}" "${TORCH_HOME}" "${XDG_CACHE_HOME}" "${TMPDIR}" \ + "${FLASHINFER_WORKSPACE_BASE}/.cache/flashinfer" + +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 VLLM_WORKER_MULTIPROC_METHOD="${VLLM_WORKER_MULTIPROC_METHOD:-spawn}" +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 VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK="${VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK:-1}" +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 MEGATRON_ALLOW_MOE_TP_WITHOUT_SP="${MEGATRON_ALLOW_MOE_TP_WITHOUT_SP:-1}" +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/models/qwen3-30B-A3B.sh b/scripts/models/qwen3-30B-A3B.sh index dfab8682..07590c08 100644 --- a/scripts/models/qwen3-30B-A3B.sh +++ b/scripts/models/qwen3-30B-A3B.sh @@ -44,6 +44,9 @@ MODEL_ARGS=( --moe-grouped-gemm --moe-token-drop-policy probs --moe-router-dtype fp32 - --moe-permute-fusion --moe-aux-loss-coeff 0 -) \ No newline at end of file +) + +if [[ "${VIME_NO_MOE_PERMUTE_FUSION:-0}" != "1" ]]; then + MODEL_ARGS+=(--moe-permute-fusion) +fi diff --git a/scripts/run-qwen3-30B-A3B.sh b/scripts/run-qwen3-30B-A3B.sh index 83dc0c6d..003d1b13 100644 --- a/scripts/run-qwen3-30B-A3B.sh +++ b/scripts/run-qwen3-30B-A3B.sh @@ -1,20 +1,66 @@ #!/bin/bash +WORKSPACE_ROOT=${WORKSPACE_ROOT:-/workspace} +VIME_PYTHON_ENV=${VIME_PYTHON_ENV:-${WORKSPACE_ROOT}/vime-rlk-env} +export HF_HOME="${HF_HOME:-${WORKSPACE_ROOT}/.cache/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export PIP_CACHE_DIR="${PIP_CACHE_DIR:-${WORKSPACE_ROOT}/.cache/pip}" +export TORCH_HOME="${TORCH_HOME:-${WORKSPACE_ROOT}/.cache/torch}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${WORKSPACE_ROOT}/.cache}" +export TMPDIR="${TMPDIR:-${WORKSPACE_ROOT}/.cache/tmp}" +export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${WORKSPACE_ROOT}}" +export FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST:-9.0}" +export VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK="${VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK:-1}" +mkdir -p "${HF_HOME}" "${HUGGINGFACE_HUB_CACHE}" "${TRANSFORMERS_CACHE}" \ + "${PIP_CACHE_DIR}" "${TORCH_HOME}" "${XDG_CACHE_HOME}" "${TMPDIR}" \ + "${FLASHINFER_WORKSPACE_BASE}/.cache/flashinfer" +if [[ -d "${VIME_PYTHON_ENV}/bin" ]]; then + export PATH="${VIME_PYTHON_ENV}/bin:${PATH}" +fi + # for rerun the task -pkill -9 vllm -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python -pkill -9 redis +if [[ "${VIME_SKIP_PROCESS_CLEANUP:-0}" != "1" ]]; then + pkill -9 -f "vllm serve" + sleep 3 + ray stop --force + pkill -9 ray + pkill -9 python + sleep 3 + pkill -9 ray + pkill -9 python + pkill -9 redis +fi set -ex # will prevent ray from buffering stdout/stderr export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-max_split_size_mb:256}" +export CUDA_MODULE_LOADING="${CUDA_MODULE_LOADING:-LAZY}" +if [[ -z "${CUDA_HOME:-}" ]]; then + if [[ -d "${VIME_PYTHON_ENV}/lib/python3.11/site-packages/nvidia/cu13" ]]; then + export CUDA_HOME="${VIME_PYTHON_ENV}/lib/python3.11/site-packages/nvidia/cu13" + elif [[ -d /usr/local/lib/python3.11/dist-packages/nvidia/cu13 ]]; then + export CUDA_HOME=/usr/local/lib/python3.11/dist-packages/nvidia/cu13 + else + export CUDA_HOME=/usr/local/cuda + fi +fi +export PATH="${CUDA_HOME}/bin:${PATH}" +if [[ -d "${VIME_PYTHON_ENV}/lib/python3.11/site-packages/nvidia/cudnn" ]]; then + CUDNN_HOME="${VIME_PYTHON_ENV}/lib/python3.11/site-packages/nvidia/cudnn" +else + CUDNN_HOME="/usr/local/lib/python3.11/dist-packages/nvidia/cudnn" +fi +TORCH_LIB_DIR="${VIME_PYTHON_ENV}/lib/python3.11/site-packages/torch/lib" +if [[ -d "${TORCH_LIB_DIR}" ]]; then + export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${CUDA_HOME}/lib:${CUDA_HOME}/lib64:${CUDNN_HOME}/lib:${LD_LIBRARY_PATH:-}" +else + export LD_LIBRARY_PATH="${CUDA_HOME}/lib:${CUDA_HOME}/lib64:${CUDNN_HOME}/lib:${LD_LIBRARY_PATH:-}" +fi +export CPATH="${CUDA_HOME}/include:${CUDNN_HOME}/include:${CPATH:-}" +export LIBRARY_PATH="${CUDA_HOME}/lib:${CUDA_HOME}/lib64:${CUDNN_HOME}/lib:${LIBRARY_PATH:-}" NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) if [ "$NVLINK_COUNT" -gt 0 ]; then @@ -22,51 +68,224 @@ if [ "$NVLINK_COUNT" -gt 0 ]; then else HAS_NVLINK=0 fi +NCCL_NVLS_ENABLE=${NCCL_NVLS_ENABLE:-0} +NCCL_CUMEM_ENABLE=${NCCL_CUMEM_ENABLE:-0} echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" +echo "NCCL_NVLS_ENABLE: $NCCL_NVLS_ENABLE" +echo "NCCL_CUMEM_ENABLE: $NCCL_CUMEM_ENABLE" + +if command -v nvidia-smi >/dev/null 2>&1; then + DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ') + DETECTED_GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -n 1) +else + DETECTED_GPUS=0 + DETECTED_GPU_NAME="unknown" +fi + +validate_positive_int() { + local name="$1" + local value="$2" + if ! [[ "$value" =~ ^[0-9]+$ ]] || [ "$value" -le 0 ]; then + echo "${name} must be a positive integer, got '${value}'" >&2 + exit 1 + fi +} + +validate_at_most_num_gpus() { + local name="$1" + local value="$2" + if [ "$value" -gt "$NUM_GPUS" ]; then + echo "${name}=${value} cannot exceed NUM_GPUS=${NUM_GPUS}" >&2 + exit 1 + fi +} + +validate_divides_num_gpus() { + local name="$1" + local value="$2" + if [ $((NUM_GPUS % value)) -ne 0 ]; then + echo "${name}=${value} must divide NUM_GPUS=${NUM_GPUS}" >&2 + exit 1 + fi +} + +NUM_GPUS=${NUM_GPUS:-2} +validate_positive_int "NUM_GPUS" "$NUM_GPUS" +if [ "$DETECTED_GPUS" -gt 0 ] && [ "$NUM_GPUS" -gt "$DETECTED_GPUS" ]; then + echo "Requested NUM_GPUS=$NUM_GPUS but only detected $DETECTED_GPUS GPUs" >&2 + exit 1 +fi +echo "BENCHMARK_GPU: ${DETECTED_GPU_NAME}" +echo "NUM_GPUS: $NUM_GPUS" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_ROOT="$(cd -- "${SCRIPT_DIR}/.." &>/dev/null && pwd)" +MEGATRON_ROOT=${MEGATRON_ROOT:-${WORKSPACE_ROOT}/Megatron-LM} +QWEN3_30B_A3B_HF_DIR=${QWEN3_30B_A3B_HF_DIR:-${WORKSPACE_ROOT}/Qwen3-30B-A3B} +QWEN3_30B_A3B_TORCH_DIST_DIR=${QWEN3_30B_A3B_TORCH_DIST_DIR:-${WORKSPACE_ROOT}/Qwen3-30B-A3B_torch_dist} +DAPO_MATH_17K_DIR=${DAPO_MATH_17K_DIR:-${WORKSPACE_ROOT}/dapo-math-17k} +AIME_2024_DIR=${AIME_2024_DIR:-${WORKSPACE_ROOT}/aime-2024} +VIME_NO_MOE_PERMUTE_FUSION=${VIME_NO_MOE_PERMUTE_FUSION:-1} source "${SCRIPT_DIR}/models/qwen3-30B-A3B.sh" +MEGATRON_TP=${MEGATRON_TP:-2} +MEGATRON_EP=${MEGATRON_EP:-${NUM_GPUS}} +MEGATRON_CP=${MEGATRON_CP:-1} +MAX_TOKENS_PER_GPU=${MAX_TOKENS_PER_GPU:-4096} +NUM_ROLLOUT=${NUM_ROLLOUT:-24} +ROLLOUT_BATCH_SIZE=${ROLLOUT_BATCH_SIZE:-2} +N_SAMPLES_PER_PROMPT=${N_SAMPLES_PER_PROMPT:-2} +ROLLOUT_MAX_RESPONSE_LEN=${ROLLOUT_MAX_RESPONSE_LEN:-1024} +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT))} +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:-} +VLLM_WORKER_MULTIPROC_METHOD=${VLLM_WORKER_MULTIPROC_METHOD:-spawn} +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} +VIME_VLLM_ENFORCE_EAGER=${VIME_VLLM_ENFORCE_EAGER:-1} +VIME_LOAD_DEBUG_ROLLOUT_DATA=${VIME_LOAD_DEBUG_ROLLOUT_DATA:-} +VIME_NO_GRAD_ACCUM_FUSION=${VIME_NO_GRAD_ACCUM_FUSION:-1} +VIME_NO_MASKED_SOFTMAX_FUSION=${VIME_NO_MASKED_SOFTMAX_FUSION:-1} +VIME_TRANSFORMER_IMPL=${VIME_TRANSFORMER_IMPL:-local} +VIME_NO_ROPE_FUSION=${VIME_NO_ROPE_FUSION:-1} +VIME_NO_PERSIST_LAYER_NORM=${VIME_NO_PERSIST_LAYER_NORM:-1} +VIME_SEQUENCE_PARALLEL=${VIME_SEQUENCE_PARALLEL:-0} +MEGATRON_ALLOW_MOE_TP_WITHOUT_SP=${MEGATRON_ALLOW_MOE_TP_WITHOUT_SP:-0} +VIME_USE_DISTRIBUTED_OPTIMIZER=${VIME_USE_DISTRIBUTED_OPTIMIZER:-1} +VIME_USE_PRECISION_AWARE_OPTIMIZER=${VIME_USE_PRECISION_AWARE_OPTIMIZER:-1} +VIME_OPTIMIZER_CPU_OFFLOAD=${VIME_OPTIMIZER_CPU_OFFLOAD:-1} +VIME_USE_FP32_GRAD_BUFFER=${VIME_USE_FP32_GRAD_BUFFER:-1} +VIME_GRAD_REDUCE_IN_BF16=${VIME_GRAD_REDUCE_IN_BF16:-0} +VIME_TRAIN_MEMORY_MARGIN_BYTES=${VIME_TRAIN_MEMORY_MARGIN_BYTES:-1073741824} +VIME_UPDATE_WEIGHT_BUFFER_SIZE=${VIME_UPDATE_WEIGHT_BUFFER_SIZE:-} +VIME_DDP_BUCKET_SIZE=${VIME_DDP_BUCKET_SIZE:-} +VIME_DDP_NUM_BUCKETS=${VIME_DDP_NUM_BUCKETS:-} +VIME_ONLY_TRAIN_PARAMS_NAME_LIST=${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-} +VIME_SYNC_TRAINABLE_WEIGHTS_ONLY=${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY:-0} +VIME_USE_KL_LOSS=${VIME_USE_KL_LOSS:-1} +VIME_USE_ROLLOUT_LOGPROBS=${VIME_USE_ROLLOUT_LOGPROBS:-0} +if [[ "${VIME_RL_KERNEL:-0}" == "1" ]]; then + VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=${VIME_RL_KERNEL_LINEAR_LOGP_BACKEND:-cuda} + VIME_RL_KERNEL_CUDA_EVENT_TIMER=${VIME_RL_KERNEL_CUDA_EVENT_TIMER:-1} + RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-1} + if [[ -z "${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-}" && "${VIME_ONLY_TRAIN_PARAMS_NAME_LIST}" == *"output_layer"* ]]; then + VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN=1 + fi +fi + +validate_positive_int "MEGATRON_TP" "$MEGATRON_TP" +validate_positive_int "MEGATRON_EP" "$MEGATRON_EP" +validate_positive_int "MEGATRON_CP" "$MEGATRON_CP" +validate_positive_int "MAX_TOKENS_PER_GPU" "$MAX_TOKENS_PER_GPU" +validate_positive_int "NUM_ROLLOUT" "$NUM_ROLLOUT" +validate_positive_int "ROLLOUT_BATCH_SIZE" "$ROLLOUT_BATCH_SIZE" +validate_positive_int "N_SAMPLES_PER_PROMPT" "$N_SAMPLES_PER_PROMPT" +validate_positive_int "ROLLOUT_MAX_RESPONSE_LEN" "$ROLLOUT_MAX_RESPONSE_LEN" +validate_positive_int "GLOBAL_BATCH_SIZE" "$GLOBAL_BATCH_SIZE" +validate_positive_int "ROLLOUT_NUM_GPUS_PER_ENGINE" "$ROLLOUT_NUM_GPUS_PER_ENGINE" +if [[ -n "${VLLM_MAX_MODEL_LEN}" ]]; then + validate_positive_int "VLLM_MAX_MODEL_LEN" "$VLLM_MAX_MODEL_LEN" +fi +if [[ -n "${VLLM_MAX_NUM_SEQS}" ]]; then + validate_positive_int "VLLM_MAX_NUM_SEQS" "$VLLM_MAX_NUM_SEQS" +fi +validate_at_most_num_gpus "MEGATRON_TP" "$MEGATRON_TP" +validate_at_most_num_gpus "MEGATRON_EP" "$MEGATRON_EP" +validate_at_most_num_gpus "ROLLOUT_NUM_GPUS_PER_ENGINE" "$ROLLOUT_NUM_GPUS_PER_ENGINE" +validate_divides_num_gpus "MEGATRON_TP" "$MEGATRON_TP" +validate_divides_num_gpus "MEGATRON_EP" "$MEGATRON_EP" +validate_divides_num_gpus "ROLLOUT_NUM_GPUS_PER_ENGINE" "$ROLLOUT_NUM_GPUS_PER_ENGINE" + +echo "MEGATRON_TP: $MEGATRON_TP" +echo "MEGATRON_EP: $MEGATRON_EP" +echo "MEGATRON_CP: $MEGATRON_CP" +echo "ROLLOUT_NUM_GPUS_PER_ENGINE: $ROLLOUT_NUM_GPUS_PER_ENGINE" +echo "NUM_ROLLOUT: $NUM_ROLLOUT" +echo "ROLLOUT_BATCH_SIZE: $ROLLOUT_BATCH_SIZE" +echo "N_SAMPLES_PER_PROMPT: $N_SAMPLES_PER_PROMPT" +echo "GLOBAL_BATCH_SIZE: $GLOBAL_BATCH_SIZE" +echo "MAX_TOKENS_PER_GPU: $MAX_TOKENS_PER_GPU" +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 "VLLM_WORKER_MULTIPROC_METHOD: ${VLLM_WORKER_MULTIPROC_METHOD}" +echo "WORKSPACE_ROOT: $WORKSPACE_ROOT" +echo "MEGATRON_ROOT: $MEGATRON_ROOT" +echo "QWEN3_30B_A3B_HF_DIR: $QWEN3_30B_A3B_HF_DIR" +echo "QWEN3_30B_A3B_TORCH_DIST_DIR: $QWEN3_30B_A3B_TORCH_DIST_DIR" +echo "DAPO_MATH_17K_DIR: $DAPO_MATH_17K_DIR" +echo "AIME_2024_DIR: $AIME_2024_DIR" +echo "VIME_CKPT_DIR: $VIME_CKPT_DIR" +echo "VIME_LOAD_DEBUG_ROLLOUT_DATA: ${VIME_LOAD_DEBUG_ROLLOUT_DATA:-}" +echo "VIME_USE_DISTRIBUTED_OPTIMIZER: $VIME_USE_DISTRIBUTED_OPTIMIZER" +echo "VIME_USE_PRECISION_AWARE_OPTIMIZER: $VIME_USE_PRECISION_AWARE_OPTIMIZER" +echo "VIME_OPTIMIZER_CPU_OFFLOAD: $VIME_OPTIMIZER_CPU_OFFLOAD" +echo "VIME_USE_FP32_GRAD_BUFFER: $VIME_USE_FP32_GRAD_BUFFER" +echo "VIME_GRAD_REDUCE_IN_BF16: $VIME_GRAD_REDUCE_IN_BF16" +echo "VIME_TRAIN_MEMORY_MARGIN_BYTES: $VIME_TRAIN_MEMORY_MARGIN_BYTES" +echo "VIME_UPDATE_WEIGHT_BUFFER_SIZE: ${VIME_UPDATE_WEIGHT_BUFFER_SIZE:-}" +echo "VIME_TRANSFORMER_IMPL: $VIME_TRANSFORMER_IMPL" +echo "VIME_ONLY_TRAIN_PARAMS_NAME_LIST: ${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-}" +echo "VIME_SYNC_TRAINABLE_WEIGHTS_ONLY: $VIME_SYNC_TRAINABLE_WEIGHTS_ONLY" +echo "VIME_USE_KL_LOSS: $VIME_USE_KL_LOSS" +echo "VIME_USE_ROLLOUT_LOGPROBS: $VIME_USE_ROLLOUT_LOGPROBS" +echo "VIME_RL_KERNEL_LINEAR_LOGP_BACKEND: ${VIME_RL_KERNEL_LINEAR_LOGP_BACKEND:-}" +echo "VIME_RL_KERNEL_CUDA_EVENT_TIMER: ${VIME_RL_KERNEL_CUDA_EVENT_TIMER:-}" +echo "VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN: ${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-}" +echo "RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16: ${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-}" + CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-30B-A3B + --hf-checkpoint "${QWEN3_30B_A3B_HF_DIR}" #--hf-checkpoint /root/Qwen3-30B-A3B-FP8 - --ref-load /root/Qwen3-30B-A3B_torch_dist - --load /root/Qwen3-30B-A3B_vime/ - --save /root/Qwen3-30B-A3B_vime/ - --save-interval 20 + --ref-load "${QWEN3_30B_A3B_TORCH_DIST_DIR}" + --load "${VIME_CKPT_DIR}/" ) +if [[ "${VIME_DISABLE_SAVE:-0}" != "1" ]]; then + CKPT_ARGS+=( + --save "${VIME_CKPT_DIR}/" + --save-interval "${VIME_SAVE_INTERVAL:-20}" + ) +fi ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data "${DAPO_MATH_17K_DIR}/dapo-math-17k.jsonl" --input-key prompt --label-key label --apply-chat-template --rollout-shuffle --rm-type deepscaler - --num-rollout 3000 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN}" --rollout-temperature 1 - --global-batch-size 256 + --global-batch-size "${GLOBAL_BATCH_SIZE}" --balance-data ) EVAL_ARGS=( --eval-interval 20 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --eval-prompt-data aime "${AIME_2024_DIR}/aime-2024.jsonl" --n-samples-per-eval-prompt 16 --eval-max-response-len 16384 --eval-top-p 1 ) +if [[ "${VIME_SKIP_EVAL_BEFORE_TRAIN:-0}" == "1" ]]; then + EVAL_ARGS+=(--skip-eval-before-train) +fi PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel + --tensor-model-parallel-size "${MEGATRON_TP}" --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 + --context-parallel-size "${MEGATRON_CP}" + --expert-model-parallel-size "${MEGATRON_EP}" --expert-tensor-parallel-size 1 --recompute-granularity full @@ -75,18 +294,36 @@ PERF_ARGS=( # --micro-batch-size 1 --use-dynamic-batch-size - --max-tokens-per-gpu 20480 + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}" ) +if [[ "${VIME_SEQUENCE_PARALLEL:-0}" == "1" ]]; then + PERF_ARGS+=(--sequence-parallel) +fi +if [[ "${VIME_NO_GRAD_ACCUM_FUSION:-0}" == "1" ]]; then + PERF_ARGS+=(--no-gradient-accumulation-fusion) +fi +if [[ "${VIME_NO_MASKED_SOFTMAX_FUSION:-0}" == "1" ]]; then + PERF_ARGS+=(--no-masked-softmax-fusion) +fi +PERF_ARGS+=(--transformer-impl "${VIME_TRANSFORMER_IMPL}") +if [[ "${VIME_NO_ROPE_FUSION:-0}" == "1" ]]; then + PERF_ARGS+=(--no-rope-fusion) +fi +if [[ "${VIME_NO_PERSIST_LAYER_NORM:-0}" == "1" ]]; then + PERF_ARGS+=(--no-persist-layer-norm) +fi GRPO_ARGS=( --advantage-estimator grpo - --use-kl-loss --kl-loss-coef 0.00 --kl-loss-type low_var_kl --entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 ) +if [[ "${VIME_USE_KL_LOSS:-1}" == "1" ]]; then + GRPO_ARGS+=(--use-kl-loss) +fi OPTIMIZER_ARGS=( --optimizer adam @@ -95,11 +332,13 @@ OPTIMIZER_ARGS=( --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer ) +if [[ "${VIME_OPTIMIZER_CPU_OFFLOAD:-1}" == "1" ]]; then + OPTIMIZER_ARGS+=(--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d) +fi +if [[ "${VIME_USE_PRECISION_AWARE_OPTIMIZER:-1}" == "1" ]]; then + OPTIMIZER_ARGS+=(--use-precision-aware-optimizer) +fi WANDB_ARGS=( #--use-wandb @@ -108,33 +347,139 @@ WANDB_ARGS=( # --wandb-key ${WANDB_KEY} ) +TB_ARGS=() +if [[ "${VIME_TENSORBOARD:-0}" == "1" ]]; then + export TENSORBOARD_DIR="${TENSORBOARD_DIR:-${VIME_ROOT}/tensorboard_log/${TB_EXPERIMENT_NAME:-qwen3-30B-A3B}}" + TB_ARGS+=(--use-tensorboard) + TB_ARGS+=(--tb-project-name "${TB_PROJECT_NAME:-vime-rlk}") + TB_ARGS+=(--tb-experiment-name "${TB_EXPERIMENT_NAME:-qwen3-30B-A3B}") +fi + VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE}" + --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}" + --vllm-enable-expert-parallel ) +if [[ -n "${VLLM_MAX_MODEL_LEN}" ]]; then + VLLM_ARGS+=(--vllm-max-model-len "${VLLM_MAX_MODEL_LEN}") +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 + VLLM_ARGS+=(--vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256)) +fi MISC_ARGS=( # default dropout in megatron is 0.1 --attention-dropout 0.0 --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 + --train-memory-margin-bytes "${VIME_TRAIN_MEMORY_MARGIN_BYTES}" # need to comment this when using model with MLA --attention-backend flash ) +if [[ "${VIME_USE_FP32_GRAD_BUFFER:-1}" == "1" ]]; then + MISC_ARGS+=(--accumulate-allreduce-grads-in-fp32) +fi +if [[ "${VIME_GRAD_REDUCE_IN_BF16:-0}" == "1" ]]; then + MISC_ARGS+=(--grad-reduce-in-bf16) +fi +if [[ "${VIME_USE_ROLLOUT_LOGPROBS:-0}" == "1" ]]; then + MISC_ARGS+=(--use-rollout-logprobs) +fi +if [[ -n "${VIME_DDP_BUCKET_SIZE}" ]]; then + MISC_ARGS+=(--ddp-bucket-size "${VIME_DDP_BUCKET_SIZE}") +fi +if [[ -n "${VIME_UPDATE_WEIGHT_BUFFER_SIZE}" ]]; then + MISC_ARGS+=(--update-weight-buffer-size "${VIME_UPDATE_WEIGHT_BUFFER_SIZE}") +fi +if [[ -n "${VIME_DDP_NUM_BUCKETS}" ]]; then + MISC_ARGS+=(--ddp-num-buckets "${VIME_DDP_NUM_BUCKETS}") +fi +if [[ -n "${VIME_LOAD_DEBUG_ROLLOUT_DATA}" ]]; then + MISC_ARGS+=(--load-debug-rollout-data "${VIME_LOAD_DEBUG_ROLLOUT_DATA}") +fi +if [[ -n "${VIME_ONLY_TRAIN_PARAMS_NAME_LIST}" ]]; then + IFS=',' read -ra _ONLY_TRAIN_PATTERNS <<< "${VIME_ONLY_TRAIN_PARAMS_NAME_LIST}" + MISC_ARGS+=(--only-train-params-name-list) + for _pattern in "${_ONLY_TRAIN_PATTERNS[@]}"; do + if [[ -n "${_pattern}" ]]; then + MISC_ARGS+=("${_pattern}") + fi + done +fi + +RLK_ARGS=() +if [[ "${VIME_RL_KERNEL:-0}" == "1" ]]; then + RLK_ARGS+=(--enable-rl-kernel --rl-kernel-ops "${VIME_RL_KERNEL_OPS:-linear_logp}") + if [[ "${VIME_RL_KERNEL_STRICT:-0}" == "1" ]]; then + RLK_ARGS+=(--rl-kernel-strict) + fi +fi # launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +cd "${VIME_ROOT}" +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="{ \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"PYTHONPATH\": \"${VIME_ROOT}:${MEGATRON_ROOT}\", + \"PATH\": \"${PATH}\", + \"HF_HOME\": \"${HF_HOME}\", + \"HUGGINGFACE_HUB_CACHE\": \"${HUGGINGFACE_HUB_CACHE}\", + \"TRANSFORMERS_CACHE\": \"${TRANSFORMERS_CACHE}\", + \"PIP_CACHE_DIR\": \"${PIP_CACHE_DIR}\", + \"TORCH_HOME\": \"${TORCH_HOME}\", + \"XDG_CACHE_HOME\": \"${XDG_CACHE_HOME}\", + \"TMPDIR\": \"${TMPDIR}\", + \"FLASHINFER_WORKSPACE_BASE\": \"${FLASHINFER_WORKSPACE_BASE}\", + \"FLASHINFER_CUDA_ARCH_LIST\": \"${FLASHINFER_CUDA_ARCH_LIST}\", + \"CUDA_HOME\": \"${CUDA_HOME:-}\", + \"LD_LIBRARY_PATH\": \"${LD_LIBRARY_PATH:-}\", + \"CPATH\": \"${CPATH:-}\", + \"LIBRARY_PATH\": \"${LIBRARY_PATH:-}\", + \"PYTORCH_CUDA_ALLOC_CONF\": \"${PYTORCH_CUDA_ALLOC_CONF}\", + \"CUDA_MODULE_LOADING\": \"${CUDA_MODULE_LOADING}\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + \"NCCL_NVLS_ENABLE\": \"${NCCL_NVLS_ENABLE}\", + \"NCCL_CUMEM_ENABLE\": \"${NCCL_CUMEM_ENABLE}\", + \"VIME_USE_DISTRIBUTED_OPTIMIZER\": \"${VIME_USE_DISTRIBUTED_OPTIMIZER}\", + \"VIME_CPU_MOE_CKPT_MERGE\": \"${VIME_CPU_MOE_CKPT_MERGE:-1}\", + \"VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK\": \"${VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK}\", + \"VIME_SYNC_TRAINABLE_WEIGHTS_ONLY\": \"${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY}\", + \"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_VLLM_FAULTHANDLER\": \"${VIME_VLLM_FAULTHANDLER:-0}\", + \"VLLM_WORKER_MULTIPROC_METHOD\": \"${VLLM_WORKER_MULTIPROC_METHOD}\", + \"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:-}\" } }" @@ -142,7 +487,7 @@ ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ -- python3 train.py \ --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ + --actor-num-gpus-per-node ${NUM_GPUS} \ --colocate \ ${MODEL_ARGS[@]} \ ${CKPT_ARGS[@]} \ @@ -150,7 +495,9 @@ ray job submit --address="http://127.0.0.1:8265" \ ${OPTIMIZER_ARGS[@]} \ ${GRPO_ARGS[@]} \ ${WANDB_ARGS[@]} \ + ${TB_ARGS[@]} \ ${PERF_ARGS[@]} \ ${EVAL_ARGS[@]} \ ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} + ${MISC_ARGS[@]} \ + ${RLK_ARGS[@]} diff --git a/tests/_unit_stubs.py b/tests/_unit_stubs.py index 2ba863de..7fe27a24 100644 --- a/tests/_unit_stubs.py +++ b/tests/_unit_stubs.py @@ -57,7 +57,8 @@ def install_rollout_optional_stubs() -> None: """Stub rollout-side optional imports when not installed.""" ensure_ray_stub() - install_vllm_router_stub() + if not real_module_available("vllm_router"): + sys.modules["vllm_router"] = types.ModuleType("vllm_router") if not real_module_available("PIL"): pil = types.ModuleType("PIL") @@ -96,48 +97,6 @@ def _raise_os_error(*args, **kwargs): sys.modules["pylatexenc"] = pylatexenc sys.modules["pylatexenc.latex2text"] = latex2text - install_wandb_stub() - - -def install_vllm_router_stub() -> None: - if real_module_available("vllm_router"): - return - - class RouterArgs: - @classmethod - def add_cli_args(cls, parser, *args, **kwargs): # noqa: ARG003 - return parser - - @classmethod - def from_cli_args(cls, args, *unused_args, **unused_kwargs): # noqa: ARG003 - return types.SimpleNamespace() - - router_mod = types.ModuleType("vllm_router") - router_mod.__path__ = [] - launch_router_mod = types.ModuleType("vllm_router.launch_router") - router_args_mod = types.ModuleType("vllm_router.router_args") - launch_router_mod.RouterArgs = RouterArgs - router_args_mod.RouterArgs = RouterArgs - router_mod.launch_router = launch_router_mod - router_mod.router_args = router_args_mod - sys.modules["vllm_router"] = router_mod - sys.modules["vllm_router.launch_router"] = launch_router_mod - sys.modules["vllm_router.router_args"] = router_args_mod - - -def install_wandb_stub() -> None: - if real_module_available("wandb"): - return - wandb_mod = types.ModuleType("wandb") - wandb_mod.run = None - wandb_mod.log = MagicMock() - wandb_mod.finish = MagicMock() - wandb_mod.login = MagicMock() - wandb_mod.init = MagicMock() - wandb_mod.Settings = MagicMock() - wandb_mod.util = types.SimpleNamespace(generate_id=lambda: "unit-test") - sys.modules["wandb"] = wandb_mod - def save_sys_modules(names: Iterable[str]) -> dict[str, Any]: return {k: sys.modules.get(k) for k in names} @@ -237,57 +196,14 @@ def add_cli_args(cls, parser): # noqa: ARG003 arg_utils.AsyncEngineArgs = AsyncEngineArgs engine_mod.arg_utils = arg_utils - system_utils_mod = types.ModuleType("vllm.utils.system_utils") - system_utils_mod.kill_process_tree = lambda pid, include_parent=True: None # noqa: ARG005 - utils_mod.system_utils = system_utils_mod - - # vllm.entrypoints stubs (used by arguments.add_vllm_arguments and vllm_engine._vllm_server_field_names) - entrypoints_mod = types.ModuleType("vllm.entrypoints") - entrypoints_mod.__path__ = [] - openai_mod = types.ModuleType("vllm.entrypoints.openai") - openai_mod.__path__ = [] - cli_args_mod = types.ModuleType("vllm.entrypoints.openai.cli_args") - - import dataclasses as _dc - - @_dc.dataclass - class FrontendArgs: - @classmethod - def add_cli_args(cls, parser): # noqa: ARG003 - return parser - - cli_args_mod.FrontendArgs = FrontendArgs - cli_args_mod.make_arg_parser = lambda parser=None: parser - cli_args_mod.validate_parsed_serve_args = lambda args: args - openai_mod.cli_args = cli_args_mod - entrypoints_mod.openai = openai_mod - vllm_mod.entrypoints = entrypoints_mod - - cli_mod = types.ModuleType("vllm.entrypoints.cli") - cli_mod.__path__ = [] - serve_mod = types.ModuleType("vllm.entrypoints.cli.serve") - - class ServeSubcommand: - pass - - serve_mod.ServeSubcommand = ServeSubcommand - cli_mod.serve = serve_mod - entrypoints_mod.cli = cli_mod - vllm_mod.engine = engine_mod vllm_mod.utils = utils_mod sys.modules["vllm"] = vllm_mod sys.modules["vllm.utils"] = utils_mod sys.modules["vllm.utils.argparse_utils"] = argparse_utils - sys.modules["vllm.utils.system_utils"] = system_utils_mod sys.modules["vllm.engine"] = engine_mod sys.modules["vllm.engine.arg_utils"] = arg_utils - sys.modules["vllm.entrypoints"] = entrypoints_mod - sys.modules["vllm.entrypoints.openai"] = openai_mod - sys.modules["vllm.entrypoints.openai.cli_args"] = cli_args_mod - sys.modules["vllm.entrypoints.cli"] = cli_mod - sys.modules["vllm.entrypoints.cli.serve"] = serve_mod def install_triton_stub() -> None: @@ -307,4 +223,5 @@ def install_triton_stub() -> None: def install_vime_distributed_utils_stub() -> None: vime_utils = types.ModuleType("vime.utils.distributed_utils") vime_utils.get_gloo_group = MagicMock(return_value="gloo") + vime_utils.distributed_masked_whiten = MagicMock(side_effect=lambda values, *args, **kwargs: values) sys.modules.setdefault("vime.utils.distributed_utils", vime_utils) diff --git a/tests/test_cuda_event_timer.py b/tests/test_cuda_event_timer.py new file mode 100644 index 00000000..d0c3e8e7 --- /dev/null +++ b/tests/test_cuda_event_timer.py @@ -0,0 +1,39 @@ +from vime.backends.megatron_utils.cuda_event_timer import CudaEventTimerQueue + + +class _FakeEvent: + def __init__(self, *, elapsed_ms: float = 0.0, ready: bool = True) -> None: + self._elapsed_ms = float(elapsed_ms) + self.ready = ready + + def query(self) -> bool: + return self.ready + + def elapsed_time(self, other) -> float: + return float(other._elapsed_ms) + + +def test_cuda_event_timer_queue_flushes_when_event_becomes_ready(): + queue = CudaEventTimerQueue() + observed = [] + start_event = _FakeEvent() + end_event = _FakeEvent(elapsed_ms=12.5, ready=False) + + queue.enqueue(start_event, end_event, observed.append) + assert observed == [] + + end_event.ready = True + queue.flush_ready() + + assert observed == [0.0125] + + +def test_cuda_event_timer_queue_clear_discards_pending_events(): + queue = CudaEventTimerQueue() + observed = [] + + queue.enqueue(_FakeEvent(), _FakeEvent(elapsed_ms=9.0, ready=False), observed.append) + queue.clear() + queue.flush_ready() + + assert observed == [] diff --git a/tests/test_rl_kernel_args.py b/tests/test_rl_kernel_args.py new file mode 100644 index 00000000..cd3916c4 --- /dev/null +++ b/tests/test_rl_kernel_args.py @@ -0,0 +1,83 @@ +from argparse import Namespace + +import pytest + +from vime.utils.rl_kernel import is_rl_kernel_op_enabled, normalize_rl_kernel_args, parse_rl_kernel_ops + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_parse_rl_kernel_ops_defaults_to_linear_logp(): + assert parse_rl_kernel_ops(None) == ("linear_logp",) + assert parse_rl_kernel_ops("") == ("linear_logp",) + + +@pytest.mark.unit +def test_parse_rl_kernel_ops_deduplicates_comma_and_space_separated_values(): + assert parse_rl_kernel_ops("linear_logp, linear_logp") == ("linear_logp",) + + +@pytest.mark.unit +def test_parse_rl_kernel_ops_rejects_unknown_ops(): + with pytest.raises(ValueError, match="Unsupported RL-Kernel op"): + parse_rl_kernel_ops("linear_logp,moe") + + +@pytest.mark.unit +def test_normalize_rl_kernel_args_keeps_default_disabled(): + args = Namespace(enable_rl_kernel=False, rl_kernel_ops="linear_logp", rl_kernel_strict=False) + + normalize_rl_kernel_args(args) + + assert args.enable_rl_kernel is False + assert args.rl_kernel_ops == ("linear_logp",) + assert is_rl_kernel_op_enabled(args, "linear_logp") is False + + +@pytest.mark.unit +def test_normalize_rl_kernel_args_accepts_env_enable(monkeypatch): + monkeypatch.setenv("VIME_RL_KERNEL", "1") + args = Namespace(enable_rl_kernel=False, rl_kernel_ops="linear_logp", rl_kernel_strict=False) + + normalize_rl_kernel_args(args) + + assert args.enable_rl_kernel is True + assert args.rl_kernel_ops == ("linear_logp",) + assert is_rl_kernel_op_enabled(args, "linear_logp") is True + + +@pytest.mark.unit +def test_normalize_rl_kernel_args_rejects_non_linear_logp_ops(): + args = Namespace(enable_rl_kernel=True, rl_kernel_ops="logp", rl_kernel_strict=False) + + with pytest.raises(ValueError, match="Unsupported RL-Kernel op"): + normalize_rl_kernel_args(args) + + +@pytest.mark.unit +@pytest.mark.parametrize("op", ["ratio_kl", "grpo_loss", "sampling"]) +def test_normalize_rl_kernel_args_rejects_out_of_scope_ops(op): + args = Namespace(enable_rl_kernel=True, rl_kernel_ops=op, rl_kernel_strict=False) + + with pytest.raises(ValueError, match="Unsupported RL-Kernel op"): + normalize_rl_kernel_args(args) + + +@pytest.mark.unit +def test_normalize_rl_kernel_args_accepts_linear_logp(): + args = Namespace(enable_rl_kernel=True, rl_kernel_ops="linear_logp", rl_kernel_strict=False) + + normalize_rl_kernel_args(args) + + assert args.rl_kernel_ops == ("linear_logp",) + assert is_rl_kernel_op_enabled(args, "linear_logp") is True + + +@pytest.mark.unit +def test_normalize_rl_kernel_args_rejects_bad_env_bool(monkeypatch): + monkeypatch.setenv("VIME_RL_KERNEL", "maybe") + args = Namespace(enable_rl_kernel=False, rl_kernel_ops="linear_logp", rl_kernel_strict=False) + + with pytest.raises(ValueError, match="VIME_RL_KERNEL"): + normalize_rl_kernel_args(args) diff --git a/tests/test_rl_kernel_linear_logp_integration.py b/tests/test_rl_kernel_linear_logp_integration.py new file mode 100644 index 00000000..c3f2d878 --- /dev/null +++ b/tests/test_rl_kernel_linear_logp_integration.py @@ -0,0 +1,476 @@ +from __future__ import annotations + +import builtins +import sys +import types +from argparse import Namespace +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + +_tests_root = Path(__file__).resolve().parent +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs + +_unit_stubs.install_megatron_mpu_stub() +_unit_stubs.install_vime_distributed_utils_stub() + +from megatron.core import mpu # noqa: E402 + +from vime.backends.megatron_utils import loss as loss_mod # noqa: E402 +from vime.backends.megatron_utils import rl_kernel as rlk_mod # noqa: E402 + +NUM_GPUS = 0 + + +class _FakeLinearLogpOp: + calls: list[dict] = [] + + def __call__( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + target_ids: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + type(self).calls.append( + { + "hidden_shape": tuple(hidden.shape), + "weight_shape": tuple(weight.shape), + "target_shape": tuple(target_ids.shape), + "bias": bias is not None, + "kwargs": kwargs, + "hidden_requires_grad": hidden.requires_grad, + "hidden_dtype": hidden.dtype, + } + ) + logits = F.linear(hidden.float(), weight.float(), None if bias is None else bias.float()) + return torch.gather(torch.log_softmax(logits, dim=-1), -1, target_ids.long().unsqueeze(-1)).squeeze(-1) + + +class _FakeLegacyLinearLogpOp: + def __call__( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + target_ids: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + logits = F.linear(hidden.float(), weight.float(), None if bias is None else bias.float()) + return torch.gather(torch.log_softmax(logits, dim=-1), -1, target_ids.long().unsqueeze(-1)).squeeze(-1) + + +def _reset_rl_kernel_state(): + rlk_mod._LOGP_OP = None + rlk_mod._LOGP_OP_LOAD_ERROR = None + rlk_mod._LINEAR_LOGP_OP = None + rlk_mod._LINEAR_LOGP_OP_LOAD_ERROR = None + rlk_mod._LINEAR_LOGP_SAVE_PROBS_CAST_LOGGED = False + rlk_mod._WARNED_FALLBACK_REASONS.clear() + rlk_mod._FALLBACK_COUNTS.clear() + rlk_mod._FALLBACK_COUNTS.update({"logp": 0, "linear_logp": 0}) + rlk_mod.reset_rl_kernel_runtime_counters() + _FakeLinearLogpOp.calls.clear() + + +def _install_fake_rl_engine(monkeypatch): + rl_engine = types.ModuleType("rl_engine") + kernels = types.ModuleType("rl_engine.kernels") + registry = types.ModuleType("rl_engine.kernels.registry") + registry.kernel_registry = types.SimpleNamespace(get_op=lambda name: _FakeLinearLogpOp()) + monkeypatch.setitem(sys.modules, "rl_engine", rl_engine) + monkeypatch.setitem(sys.modules, "rl_engine.kernels", kernels) + monkeypatch.setitem(sys.modules, "rl_engine.kernels.registry", registry) + + +def _install_fake_legacy_rl_engine(monkeypatch): + rl_engine = types.ModuleType("rl_engine") + kernels = types.ModuleType("rl_engine.kernels") + registry = types.ModuleType("rl_engine.kernels.registry") + registry.kernel_registry = types.SimpleNamespace(get_op=lambda name: _FakeLegacyLinearLogpOp()) + monkeypatch.setitem(sys.modules, "rl_engine", rl_engine) + monkeypatch.setitem(sys.modules, "rl_engine.kernels", kernels) + monkeypatch.setitem(sys.modules, "rl_engine.kernels.registry", registry) + + +def _make_args(**overrides) -> Namespace: + values = { + "enable_rl_kernel": True, + "rl_kernel_ops": ("linear_logp",), + "rl_kernel_strict": False, + "allgather_cp": False, + "qkv_format": "thd", + "rollout_temperature": 1.0, + "log_probs_chunk_size": -1, + "entropy_coef": 0.0, + "sequence_parallel": False, + "padded_vocab_size": None, + "vocab_size": None, + } + values.update(overrides) + return Namespace(**values) + + +@pytest.fixture(autouse=True) +def reset_parallelism(): + _reset_rl_kernel_state() + mpu.get_tensor_model_parallel_world_size.return_value = 1 + mpu.get_tensor_model_parallel_rank.return_value = 0 + mpu.get_tensor_model_parallel_group.return_value = None + mpu.get_context_parallel_world_size.return_value = 1 + mpu.get_context_parallel_rank.return_value = 0 + mpu.get_virtual_pipeline_model_parallel_world_size.return_value = None + mpu.is_pipeline_last_stage.return_value = True + yield + _reset_rl_kernel_state() + + +def _reference_logp(hidden: torch.Tensor, weight: torch.Tensor, target: torch.Tensor, bias: torch.Tensor | None): + logits = F.linear(hidden.float(), weight.float(), None if bias is None else bias.float()) + return torch.gather(torch.log_softmax(logits, dim=-1), -1, target.long().unsqueeze(-1)).squeeze(-1) + + +def _cpu_calculate_log_probs_and_entropy( + logits: torch.Tensor, + tokens: torch.Tensor, + tp_group, + *, + with_entropy: bool, + chunk_size: int, +): + del tp_group, chunk_size + log_probs = torch.log_softmax(logits.float(), dim=-1) + selected = torch.gather(log_probs, -1, tokens.long().unsqueeze(-1)).squeeze(-1) + entropy = None + if with_entropy: + probs = torch.softmax(logits.float(), dim=-1) + entropy = -(probs * log_probs).sum(dim=-1) + return selected, entropy + + +@pytest.mark.unit +def test_maybe_compute_linear_logp_passes_tensor_parallel_metadata(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args() + torch.manual_seed(1) + hidden = torch.randn(6, 5) + weight = torch.randn(8, 5) + bias = torch.randn(8) + target = torch.randint(0, 8, (6,)) + context = rlk_mod.LinearLogpContext( + lm_head_weight=weight, + bias=bias, + tp_group="tp", + vocab_start_index=16, + global_vocab_size=32, + ) + + actual = rlk_mod.maybe_compute_linear_logp(hidden, target, context=context, args=args, with_entropy=False) + + torch.testing.assert_close(actual, _reference_logp(hidden, weight, target, bias)) + assert _FakeLinearLogpOp.calls == [ + { + "hidden_shape": (6, 5), + "weight_shape": (8, 5), + "target_shape": (6,), + "bias": True, + "kwargs": { + "tp_group": "tp", + "vocab_start_index": 16, + "global_vocab_size": 32, + }, + "hidden_requires_grad": False, + } + ] + counters = rlk_mod.get_rl_kernel_runtime_counters() + assert counters["linear_logp_call_count"] == 1.0 + assert counters["linear_logp_token_count"] == 6.0 + assert counters["linear_logp_dispatch_elapsed_s"] >= 0.0 + + +@pytest.mark.unit +def test_linear_logp_runtime_counter_delta_tracks_since_last_read(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args() + hidden = torch.randn(4, 3) + weight = torch.randn(5, 3) + target = torch.randint(0, 5, (4,)) + context = rlk_mod.LinearLogpContext(lm_head_weight=weight, bias=None, tp_group=None) + + rlk_mod.maybe_compute_linear_logp(hidden[:2], target[:2], context=context, args=args, with_entropy=False) + first_delta = rlk_mod.get_rl_kernel_runtime_counter_delta() + + rlk_mod.maybe_compute_linear_logp(hidden[2:], target[2:], context=context, args=args, with_entropy=False) + second_delta = rlk_mod.get_rl_kernel_runtime_counter_delta() + totals = rlk_mod.get_rl_kernel_runtime_counters() + + assert first_delta["linear_logp_call_count"] == 1.0 + assert first_delta["linear_logp_token_count"] == 2.0 + assert second_delta["linear_logp_call_count"] == 1.0 + assert second_delta["linear_logp_token_count"] == 2.0 + assert totals["linear_logp_call_count"] == 2.0 + assert totals["linear_logp_token_count"] == 4.0 + + +@pytest.mark.unit +def test_linear_logp_detaches_hidden_for_output_layer_only_training(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args(only_train_params_name_list=("output_layer",)) + hidden = torch.randn(4, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + target = torch.randint(0, 5, (4,)) + context = rlk_mod.LinearLogpContext(lm_head_weight=weight, bias=None, tp_group=None) + + rlk_mod.maybe_compute_linear_logp(hidden, target, context=context, args=args, with_entropy=False) + + assert _FakeLinearLogpOp.calls[-1]["hidden_requires_grad"] is False + + +@pytest.mark.unit +def test_linear_logp_matches_vime_response_slicing_from_hidden_states(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args() + vocab_size = 23 + hidden_size = 7 + total_lengths = [5, 4, 6] + response_lengths = [2, 3, 4] + torch.manual_seed(2) + unconcat_tokens = [torch.randint(0, vocab_size, (length,), dtype=torch.long) for length in total_lengths] + hidden = torch.randn(sum(total_lengths), 1, hidden_size) + weight = torch.randn(vocab_size, hidden_size) + bias = torch.randn(vocab_size) + context = rlk_mod.LinearLogpContext(lm_head_weight=weight, bias=bias, tp_group=None) + + _, result = loss_mod.get_log_probs_and_entropy( + hidden, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=False, + rl_kernel_linear_logp_context=context, + ) + + full_tokens = torch.zeros(sum(total_lengths), dtype=torch.long) + offset = 0 + for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False): + full_tokens[offset : offset + total_length - 1] = tokens[1:total_length] + offset += total_length + full_logp = _reference_logp(hidden.squeeze(1), weight, full_tokens, bias) + + expected = [] + offset = 0 + for total_length, response_length in zip(total_lengths, response_lengths, strict=False): + end = offset + total_length + start = end - response_length + expected.append(full_logp[start - 1 : end - 1]) + offset += total_length + + assert len(_FakeLinearLogpOp.calls) == 1 + for actual, expected_item in zip(result["log_probs"], expected, strict=True): + torch.testing.assert_close(actual, expected_item, rtol=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_linear_logp_materializes_logits_fallback_when_optional_package_missing(monkeypatch): + monkeypatch.setattr(loss_mod, "calculate_log_probs_and_entropy", _cpu_calculate_log_probs_and_entropy) + for name in list(sys.modules): + if name == "rl_engine" or name.startswith("rl_engine."): + monkeypatch.delitem(sys.modules, name, raising=False) + + original_import = builtins.__import__ + + def fail_rl_engine_import(name, *args, **kwargs): + if name == "rl_engine" or name.startswith("rl_engine."): + raise ModuleNotFoundError("No module named 'rl_engine'") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fail_rl_engine_import) + args = _make_args() + vocab_size = 19 + hidden_size = 5 + total_lengths = [4, 5] + response_lengths = [2, 3] + torch.manual_seed(3) + unconcat_tokens = [torch.randint(0, vocab_size, (length,), dtype=torch.long) for length in total_lengths] + hidden = torch.randn(1, sum(total_lengths), hidden_size) + weight = torch.randn(vocab_size, hidden_size) + bias = torch.randn(vocab_size) + context = rlk_mod.LinearLogpContext(lm_head_weight=weight, bias=bias, tp_group=None) + + _, result = loss_mod.get_log_probs_and_entropy( + hidden, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=False, + rl_kernel_linear_logp_context=context, + ) + + logits = F.linear(hidden.squeeze(0).float(), weight.float(), bias.float()) + expected = [] + offset = 0 + for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False): + end = offset + total_length + start = end - response_length + target = tokens[-response_length:] + expected.append( + torch.gather(torch.log_softmax(logits[start - 1 : end - 1], dim=-1), -1, target.unsqueeze(-1)).squeeze(-1) + ) + offset += total_length + + assert rlk_mod.get_rl_kernel_fallback_count("linear_logp") == 1 + for actual, expected_item in zip(result["log_probs"], expected, strict=True): + torch.testing.assert_close(actual, expected_item, rtol=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_linear_logp_falls_back_when_op_lacks_tp_interface(monkeypatch): + _install_fake_legacy_rl_engine(monkeypatch) + args = _make_args() + hidden = torch.randn(3, 4) + weight = torch.randn(6, 4) + target = torch.randint(0, 6, (3,)) + context = rlk_mod.LinearLogpContext( + lm_head_weight=weight, + bias=None, + tp_group="tp_group", + vocab_start_index=6, + global_vocab_size=12, + ) + + actual = rlk_mod.maybe_compute_linear_logp(hidden, target, context=context, args=args, with_entropy=False) + + assert actual is None + assert rlk_mod.get_rl_kernel_fallback_count("linear_logp") == 1 + + +@pytest.mark.unit +def test_linear_logp_context_from_model_uses_tp_vocab_offsets(): + mpu.get_tensor_model_parallel_world_size.return_value = 4 + mpu.get_tensor_model_parallel_rank.return_value = 2 + mpu.get_tensor_model_parallel_group.return_value = "tp_group" + + output_layer = types.SimpleNamespace( + weight=torch.empty(8, 4), + bias=None, + sequence_parallel=True, + ) + model = types.SimpleNamespace(output_layer=output_layer, post_process=True) + args = _make_args(padded_vocab_size=32, sequence_parallel=False) + + context = rlk_mod.get_linear_logp_context_from_model(args, model) + + assert context is not None + assert context.lm_head_weight is output_layer.weight + assert context.tp_group == "tp_group" + assert context.vocab_start_index == 16 + assert context.global_vocab_size == 32 + assert context.sequence_parallel is True + + +@pytest.mark.unit +def test_linear_logp_context_prefers_output_layer_weight_for_untied_pp1_model(): + mpu.get_tensor_model_parallel_world_size.return_value = 1 + mpu.get_tensor_model_parallel_rank.return_value = 0 + mpu.get_tensor_model_parallel_group.return_value = None + + output_weight = torch.empty(8, 4) + embedding_weight = torch.empty(8, 4) + output_layer = types.SimpleNamespace(weight=output_weight, bias=None) + model = types.SimpleNamespace( + output_layer=output_layer, + post_process=True, + pre_process=True, + shared_embedding_or_output_weight=lambda: embedding_weight, + ) + args = _make_args() + + context = rlk_mod.get_linear_logp_context_from_model(args, model) + + assert context is not None + assert context.lm_head_weight is output_weight + + +@pytest.mark.unit +def test_linear_logp_context_uses_shared_weight_when_output_layer_weight_is_missing(): + mpu.get_tensor_model_parallel_world_size.return_value = 1 + mpu.get_tensor_model_parallel_rank.return_value = 0 + mpu.get_tensor_model_parallel_group.return_value = None + + embedding_weight = torch.empty(8, 4) + output_layer = types.SimpleNamespace(weight=None, bias=None) + model = types.SimpleNamespace( + output_layer=output_layer, + post_process=True, + pre_process=True, + shared_embedding_or_output_weight=lambda: embedding_weight, + ) + args = _make_args() + + context = rlk_mod.get_linear_logp_context_from_model(args, model) + + assert context is not None + assert context.lm_head_weight is embedding_weight + + +@pytest.mark.unit +def test_linear_logp_context_uses_covered_padded_vocab_when_padded_vocab_size_missing(): + mpu.get_tensor_model_parallel_world_size.return_value = 4 + mpu.get_tensor_model_parallel_rank.return_value = 1 + mpu.get_tensor_model_parallel_group.return_value = "tp_group" + + output_layer = types.SimpleNamespace(weight=torch.empty(8, 4), bias=None) + model = types.SimpleNamespace(output_layer=output_layer, post_process=True) + args = _make_args(padded_vocab_size=None, vocab_size=30) + + context = rlk_mod.get_linear_logp_context_from_model(args, model) + + assert context is not None + assert context.vocab_start_index == 8 + assert context.global_vocab_size == 32 + + +@pytest.mark.unit +def test_return_hidden_states_for_linear_logp_restores_post_process_flag(): + args = _make_args() + model = types.SimpleNamespace(post_process=True) + context = rlk_mod.LinearLogpContext( + lm_head_weight=torch.empty(4, 3), + bias=None, + tp_group=None, + ) + + with rlk_mod.return_hidden_states_for_linear_logp(args, model, context) as enabled: + assert enabled is True + assert model.post_process is False + + assert model.post_process is True + + +@pytest.mark.unit +def test_policy_loss_only_skips_entropy_when_linear_logp_context_is_active(monkeypatch): + args = _make_args(enable_rl_kernel=True, rl_kernel_ops=("linear_logp",), entropy_coef=0.0) + context = rlk_mod.LinearLogpContext( + lm_head_weight=torch.empty(4, 3), + bias=None, + tp_group=None, + ) + + monkeypatch.delenv("VIME_SKIP_ZERO_ENTROPY_METRIC", raising=False) + assert loss_mod._policy_loss_needs_entropy(args, None) is True + assert loss_mod._policy_loss_needs_entropy(args, context) is False + + monkeypatch.setenv("VIME_SKIP_ZERO_ENTROPY_METRIC", "1") + assert loss_mod._policy_loss_needs_entropy(args, None) is False + assert loss_mod._policy_loss_needs_entropy(args, context) is False + + args.entropy_coef = 0.01 + assert loss_mod._policy_loss_needs_entropy(args, None) is True + assert loss_mod._policy_loss_needs_entropy(args, context) is True diff --git a/tests/test_rl_kernel_logp_integration.py b/tests/test_rl_kernel_logp_integration.py new file mode 100644 index 00000000..5d444768 --- /dev/null +++ b/tests/test_rl_kernel_logp_integration.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import sys +import types +from argparse import Namespace +from pathlib import Path + +import pytest +import torch + +_tests_root = Path(__file__).resolve().parent +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs + +_unit_stubs.install_megatron_mpu_stub() +_unit_stubs.install_vime_distributed_utils_stub() + +from megatron.core import mpu # noqa: E402 + +from vime.backends.megatron_utils import loss as loss_mod # noqa: E402 +from vime.backends.megatron_utils import rl_kernel as rlk_mod # noqa: E402 + + +NUM_GPUS = 0 + + +class _FakeLogpOp: + calls = 0 + + def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + type(self).calls += 1 + return torch.gather(torch.log_softmax(logits.float(), dim=-1), -1, token_ids.long().unsqueeze(-1)).squeeze(-1) + + +def _reset_rl_kernel_state(): + rlk_mod._LOGP_OP = None + rlk_mod._LOGP_OP_LOAD_ERROR = None + rlk_mod._LINEAR_LOGP_OP = None + rlk_mod._LINEAR_LOGP_OP_LOAD_ERROR = None + rlk_mod._WARNED_FALLBACK_REASONS.clear() + rlk_mod._FALLBACK_COUNTS.clear() + rlk_mod._FALLBACK_COUNTS.update({"logp": 0, "linear_logp": 0}) + _FakeLogpOp.calls = 0 + + +def _install_fake_rl_engine(monkeypatch): + rl_engine = types.ModuleType("rl_engine") + kernels = types.ModuleType("rl_engine.kernels") + registry = types.ModuleType("rl_engine.kernels.registry") + registry.kernel_registry = types.SimpleNamespace(get_op=lambda name: _FakeLogpOp()) + monkeypatch.setitem(sys.modules, "rl_engine", rl_engine) + monkeypatch.setitem(sys.modules, "rl_engine.kernels", kernels) + monkeypatch.setitem(sys.modules, "rl_engine.kernels.registry", registry) + + +def _make_args(**overrides) -> Namespace: + values = { + "enable_rl_kernel": True, + "rl_kernel_ops": ("logp",), + "rl_kernel_strict": False, + "allgather_cp": False, + "qkv_format": "thd", + "rollout_temperature": 1.0, + "log_probs_chunk_size": -1, + } + values.update(overrides) + return Namespace(**values) + + +@pytest.fixture(autouse=True) +def reset_parallelism(): + _reset_rl_kernel_state() + mpu.get_tensor_model_parallel_world_size.return_value = 1 + mpu.get_tensor_model_parallel_group.return_value = None + mpu.get_context_parallel_world_size.return_value = 1 + mpu.get_context_parallel_rank.return_value = 0 + yield + _reset_rl_kernel_state() + + +@pytest.mark.unit +def test_rl_kernel_logp_matches_vime_response_slicing(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args() + vocab_size = 257 + total_lengths = [9, 7, 11] + response_lengths = [4, 3, 6] + torch.manual_seed(17) + unconcat_tokens = [torch.randint(0, vocab_size, (length,), dtype=torch.long) for length in total_lengths] + logits = torch.randn(1, sum(total_lengths), vocab_size, dtype=torch.float32) + + with torch.no_grad(): + _, result = loss_mod.get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=False, + ) + + expected = [] + offset = 0 + log_probs = torch.log_softmax(logits.squeeze(0), dim=-1) + for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False): + start = offset + total_length - response_length + end = offset + total_length + shifted_tokens = tokens[-response_length:] + expected.append(torch.gather(log_probs[start - 1 : end - 1], -1, shifted_tokens.unsqueeze(-1)).squeeze(-1)) + offset += total_length + + assert _FakeLogpOp.calls == 1 + assert len(result["log_probs"]) == len(expected) + for actual, expected_item in zip(result["log_probs"], expected, strict=True): + torch.testing.assert_close(actual, expected_item, rtol=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_rl_kernel_logp_falls_back_when_entropy_requested(monkeypatch): + _install_fake_rl_engine(monkeypatch) + args = _make_args() + logits = torch.randn(8, 16, dtype=torch.float32) + tokens = torch.randint(0, 16, (8,), dtype=torch.long) + + with torch.no_grad(): + actual = rlk_mod.maybe_compute_logp(logits, tokens, args=args, with_entropy=True) + + assert actual is None + assert _FakeLogpOp.calls == 0 + + +@pytest.mark.unit +def test_rl_kernel_logp_falls_back_for_tensor_parallel_vocab(monkeypatch): + _install_fake_rl_engine(monkeypatch) + mpu.get_tensor_model_parallel_world_size.return_value = 2 + args = _make_args() + logits = torch.randn(8, 16, dtype=torch.float32) + tokens = torch.randint(0, 16, (8,), dtype=torch.long) + + with torch.no_grad(): + actual = rlk_mod.maybe_compute_logp(logits, tokens, args=args, with_entropy=False) + + assert actual is None + assert _FakeLogpOp.calls == 0 + + +@pytest.mark.unit +def test_rl_kernel_logp_strict_mode_raises_on_unsupported_parallelism(monkeypatch): + _install_fake_rl_engine(monkeypatch) + mpu.get_tensor_model_parallel_world_size.return_value = 2 + args = _make_args(rl_kernel_strict=True) + logits = torch.randn(8, 16, dtype=torch.float32) + tokens = torch.randint(0, 16, (8,), dtype=torch.long) + + with pytest.raises(RuntimeError, match="tensor-parallel vocab shards"): + with torch.no_grad(): + rlk_mod.maybe_compute_logp(logits, tokens, args=args, with_entropy=False) + + +@pytest.mark.unit +def test_rl_kernel_logp_falls_back_when_optional_package_missing(monkeypatch): + import builtins + + for name in list(sys.modules): + if name == "rl_engine" or name.startswith("rl_engine."): + monkeypatch.delitem(sys.modules, name, raising=False) + + original_import = builtins.__import__ + + def fail_rl_engine_import(name, *args, **kwargs): + if name == "rl_engine" or name.startswith("rl_engine."): + raise ModuleNotFoundError("No module named 'rl_engine'") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fail_rl_engine_import) + args = _make_args() + logits = torch.randn(8, 16, dtype=torch.float32) + tokens = torch.randint(0, 16, (8,), dtype=torch.long) + + with torch.no_grad(): + actual = rlk_mod.maybe_compute_logp(logits, tokens, args=args, with_entropy=False) + + assert actual is None 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..e3bbff85 --- /dev/null +++ b/vime-RLK-final-metrics-config-matrix.md @@ -0,0 +1,167 @@ +# vime + RL-Kernel linear_logp 8 卡训推 12 轮配置矩阵 + +日期：2026-07-07 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` + +## 当前实验进展 + +截至 2026-07-07 15:40 UTC： + +- `RL-Kernel`：已切到 PR 211 对应分支；`vime`：已切到 PR 3 对应分支。 +- 已完成 8 卡 12 轮 candidate：T01、T03、T06，均命中 `FusedLinearLogpSM90Op` fused-tile full-gradient fast path，fallback=0。 +- 已完成 8 卡 12 轮 baseline：T06。T06 candidate 单算子稳定窗口显著快于 baseline，且 peak reserved 更低。 +- 已完成 8 卡 12 轮 baseline：T01、T03、T06。T01/T03/T06 都使用 `TRACE_MODE=none`，用于和 no-trace candidate 对齐。 +- T07 no-trace candidate 已跑过首轮并确认 fast path/fallback=0，随后按要求停止并重启为 trace run；后续 no-trace 重启也已按用户要求停止，未纳入正式指标。 +- T07 trace candidate 使用 `TRACE_MODE=all TRACE_ROLLOUTS=3`，已完成 step 2/11，rollout 3 的 rollout-generate trace 已触发 `cudaProfilerStart`；因 profiler stop/落盘阶段未及时返回，已按后续测试优先级主动中断，trace 文件未落盘。 +- 正式 metrics run 使用 `TRACE_MODE=none` 降低 profiler overhead；trace run 只作为诊断，不混入 baseline/candidate timing 对比。 + +## 固定环境 + +| 项 | 值 | +| --- | --- | +| 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 必须先跑通 | candidate success; baseline success | +| 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 | candidate success; baseline success on attempt 2 | +| 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 | candidate success; baseline success | +| T07 | OOM probe | `RBS=4; NSP=2; GBS=8; MAX_TOKENS=8192; RESP=3968; VLLM_MEM=0.46; VLLM_MAX_MODEL_LEN=4096` | T06 成功且显存余量足够时跑 | candidate trace/no-trace attempts interrupted; not included in metrics | +| 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 | success | success | 20.60 | 19.89 | 104.18 | 105.76 | 17.56 | 20.02 | 16.95 | 18.88 | 24.34 | 24.16 | 10.53 | 10.60 | 0.0000 | 0.0000 | 0.02668 | 0.02537 | pass | pass | +| T03 | success | success | 24.14 | 22.49 | 155.63 | 156.55 | 17.78 | 20.19 | 16.66 | 19.32 | 74.56 | 72.68 | 10.30 | 10.57 | 0.0278 | 0.0278 | 0.02264 | 0.02395 | pass | pass | +| T05 | | | | | | | | | | | | | | | | | | | | | +| T06 | success | success | 49.26 | 46.23 | 232.20 | 228.40 | 21.84 | 21.92 | 20.54 | 21.62 | 146.76 | 143.07 | 20.33 | 21.05 | 0.1111 | 0.0972 | 0.02034 | 0.02224 | pass | pass | +| T07 | | trace interrupted step 2/11 | | 59.06 | | | | | | | | | | | | | | | | pass until interrupted | +| 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 | success | success | 796.44 | 796.44 | 796.44 | 796.44 | 5.14* | 3.51 | 1.46x | 15.24* | 12.37 | 1.23x | 4.81* | 3.88 | 996.13 | 995.77 | 4056 | 3112 | 0 | +| T03 | success | success | 1820.44 | 1820.44 | 1820.44 | 1820.44 | 7.78* | 3.36 | 2.32x | 18.56* | 10.37 | 1.79x | 6.61* | 3.06 | 2149.72 | 2013.54 | 6684 | 4862 | 0 | +| T05 | | | | | | | | | | | | | | | | | | | | +| T06 | success | success | 6769.78 | 6826.67 | 6769.78 | 6826.67 | 14.52* | 7.62* | 1.91x | 33.96* | 18.50* | 1.84x | 8.37* | 5.34* | 7724.49 | 7529.68 | 32342 | 26710 | 0 | +| T07 | | trace interrupted step 2/11 | | | | | | | | | | | | | | 10046.53 until interrupted | | 38922 until interrupted | 0 until interrupted | +| T08 | | | | | | | | | | | | | | | | | | | | + +解读： + +- `save-logits` 必须命中 fused-tile full-gradient fast path，fallback=0。 +- 主指标是 `fwd+bwd CUDA ms`；`dispatch ms` 只作辅助。 +- 如果 candidate 能跑通而 baseline OOM，单独记录为 candidate 最大不 OOM 优势。 +- `*`：baseline 单算子有偶发 native logprob spike，因此单算子主比较使用剔除 spike 的稳定窗口。T01 baseline step 10 出现 407/558 ms fwd/fwd+bwd spike，表中使用 3-9、11 稳定窗口；T01 3-11 median 为 5.25/15.22 ms。T03 baseline step 10 出现 391/567 ms spike，表中使用 3-9、11 稳定窗口；T03 3-11 median 为 8.05/18.43 ms。T06 baseline step 11 出现 4-5 秒级 spike，表中使用 3-10 稳定窗口；T06 3-11 median 为 14.75/34.31 ms。 + +## T06 candidate vs baseline 结论 + +| 指标 | T06 baseline | T06 candidate | 结论 | +| --- | ---: | ---: | --- | +| run status | success | success | 都完整完成 12 轮 | +| rollout time s, mean 3-11 | 146.76 | 143.07 | candidate 更快 | +| tokens/GPU/s, mean 3-11 | 20.33 | 21.05 | candidate +3.5% | +| step time s, mean 3-11 | 232.20 | 228.40 | candidate 更快 | +| train time s, mean 3-11 | 21.84 | 21.92 | 基本持平 | +| peak reserved GB, max 0-11 | 49.26 | 46.23 | candidate 少约 3.03 GB | +| tokens/call, mean 3-11 | 6769.78 | 6826.67 | 接近 | +| forward CUDA ms, stable 3-10 | 14.52 | 7.62 | candidate 约 1.91x 快 | +| fwd+bwd CUDA ms, stable 3-10 | 33.96 | 18.50 | candidate 约 1.84x 快 | +| dispatch ms, stable 3-10 | 8.37 | 5.34 | candidate 更快 | +| abs diff, mean 3-11 | 0.02034 | 0.02224 | 同一量级 | + +## 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 | `/workspace/logs/8gpu_T01_full_cuda_20260707_103211/ray_job_cuda.log` | +| T01 full baseline | `/workspace/logs/8gpu_T01_full_baseline_20260707_142633/ray_job_baseline.log` | +| T03 full baseline attempt 1 failed | `/workspace/logs/8gpu_T03_full_baseline_20260707_145557/ray_job_baseline.log` | +| T03 full baseline attempt 2 success | `/workspace/logs/8gpu_T03_full_baseline_20260707_150229/ray_job_baseline.log` | +| T03 full save-logits | `/workspace/logs/8gpu_T03_full_cuda_20260707_105946/ray_job_cuda.log` | +| T06 full baseline | `/workspace/logs/8gpu_T06_full_baseline_20260707_124934/ray_job_baseline.log` | +| T06 full save-logits | `/workspace/logs/8gpu_T06_full_cuda_20260707_115418/ray_job_cuda.log` | +| T07 full save-logits | `/workspace/logs/8gpu_T07_full_cuda_20260707_135416/ray_job_cuda.log` | +| T07 full save-logits no-trace rerun interrupted | `/workspace/logs/8gpu_T07_full_cuda_20260707_154241/ray_job_cuda.log` | +| T07 trace dir | `/workspace/nsys_traces/8gpu_T07_full_cuda_20260707_135416/` | +| 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-RLK-long-run-practice-report.md b/vime-RLK-long-run-practice-report.md new file mode 100644 index 00000000..43ee96a6 --- /dev/null +++ b/vime-RLK-long-run-practice-report.md @@ -0,0 +1,2200 @@ +# vime + RL-Kernel 8 卡长跑实践记录 + +日期：2026-07-07 UTC + +本文记录本轮围绕 `linear_logp` 的 8xH100 长时间验证经验。范围包括环境搭建、分支和运行方式、T01/T03/T06 已完成结果、T07 trace 尝试状态、调试中做过的代码级改动，以及 baseline 与 candidate 算子实现的差异。 + +本文只纳入 `TRACE_MODE=none` 的完整 12 轮结果做性能比较。trace run 只用于诊断，不混入 timing 表。 + +## 实验目标 + +本轮目标不是单独跑一个 microbenchmark，而是在 vime 的完整训推链路里验证 RL-Kernel `linear_logp`： + +- 模型：Qwen3-30B-A3B。 +- 机器：8x H100 80GB。 +- vime 分支：`pr-3`。 +- RL-Kernel 分支：PR 211 对应分支。 +- 训练并行：TP=2、PP=1、CP=1、EP=8。 +- rollout：完整 vLLM rollout，`NUM_ROLLOUT=12`，不使用 debug rollout。 +- 对比： + - baseline：Megatron/vime 原生 output layer + 原生 selected logprob。 + - candidate：RL-Kernel `FusedLinearLogpSM90Op`，`save-logits` / fused-tile full-gradient path。 + +本文里的 candidate 特指 `save-logits` / full-gradient fused-tile 路径。它对应 `TRAIN_SCOPE=full`、`RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0`、`RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=1`，日志应出现 `Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path.`。它不是 output-layer-only 的 `save-prob` 路径，也不是一个单独叫 `save-logp` 的模式。 + +最终要回答两个问题： + +1. candidate 在真实 8 卡训推流程中是否稳定完成 12 轮，是否 fallback=0。 +2. candidate 的 `linear_logp` 单算子是否比 baseline 更快、显存更低，并解释为什么。 + +## 工作区和环境 + +用户要求“环境尽可能装在 workspace”，本轮按这个原则处理： + +| 项 | 路径或设置 | +| --- | --- | +| workspace | `/workspace` | +| vime | `/workspace/vime` | +| RL-Kernel | `/workspace/RL-Kernel` | +| Megatron-LM | `/workspace/Megatron-LM` | +| Python env | `/workspace/vime-rlk-env` | +| HF cache | `/workspace/.cache/huggingface` | +| pip cache | `/workspace/.cache/pip` | +| torch cache | `/workspace/.cache/torch` | +| temp dir | `/workspace/.cache/tmp` | +| FlashInfer workspace | `/workspace/.cache/flashinfer` | +| logs | `/workspace/logs` | +| traces | `/workspace/nsys_traces` | + +脚本里已经把这些 cache 环境变量写入 workspace： + +```bash +export HF_HOME="${WORKSPACE_ROOT}/.cache/huggingface" +export HUGGINGFACE_HUB_CACHE="${HF_HOME}/hub" +export TRANSFORMERS_CACHE="${HF_HOME}/transformers" +export PIP_CACHE_DIR="${WORKSPACE_ROOT}/.cache/pip" +export TORCH_HOME="${WORKSPACE_ROOT}/.cache/torch" +export XDG_CACHE_HOME="${WORKSPACE_ROOT}/.cache" +export TMPDIR="${WORKSPACE_ROOT}/.cache/tmp" +export FLASHINFER_WORKSPACE_BASE="${WORKSPACE_ROOT}" +export FLASHINFER_CUDA_ARCH_LIST="9.0" +``` + +这点很重要：Qwen3-30B-A3B、vLLM、FlashInfer、torch extension 会产生较多缓存和临时文件，如果落到默认 HOME 或系统临时目录，长跑中更容易遇到空间、权限或不可复现问题。 + +## 运行入口 + +主入口是： + +```bash +cd /workspace/vime +WORKSPACE_ROOT=/workspace \ +VIME_PYTHON_ENV=/workspace/vime-rlk-env \ +TRACE_MODE=none \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T01 cuda +``` + +其中： + +- `MODE=cuda` 表示 candidate，开启 RL-Kernel `linear_logp`。 +- `MODE=baseline` 表示关闭 RL-Kernel，打开 baseline timer。 +- `TRACE_MODE=none` 是正式性能数据口径。 +- `TRACE_MODE=all TRACE_ROLLOUTS=3` 只用于单独抓 trace。 + +candidate full-gradient 的关键环境变量： + +```bash +export VIME_RL_KERNEL=1 +export VIME_RL_KERNEL_OPS=linear_logp +export VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=cuda +export VIME_RL_KERNEL_CUDA_EVENT_TIMER=1 +export VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN=0 +export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0 +export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=1 +``` + +baseline 的关键环境变量： + +```bash +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 +``` + +清理残留进程的实际经验： + +```bash +/workspace/vime-rlk-env/bin/ray stop --force +pkill -f 'vllm|ray::|redis-server|train' +nvidia-smi +``` + +T03 baseline 第一次失败后，必须清理 Ray/vLLM/redis 残留再重跑，否则下一次 run 容易继承坏状态。 + +## 完整操作步骤 + +这一节按“从空闲机器到拿到一组可用 metrics”的顺序写。命令默认在 `/workspace` 下执行，除非特别说明。 + +### 1. 确认分支和工作树 + +先确认三个仓库的分支状态。不要在未确认分支时直接开跑，因为 vime PR、RL-Kernel PR 和 Megatron 本地兼容补丁缺一项都可能导致数据不可复现。 + +```bash +git -C /workspace/vime status --short --branch +git -C /workspace/RL-Kernel status --short --branch +git -C /workspace/Megatron-LM status --short --branch +``` + +本轮期望状态： + +```text +/workspace/vime branch: pr-3 +/workspace/RL-Kernel branch: pr-211 / origin/pr/211 +/workspace/Megatron-LM local compatibility patches present +``` + +如果分支不对，按下面方式切换： + +```bash +cd /workspace/vime +gh pr checkout 3 + +cd /workspace/RL-Kernel +gh pr checkout 211 +``` + +注意事项： + +- 不要把 GitHub token 写入文档、日志或 shell history。 +- `/workspace/Megatron-LM` 当前是 detached/local patch 状态，不要用 `git reset --hard` 清掉；这些补丁用于本环境跑通 Qwen3-MoE。 +- 如果 `git status` 里已有用户改动，提交 PR 前单独确认，不要为了跑实验而回滚。 + +### 2. 激活 workspace 内环境 + +本轮使用 `/workspace/vime-rlk-env`，并让 cache、临时目录、FlashInfer 编译产物尽量落到 `/workspace`。 + +```bash +export WORKSPACE_ROOT=/workspace +export VIME_PYTHON_ENV=/workspace/vime-rlk-env +export PATH="${VIME_PYTHON_ENV}/bin:${PATH}" + +export HF_HOME="${WORKSPACE_ROOT}/.cache/huggingface" +export HUGGINGFACE_HUB_CACHE="${HF_HOME}/hub" +export TRANSFORMERS_CACHE="${HF_HOME}/transformers" +export PIP_CACHE_DIR="${WORKSPACE_ROOT}/.cache/pip" +export TORCH_HOME="${WORKSPACE_ROOT}/.cache/torch" +export XDG_CACHE_HOME="${WORKSPACE_ROOT}/.cache" +export TMPDIR="${WORKSPACE_ROOT}/.cache/tmp" +export FLASHINFER_WORKSPACE_BASE="${WORKSPACE_ROOT}" +export FLASHINFER_CUDA_ARCH_LIST=9.0 + +mkdir -p \ + "${HF_HOME}" \ + "${HUGGINGFACE_HUB_CACHE}" \ + "${TRANSFORMERS_CACHE}" \ + "${PIP_CACHE_DIR}" \ + "${TORCH_HOME}" \ + "${XDG_CACHE_HOME}" \ + "${TMPDIR}" \ + "${FLASHINFER_WORKSPACE_BASE}/.cache/flashinfer" +``` + +验证 Python 和包路径： + +```bash +which python +python - <<'PY' +import sys +print(sys.executable) +print(sys.version) +PY +``` + +期望 `sys.executable` 是： + +```text +/workspace/vime-rlk-env/bin/python +``` + +### 3. 确认 RL-Kernel CUDA extension 可用 + +candidate 必须命中 `FusedLinearLogpSM90Op`，所以开跑前先确认 extension 里有 SM90 forward/backward 符号。 + +```bash +cd /workspace/RL-Kernel +python - <<'PY' +from rl_engine.kernels.ops.cuda.loss.linear_logp import _C +print("fused_fwd", hasattr(_C, "fused_linear_logp_sm90")) +print("fused_bwd", hasattr(_C, "fused_linear_logp_sm90_backward")) +print("tp_global_fwd", hasattr(_C, "fused_linear_logp_sm90_global_target")) +print("save_probs_fwd", hasattr(_C, "linear_logp_probs_bf16_forward")) +PY +``` + +至少需要： + +```text +fused_fwd True +fused_bwd True +``` + +如果这里失败，先不要跑 8 卡。通常处理顺序是： + +```bash +cd /workspace/RL-Kernel +KERNEL_ALIGN_FORCE_SM90=1 pip install -e . +``` + +然后重新执行上面的 Python 检查。 + +### 4. 清理上一轮残留进程 + +每次完整 8 卡 run 之前都要清理 Ray、vLLM、redis 和训练进程。尤其是 EngineDead、OOM、中断 trace 之后，残留进程会导致下一轮端口冲突、GPU 显存残留或 Ray object store 状态污染。 + +```bash +/workspace/vime-rlk-env/bin/ray stop --force || true +pkill -f 'vllm|ray::|redis-server|train' || true +sleep 5 +nvidia-smi +``` + +理想状态是每张 H100 只剩很低的 context 占用，例如几 MB。如果某张卡还有几十 GB，说明仍有进程没杀掉，需要用： + +```bash +nvidia-smi +ps -ef | rg 'vllm|ray|python|train' +``` + +定位 PID 后再处理。 + +### 5. 确认 benchmark 脚本的配置展开 + +入口脚本是： + +```text +/workspace/vime/scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh +``` + +脚本先根据 `CONFIG_ID` 展开 T01/T03/T06/T07 的关键配置。T07 当前配置在脚本里等价于： + +```bash +ROLLOUT_BATCH_SIZE=4 +N_SAMPLES_PER_PROMPT=2 +GLOBAL_BATCH_SIZE=8 +MAX_TOKENS_PER_GPU=8192 +ROLLOUT_MAX_RESPONSE_LEN=3968 +VLLM_GPU_MEMORY_UTILIZATION=0.46 +VLLM_MAX_MODEL_LEN=4096 +``` + +并统一设置： + +```bash +NUM_GPUS=8 +MEGATRON_TP=2 +MEGATRON_EP=8 +MEGATRON_CP=1 +ROLLOUT_NUM_GPUS_PER_ENGINE=8 +NUM_ROLLOUT=12 +VIME_DISABLE_SAVE=1 +VIME_SKIP_EVAL_BEFORE_TRAIN=1 +VIME_USE_KL_LOSS=0 +VIME_SKIP_ZERO_ENTROPY_METRIC=1 +VIME_VLLM_ENFORCE_EAGER=1 +VLLM_WORKER_MULTIPROC_METHOD=spawn +VIME_TRAIN_MEMORY_MARGIN_BYTES=1073741824 +VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK=1 +MEGATRON_ALLOW_MOE_TP_WITHOUT_SP=1 +CUDA_MODULE_LOADING=EAGER +``` + +这些环境变量不是装饰项： + +- `VIME_SKIP_ZERO_ENTROPY_METRIC=1` 让 `entropy_coef=0` 时不计算无用 entropy，避免 candidate 因 entropy fallback。 +- `VLLM_WORKER_MULTIPROC_METHOD=spawn` 避免 vLLM worker 继承复杂父进程 CUDA 状态。 +- `VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK=1` 只绕过 Megatron validation 阶段的 capability check，runtime 仍保留 grouped GEMM。 +- `MEGATRON_ALLOW_MOE_TP_WITHOUT_SP=1` 是本地 Megatron 兼容补丁的开关。 + +### 6. candidate 与 baseline 在脚本里的差异 + +脚本根据第二个参数 `MODE` 设置两套路径。 + +candidate： + +```bash +MODE=cuda +VIME_RL_KERNEL=1 +VIME_BASELINE_LINEAR_LOGP_TIMER=0 +VIME_BASELINE_CUDA_EVENT_TIMER=0 +VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=cuda +VIME_RL_KERNEL_CUDA_EVENT_TIMER=1 +VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN=0 +RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0 +RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=1 +``` + +baseline： + +```bash +MODE=baseline +VIME_RL_KERNEL=0 +VIME_BASELINE_LINEAR_LOGP_TIMER=1 +VIME_BASELINE_CUDA_EVENT_TIMER=1 +RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0 +RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=0 +``` + +full-gradient 主结果都使用： + +```bash +TRAIN_SCOPE=full +``` + +`output_layer` scope 只用于可选补充，不作为本轮主结论。 + +### 7. 启动正式 no-trace run + +正式性能 run 必须用： + +```bash +TRACE_MODE=none +``` + +T07 candidate 本轮重跑命令如下： + +```bash +cd /workspace/vime +WORKSPACE_ROOT=/workspace \ +RUN_ROOT=/workspace \ +VIME_PYTHON_ENV=/workspace/vime-rlk-env \ +TRACE_MODE=none \ +TRAIN_SCOPE=full \ +RUN_NAME=8gpu_T07_full_cuda_20260707_155623 \ +VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T07 cuda +``` + +预期日志路径： + +```text +/workspace/logs/8gpu_T07_full_cuda_20260707_155623/ray_job_cuda.log +``` + +为什么设置 `VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728`： + +- 这个值是 128MiB。 +- 8 卡长跑中 update weight 是固定开销之一。 +- 分桶太大容易让显存峰值和 Ray object 传输抖动变大。 +- 分桶日志可以帮助区分“算子慢”和“权重同步慢”。 + +### 8. 启动后第一轮必须检查的日志 + +启动后先确认 Ray job 进入训练，而不是卡在环境安装或 vLLM 启动。 + +```bash +tail -f /workspace/logs/8gpu_T07_full_cuda_20260707_155623/ray_job_cuda.log +``` + +另开一个 shell 用 `rg` 查关键行： + +```bash +LOG=/workspace/logs/8gpu_T07_full_cuda_20260707_155623/ray_job_cuda.log +rg -n \ + "Using RL-Kernel linear_logp op|fused-tile|rl_kernel_fallback_count|perf [0-9]+:|step [0-9]+:|train_rollout_logprob_abs_diff|CUDA out of memory|EngineDead|ConnectionRefusedError" \ + "$LOG" +``` + +必须看到： + +```text +Using RL-Kernel linear_logp op: FusedLinearLogpSM90Op +Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path. +``` + +并且每轮 `fallback` 必须保持 0。若出现 fallback 非 0，这组 candidate 数据不能作为宣传结果，需要先 debug。 + +### 9. 代码级路径：candidate 怎么绕开完整 logits + +T07 candidate run 的关键代码路径如下。 + +第一步，vime 根据参数启用 RL-Kernel： + +```text +/workspace/vime/scripts/run-qwen3-30B-A3B.sh +``` + +脚本把参数传进 Ray runtime env： + +```bash +RLK_ARGS+=(--enable-rl-kernel --rl-kernel-ops "${VIME_RL_KERNEL_OPS:-linear_logp}") +``` + +第二步，Megatron model 构建后，vime 从 model 里取 output layer 权重和 TP 信息： + +```text +/workspace/vime/vime/backends/megatron_utils/rl_kernel.py +get_linear_logp_context_from_model() +``` + +核心字段： + +```python +LinearLogpContext( + lm_head_weight=weight, + bias=bias, + tp_group=tp_group, + vocab_start_index=vocab_start_index, + global_vocab_size=global_vocab_size, + sequence_parallel=..., +) +``` + +第三步，前向时临时让 Megatron 返回 hidden states，而不是让 output layer 直接产出 logits： + +```python +@contextmanager +def return_hidden_states_for_linear_logp(args, model, context): + old_post_process = module.post_process + module.post_process = False + try: + yield True + finally: + module.post_process = old_post_process +``` + +第四步，loss 侧把 hidden states 展平，构建 shifted target token，然后调用： + +```text +/workspace/vime/vime/backends/megatron_utils/loss.py +get_log_probs_and_entropy() +``` + +关键分支： + +```python +if linear_logp_context is not None: + log_prob_full = maybe_compute_linear_logp( + logits, + full_tokens, + context=linear_logp_context, + args=args, + with_entropy=with_entropy, + ) +``` + +这里变量名仍叫 `logits`，但在 candidate 路径它实际是 hidden states；如果 RL-Kernel 不能使用，才会调用 `_materialize_linear_logits()` 退回完整 logits。 + +第五步，真正调用 RL-Kernel： + +```text +/workspace/vime/vime/backends/megatron_utils/rl_kernel.py +maybe_compute_linear_logp() +``` + +核心调用： + +```python +log_prob = op( + hidden_states, + weight, + target_ids.long(), + bias, + tp_group=context.tp_group, + vocab_start_index=context.vocab_start_index, + global_vocab_size=context.global_vocab_size, +) +``` + +第六步，RL-Kernel 选择 SM90 fused path： + +```text +/workspace/RL-Kernel/rl_engine/kernels/ops/cuda/loss/linear_logp.py +FusedLinearLogpSM90Op.apply() +``` + +T07 full-gradient 期望命中： + +```python +return _TensorParallelLinearLogpFusedTileBF16Function.apply( + hidden, + lm_head_weight, + target_ids, + int(vocab_start_index), + None if global_vocab_size is None else int(global_vocab_size), + tp_group, +) +``` + +第七步，CUDA extension forward： + +```text +/workspace/RL-Kernel/csrc/cuda/fused_linear_logp_sm90.cu +fused_linear_logp_sm90_forward_impl() +``` + +这个函数做的事： + +- 检查 hidden/weight 是 CUDA bf16 contiguous tensor。 +- 用 TMA descriptor 读取 hidden tile 和 weight tile。 +- 在 tile GEMM 中累计 local max、sum exp、target logit。 +- 返回 `[N]` 级别的 `out_value` 和 `lse`，不返回 `[N, V]` logits。 + +第八步，CUDA extension backward： + +```text +/workspace/RL-Kernel/csrc/cuda/fused_linear_logp_sm90.cu +fused_linear_logp_sm90_backward() +``` + +T07 full-gradient 需要 `compute_grad_hidden` 和 `compute_grad_weight`，因此重点看 full fused tile branch： + +```cpp +if ((full_fused_tile_mode == "tile_cublas" || full_fused_tile_mode == "tile" || + full_fused_tile_mode == "streaming" || full_fused_tile_mode == "tiled") && + (compute_grad_hidden || compute_grad_weight) && !compute_grad_bias && + !bias.has_value() && hidden.scalar_type() == at::kBFloat16 && + weight.scalar_type() == at::kBFloat16 && D % BK == 0) { + ... +} +``` + +### 10. 代码级路径：baseline 为什么更重 + +baseline 不启用 `linear_logp_context`，所以 `get_log_probs_and_entropy()` 进入原生路径： + +```python +if linear_logp_context is None: + logits = logits.contiguous() + log_prob_full, entropy_full = calculate_log_probs_and_entropy( + logits, + full_tokens, + tp_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + ) +``` + +对应： + +```text +/workspace/vime/vime/utils/ppo_utils.py +calculate_log_probs_and_entropy() +``` + +核心逻辑： + +```python +log_prob = compute_log_probs(logits.clone(), tokens, tp_group) +``` + +这就是 baseline 的主要额外成本： + +- output layer 已经物化 `[T, V]` logits。 +- logprob 又 clone 一份 logits。 +- softmax/logsumexp/gather 在完整 vocab 维度上做。 +- backward 也围绕完整 logits 图回传。 + +baseline 计时 hook 在： + +```text +/workspace/vime/vime/backends/megatron_utils/model.py +_probe_baseline_output_layer_forward() +``` + +它记录 output layer 的 forward 和 forward+backward CUDA event，再与 native logprob timer 相加，形成： + +```text +train/baseline_linear_logp_forward_cuda_event_elapsed_s_delta +train/baseline_linear_logp_forward_backward_cuda_event_elapsed_s_delta +train/baseline_linear_logp_dispatch_elapsed_s_delta +``` + +### 11. 每轮判定标准 + +一轮可接受 candidate 结果至少满足： + +```text +run_status: success +rl_kernel_fallback_count_delta: 0 +rl_kernel_linear_logp_call_count_delta > 0 +rl_kernel_linear_logp_token_count_delta > 0 +train_rollout_logprob_abs_diff: finite, same order as previous runs +loss/reward: finite +``` + +T01/T03/T06 已完成结果的 abs diff 参考： + +```text +T01 candidate abs diff mean 3-11: 0.02537 +T03 candidate abs diff mean 3-11: 0.02395 +T06 candidate abs diff mean 3-11: 0.02224 +``` + +T07 如果 abs diff 是同一量级，可以认为 correctness 没有明显异常；如果突然变成很大或 NaN，要先查 rollout/train logprob 对齐和 target 构造。 + +### 12. metrics 提取步骤 + +正式统计使用 rollout 3-11 均值，避开前 3 轮 warmup。 + +T07 完成后先查这些行： + +```bash +LOG=/workspace/logs/8gpu_T07_full_cuda_20260707_155623/ray_job_cuda.log + +rg -n \ + "perf [3-9]:|perf 1[01]:|step [3-9]:|step 1[01]:|rl_kernel_linear_logp|train_rollout_logprob_abs_diff|raw_reward|peak_reserved|fallback" \ + "$LOG" +``` + +需要写入最终矩阵的字段： + +```text +run_status +peak_vram_gb +step_time_s mean 3-11 +train_time_s mean 3-11 +actor_train_time_s mean 3-11 +rollout_time_s mean 3-11 +tokens_per_gpu_per_sec mean 3-11 +raw_reward mean 3-11 +train_rollout_logprob_abs_diff mean 3-11 +rl_kernel_linear_logp_tokens_per_call_delta mean 3-11 +rl_kernel_linear_logp_forward_cuda_event_elapsed_s_delta mean 3-11 +rl_kernel_linear_logp_forward_backward_cuda_event_elapsed_s_delta mean 3-11 +rl_kernel_linear_logp_dispatch_elapsed_s_delta mean 3-11 +peak_alloc_delta MB +peak_reserved_delta MB +fallback +``` + +如果 T07 candidate 成功，再决定是否补 baseline。baseline 命令为： + +```bash +cd /workspace/vime +WORKSPACE_ROOT=/workspace \ +RUN_ROOT=/workspace \ +VIME_PYTHON_ENV=/workspace/vime-rlk-env \ +TRACE_MODE=none \ +TRAIN_SCOPE=full \ +RUN_NAME=8gpu_T07_full_baseline_YYYYMMDD_HHMMSS \ +VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T07 baseline +``` + +baseline 完成后再用同样 rollout 3-11 口径更新 `vime-RLK-final-metrics-config-matrix.md`。 + +### 13. T07 重新启动前的当前阻塞状态 + +2026-07-07 15:59 UTC 准备按以下命令重新启动 T07 candidate no-trace： + +```bash +cd /workspace/vime +WORKSPACE_ROOT=/workspace \ +RUN_ROOT=/workspace \ +VIME_PYTHON_ENV=/workspace/vime-rlk-env \ +TRACE_MODE=none \ +TRAIN_SCOPE=full \ +RUN_NAME=8gpu_T07_full_cuda_20260707_155623 \ +VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T07 cuda +``` + +启动前检查发现每卡仍有约 41GB 显存占用： + +```text +GPU0 41569 MiB +GPU1 41713 MiB +GPU2 41713 MiB +GPU3 41713 MiB +GPU4 41713 MiB +GPU5 41713 MiB +GPU6 41713 MiB +GPU7 40993 MiB +``` + +`nvidia-smi` 主进程表和 `nvidia-smi pmon` 没列出进程，但 NVML compute-app query 能看到宿主侧不可见 PID： + +```text +2195449 [Not Found] GPU0 41560 MiB +2195450 [Not Found] GPU1 41704 MiB +2195451 [Not Found] GPU2 41704 MiB +2195452 [Not Found] GPU3 41704 MiB +2195453 [Not Found] GPU4 41704 MiB +2195454 [Not Found] GPU5 41704 MiB +2195455 [Not Found] GPU6 41704 MiB +2195456 [Not Found] GPU7 40984 MiB +``` + +在当前容器里 `/proc/2195449` 到 `/proc/2195456` 不存在，`kill -9` 返回 `No such process`，说明这些 PID 不在当前 PID namespace。`nvidia-smi --gpu-reset -i 0,1,2,3,4,5,6,7` 返回 `Not Supported`。因此 T07 暂时不能安全启动；否则会在已有 41GB 占用上叠加训练/rollout 显存，基本确定 OOM。 + +处理建议： + +1. 在宿主机或拥有宿主 PID namespace 的管理端 kill `2195449-2195456`。 +2. 或释放对应占用 GPU 的外部容器/作业。 +3. 释放后确认 `nvidia-smi --query-gpu=index,memory.used` 每卡回到几 MB。 +4. 再执行上面的 T07 candidate no-trace 命令。 + +### 14. 异常处理步骤 + +如果 T07 candidate OOM： + +1. 保存 log，不覆盖。 +2. 记录最后一轮 peak reserved 和 OOM 栈。 +3. 清理 Ray/vLLM 进程。 +4. 回退 T08：`RESP=3584`，`VLLM_MEM=0.45`。 +5. 不补 T07 baseline，除非需要证明 baseline 也 OOM。 + +如果 T07 candidate fallback 非 0： + +1. 查是否缺少 `Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path.`。 +2. 查 hidden/weight dtype 是否仍是 bf16。 +3. 查 bias 是否变成非 None。 +4. 查 `with_entropy` 是否被打开。 +5. 查 `VIME_SKIP_ZERO_ENTROPY_METRIC=1` 是否传进 Ray runtime env。 +6. 查 `RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=1` 是否生效。 + +如果 vLLM wakeup 或 onload weights 失败： + +1. 先看是否有 `ConnectionRefusedError`、`EngineDead`、`APIServer` 退出。 +2. 不要直接重跑同一个 shell；先执行清理命令。 +3. 检查端口和 GPU 进程。 +4. 重跑同一配置，失败 run 单独标记为 startup failure，不纳入 metrics。 + +如果 candidate 完整 step 慢于 baseline： + +1. 先比较单算子 fwd/fwd+bwd，而不是 step time。 +2. 查 rollout time、weight-sync bucket time、Ray object store 和 vLLM wakeup。 +3. 查 tokens/call 是否太小。 +4. 查是否开了 trace。 +5. 查 baseline 是否包含 spike，不能用异常 spike 夸大或误判。 + +## RL 框架级工作流程 + +这一节按 vime 本轮 8 卡 colocate 训练的实际代码路径讲清楚 RL 框架如何工作。重点不是命令行参数，而是 Ray、vLLM、Megatron、loss、权重同步之间的数据和控制流。 + +本轮入口脚本最终提交的 Ray job 是： + +```bash +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node ${NUM_GPUS} \ + --colocate \ + ... +``` + +所以本轮主流程在： + +```text +/workspace/vime/train.py +``` + +不是 `train_async.py`。`train_async.py` 明确 `assert not args.colocate`，而本轮使用 `--colocate`，训练和 rollout 共享同一组 8 张 H100，通过 sleep/wake 和 weight sync 协调显存。 + +### 总体循环 + +`train.py` 的主循环可以抽象成： + +```python +def train(args): + pgs = create_placement_groups(args) + rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + actor_model, critic_model = create_training_models(args, pgs, rollout_manager) + + if args.offload_rollout: + rollout_manager.onload_weights() + + actor_model.update_weights() + + if args.offload_rollout: + rollout_manager.onload_kv() + + for rollout_id in range(args.start_rollout_id, args.num_rollout): + rollout_data_ref = rollout_manager.generate.remote(rollout_id) + + if args.offload_rollout: + rollout_manager.offload() + + actor_model.async_train(rollout_id, rollout_data_ref) + + if args.offload_rollout: + rollout_manager.onload_weights() + + actor_model.update_weights() + + if args.offload_rollout: + rollout_manager.onload_kv() +``` + +真实代码里还有 eval、save、critic、global dataset、periodic action，但 T01/T03/T06/T07 这条主线就是： + +```text +启动 Ray/vLLM/Megatron +-> 初始 Megatron actor weights 推给 vLLM +-> vLLM rollout 生成 samples +-> samples 转成 Megatron train batch +-> Megatron actor train +-> Megatron 新权重同步到 vLLM +-> 下一轮 rollout +``` + +### Ray placement group：谁占哪张 GPU + +代码位置： + +```text +/workspace/vime/vime/ray/placement_group.py +create_placement_groups() +create_rollout_manager() +create_training_models() +``` + +本轮 `--colocate` 为 True，因此： + +```python +elif args.colocate: + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + rollout_offset = 0 +``` + +含义： + +- Ray 只创建一个包含 8 个 GPU bundle 的 placement group。 +- actor 和 rollout 都从同一个 placement group 里取资源。 +- rollout offset 是 0，说明 vLLM engine 和 Megatron actor 逻辑上共享同一批 GPU。 + +`_create_placement_group()` 做两件关键事： + +1. 创建 `bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]`。 +2. 用临时 `InfoActor` 查询每个 bundle 实际落在哪个 node/GPU，然后按 node 和 GPU ID 排序，得到稳定的 `pg_reordered_bundle_indices` 和 `pg_reordered_gpu_ids`。 + +这一步重要是因为 Ray 的 bundle index 不一定天然等于物理 GPU ID。后续 vLLM 的 `base_gpu_id`、Megatron rank 和 colocate weight sync 都依赖这个排序。 + +### RolloutManager：vLLM rollout 的控制面 + +代码位置： + +```text +/workspace/vime/vime/ray/rollout.py +RolloutManager +RolloutServer +ServerGroup +``` + +`create_rollout_manager()` 创建一个 Ray actor： + +```python +rollout_manager = RolloutManager.options( + num_cpus=1, + num_gpus=0, +).remote(args, pg) +``` + +`RolloutManager.__init__()` 做的事： + +1. 加载数据源： + + ```python + data_source_cls = load_function(self.args.data_source_path) + self.data_source = data_source_cls(args) + ``` + +2. 加载 rollout 函数： + + ```python + self.generate_rollout = load_function(self.args.rollout_function_path) + self.eval_generate_rollout = load_function(self.args.eval_function_path) + ``` + +3. 如果不是 debug train-only，就启动 vLLM servers： + + ```python + self.servers = start_rollout_servers(args, pg) + ``` + +4. 创建 rollout engine lock： + + ```python + self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() + ``` + +这个 lock 在权重同步时用来协调 rollout engines，不让 vLLM 在更新权重时同时生成。 + +### ServerGroup：如何启动 vLLM engine + +代码位置： + +```text +/workspace/vime/vime/ray/rollout.py +ServerGroup.start_engines() +``` + +每个 `ServerGroup` 表示一组同构 vLLM engine。T01/T03/T06/T07 这种单模型单 engine 配置通常只有一个主要 server group。 + +启动逻辑： + +```python +RolloutRayActor = ray.remote(VLLMEngine) + +rollout_engine = RolloutRayActor.options( + num_cpus=num_cpus, + num_gpus=0.2, + scheduling_strategy=PlacementGroupSchedulingStrategy(...), + runtime_env={"env_vars": env_vars}, +).remote( + self.args, + rank=global_rank, + worker_type=self.worker_type, + base_gpu_id=base_gpu_id, + vllm_overrides=self.vllm_overrides, + num_gpus_per_engine=self.num_gpus_per_engine, +) +``` + +这里 Ray actor 只申请 `num_gpus=0.2`，不是因为 vLLM 只用 0.2 张卡，而是因为真正的 vLLM server 是 actor 里再 spawn 出来的子进程。Ray 资源只是占位和调度，实际 CUDA 可见设备通过 `CUDA_VISIBLE_DEVICES` 控制。 + +`base_gpu_id` 来自 placement group 的排序结果，表示该 vLLM engine 从哪张物理 GPU 开始取连续设备。 + +### VLLMEngine：Ray actor 到 vLLM server 子进程 + +代码位置： + +```text +/workspace/vime/vime/backends/vllm_utils/vllm_engine.py +VLLMEngine.init() +launch_server_process() +_build_subprocess_env() +_run_vllm_server() +``` + +`VLLMEngine.init()` 先计算 server args： + +```python +server_args_dict, external_engine_need_check_fields = _compute_server_args(...) +``` + +然后普通本地模式走： + +```python +self._init_normal(server_args_dict) +``` + +`_init_normal()` 里真正启动 vLLM： + +```python +self.process = launch_server_process(server_args_dict) +``` + +`launch_server_process()` 做三件关键事： + +1. 构造子进程环境： + + ```python + env = _build_subprocess_env(server_args_dict) + ``` + +2. 强制 multiprocessing spawn： + + ```python + multiprocessing.set_start_method("spawn", force=True) + p = multiprocessing.Process(target=_run_vllm_server, args=(kwargs, env)) + p.start() + ``` + +3. node rank 0 等 `/health`： + + ```python + _wait_server_healthy(base_url=..., is_process_alive=lambda: p.is_alive()) + ``` + +`_build_subprocess_env()` 是本轮 debug 的关键点之一： + +```python +env["CUDA_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"] +env.setdefault("VLLM_SERVER_DEV_MODE", "1") +env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +env.setdefault("NCCL_CUMEM_ENABLE", "0") +``` + +含义： + +- Ray actor 自己不直接决定 vLLM 用哪些 GPU，vLLM server 子进程通过 `CUDA_VISIBLE_DEVICES` 绑定。 +- `VLLM_WORKER_MULTIPROC_METHOD=spawn` 避免 fork 继承父进程复杂 CUDA/Ray 状态。 +- colocate 模式会把 vime root 补进 `PYTHONPATH`，并允许 vLLM IPC weight update 的序列化： + + ```python + env.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + ``` + +`_run_vllm_server()` 直接调用 vLLM OpenAI server 入口： + +```python +from vllm.entrypoints.cli.serve import ServeSubcommand +... +ServeSubcommand.cmd(args) +``` + +因此 vime 不是手写推理内核，而是把 vLLM 当作一个 HTTP server 管理。 + +### vLLM sleep/wake：colocate 显存协调 + +代码位置： + +```text +/workspace/vime/vime/backends/vllm_utils/vllm_engine.py +release_memory_occupation() +resume_memory_occupation() +``` + +vLLM offload： + +```python +def release_memory_occupation(self, level: int = 2): + self.flush_cache() + response = requests.post(f"http://{self.server_host}:{self.server_port}/sleep", params={"level": level}) +``` + +vLLM onload： + +```python +def resume_memory_occupation(self, tags: list[str] = None): + tags = _normalize_vllm_wake_tags(tags) + response = requests.post(f"http://{self.server_host}:{self.server_port}/wake_up", params=wake_params) +``` + +本轮主要用两个 tag： + +```text +weights +kv_cache +``` + +`train.py` 里可以看到顺序： + +```python +if args.offload_rollout: + rollout_manager.onload_weights() + +actor_model.update_weights() + +if args.offload_rollout: + rollout_manager.onload_kv() +``` + +这表示： + +1. 先 wake vLLM weights，让它能接收新权重。 +2. Megatron actor 把训练后的权重同步到 vLLM。 +3. 再 wake KV cache / CUDA graph 等生成所需状态。 + +训练时则反向： + +```python +rollout_data_ref = rollout_manager.generate.remote(rollout_id) + +if args.offload_rollout: + rollout_manager.offload() + +actor_model.async_train(...) +``` + +也就是 rollout 生成结束后先让 vLLM sleep，释放显存给 Megatron train。 + +### rollout 数据如何变成训练 batch + +代码位置： + +```text +/workspace/vime/vime/ray/rollout.py +RolloutManager.generate() +RolloutManager._get_rollout_data() +RolloutManager._convert_samples_to_train_data() +RolloutManager._split_train_data_by_dp() +``` + +`RolloutManager.generate()` 的主逻辑： + +```python +data, metrics = self._get_rollout_data(rollout_id=rollout_id) +data = self._convert_samples_to_train_data(data) +return self._split_train_data_by_dp(data) +``` + +`_get_rollout_data()` 有两种来源： + +1. `load_debug_rollout_data`：从磁盘读已保存 rollout。 +2. 正常路径： + + ```python + data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False) + metrics = data.metrics + data = data.samples + ``` + +`generate_rollout` 是用户配置的 rollout function，本轮 Qwen3/vLLM 路径会通过 vLLM HTTP server 生成 response，并形成 `Sample`。 + +`_convert_samples_to_train_data()` 把 `Sample` 列表转为训练所需字段： + +```python +train_data = { + "tokens": [sample.tokens for sample in samples], + "response_lengths": [sample.response_length for sample in samples], + "rewards": rewards, + "raw_reward": raw_rewards, + "truncated": ..., + "sample_indices": ..., + "rollout_ids": ..., + "loss_masks": ..., +} +``` + +几个字段的 RL 含义： + +- `tokens`：prompt + response 的完整 token 序列。 +- `response_lengths`：只在 response token 上计算 policy loss。 +- `loss_masks`：哪些 response token 参与 loss。 +- `rewards/raw_reward`：规则 RM 或外部 RM 产生的奖励。 +- `rollout_log_probs`：如果 rollout 侧带回生成时 logprob，可用于 off-policy correction 或 mismatch metric。 +- `rollout_mask_sums`：同一个 rollout 拆成多个 training samples 时，用于保持“每个 rollout 算一次”的归一化口径。 + +奖励后处理在： + +```python +raw_rewards, rewards = self._post_process_rewards(samples) +``` + +例如 GRPO/GSPO 下可以做 group normalization： + +```python +rewards = rewards.reshape(-1, self.args.n_samples_per_prompt) +rewards = rewards - mean +rewards = rewards / (std + 1e-6) +``` + +最后 `_split_train_data_by_dp()` 根据 DP size 和动态 batch schedule 切分给每个 DP rank： + +```python +partitions, micro_batch_indices, num_microbatches, global_batch_sizes = build_dp_schedule(...) +... +rollout_data_refs.append(Box(ray.put(rollout_data))) +``` + +这里返回的是一组 Ray object refs，每个 DP rank 一个 `Box(ray.put(...))`。训练 actor 后面会按自己的 DP rank 取对应切片。 + +### Megatron train actors 如何创建 + +代码位置： + +```text +/workspace/vime/vime/ray/actor_group.py +RayTrainGroup + +/workspace/vime/vime/backends/megatron_utils/actor.py +MegatronTrainRayActor +``` + +`create_training_models()` 创建 actor train group： + +```python +actor_model = allocate_train_group( + args=actor_args, + num_nodes=args.actor_num_nodes, + num_gpus_per_node=args.actor_num_gpus_per_node, + pg=pgs["actor"], +) +``` + +`RayTrainGroup._allocate_gpus_for_actor()` 中： + +```python +TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(MegatronTrainRayActor) +... +actor = TrainRayActor.options( + num_cpus=num_gpus_per_actor, + num_gpus=num_gpus_per_actor, + scheduling_strategy=PlacementGroupSchedulingStrategy(...), +).remote(world_size, rank, master_addr, master_port) +``` + +每个 GPU 一个 Megatron train actor。Ray 层给 actor 放到对应 bundle 上，Megatron 内部再用 `rank/world_size/master_addr/master_port` 初始化 torch distributed。 + +传给 train actor 的关键环境变量来自 `scripts/run-qwen3-30B-A3B.sh` 的 Ray runtime env，例如： + +```text +PYTHONPATH +HF_HOME / TRANSFORMERS_CACHE / TMPDIR +VIME_RL_KERNEL_LINEAR_LOGP_BACKEND +VIME_RL_KERNEL_CUDA_EVENT_TIMER +VIME_SKIP_ZERO_ENTROPY_METRIC +MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ +MEGATRON_ALLOW_MOE_TP_WITHOUT_SP +``` + +`MegatronTrainRayActor.init()` 做的事： + +1. 初始化 torch distributed / Megatron： + + ```python + monkey_patch_torch_dist() + super().init(args, role, ...) + init(args) + ``` + +2. 每个 local GPU 依次读 HF config/tokenizer，避免并发写 cache： + + ```python + for i in range(args.num_gpus_per_node): + if i == dist.get_rank() % args.num_gpus_per_node: + self.hf_config = AutoConfig.from_pretrained(...) + self.tokenizer = AutoTokenizer.from_pretrained(...) + dist.barrier(group=get_gloo_group()) + ``` + +3. 构建 Megatron model/optimizer/scheduler： + + ```python + self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(...) + ``` + +4. 记录训练并行配置，给 rollout DP split 使用： + + ```python + self.train_parallel_config = { + "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False), + "cp_size": mpu.get_context_parallel_world_size(), + "vpp_size": vpp_size, + "microbatch_group_size_per_vp_stage": microbatch_group_size_per_vp_stage, + } + ``` + +5. 创建权重备份器： + + ```python + self.weights_backuper = TensorBackuper.create(...) + self.weights_backuper.backup("actor") + ``` + +6. 根据 colocate 选择权重同步实现： + + ```python + if self.args.colocate: + update_weight_cls = UpdateWeightFromTensor + else: + update_weight_cls = UpdateWeightFromDistributed + self.weight_updater = update_weight_cls(...) + ``` + +本轮是 colocate，所以走 `UpdateWeightFromTensor`。 + +### Megatron actor train：从 Ray object 到 GPU tensor + +代码位置： + +```text +/workspace/vime/vime/backends/megatron_utils/actor.py +MegatronTrainRayActor.train() +MegatronTrainRayActor._get_rollout_data() +MegatronTrainRayActor.train_actor() +``` + +`RayTrainGroup.async_train()` 会对每个 train actor 调： + +```python +actor.train.remote(rollout_id, rollout_data_ref, external_data=...) +``` + +`MegatronTrainRayActor.train()` 的结构： + +```python +if self.args.offload_train: + self.wake_up() + +rollout_data = self._get_rollout_data(rollout_data_ref) + +if self.role == "critic": + result = self.train_critic(...) +else: + self.train_actor(...) + +if self.args.offload_train: + del rollout_data + self.sleep() +``` + +`_get_rollout_data()` 把 CPU/Ray 数据搬到当前 rank 的 GPU： + +```python +rollout_data = process_rollout_data(...) +rollout_data["tokens"] = [ + torch.tensor(t, dtype=torch.long, device=torch.cuda.current_device()) + for t in rollout_data["tokens"] +] +rollout_data["loss_masks"] = [ + torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) + for t in rollout_data["loss_masks"] +] +``` + +如果有 `rollout_log_probs` 或 `teacher_log_probs`，还会按 CP/qkv layout 切到当前 rank 需要的 response 片段： + +```python +slice_log_prob_with_cp(log_prob, total_length, response_length, ...) +``` + +### actor train 内部的 RL 计算顺序 + +代码位置： + +```text +/workspace/vime/vime/backends/megatron_utils/actor.py +MegatronTrainRayActor.train_actor() +``` + +主流程： + +```python +data_iterator = get_data_iterator(rollout_data) +num_microbatches = rollout_data["num_microbatches"] +global_batch_sizes = rollout_data["global_batch_sizes"] + +if self.args.compute_advantages_and_returns: + # 可选：ref / teacher / old_actor logprob + rollout_data.update(self.compute_log_prob(...)) + + # critic values or external values + ... + + compute_advantages_and_returns(self.args, rollout_data) + +log_rollout_data(...) + +train( + rollout_id, + self.model, + self.optimizer, + self.opt_param_scheduler, + data_iterator, + num_microbatches, + global_batch_sizes, +) + +self.weights_backuper.backup("actor") +``` + +对 PPO/GRPO 类训练来说，关键概念是： + +- rollout 阶段拿到 sample/reward。 +- train 阶段重新计算当前 actor 对这些 response token 的 logprob。 +- 结合 reward/advantage 计算 policy loss。 +- 反向更新 Megatron actor。 +- 更新后的 actor 权重再同步回 vLLM，供下一轮 rollout 使用。 + +本轮 `kl_loss_coef=0`、`entropy_coef=0`、`VIME_SKIP_ZERO_ENTROPY_METRIC=1`，所以主关注点变成 actor policy loss 所需的 selected-token logprob，这正是 `linear_logp` 的位置。 + +### Megatron pipeline train step 与 loss_function + +代码位置： + +```text +/workspace/vime/vime/backends/megatron_utils/model.py +train() +train_one_step() +``` + +`train()` 会按 rollout 内的 step 切分调用 `train_one_step()`。`train_one_step()` 定义了给 Megatron pipeline engine 的 `forward_step()`： + +```python +def forward_step(data_iterator, model, return_schedule_plan=False): + batch = get_batch(...) + + linear_logp_context = None + if _train_should_return_hidden_for_linear_logp(args, return_schedule_plan=return_schedule_plan): + linear_logp_context = get_linear_logp_context_from_model(args, model) + + with _probe_baseline_output_layer_forward(args, model): + with return_hidden_states_for_linear_logp(args, model, linear_logp_context): + output_tensor = model(**forward_kwargs) + + return output_tensor, partial( + loss_function, + args, + batch, + num_microbatches, + step_global_batch_size, + rl_kernel_linear_logp_context=linear_logp_context, + ) +``` + +这个函数是 baseline/candidate 分叉的核心： + +- baseline：`linear_logp_context is None`，Megatron model 正常返回 logits，后续 loss 走 `calculate_log_probs_and_entropy(logits, tokens, ...)`。 +- candidate：`linear_logp_context` 非空，`return_hidden_states_for_linear_logp()` 临时让 Megatron 返回 hidden states，loss 里调用 RL-Kernel `linear_logp`，不在 PyTorch 层物化完整 logits。 + +Megatron 的 forward/backward 由： + +```python +forward_backward_func = get_forward_backward_func() +losses_reduced = forward_backward_func( + forward_step_func=..., + data_iterator=data_iterator, + model=model, + num_microbatches=num_microbatches, + ... + forward_only=False, +) +``` + +执行。也就是说，RL-Kernel `linear_logp` 并不是绕过 Megatron 训练；它只是替换 actor loss 中“hidden/output_layer -> selected logprob”这一段，仍然在 Megatron pipeline/DDP/optimizer 框架内参与 autograd。 + +### 权重同步：为什么训练后必须 update_weights + +RL 训练里有两个 actor 副本： + +```text +Megatron actor: 训练副本，负责 forward/backward/optimizer.step +vLLM actor: rollout 副本，负责高吞吐生成 response +``` + +训练后 Megatron actor 权重变了。如果不把新权重同步到 vLLM，下一轮 rollout 仍然用旧策略生成，训练就会偏离 on-policy 目标。 + +`train.py` 因此每轮 train 后调用： + +```python +actor_model.update_weights() +``` + +它最终广播到每个 Megatron rank： + +```python +RayTrainGroup.update_weights() +-> MegatronTrainRayActor.update_weights() +-> self.weight_updater.update_weights() +``` + +本轮 colocate 下 `self.weight_updater` 是： + +```text +/workspace/vime/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +UpdateWeightFromTensor +``` + +### colocate 权重同步的完整数据流 + +代码位置： + +```text +/workspace/vime/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +UpdateWeightFromTensor.update_weights() +UpdateWeightFromTensor._send_hf_params() +_send_to_colocated_engine() + +/workspace/vime/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +HfWeightIteratorDirect.get_hf_weight_chunks() + +/workspace/vime/vime/backends/megatron_utils/update_weight/common.py +named_params_and_buffers() +all_gather_params_async() + +/workspace/vime/vime/backends/vllm_utils/vllm_engine.py +VLLMEngine.update_weights_from_tensor() +``` + +整体流程： + +```text +Megatron sharded params +-> collect global param metadata +-> PP/EP broadcast +-> TP all-gather full param +-> convert Megatron names/layout to HF/vLLM names/layout +-> build CUDA IPC handles +-> Gloo gather IPC payloads to vLLM slot leader +-> Ray call VLLMEngine.update_weights_from_tensor() +-> HTTP POST /update_weights to vLLM server +``` + +`UpdateWeightFromTensor.update_weights()` 先让 vLLM 暂停生成并清 cache： + +```python +if rank == 0: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) +``` + +然后每个 colocated engine 进入 vLLM weight update mode： + +```python +if self._ipc_engine is not None and rank == self._ipc_gather_src: + ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True)) +``` + +接着从 Megatron 取 actor 权重： + +```python +megatron_local_weights = self.weights_getter() +``` + +本轮 `weights_getter` 来自： + +```python +weights_getter=lambda: self.weights_backuper.get("actor") +``` + +也就是刚训练完并 `backup("actor")` 的 actor 参数。 + +### Megatron 参数如何变成 HF/vLLM 参数 + +`HfWeightIteratorDirect.get_hf_weight_chunks()` 是核心： + +```python +for bucket_idx, megatron_local_param_infos in enumerate(self.megatron_local_param_info_buckets, start=1): + megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights) + hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos) + yield hf_named_tensors +``` + +为什么要分 bucket： + +- Qwen3-30B-A3B 参数很大。 +- 一次性 all-gather + 转 HF + IPC 可能打爆显存。 +- `VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728` 把每个 bucket 控制在 128MiB 量级。 +- 本轮日志中的 `[weight-sync] bucket ...` 就来自这里。 + +`_get_megatron_full_params()` 做多级并行收集： + +1. 参数只在 `info.src_rank` 上真实存在，其它 rank 创建 empty tensor。 +2. 如果 PP>1，跨 pipeline parallel group broadcast。 +3. 如果 EP>1，expert 参数跨 expert parallel group broadcast。 +4. 恢复 tensor parallel attrs。 +5. 调 `all_gather_params_async()` 跨 TP/ETP all-gather 成完整权重。 + +```python +gathered_params = all_gather_params_async(list(zip(megatron_local_param_infos, params, strict=False))) +``` + +然后转换成 HF/vLLM 命名： + +```python +hf_named_tensors.extend( + convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config) +) +``` + +本轮对 Qwen3MoE 做了两个关键修正： + +1. layernorm 名字兼容： + + ```python + rest in {"self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"} + rest in {"mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight", "post_attention_layernorm.weight"} + ``` + +2. grouped expert 权重拆成 per-expert： + + ```python + if rest == "mlp.experts.weight1": + expert_tensors = param.view(num_local_experts, args.hidden_size, -1).transpose(-1, -2) + target = "linear_fc1" + elif rest == "mlp.experts.weight2": + expert_tensors = param.view(num_local_experts, -1, args.hidden_size).transpose(-1, -2) + target = "linear_fc2" + ``` + +原因是 Megatron 训练侧的 grouped MoE 参数布局和 vLLM/HF 推理侧的 per-expert 参数布局不同。如果这里没拆对，vLLM 能收到权重，但专家层语义会错。 + +### CUDA IPC 到 vLLM + +HF named tensors 准备好后： + +```python +refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) +ray.get(refs) +``` + +colocate 路径进入： + +```python +_send_to_colocated_engine( + hf_named_tensors, + ipc_engine=self._ipc_engine, + ipc_gather_src=self._ipc_gather_src, + ipc_gather_group=self._ipc_gather_group, + weight_version=self.weight_version, +) +``` + +如果一个 vLLM engine 使用多个 GPU slot，先在 Gloo group 内 gather 每个 rank 的 IPC payload： + +```python +dist.gather_object(payload, object_gather_list=gathered_payloads, dst=ipc_gather_src, group=ipc_gather_group) +``` + +slot leader 合并 payload 后调用 vLLM engine： + +```python +ipc_engine.update_weights_from_tensor.remote(**merged, weight_version=str(weight_version)) +``` + +`VLLMEngine.update_weights_from_tensor()` 再通过 HTTP 调 vLLM server： + +```python +payload = {"names": names, "dtype_names": dtype_names, "shapes": shapes} +payload["ipc_handles_pickled"] = base64.b64encode(cloudpickle.dumps(ipc_handles)).decode("utf-8") +result = self._make_request("update_weights", {"update_info": payload}) +self._weight_version = str(weight_version) +``` + +也就是说，真正的大 tensor 不通过 Ray object store 拷贝；Ray/HTTP 传的是 CUDA IPC handle 和 metadata，vLLM 进程打开 handle 后读取同 GPU 上的 tensor。 + +每个 bucket 完成后释放 IPC cache： + +```python +del long_lived_tensors, hf_named_tensors +torch.cuda.ipc_collect() +``` + +所有 bucket 完成后退出 vLLM weight update mode： + +```python +if self._ipc_engine is not None and rank == self._ipc_gather_src: + ray.get(self._ipc_engine.finish_weight_update.remote()) +``` + +最后恢复生成： + +```python +if rank == 0: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) +``` + +### 分布式权重同步分支 + +本轮是 colocate，所以主路径是 `UpdateWeightFromTensor`。但代码还支持非 colocate： + +```text +/workspace/vime/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +UpdateWeightFromDistributed +``` + +这个分支的思想是： + +```text +Megatron trainer rank 0 + vLLM engine GPUs +-> 建 NCCLWeightTransferEngine group +-> Ray 传 metadata +-> NCCL broadcast tensor 到远端 engine +``` + +`UpdateWeightFromTensor` 里也有 mixed colocate/distributed 支持： + +```python +self.use_distribute = len(rollout_engines) > colocate_engine_nums +``` + +如果 rollout engines 有一部分不在 actor GPU 范围内，就 colocated engine 走 IPC，剩余 engine 走 distributed NCCL。 + +### 为什么 T03 baseline 会卡在 wake_up weights + +T03 baseline 第一次失败在： + +```text +/wake_up?tags=weights +ConnectionRefusedError +``` + +框架级解释： + +1. `train.py` 在 train 后准备更新 rollout 权重。 +2. colocate/offload 模式下先 `rollout_manager.onload_weights()`。 +3. `RolloutServer.onload_weights()` 会对需要 offload 的 server group 调： + + ```python + engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + ``` + +4. `VLLMEngine.resume_memory_occupation()` HTTP POST 到 vLLM server： + + ```python + POST /wake_up?tags=weights + ``` + +5. 如果 vLLM APIServer/EngineCore 已退出，15000 端口无人监听，就会 `ConnectionRefusedError`。 + +所以这个错误不是 Megatron loss 或 `linear_logp` 本身失败，而是 rollout server 生命周期/残留状态问题。清理 Ray/vLLM/redis 后重跑成功，也符合这个判断。 + +### 框架级排错顺序 + +如果后续继续跑 T07 或其它配置，建议按层次排错： + +1. Ray 层： + - `ray status` + - Ray job 是否启动。 + - placement group 是否 ready。 + - actor 是否 `DEAD` 或 `PENDING`。 + +2. vLLM 层： + - `/health` 是否 200。 + - `VLLMEngine` 是否成功 spawn server process。 + - `CUDA_VISIBLE_DEVICES` 是否对应预期 GPU。 + - sleep/wake 的 `/sleep`、`/wake_up?tags=weights`、`/wake_up?tags=kv_cache` 是否成功。 + +3. rollout 数据层： + - `RolloutManager.generate()` 是否返回 `Sample`。 + - `tokens/response_lengths/loss_masks/rewards` 是否长度一致。 + - `build_dp_schedule()` 是否满足 `RBS*NSP >= DP` 和 `GBS <= RBS*NSP`。 + +4. Megatron train 层： + - `MegatronTrainRayActor._get_rollout_data()` 是否能把数据搬上 GPU。 + - `train_actor()` 是否能算 advantage/logprob/loss。 + - `train_one_step()` 是否有 finite loss/grad norm。 + +5. RL-Kernel 算子层： + - 是否出现 `Using RL-Kernel linear_logp op: FusedLinearLogpSM90Op`。 + - 是否出现 `Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path.`。 + - `rl_kernel_fallback_count_delta` 是否为 0。 + +6. 权重同步层： + - `[weight-sync] bucket ... begin/all-gather/HF convert/IPC payload/update returned` 是否连续出现。 + - vLLM `start_weight_update` / `finish_weight_update` 是否成功。 + - `weight_version` 是否随 rollout 增长。 + +7. 显存/offload 层： + - rollout 生成后 vLLM 是否 offload。 + - train 前 Megatron 是否 wake。 + - train 后 vLLM 是否只先 wake weights，再 update weights，再 wake kv。 + - `actor_train_peak_reserved_delta_mb` 和 `peak_vram_gb` 是否异常上升。 + +这个分层视角比只看命令行更适合理解 RL 框架：vime 把“生成”和“训练”拆成两个执行系统，vLLM 负责高吞吐 rollout，Megatron 负责大模型训练，中间通过 Ray object refs 传样本、通过 CUDA IPC/NCCL 同步权重。 + +## 已完成结果 + +T01、T03、T06 都完成了 candidate 与 baseline 的完整 12 轮 no-trace 对比。T07 做过 trace 尝试和一次后续 no-trace 重启，但按用户后续要求已经停止，不纳入正式指标。 + +| 配置 | baseline log | candidate log | +| --- | --- | --- | +| T01 | `/workspace/logs/8gpu_T01_full_baseline_20260707_142633/ray_job_baseline.log` | `/workspace/logs/8gpu_T01_full_cuda_20260707_103211/ray_job_cuda.log` | +| T03 | `/workspace/logs/8gpu_T03_full_baseline_20260707_150229/ray_job_baseline.log` | `/workspace/logs/8gpu_T03_full_cuda_20260707_105946/ray_job_cuda.log` | +| T06 | `/workspace/logs/8gpu_T06_full_baseline_20260707_124934/ray_job_baseline.log` | `/workspace/logs/8gpu_T06_full_cuda_20260707_115418/ray_job_cuda.log` | + +完整训推总览： + +| 配置 | baseline status | candidate status | baseline step s | candidate step s | baseline rollout s | candidate rollout s | baseline peak reserved GB | candidate peak reserved GB | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| T01 | success | success | 104.18 | 105.76 | 24.34 | 24.16 | 20.60 | 19.89 | +| T03 | success | success | 155.63 | 156.55 | 74.56 | 72.68 | 24.14 | 22.49 | +| T06 | success | success | 232.20 | 228.40 | 146.76 | 143.07 | 49.26 | 46.23 | + +单算子主指标： + +| 配置 | tokens/call baseline | tokens/call candidate | fwd ms baseline | fwd ms candidate | fwd speedup | fwd+bwd ms baseline | fwd+bwd ms candidate | fwd+bwd speedup | reserved delta baseline | reserved delta candidate | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| T01 | 796.44 | 796.44 | 5.14 | 3.51 | 1.46x | 15.24 | 12.37 | 1.23x | 4056 MB | 3112 MB | +| T03 | 1820.44 | 1820.44 | 7.78 | 3.36 | 2.32x | 18.56 | 10.37 | 1.79x | 6684 MB | 4862 MB | +| T06 | 6769.78 | 6826.67 | 14.52 | 7.62 | 1.91x | 33.96 | 18.50 | 1.84x | 32342 MB | 26710 MB | + +口径说明： + +- 表中单算子指标来自 rollout 3-11 的稳定窗口。 +- baseline 有偶发 native logprob spike，因此 T01/T03/T06 的 baseline 单算子主表剔除了已确认的 spike 轮次。 +- T01 是 smoke 档，整体 step time candidate 没有胜出，因为 rollout、weight sync、调度和日志等固定开销占比更大；但单算子仍快。 +- T03/T06 放大了每次 `linear_logp` 的 token 数，candidate 的收益开始稳定体现。 +- T06 candidate 比 baseline 少约 3.03GB full-run peak reserved；单算子 memory probe 的 reserved delta 少约 5.5GB。 + +## 配置矩阵经验 + +本轮实际跑通的关键配置： + +| 配置 | RBS | NSP | GBS | MAX_TOKENS | RESP | VLLM_MEM | VLLM_MAX_MODEL_LEN | 定位 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | +| T01 | 2 | 2 | 4 | 2048 | 512 | 0.40 | 2048 | smoke | +| T03 | 2 | 2 | 4 | 4096 | 1536 | 0.45 | 4096 | target-low | +| T06 | 4 | 2 | 8 | 8192 | 3072 | 0.45 | 4096 | max-token-call probe | +| T07 | 4 | 2 | 8 | 8192 | 3968 | 0.46 | 4096 | trace/interrupted，不纳入正式数据 | + +经验： + +- 先跑 candidate，再补 baseline。candidate 如果 OOM，baseline 大概率没有必要先跑。 +- 想放大 `linear_logp` 收益，优先提高 `RESP` 和 `MAX_TOKENS_PER_GPU`，再提高 `GBS`/`NSP`。 +- T06 是当前最清晰的宣传档：tokens/call 达到 6.8k，fwd+bwd 约 1.84x，完整 step 也有小幅优势。 +- T07 trace 开销和 profiler 落盘行为会影响训练推进，不能拿 trace run 与 no-trace baseline 直接比较。 + +## baseline 算子路径 + +baseline 的本质是两段式： + +1. Megatron output layer 先计算完整 vocab logits。 +2. vime 再对完整 logits 做 selected-token logprob。 + +对应代码位置： + +- `/workspace/Megatron-LM/megatron/core/models/gpt/gpt_model.py`：GPT model 使用 `LinearCrossEntropyModule`/output layer。 +- `/workspace/vime/vime/backends/megatron_utils/model.py`：`_probe_baseline_output_layer_forward()` 在 output layer 上挂 hook，记录 forward 和 forward+backward CUDA event。 +- `/workspace/vime/vime/backends/megatron_utils/loss.py`：`get_log_probs_and_entropy()` 在 baseline 路径下调用原生 logprob。 +- `/workspace/vime/vime/utils/ppo_utils.py`：`calculate_log_probs_and_entropy()` 里会对 `logits.clone()` 调 `compute_log_probs()`。 + +baseline 关键逻辑可以简化成： + +```python +# loss.py +if linear_logp_context is None: + logits = logits.contiguous() + log_prob_full, entropy_full = calculate_log_probs_and_entropy( + logits, + full_tokens, + tp_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + ) +``` + +`calculate_log_probs_and_entropy()` 里又会做： + +```python +log_prob = compute_log_probs(logits.clone(), tokens, tp_group) +``` + +对 Qwen3-30B-A3B 这类 vocab 很大的模型，这意味着 baseline 的 hot path 会显式处理 `[T, V]` 形状的 logits： + +- `T` 是 packed tokens，本轮 T06 大约 6.8k tokens/call。 +- `V` 是 padded vocab，本轮 TP=2 时每 rank 是本地 vocab shard，全局 vocab 仍很大。 +- output layer 要先产生完整 logits。 +- logprob 还要对完整 logits 做 softmax/logsumexp/gather 相关操作。 +- 即使最终 PPO loss 只需要 selected token 的 logprob，baseline 仍为所有 vocab token 支付了额外 HBM 读写和中间 tensor 成本。 + +baseline 计时是这样拼成 `baseline_linear_logp_*` 的： + +- `output_layer_*` 记录 output layer 计算完整 logits 的耗时和 CUDA event。 +- `native_logprob_*` 记录 `calculate_log_probs_and_entropy()` 的耗时和 CUDA event。 +- `baseline_linear_logp_*` 把二者加总，形成与 candidate `linear_logp` 可比较的总成本。 + +这也解释了 baseline spike 的来源：output layer 和 native logprob 是两个独立阶段，任一阶段发生调度、clone、allocator、softmax 或 TP 通信抖动，都会放大到 `baseline_linear_logp_*` 上。本轮看到 T01/T03/T06 baseline 均有 native logprob spike，因此宣传图使用稳定窗口，不使用 spike 轮次夸大结论。 + +## candidate 算子路径 + +candidate 的路径是直接用 RL-Kernel 在 hidden state 上计算 selected logprob，避免完整 logits 作为 Python/PyTorch 层面的中间结果。 + +对应代码位置： + +- `/workspace/vime/vime/backends/megatron_utils/rl_kernel.py` + - `_get_linear_logp_op()` 根据 `VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=cuda` 加载 `FusedLinearLogpSM90Op`。 + - `get_linear_logp_context_from_model()` 从 Megatron model 取 `lm_head_weight`、bias、TP group、`vocab_start_index`、`global_vocab_size`。 + - `return_hidden_states_for_linear_logp()` 临时把 Megatron `post_process=False`，让模型返回 hidden states 而不是 logits。 + - `maybe_compute_linear_logp()` 直接调用 RL-Kernel op。 +- `/workspace/RL-Kernel/rl_engine/kernels/ops/cuda/loss/linear_logp.py` + - `FusedLinearLogpSM90Op.apply()` 选择 SM90 tensor-parallel fast path。 + - `_TensorParallelLinearLogpFusedTileBF16Function` 负责 TP 下 forward/backward。 +- `/workspace/RL-Kernel/csrc/cuda/fused_linear_logp_sm90.cu` + - `fused_linear_logp_sm90_forward_impl()` 是 fused forward CUDA 入口。 + - `fused_linear_logp_sm90_backward()` 是 fused backward CUDA 入口。 + +candidate 关键逻辑可以简化成： + +```python +# vime/backends/megatron_utils/rl_kernel.py +op = FusedLinearLogpSM90Op() +log_prob = op( + hidden_states, + lm_head_weight, + target_ids.long(), + bias, + tp_group=context.tp_group, + vocab_start_index=context.vocab_start_index, + global_vocab_size=context.global_vocab_size, +) +``` + +RL-Kernel 里 `FusedLinearLogpSM90Op` 的 docstring 已经说明了核心目标： + +```python +Computes log_softmax(hidden @ W^T + b)[target] without materializing the [N, V] logits. +``` + +本轮 candidate 命中的 fast path 是日志里的： + +```text +Using fused-tile bf16 full-gradient tensor-parallel linear_logp fast path. +``` + +命中条件包括： + +- SM90/Hopper 设备。 +- hidden 和 weight 是 bf16。 +- TP vocab shard 场景可识别。 +- bias 为 None。 +- full-gradient 训练需要 hidden 或 weight 的梯度。 +- `RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=1`。 + +CUDA forward 的关键实现点： + +- 用 TMA descriptor 读取 hidden tile 和 weight tile。 +- 用 WGMMA/Tensor Core 做 tile GEMM。 +- 在 tile 内累计每个 token 的局部 max、sum exp、target logit。 +- forward 返回的是 `[N]` 的 selected logprob 或 target logit/lse，不返回 `[N, V]` logits。 +- TP 场景下各 rank 先算 local shard 的 `local_target_logit` 和 `local_lse`，再合并成 global selected logprob。 + +CUDA backward 的关键实现点： + +- full-gradient path 以 vocab tile 为单位重算或生成局部 `dlogits`。 +- 对 `grad_hidden` 和 `grad_weight` 用 tile GEMM 累积。 +- 避免 Python 层 chunk loop 和大量小 matmul dispatch。 +- bf16 输入下尽量让 GEMM 输入保持 Tensor Core 友好的布局和 dtype。 + +## 为什么 candidate 比 baseline 好 + +这不是简单的“少一个 kernel launch”，主要是计算图和数据流都变了。 + +### 1. 避免完整 logits 物化 + +baseline 为 selected logprob 先生成完整 `[T, V]` logits。以 T06 约 6.8k tokens/call 为例，即使 TP=2 后每 rank 只看本地 vocab shard，这个矩阵仍然很大。完整 logits 会带来： + +- output layer 写出大矩阵。 +- native logprob 再读入大矩阵。 +- `logits.clone()` 带来额外读写。 +- softmax/logsumexp 相关中间结果继续访问同一大矩阵。 +- backward 还要为 logits 梯度走一遍大矩阵。 + +candidate 在 CUDA kernel 内直接围绕 selected target 计算 logsumexp 和 target logit，不把完整 logits 暴露为框架层 tensor。最终需要保存的是 `[N]` 级别的 lse/logp，以及 backward 所需的少量状态。显存峰值和 HBM traffic 都下降。 + +### 2. 把 linear + logprob 融合成同一个算子语义 + +baseline 的算子边界是： + +```text +hidden -> output_layer -> logits -> calculate_log_probs_and_entropy -> selected logprob +``` + +candidate 的算子边界是： + +```text +hidden + lm_head_weight + target_ids -> selected logprob +``` + +这个边界变化让 CUDA 实现可以在 tile 内边算 GEMM，边维护 max/sum/target 统计量，不需要先完成整张 logits 矩阵再进入下一步。 + +### 3. TP vocab shard 更自然 + +baseline 在 vime/Megatron 原生路径里仍然围绕 logits tensor 和后续 logprob 工具函数组织逻辑。candidate 在 `LinearLogpContext` 里明确携带： + +- `tp_group` +- `vocab_start_index` +- `global_vocab_size` + +RL-Kernel 的 TP path 让每个 rank 只在本地 vocab shard 里做 fused tile forward，然后用 TP group 合并 `local_lse` 和 owned target logit。这个结构更贴近 vocab-parallel output layer 的真实分片。 + +### 4. backward 更少 Python 调度和 allocator 压力 + +baseline 的 backward 经过 output layer autograd 和 native logprob 的组合，框架层中间 tensor 更多。candidate 的 `fused_linear_logp_sm90_backward()` 在 CUDA/C++ 里组织 vocab tile 循环，减少 Python 层 chunk loop、小 kernel、小 matmul dispatch 和 allocator 抖动。 + +这也是为什么 T01 这种小 token/call 配置里，candidate 单算子仍快，但完整 step 不一定赢；一旦 T03/T06 把 token/call 放大，kernel hot path 占比上升，candidate 的收益就更明显。 + +### 5. 实测支持这个解释 + +| 配置 | tokens/call | fwd speedup | fwd+bwd speedup | full-run reserved saving | +| --- | ---: | ---: | ---: | ---: | +| T01 | 796 | 1.46x | 1.23x | 0.71GB | +| T03 | 1820 | 2.32x | 1.79x | 1.65GB | +| T06 | 6827 | 1.91x | 1.84x | 3.03GB | + +随着 tokens/call 增大，candidate 避免 `[T, V]` 中间结果的收益变得更实在。T06 的 full-run peak reserved 从 49.26GB 降到 46.23GB，和单算子 memory probe 的趋势一致。 + +## 跑通 8 卡时做过的代码级修正 + +这些修改不是都直接影响 `linear_logp` 性能，但它们决定了长跑能否稳定完成。 + +### workspace-local cache + +`scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh` 和 `scripts/run-qwen3-30B-A3B.sh` 都加入了 workspace-local cache 变量。这样 Ray worker、vLLM worker、FlashInfer 编译缓存和 torch cache 不会散落到系统目录。 + +### vLLM spawn 和 faulthandler + +`scripts/run-qwen3-30B-A3B.sh` 和 `vime/backends/vllm_utils/vllm_engine.py` 里加入： + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn +VIME_VLLM_FAULTHANDLER=1 +``` + +实践原因： + +- vLLM engine 在 Ray worker 内启动，fork 方式更容易继承复杂 CUDA/Ray 状态。 +- spawn 更干净，出错时配合 faulthandler 更容易看到 Python 栈。 +- T03 baseline 第一次失败表现为 vLLM APIServer/EngineCore 退出，onload weights 请求 `/wake_up?tags=weights` 时 `ConnectionRefusedError`。清理后重跑成功，因此该失败不计入性能。 + +### Megatron grouped GEMM validation bypass + +`vime/backends/megatron_utils/arguments.py` 加入： + +```python +skip_grouped_gemm_capability_check = ( + os.environ.get("VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK", "0") == "1" + and getattr(args, "moe_grouped_gemm", False) +) +``` + +做法是 validation 时临时把 `args.moe_grouped_gemm=False`，跑完 Megatron validation 后恢复为 True。原因是当前环境里 Megatron 的 capability check 会阻止配置通过，但 runtime 实际需要保留 grouped GEMM 设置。 + +### Qwen3MoE layernorm 权重映射 + +`vime/backends/megatron_utils/megatron_to_hf/qwen3moe.py` 扩展了 layernorm 名称映射： + +```python +rest in {"self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"} +rest in {"mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight", "post_attention_layernorm.weight"} +``` + +实践原因：不同 Megatron/MoE 代码路径导出的 layernorm key 不完全一致，如果不兼容这些名字，Megatron -> HF/vLLM 权重同步容易漏映射。 + +### grouped MLP expert 权重拆分 + +`vime/backends/megatron_utils/update_weight/common.py` 加入 `_iter_grouped_mlp_expert_weights()`，把 grouped expert 权重拆成 vLLM/HF 侧期望的 per-expert 权重名： + +```python +if rest == "mlp.experts.weight1": + expert_tensors = param.view(num_local_experts, args.hidden_size, -1).transpose(-1, -2) + partition_dim = 0 + target = "linear_fc1" +elif rest == "mlp.experts.weight2": + expert_tensors = param.view(num_local_experts, -1, args.hidden_size).transpose(-1, -2) + partition_dim = 1 + target = "linear_fc2" +``` + +并保留 tensor parallel attrs： + +```python +tensor.tensor_model_parallel = getattr(source, "tensor_model_parallel", False) +tensor.partition_dim = partition_dim +tensor.partition_stride = 1 +tensor.parallel_mode = getattr(source, "parallel_mode", None) +``` + +实践原因：Qwen3 MoE 的训练侧 grouped expert 参数和推理侧 per-expert 参数命名/布局不一致，权重同步必须在 vime 侧做结构化转换，不能靠字符串硬凑。 + +### update weight 分桶和计时 + +`hf_weight_iterator_direct.py` 和 `update_weight_from_tensor.py` 增加了分桶计时日志，并支持： + +```bash +VIME_UPDATE_WEIGHT_BUFFER_SIZE=134217728 +``` + +实践原因： + +- 8 卡长跑里 weight sync 是 step time 的固定成本来源之一。 +- 分桶过大可能压显存或造成 Ray object 传输抖动。 +- 分桶计时能区分“kernel 慢”和“权重同步慢”，否则容易把 full step 差异误归因到 `linear_logp`。 + +### Megatron-LM 本地兼容补丁 + +`/workspace/Megatron-LM` 当前有本地兼容补丁，主要用于让本环境的 Qwen3-MoE + packed THD + fallback attention 能跑通： + +- `megatron/core/transformer/dot_product_attention.py`：给 non-TE DotProductAttention 增加 packed THD fallback。 +- `megatron/core/transformer/moe/moe_layer.py`：通过 `MEGATRON_ALLOW_MOE_TP_WITHOUT_SP=1` 放过 MoE TP without SP 的训练校验。 +- `megatron/core/transformer/moe/moe_utils.py`：TE import 失败时补 `te_general_gemm = None`，避免后续空引用。 +- 其他文件有本地修改，提交 PR 前需要确认哪些属于本轮 vime/RL-Kernel 变更范围。 + +这类 Megatron 本地补丁要单独看待：它们保证环境可运行，但不应在宣传时被描述成 RL-Kernel 算子收益来源。 + +## trace 经验 + +T07 做过一次 trace run： + +```bash +TRACE_MODE=all TRACE_ROLLOUTS=3 \ +scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh T07 cuda +``` + +实际状态： + +- 运行到 step 2/11。 +- rollout 3 的 rollout-generate trace 已触发 `cudaProfilerStart`。 +- profiler stop/落盘阶段没有及时返回。 +- 后续为了先跑 T01/T03 baseline，已主动中断。 +- trace 文件未形成可用产物。 + +结论： + +- trace 会改变 timing，不能与 no-trace baseline/candidate 直接比较。 +- 需要 trace 时应单独开 run，最好只抓一个明确 rollout/actor train 窗口。 +- 正式指标必须使用 `TRACE_MODE=none`。 + +## debug 经验 + +### T03 baseline EngineDead + +失败日志： + +```text +/workspace/logs/8gpu_T03_full_baseline_20260707_145557/ray_job_baseline.log +``` + +症状： + +- vLLM APIServer/EngineCore 退出。 +- onload weights 调 `/wake_up?tags=weights` 连接 15000 端口失败。 +- 抛出 `ConnectionRefusedError` / `requests.exceptions.ConnectionError`。 + +处理： + +1. `ray stop --force`。 +2. 杀掉残留 vLLM/train/redis 进程。 +3. 确认 `nvidia-smi` 每卡只剩少量 context 占用。 +4. 重跑 T03 baseline，成功完成 12 轮。 + +判断：这是启动/残留状态问题，不是 baseline 算子真实性能失败。 + +### baseline spike + +T01/T03/T06 baseline 都观察到 native logprob 相关 spike： + +- T01 baseline step 10 有 400ms+ 级 forward spike。 +- T03 baseline step 10 有 390ms+ 级 forward spike。 +- T06 baseline step 11 有数秒级 spike。 + +处理方式： + +- 主表使用稳定窗口，剔除 spike 轮。 +- 同时在最终矩阵中保留 spike 说明，避免选择性隐瞒。 +- 对宣传图使用稳定窗口数据，不用 spike 夸大 speedup。 + +### candidate 低于 baseline 时怎么看 + +如果后续某个配置出现 candidate 完整 step 不如 baseline，不要立即判定算子失败，需要拆开看： + +1. `train/rl_kernel_linear_logp_forward_backward_cuda_event_elapsed_s_delta` 是否仍优于 baseline。 +2. `fallback` 是否为 0。 +3. rollout time、weight sync time、vLLM wakeup、Ray object store 是否变化。 +4. 是否 trace/profiler 开着。 +5. 是否 tokens/call 太小，kernel hot path 占比被固定开销淹没。 + +T01 就是典型例子：candidate 单算子更快，但完整 step 略慢，原因更可能是小配置下固定开销占比太高，而不是 `linear_logp` 算子不行。 + +## 可宣传结论 + +可以对外使用的结论应限定在 T01/T03/T06 已完成 no-trace 数据上： + +- 8xH100、Qwen3-30B-A3B、完整 12 rollout 链路中，candidate T01/T03/T06 均成功完成。 +- candidate 均命中 `FusedLinearLogpSM90Op` fused-tile bf16 full-gradient tensor-parallel fast path。 +- fallback=0。 +- 单算子 forward 最多 2.32x，forward+backward 最多 1.84x。 +- T06 单算子 forward+backward 从 33.96ms 降到 18.50ms。 +- T06 full-run peak reserved 从 49.26GB 降到 46.23GB。 +- T06 rollout time 从 146.76s 降到 143.07s，tokens/GPU/s 从 20.33 提到 21.05。 + +配套可视化： + +```text +/workspace/vime/vime-RLK-single-op-performance-visualization.svg +``` + +这张图只使用 T01/T03/T06 的 no-trace 稳定窗口数据，展示 forward、forward+backward 和 single-op reserved memory delta 三组指标，并在底部画出 baseline 与 candidate 的算子数据流差异。 + +建议宣传表达： + +```text +On Qwen3-30B-A3B 8xH100 full rollout training, RL-Kernel fused linear_logp +hits the SM90 tensor-parallel fast path with zero fallback and cuts the +linear_logp fwd+bwd CUDA time by up to 1.84x, while reducing peak reserved +memory by 3.03GB in the T06 full-run setting. +``` + +中文表达： + +```text +在 Qwen3-30B-A3B 8xH100 完整训推链路中，RL-Kernel fused linear_logp +稳定命中 SM90 tensor-parallel fast path，fallback=0；T06 配置下单算子 +forward+backward 从 33.96ms 降至 18.50ms，约 1.84x，加上 full-run +peak reserved 显存减少约 3.03GB。 +``` + +## 后续工作 + +1. 暂停 T07，除非用户明确要求继续。 +2. 如果继续 T07，先跑 no-trace candidate 完整 12 轮，再决定是否补 baseline。 +3. 若 T07 OOM，回退 T08。 +4. 提 PR 前确认 Megatron-LM 本地补丁是否进入本轮范围；当前用户关心的是 vime 和 RL-Kernel。 +5. 最终矩阵里 T07 必须保持“中断/不计入指标”，不能写成 running。 +6. 发布宣传图时注明 stable window 和 no-trace 口径。 diff --git a/vime-RLK-single-op-performance-visualization.svg b/vime-RLK-single-op-performance-visualization.svg new file mode 100644 index 00000000..f12cd328 --- /dev/null +++ b/vime-RLK-single-op-performance-visualization.svg @@ -0,0 +1,250 @@ + + RL-Kernel linear_logp single-operator performance on Qwen3-30B-A3B 8xH100 + Comparison of baseline native output layer plus native logprob against RL-Kernel fused linear_logp candidate for T01, T03, and T06. + + + + + + + + + + + + + + + + + RL-Kernel fused linear_logp cuts single-op cost + Qwen3-30B-A3B · 8xH100 · TP=2 · full-gradient · 12 rollout · no trace · stable window + + + baseline: native output layer + native logprob + + candidate: save-logits full-gradient fused-tile linear_logp + + + + + 0 + FALLBACKS + + + + 2.32x + MAX FWD + + + + 1.84x + MAX FWD+BWD + + + + + + + Forward CUDA time + Lower is better. Baseline output layer + native logprob vs fused candidate. + + + + + + 0 + 5 + 10 + 15 ms + + T01 + 796 tok/call + + + 5.14 ms + 3.51 ms + 1.46x + + T03 + 1.82k tok/call + + + 7.78 ms + 3.36 ms + 2.32x + + T06 + 6.83k tok/call + + + 14.52 ms + 7.62 ms + 1.91x + + + + + + + Forward + backward CUDA time + Full-gradient actor train path. Candidate keeps the fused tile path in backward. + + + + + + + 0 + 10 + 20 + 30 + 36 ms + + T01 + smoke + + + 15.24 ms + 12.37 ms + 1.23x + + T03 + target-low + + + 18.56 ms + 10.37 ms + 1.79x + + T06 + max token call + + + 33.96 ms + 18.50 ms + 1.84x + + + + + + + Reserved memory delta + Single-op memory probe, MB converted to GiB. Lower is better. + + + + + + 0 + 10 + 20 + 32 GiB + + T01 + small logits + + + 3.96 GiB + 3.04 GiB + + T03 + larger tokens + + + 6.53 GiB + 4.75 GiB + + T06 + large token call + + + 31.58 GiB + 26.08 GiB + + + T06 saves 5.50 GiB + + + + + + + What changed in the operator path + The candidate changes the dataflow, not just the timing wrapper. + + baseline + + hidden states + [T, hidden] + + + native output layer + materializes full logits + + + [T, V] logits + large HBM read/write + + + native logprob + clone + softmax/logsumexp + + gather selected token + + + logprob + [T] + + candidate + + hidden + W + bf16, TP shard + + + SM90 fused-tile linear_logp + TMA/WGMMA tile GEMM + local lse + target logit + + + TP merge + global selected logprob + + + logprob + [T] + + + Key reason: + candidate computes selected-token logprob without framework-level full logits materialization, reducing HBM traffic, allocator pressure, and Python-side dispatch. + + + Source: /workspace/vime/vime-RLK-final-metrics-config-matrix.md · T01/T03/T06 full-gradient runs · baseline spike rounds excluded from stable-window single-op averages. + Candidate path: save-logits / RL-Kernel FusedLinearLogpSM90Op, fused-tile bf16 full-gradient tensor-parallel fast path, fallback=0. + diff --git a/vime-RLK.md b/vime-RLK.md new file mode 100644 index 00000000..e50cf7ca --- /dev/null +++ b/vime-RLK.md @@ -0,0 +1,321 @@ +# vime + RL-Kernel linear_logp 2xH100 性能预验证 + +## 0. 我们要做什么 + +本轮先在 2xH100 上做 A/B 性能预验证，不做 smoke-only。只有 2 卡已经明显优于 vime 原生路径，才扩大到 8 卡主宣传 benchmark。 + +```text +baseline: RL-Align/vime#2, RL-Kernel off, Qwen3-30B-A3B, 2xH100 colocate +candidate: RL-Align/vime#2 + RL-Align/RL-Kernel#189, RL-Kernel linear_logp on, Qwen3-30B-A3B, 2xH100 colocate +``` + +2 卡阶段仍然使用和 8 卡一致的指标验收线： + +- `rl_kernel_fallback_count = 0` +- `raw_reward` 不低于 baseline 同量级 +- `train_rollout_logprob_abs_diff` 不持续高于 baseline +- `mean_log_probs_time_s` 或 `peak_vram_gb` 有明确下降 +- 最好能看到明显收益后再上 8 卡：建议 `mean_log_probs_time_s` 下降 >= 20% 或 `peak_vram_gb` 下降 >= 10% + +2 卡结果只作为上 8 卡前的门禁，不直接进入宣传材料；但该门禁必须放大 selected-logprob workload，能看出 RL-Kernel `linear_logp` 的真实收益。 + +## 1. 范围 + +只保留： + +- `linear_logp` +- Qwen3-30B-A3B +- TP=2 +- 2xH100 单机 colocate +- baseline 和 candidate 都必须跑 +- 指标集合与 8xH100 主 benchmark 保持一致 + +不做： + +- 8xH100 主宣传实验 +- Qwen3-4B smoke +- R3 单独对比 +- GLM-4.5 +- GB200/H200/A100 硬件对照 +- 训推一致性专项 benchmark +- MoE expert/router RL-Kernel 算子 + +## 2. 性能预验证配置 + +默认配置不是 smoke，而是 24 step 的 2 卡性能预验证。核心思路是增加 selected-logprob token 数，让 `linear_logp` 的收益不要被 rollout、update weights 等固定开销完全淹没。 + +```bash +export CUDA_VISIBLE_DEVICES=0,1 +export NUM_GPUS=2 +export MEGATRON_TP=2 +export MEGATRON_EP=2 +export MEGATRON_CP=1 +export ROLLOUT_NUM_GPUS_PER_ENGINE=2 + +export NUM_ROLLOUT=24 +export ROLLOUT_BATCH_SIZE=2 +export N_SAMPLES_PER_PROMPT=2 +export GLOBAL_BATCH_SIZE=4 +export MAX_TOKENS_PER_GPU=4096 +export ROLLOUT_MAX_RESPONSE_LEN=1024 +export VLLM_GPU_MEMORY_UTILIZATION=0.50 + +export VIME_CKPT_DIR=/root/Qwen3-30B-A3B_vime_tp2_dev +export VIME_DISABLE_SAVE=1 +export VIME_SKIP_EVAL_BEFORE_TRAIN=1 +export VIME_VLLM_ENFORCE_EAGER=1 +export VIME_NO_GRAD_ACCUM_FUSION=1 +``` + +如果 2xH100 OOM，先只做这一档降级；降级后仍然不是 smoke，因为 response len 和 step 数保持较大： + +```text +MAX_TOKENS_PER_GPU=4096 +ROLLOUT_MAX_RESPONSE_LEN=1024 +ROLLOUT_BATCH_SIZE=1 +N_SAMPLES_PER_PROMPT=2 +GLOBAL_BATCH_SIZE=2 +``` + +## 3. 拉代码 + +从官方仓库开始： + +```bash +cd /workspace +git clone https://github.com/RL-Align/RL-Kernel.git RL-Kernel +git clone https://github.com/RL-Align/vime.git vime-rlk-tp2 +``` + +RL-Kernel 使用 TP 版 `linear_logp`： + +```bash +cd /workspace/RL-Kernel +git checkout main +git pull origin main +gh pr checkout 189 +``` + +vime 使用 2xH100 开发验证 PR： + +```bash +cd /workspace/vime-rlk-tp2 +git checkout main +git pull origin main +gh pr checkout 2 +``` + +## 4. 安装 + +```bash +cd /workspace/RL-Kernel +pip install -e . +python setup.py build_ext --inplace -v + +cd /workspace/vime-rlk-tp2 +pip install -e . +``` + +## 5. 模型和数据 + +```bash +pip install -U "huggingface_hub[cli]" +huggingface-cli login + +hf download Qwen/Qwen3-30B-A3B --local-dir /root/Qwen3-30B-A3B + +hf download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +hf download --repo-type dataset zhuzilin/aime-2024 \ + --local-dir /root/aime-2024 +``` + +转换 Megatron `torch_dist` checkpoint： + +```bash +cd /workspace/vime-rlk-tp2 +source scripts/models/qwen3-30B-A3B.sh + +PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 2 \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-30B-A3B \ + --save /root/Qwen3-30B-A3B_torch_dist + +mkdir -p /root/Qwen3-30B-A3B_vime_tp2_dev +``` + +## 6. 跑 baseline + +baseline 必跑，用来代表 vime 原生路径；不要打开 RL-Kernel。 + +```bash +cd /workspace/vime-rlk-tp2 + +export CUDA_VISIBLE_DEVICES=0,1 +export NUM_GPUS=2 +export MEGATRON_TP=2 +export MEGATRON_EP=2 +export MEGATRON_CP=1 +export ROLLOUT_NUM_GPUS_PER_ENGINE=2 + +export NUM_ROLLOUT=24 +export ROLLOUT_BATCH_SIZE=2 +export N_SAMPLES_PER_PROMPT=2 +export GLOBAL_BATCH_SIZE=4 +export MAX_TOKENS_PER_GPU=4096 +export ROLLOUT_MAX_RESPONSE_LEN=1024 +export VLLM_GPU_MEMORY_UTILIZATION=0.50 + +export VIME_CKPT_DIR=/root/Qwen3-30B-A3B_vime_tp2_dev +export VIME_DISABLE_SAVE=1 +export VIME_SKIP_EVAL_BEFORE_TRAIN=1 +export VIME_VLLM_ENFORCE_EAGER=1 +export VIME_NO_GRAD_ACCUM_FUSION=1 + +unset VIME_RL_KERNEL VIME_RL_KERNEL_OPS VIME_RL_KERNEL_STRICT + +bash scripts/run-qwen3-30B-A3B.sh 2>&1 | tee /workspace/vime-rlk-tp2-baseline.log +``` + +## 7. 跑 candidate + +candidate 使用同一套 2 卡配置，只打开 RL-Kernel。 + +```bash +cd /workspace/vime-rlk-tp2 + +export CUDA_VISIBLE_DEVICES=0,1 +export NUM_GPUS=2 +export MEGATRON_TP=2 +export MEGATRON_EP=2 +export MEGATRON_CP=1 +export ROLLOUT_NUM_GPUS_PER_ENGINE=2 + +export NUM_ROLLOUT=24 +export ROLLOUT_BATCH_SIZE=2 +export N_SAMPLES_PER_PROMPT=2 +export GLOBAL_BATCH_SIZE=4 +export MAX_TOKENS_PER_GPU=4096 +export ROLLOUT_MAX_RESPONSE_LEN=1024 +export VLLM_GPU_MEMORY_UTILIZATION=0.50 + +export VIME_CKPT_DIR=/root/Qwen3-30B-A3B_vime_tp2_dev +export VIME_DISABLE_SAVE=1 +export VIME_SKIP_EVAL_BEFORE_TRAIN=1 +export VIME_VLLM_ENFORCE_EAGER=1 +export VIME_NO_GRAD_ACCUM_FUSION=1 + +export VIME_RL_KERNEL=1 +export VIME_RL_KERNEL_OPS=linear_logp +export VIME_RL_KERNEL_STRICT=1 + +bash scripts/run-qwen3-30B-A3B.sh 2>&1 | tee /workspace/vime-rlk-tp2-candidate.log +``` + +## 8. 验收线 + +每组先跑 1 次确认无错误；稳定后 baseline/candidate 各跑至少 3 次，丢弃前 5-10 step warmup 后统计。 + +candidate 必须满足： + +```text +RL-Kernel linear_logp backend 被加载 +VIME_RL_KERNEL_STRICT=1 没有触发 RuntimeError +rl_kernel_fallback_count = 0 +rl_kernel_linear_logp_call_count_delta > 0 +rl_kernel_linear_logp_token_count_delta > 0 +rl_kernel_linear_logp_dispatch_elapsed_s_delta > 0 +log_probs / loss / reward 指标为 finite +raw_reward 不低于 baseline 同量级 +train_rollout_logprob_abs_diff 不持续高于 baseline +mean_log_probs_time_s 或 peak_vram_gb 有明确下降 +``` + +2 卡上卡门槛： + +```text +每组至少 24 train step +丢弃前 5 step warmup +mean_log_probs_time_s 下降 >= 20% +或 peak_vram_gb 下降 >= 10% +或二者都有小幅但稳定下降，且 mean_step_time_s 不明显变差 +``` + +不允许： + +```text +fallback 到 vime materialized logits 路径 +target vocab shard 报错 +TP collective hang +loss/logprob NaN 或 Inf +candidate 质量指标明显劣于 baseline +rl_kernel_linear_logp_call_count_delta 长时间为 0 +rl_kernel_linear_logp_token_count_delta 只覆盖极少 token +``` + +runtime counter 解释： + +```text +*_total：当前进程累计命中的 RL-Kernel linear_logp 调用、token 和 dispatch 耗时。 +*_delta：两次 train log 之间新增的调用、token 和 dispatch 耗时；第一个 train step 会覆盖此前 ref-logprob 加本 step train-logprob。 +tokens_per_call = token_count_delta / max(call_count_delta, 1)，用于判断是否只是空调用或很小 workload。 +dispatch_elapsed_s 不做 CUDA synchronize，不作为 GPU kernel time 宣传；正式性能仍看 mean_log_probs_time_s、step time 和 profiler。 +``` + +## 9. 必须记录 + +```text +gpu_name +num_gpus +vime_commit +rl_kernel_commit +vime_pr +rl_kernel_pr +model +tp +ep +cp +rollout_batch_size +n_samples_per_prompt +global_batch_size +max_tokens_per_gpu +rollout_max_response_len +vllm_gpu_memory_utilization +selected_rl_kernel_backend +rl_kernel_fallback_count +rl_kernel_linear_logp_call_count_total +rl_kernel_linear_logp_call_count_delta +rl_kernel_linear_logp_token_count_total +rl_kernel_linear_logp_token_count_delta +rl_kernel_linear_logp_dispatch_elapsed_s_total +rl_kernel_linear_logp_dispatch_elapsed_s_delta +rl_kernel_linear_logp_tokens_per_call_total +rl_kernel_linear_logp_tokens_per_call_delta +first_successful_train_step +mean_step_time_s +p50_step_time_s +p90_step_time_s +mean_log_probs_time_s +p50_log_probs_time_s +p90_log_probs_time_s +peak_vram_gb +raw_reward_mean +train_rollout_logprob_abs_diff_mean +error_stack_if_failed +``` + +## 10. 下一步 + +2xH100 指标门禁通过后再进入正式 benchmark： + +```text +8xH100 +Qwen3-30B-A3B +baseline vs candidate +至少 3 次 run +统计 step time、logprob time、peak VRAM、raw_reward、train_rollout_logprob_abs_diff +``` + +只有 8xH100 正式 benchmark 结果可以进入宣传材料。 diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 5bb7c8a7..50b5c10f 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -17,7 +17,12 @@ from vime.utils.data import process_rollout_data from vime.utils.distributed_utils import get_gloo_group from vime.utils.logging_utils import init_tracking -from vime.utils.memory_utils import clear_memory, print_memory +from vime.utils.memory_utils import ( + clear_memory, + get_peak_memory_tracker, + print_memory, + reset_peak_memory_tracker, +) from vime.utils.misc import Box from vime.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups from vime.utils.routing_replay import RoutingReplay @@ -41,7 +46,126 @@ logger = logging.getLogger(__name__) +def _env_flag(name: str) -> bool: + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} + + +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 + + +def _collect_actor_train_memory_metrics() -> dict[str, float]: + device = torch.cuda.current_device() + stats = get_peak_memory_tracker("actor_train", device=device) + metrics_tensor = torch.tensor( + [ + float(stats["alloc_before"]), + float(stats["reserved_before"]), + float(stats["alloc_after"]), + float(stats["reserved_after"]), + float(stats["peak_alloc"]), + float(stats["peak_reserved"]), + ], + device=device, + dtype=torch.float64, + ) + dist.all_reduce(metrics_tensor, op=dist.ReduceOp.MAX) + alloc_before, reserved_before, alloc_after, reserved_after, peak_alloc, peak_reserved = metrics_tensor.tolist() + mib = float(1024**2) + return { + "actor_train_alloc_before_mb": alloc_before / mib, + "actor_train_reserved_before_mb": reserved_before / mib, + "actor_train_alloc_after_mb": alloc_after / mib, + "actor_train_reserved_after_mb": reserved_after / mib, + "actor_train_peak_alloc_mb": peak_alloc / mib, + "actor_train_peak_reserved_mb": peak_reserved / mib, + "actor_train_peak_alloc_delta_mb": (peak_alloc - alloc_before) / mib, + "actor_train_peak_reserved_delta_mb": (peak_reserved - reserved_before) / mib, + } + + +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: + 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, + ) + 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 MegatronTrainRayActor(TrainRayActor): + def _keep_train_process_groups_during_offload(self) -> bool: + return _env_flag("VIME_KEEP_TRAIN_PROCESS_GROUPS_DURING_OFFLOAD") + + def _keep_host_cache_during_train_offload(self) -> bool: + return _env_flag("VIME_KEEP_HOST_CACHE_DURING_TRAIN_OFFLOAD") + + def _skip_memory_clear_after_train_wake_up(self) -> bool: + return _env_flag("VIME_SKIP_MEMORY_CLEAR_AFTER_TRAIN_WAKE_UP") + @with_defer(lambda: Timer().start("train_wait")) def init( self, @@ -170,7 +294,11 @@ def init( def sleep(self) -> None: assert self.args.offload_train - clear_memory(clear_host_memory=True) + keep_host_cache = self._keep_host_cache_during_train_offload() + if keep_host_cache: + logger.info("Keeping host pinned cache before train offload.") + # Reuse pinned host allocations for torch_memory_saver's large CPU backups. + clear_memory(clear_host_memory=not keep_host_cache) print_memory("before offload model") if ( self.role == "actor" @@ -179,7 +307,10 @@ def sleep(self) -> None: and hasattr(self.weight_updater, "disconnect_rollout_engines") ): self.weight_updater.disconnect_rollout_engines() - destroy_process_groups() + if self._keep_train_process_groups_during_offload(): + logger.info("Keeping train process groups alive during offload.") + else: + destroy_process_groups() torch_memory_saver.pause() @@ -192,8 +323,13 @@ def wake_up(self) -> None: torch_memory_saver.resume() - clear_memory() - reload_process_groups() + if self._skip_memory_clear_after_train_wake_up(): + logger.info("Skipping gc/empty_cache after train wake_up; synchronizing only.") + torch.cuda.synchronize() + else: + clear_memory() + if not self._keep_train_process_groups_during_offload(): + reload_process_groups() if self.role == "actor": self._switch_model("actor") print_memory("after wake_up model") @@ -523,16 +659,19 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data # Train if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" - with timer("actor_train"): - train( - rollout_id, - self.model, - self.optimizer, - self.opt_param_scheduler, - data_iterator, - num_microbatches, - global_batch_sizes, - ) + with _NsysCudaProfilerCapture(self.role, rollout_id, "actor_train"): + reset_peak_memory_tracker("actor_train") + with timer("actor_train"): + train( + rollout_id, + self.model, + self.optimizer, + self.opt_param_scheduler, + data_iterator, + num_microbatches, + global_batch_sizes, + ) + Timer().update_metrics(_collect_actor_train_memory_metrics()) self.prof.step(rollout_id=rollout_id) @@ -603,7 +742,8 @@ def update_weights(self) -> None: if reconnect_rollout_engines: self.wake_up() elif self.args.offload_train: - reload_process_groups() + if not self._keep_train_process_groups_during_offload(): + reload_process_groups() if num_new_engines > 0 or reconnect_rollout_engines: self.weight_updater.connect_rollout_engines( @@ -643,7 +783,10 @@ def update_weights(self) -> None: if reconnect_rollout_engines: self.sleep() elif self.args.offload_train: - destroy_process_groups() + if self._keep_train_process_groups_during_offload(): + logger.info("Keeping train process groups alive after update_weights.") + else: + destroy_process_groups() def load_other_checkpoint(self, model_tag: str, path: str) -> None: old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune diff --git a/vime/backends/megatron_utils/arguments.py b/vime/backends/megatron_utils/arguments.py index 59a05dca..04fa4635 100644 --- a/vime/backends/megatron_utils/arguments.py +++ b/vime/backends/megatron_utils/arguments.py @@ -1,5 +1,6 @@ import ast import logging +import os from megatron.training.arguments import parse_args as _megatron_parse_args from megatron.training.arguments import validate_args as _megatron_validate_args @@ -72,7 +73,22 @@ def _is_moe_config(hf_config): def validate_args(args): """Run megatron's own validate_args plus vime-specific megatron validations.""" - _megatron_validate_args(args) + skip_grouped_gemm_capability_check = ( + os.environ.get("VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK", "0") == "1" + and getattr(args, "moe_grouped_gemm", False) + ) + if skip_grouped_gemm_capability_check: + logger.info( + "Skipping Megatron grouped GEMM compute capability check during argument " + "validation; preserving --moe-grouped-gemm for runtime." + ) + args.moe_grouped_gemm = False + try: + _megatron_validate_args(args) + finally: + args.moe_grouped_gemm = True + else: + _megatron_validate_args(args) # always use varlen args.variable_seq_lengths = True @@ -146,7 +162,7 @@ def equal(x, y): def _set_default_megatron_args(args): # always use zero optimizer - args.use_distributed_optimizer = True + args.use_distributed_optimizer = os.environ.get("VIME_USE_DISTRIBUTED_OPTIMIZER", "1") != "0" # TODO: maybe change this after megatron has good fp8 support args.bf16 = not args.fp16 # placeholders diff --git a/vime/backends/megatron_utils/baseline_timer.py b/vime/backends/megatron_utils/baseline_timer.py new file mode 100644 index 00000000..7dc1da24 --- /dev/null +++ b/vime/backends/megatron_utils/baseline_timer.py @@ -0,0 +1,122 @@ +import os +from argparse import Namespace + +import torch + +from .cuda_event_timer import CudaEventTimerQueue + + +_COUNTER_KEYS = ( + "output_layer_call_count", + "output_layer_token_count", + "output_layer_dispatch_elapsed_s", + "output_layer_cuda_event_count", + "output_layer_cuda_event_elapsed_s", + "output_layer_forward_backward_cuda_event_count", + "output_layer_forward_backward_cuda_event_elapsed_s", + "native_logprob_call_count", + "native_logprob_token_count", + "native_logprob_dispatch_elapsed_s", + "native_logprob_cuda_event_count", + "native_logprob_cuda_event_elapsed_s", +) +_COUNTERS: dict[str, float] = dict.fromkeys(_COUNTER_KEYS, 0.0) +_LAST_SNAPSHOT: dict[str, float] = dict.fromkeys(_COUNTER_KEYS, 0.0) +_CUDA_EVENT_TIMER_QUEUE = CudaEventTimerQueue() + + +def _env_flag(name: str) -> bool: + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} + + +def baseline_linear_logp_timer_enabled(args: Namespace | None = None) -> bool: + if not _env_flag("VIME_BASELINE_LINEAR_LOGP_TIMER"): + return False + if args is not None and getattr(args, "enable_rl_kernel", False): + return False + return True + + +def baseline_cuda_event_timer_enabled(args: Namespace | None = None) -> bool: + if not baseline_linear_logp_timer_enabled(args): + return False + return _env_flag("VIME_BASELINE_CUDA_EVENT_TIMER") + + +def reset_baseline_linear_logp_runtime_counters() -> None: + _CUDA_EVENT_TIMER_QUEUE.clear() + for key in _COUNTER_KEYS: + _COUNTERS[key] = 0.0 + _LAST_SNAPSHOT[key] = 0.0 + + +def get_baseline_linear_logp_runtime_counters() -> dict[str, float]: + _CUDA_EVENT_TIMER_QUEUE.flush_ready() + return dict(_COUNTERS) + + +def get_baseline_linear_logp_runtime_counter_delta() -> dict[str, float]: + current = get_baseline_linear_logp_runtime_counters() + delta = {key: current.get(key, 0.0) - _LAST_SNAPSHOT.get(key, 0.0) for key in _COUNTER_KEYS} + _LAST_SNAPSHOT.update(current) + return delta + + +def record_baseline_output_layer_runtime(token_count: int, elapsed_s: float) -> None: + _COUNTERS["output_layer_call_count"] += 1.0 + _COUNTERS["output_layer_token_count"] += float(token_count) + _COUNTERS["output_layer_dispatch_elapsed_s"] += float(elapsed_s) + + +def record_baseline_output_layer_cuda_event_runtime(elapsed_s: float) -> None: + _COUNTERS["output_layer_cuda_event_count"] += 1.0 + _COUNTERS["output_layer_cuda_event_elapsed_s"] += float(elapsed_s) + + +def record_baseline_output_layer_forward_backward_cuda_event_runtime(elapsed_s: float) -> None: + _COUNTERS["output_layer_forward_backward_cuda_event_count"] += 1.0 + _COUNTERS["output_layer_forward_backward_cuda_event_elapsed_s"] += float(elapsed_s) + + +def record_baseline_native_logprob_runtime(token_count: int, elapsed_s: float) -> None: + _COUNTERS["native_logprob_call_count"] += 1.0 + _COUNTERS["native_logprob_token_count"] += float(token_count) + _COUNTERS["native_logprob_dispatch_elapsed_s"] += float(elapsed_s) + + +def record_baseline_native_logprob_cuda_event_runtime(elapsed_s: float) -> None: + _COUNTERS["native_logprob_cuda_event_count"] += 1.0 + _COUNTERS["native_logprob_cuda_event_elapsed_s"] += float(elapsed_s) + + +def queue_baseline_output_layer_cuda_event( + start_event: torch.cuda.Event, + end_event: torch.cuda.Event, +) -> None: + _CUDA_EVENT_TIMER_QUEUE.enqueue( + start_event, + end_event, + record_baseline_output_layer_cuda_event_runtime, + ) + + +def queue_baseline_output_layer_forward_backward_cuda_event( + start_event: torch.cuda.Event, + end_event: torch.cuda.Event, +) -> None: + _CUDA_EVENT_TIMER_QUEUE.enqueue( + start_event, + end_event, + record_baseline_output_layer_forward_backward_cuda_event_runtime, + ) + + +def queue_baseline_native_logprob_cuda_event( + start_event: torch.cuda.Event, + end_event: torch.cuda.Event, +) -> None: + _CUDA_EVENT_TIMER_QUEUE.enqueue( + start_event, + end_event, + record_baseline_native_logprob_cuda_event_runtime, + ) diff --git a/vime/backends/megatron_utils/cuda_event_timer.py b/vime/backends/megatron_utils/cuda_event_timer.py new file mode 100644 index 00000000..0afc97bc --- /dev/null +++ b/vime/backends/megatron_utils/cuda_event_timer.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from collections.abc import Callable + +import torch + + +class CudaEventTimerQueue: + """Collect CUDA event timings without synchronizing the hot path.""" + + def __init__(self) -> None: + self._pending: list[tuple[torch.cuda.Event, torch.cuda.Event, Callable[[float], None]]] = [] + + def clear(self) -> None: + self._pending.clear() + + def enqueue( + self, + start_event: torch.cuda.Event, + end_event: torch.cuda.Event, + record_runtime: Callable[[float], None], + ) -> None: + self._pending.append((start_event, end_event, record_runtime)) + self.flush_ready() + + def flush_ready(self) -> None: + if not self._pending: + return + remaining = [] + for start_event, end_event, record_runtime in self._pending: + if end_event.query(): + record_runtime(start_event.elapsed_time(end_event) / 1000.0) + else: + remaining.append((start_event, end_event, record_runtime)) + self._pending = remaining diff --git a/vime/backends/megatron_utils/data.py b/vime/backends/megatron_utils/data.py index 42c19e7e..8972322d 100644 --- a/vime/backends/megatron_utils/data.py +++ b/vime/backends/megatron_utils/data.py @@ -1,4 +1,5 @@ import logging +import os from argparse import Namespace from collections.abc import Sequence @@ -57,6 +58,7 @@ def get_batch( batch = data_iterator.get_next(keys) tokens = batch["tokens"] + num_real_sequences = len(tokens) # use 0 as the pad token id should be fine? pad_token_id = 0 pad_size = mpu.get_tensor_model_parallel_world_size() * pad_multiplier @@ -120,6 +122,13 @@ def get_batch( max_seqlen_kv=max_seqlen, qkv_format="thd", ) + if ( + os.environ.get("MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ", "0") == "1" + and cp_size == 1 + and not allgather_cp + and num_real_sequences == 1 + ): + packed_seq_params = None tokens = tokens.unsqueeze(0) else: diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 3f6ab29c..dd3cb025 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -1,3 +1,6 @@ +import logging +import os +import time from argparse import Namespace from collections.abc import Callable, Iterator from typing import Any @@ -8,6 +11,7 @@ from megatron.core import mpu from torch.utils.checkpoint import checkpoint +from vime.utils.memory_utils import update_peak_memory_tracker from vime.utils.distributed_utils import distributed_masked_whiten from vime.utils.misc import load_function from vime.utils.ppo_utils import ( @@ -23,12 +27,25 @@ ) from vime.utils.types import RolloutBatch +from .baseline_timer import ( + baseline_cuda_event_timer_enabled, + baseline_linear_logp_timer_enabled, + queue_baseline_native_logprob_cuda_event, + record_baseline_native_logprob_runtime, +) from .cp_utils import ( all_gather_with_cp, get_logits_and_tokens_offset_with_cp, get_sum_of_sample_mean, slice_log_prob_with_cp, ) +from .rl_kernel import LinearLogpContext, get_rl_kernel_fallback_count, maybe_compute_linear_logp, maybe_compute_logp + +logger = logging.getLogger(__name__) + + +def _env_flag(name: str) -> bool: + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} def get_responses( @@ -383,6 +400,70 @@ def _extract_per_sample( return log_probs_list, entropy_list +def _gather_sequence_parallel_hidden_if_needed( + hidden_states: torch.Tensor, + context: LinearLogpContext | None, +) -> torch.Tensor: + if context is None or not context.sequence_parallel: + return hidden_states + + from megatron.core import tensor_parallel + + return tensor_parallel.gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) + + +def _flatten_logprob_model_output( + output_tensor: torch.Tensor, + *, + qkv_format: str, + max_seq_lens: list[int] | None, + linear_logp_context: LinearLogpContext | None, +) -> torch.Tensor: + if qkv_format == "thd": + assert len(output_tensor.shape) == 3, f"{output_tensor.shape}" + if output_tensor.size(0) == 1: + return output_tensor.squeeze(0) + if linear_logp_context is not None and output_tensor.size(1) == 1: + return output_tensor.squeeze(1) + assert output_tensor.size(0) == 1, f"{output_tensor.shape}" + else: + assert max_seq_lens is not None + return output_tensor.view(-1, output_tensor.size(-1)) + + raise AssertionError(f"Unsupported output tensor shape: {output_tensor.shape}") + + +def _materialize_linear_logits( + hidden_states: torch.Tensor, + *, + context: LinearLogpContext, + args: Namespace, +) -> torch.Tensor: + logits = F.linear(hidden_states, context.lm_head_weight, context.bias) + rollout_temperature = getattr(args, "rollout_temperature", 1.0) + if rollout_temperature != 1.0: + logits = logits / rollout_temperature + return logits.float() + + +def _policy_loss_needs_entropy( + args: Namespace, + rl_kernel_linear_logp_context: LinearLogpContext | None, +) -> bool: + if ( + getattr(args, "entropy_coef", 0.0) == 0 + and _env_flag("VIME_SKIP_ZERO_ENTROPY_METRIC") + ): + return False + if rl_kernel_linear_logp_context is None: + return True + return not ( + getattr(args, "enable_rl_kernel", False) + and "linear_logp" in getattr(args, "rl_kernel_ops", ()) + and getattr(args, "entropy_coef", 0.0) == 0 + ) + + def get_log_probs_and_entropy( logits: torch.Tensor, *, @@ -393,6 +474,7 @@ def get_log_probs_and_entropy( with_entropy: bool = False, non_loss_data: bool = True, max_seq_lens: list[int] | None = None, + rl_kernel_linear_logp_context: LinearLogpContext | None = None, ) -> dict[str, list[torch.Tensor]]: """Compute per-token log-probabilities (and optionally entropy) on responses. @@ -406,21 +488,28 @@ def get_log_probs_and_entropy( assert non_loss_data qkv_format = args.qkv_format - assert logits.dtype == torch.float32, f"{logits.dtype}" assert len(logits.shape) == 3, f"{logits.shape}" - if qkv_format == "thd": - assert logits.size(0) == 1, f"{logits.shape}" - logits = logits.squeeze(0) + linear_logp_context = rl_kernel_linear_logp_context + if linear_logp_context is not None: + logits = _gather_sequence_parallel_hidden_if_needed(logits, linear_logp_context) else: - assert max_seq_lens is not None - logits = logits.view(-1, logits.size(-1)) + assert logits.dtype == torch.float32, f"{logits.dtype}" + + logits = _flatten_logprob_model_output( + logits, + qkv_format=qkv_format, + max_seq_lens=max_seq_lens, + linear_logp_context=linear_logp_context, + ).contiguous() + + if linear_logp_context is None: + # Apply rollout temperature scaling to logits to match rollout-time log-probs. + rollout_temperature = getattr(args, "rollout_temperature", 1.0) + if rollout_temperature != 1.0: + logits = logits / rollout_temperature + logits = logits.contiguous() - # Apply rollout temperature scaling to logits to match rollout-time log-probs. - rollout_temperature = getattr(args, "rollout_temperature", 1.0) - if rollout_temperature != 1.0: - logits = logits / rollout_temperature - logits = logits.contiguous() T = logits.size(0) device = logits.device tp_group = mpu.get_tensor_model_parallel_group() @@ -431,14 +520,96 @@ def get_log_probs_and_entropy( T, device, unconcat_tokens, total_lengths, response_lengths, qkv_format, max_seq_lens, args.allgather_cp ) - # --- compute on full [T,V] logits at once via calculate_log_probs_and_entropy --- - log_prob_full, entropy_full = calculate_log_probs_and_entropy( - logits, - full_tokens, - tp_group, - with_entropy=with_entropy, - chunk_size=chunk_size, + # --- compute on full [T,V] logits at once --- + log_prob_full = None + if linear_logp_context is not None: + log_prob_full = maybe_compute_linear_logp( + logits, + full_tokens, + context=linear_logp_context, + args=args, + with_entropy=with_entropy, + ) + if log_prob_full is None: + logits = _materialize_linear_logits(logits, context=linear_logp_context, args=args).contiguous() + else: + log_prob_full = maybe_compute_logp(logits, full_tokens, args=args, with_entropy=with_entropy) + + native_memory_probe = ( + log_prob_full is None + and linear_logp_context is None + and _env_flag("VIME_LINEAR_LOGP_MEMORY_PROBE") + and logits.is_cuda ) + if native_memory_probe: + probe_device = logits.device + torch.cuda.synchronize(probe_device) + probe_before_alloc = torch.cuda.memory_allocated(probe_device) + probe_before_reserved = torch.cuda.memory_reserved(probe_device) + update_peak_memory_tracker("actor_train", device=probe_device) + torch.cuda.reset_peak_memory_stats(probe_device) + probe_start_s = time.perf_counter() + + if log_prob_full is None: + native_runtime_timer = linear_logp_context is None and baseline_linear_logp_timer_enabled(args) + native_timer_start_s = time.perf_counter() if native_runtime_timer else None + native_event_timer = ( + linear_logp_context is None + and baseline_cuda_event_timer_enabled(args) + and logits.is_cuda + ) + native_event_start = None + if native_event_timer: + native_event_start = torch.cuda.Event(enable_timing=True) + native_event_start.record() + log_prob_full, entropy_full = calculate_log_probs_and_entropy( + logits, + full_tokens, + tp_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + ) + if native_timer_start_s is not None: + record_baseline_native_logprob_runtime(full_tokens.numel(), time.perf_counter() - native_timer_start_s) + if native_event_start is not None: + native_event_end = torch.cuda.Event(enable_timing=True) + native_event_end.record() + queue_baseline_native_logprob_cuda_event( + native_event_start, + native_event_end, + ) + else: + entropy_full = None + + if native_memory_probe: + torch.cuda.synchronize(probe_device) + probe_after_alloc = torch.cuda.memory_allocated(probe_device) + probe_after_reserved = torch.cuda.memory_reserved(probe_device) + probe_peak_alloc = torch.cuda.max_memory_allocated(probe_device) + probe_peak_reserved = torch.cuda.max_memory_reserved(probe_device) + update_peak_memory_tracker( + "actor_train", + peak_alloc=probe_peak_alloc, + peak_reserved=probe_peak_reserved, + device=probe_device, + ) + logger.info( + "Baseline native_logprob memory_probe: op=calculate_log_probs_and_entropy " + "logits_shape=%s tokens=%d with_entropy=%s alloc_before_mb=%.2f " + "peak_alloc_mb=%.2f peak_delta_mb=%.2f alloc_after_mb=%.2f " + "alloc_after_delta_mb=%.2f reserved_before_mb=%.2f reserved_after_mb=%.2f elapsed_s=%.6f", + tuple(logits.shape), + int(full_tokens.numel()), + bool(with_entropy), + probe_before_alloc / (1024**2), + probe_peak_alloc / (1024**2), + (probe_peak_alloc - probe_before_alloc) / (1024**2), + probe_after_alloc / (1024**2), + (probe_after_alloc - probe_before_alloc) / (1024**2), + probe_before_reserved / (1024**2), + probe_after_reserved / (1024**2), + time.perf_counter() - probe_start_s, + ) log_prob_full = log_prob_full.squeeze(-1) # [T, 1] -> [T] # --- extract per-sample response portions --- @@ -802,6 +973,7 @@ def policy_loss_function( batch: RolloutBatch, logits: torch.Tensor, sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], + rl_kernel_linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute policy loss (PPO/GSPO) and metrics. @@ -834,14 +1006,17 @@ def policy_loss_function( total_lengths = batch["total_lengths"] max_seq_lens = batch.get("max_seq_lens", None) + need_entropy = _policy_loss_needs_entropy(args, rl_kernel_linear_logp_context) + _, log_probs_and_entropy = get_log_probs_and_entropy( logits, args=args, unconcat_tokens=batch["unconcat_tokens"], total_lengths=total_lengths, response_lengths=response_lengths, - with_entropy=True, + with_entropy=need_entropy, max_seq_lens=max_seq_lens, + rl_kernel_linear_logp_context=rl_kernel_linear_logp_context, ) log_probs = log_probs_and_entropy["log_probs"] @@ -963,9 +1138,12 @@ def policy_loss_function( ppo_kl = sum_of_sample_mean(ppo_kl) # entropy loss - entropy = log_probs_and_entropy["entropy"] - entropy = torch.cat(entropy, dim=0) - entropy_loss = sum_of_sample_mean(entropy) + if need_entropy: + entropy = log_probs_and_entropy["entropy"] + entropy = torch.cat(entropy, dim=0) + entropy_loss = sum_of_sample_mean(entropy) + else: + entropy_loss = log_probs.new_zeros(()) loss = pg_loss - args.entropy_coef * entropy_loss @@ -1033,6 +1211,7 @@ def value_loss_function( batch: RolloutBatch, logits: torch.Tensor, sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], + rl_kernel_linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute clipped value loss and metrics. @@ -1051,6 +1230,7 @@ def value_loss_function( Tuple of `(loss, metrics)` where `loss` is a scalar tensor and `metrics` contains detached scalars "value_loss" and "value_clipfrac". """ + del rl_kernel_linear_logp_context old_values = torch.cat(batch["values"], dim=0) _, values = get_values( @@ -1091,6 +1271,7 @@ def sft_loss_function( batch: RolloutBatch, logits: torch.Tensor, sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], + rl_kernel_linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute supervised fine-tuning loss over response tokens. @@ -1119,6 +1300,7 @@ def sft_loss_function( response_lengths=response_lengths, with_entropy=False, max_seq_lens=batch.get("max_seq_lens", None), + rl_kernel_linear_logp_context=rl_kernel_linear_logp_context, ) log_probs = log_probs_and_entropy["log_probs"] @@ -1143,6 +1325,7 @@ def loss_function( num_microbatches: int, step_global_batch_size: int, logits: torch.Tensor, + rl_kernel_linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]: """Dispatch to the configured loss and rescale for Megatron integration. @@ -1195,10 +1378,22 @@ def loss_function( case _: raise ValueError(f"Unknown loss type: {args.loss_type}") + if func in {policy_loss_function, value_loss_function, sft_loss_function}: + func_args = (args, batch, logits, sum_of_sample_mean, rl_kernel_linear_logp_context) + else: + func_args = (args, batch, logits, sum_of_sample_mean) + if args.recompute_loss_function: - loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean, use_reentrant=False) + loss, log = checkpoint(func, *func_args, use_reentrant=False) else: - loss, log = func(args, batch, logits, sum_of_sample_mean) + loss, log = func(*func_args) + + if getattr(args, "enable_rl_kernel", False): + log["rl_kernel_fallback_count"] = torch.tensor( + get_rl_kernel_fallback_count(), + device=logits.device, + dtype=torch.float32, + ) # With allgather-CP, some CP ranks may have no loss-contributing tokens (e.g., all # padding). Without this, gradient doesn't flow through their attention path, so diff --git a/vime/backends/megatron_utils/megatron_to_hf/qwen2.py b/vime/backends/megatron_utils/megatron_to_hf/qwen2.py index f7b72935..10565062 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/qwen2.py +++ b/vime/backends/megatron_utils/megatron_to_hf/qwen2.py @@ -57,9 +57,13 @@ def convert_qwen2_to_hf(args, name, param): ] elif rest == "mlp.linear_fc2.weight": return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)] - elif rest == "self_attention.linear_qkv.layer_norm_weight": + elif rest in {"self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"}: return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] - elif rest == "mlp.linear_fc1.layer_norm_weight": + elif rest in { + "mlp.linear_fc1.layer_norm_weight", + "pre_mlp_layernorm.weight", + "post_attention_layernorm.weight", + }: return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] # qk norm diff --git a/vime/backends/megatron_utils/megatron_to_hf/qwen3moe.py b/vime/backends/megatron_utils/megatron_to_hf/qwen3moe.py index 9f5b5b81..d64ee88f 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/qwen3moe.py +++ b/vime/backends/megatron_utils/megatron_to_hf/qwen3moe.py @@ -97,11 +97,13 @@ def convert_qwen3moe_to_hf(args, name, param): ] elif rest == "mlp.linear_fc2.weight": return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)] - elif rest == "self_attention.linear_qkv.layer_norm_weight": + elif rest in {"self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"}: return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] - elif rest == "mlp.linear_fc1.layer_norm_weight": - return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] - elif rest == "pre_mlp_layernorm.weight": + elif rest in { + "mlp.linear_fc1.layer_norm_weight", + "pre_mlp_layernorm.weight", + "post_attention_layernorm.weight", + }: return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] elif rest == "mlp.router.weight": return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)] diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index 6b602fc1..40af29bd 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -3,8 +3,10 @@ import logging import math import os +import time from argparse import Namespace from collections.abc import Callable, Sequence +from contextlib import contextmanager from functools import partial from pathlib import Path @@ -27,14 +29,33 @@ from megatron.core.pipeline_parallel.utils import unwrap_model except ImportError: from megatron.core.utils import unwrap_model -from vime.utils import logging_utils -from vime.utils.memory_utils import clear_memory +from vime.utils import logging_utils +from vime.utils.memory_utils import clear_memory, update_peak_memory_tracker +from vime.utils.rl_kernel import is_rl_kernel_op_enabled + +from .baseline_timer import ( + baseline_cuda_event_timer_enabled, + baseline_linear_logp_timer_enabled, + get_baseline_linear_logp_runtime_counter_delta, + get_baseline_linear_logp_runtime_counters, + queue_baseline_output_layer_cuda_event, + queue_baseline_output_layer_forward_backward_cuda_event, + record_baseline_output_layer_runtime, +) from .checkpoint import load_checkpoint, save_checkpoint from .cp_utils import reduce_train_step_metrics from .data import DataIterator, get_batch -from .loss import loss_function +from .loss import get_log_probs_and_entropy, loss_function from .model_provider import get_model_provider_func +from .rl_kernel import ( + get_linear_logp_context_from_model, + get_rl_kernel_runtime_counter_delta, + get_rl_kernel_runtime_counters, + return_hidden_states_for_linear_logp, + should_use_linear_logp_model_output, + warn_linear_logp_fallback, +) logger = logging.getLogger(__name__) @@ -73,6 +94,250 @@ def wrapped_forward_step(*args, **kwargs): return wrapped_forward_step +def _env_flag(name: str) -> bool: + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} + + +def _first_tensor(value): + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + if isinstance(value, dict): + for item in value.values(): + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + + +def _detach_first_tensor(value): + if isinstance(value, torch.Tensor): + return value.detach(), True + if isinstance(value, tuple): + new_items = [] + changed = False + for item in value: + if not changed: + new_item, changed = _detach_first_tensor(item) + new_items.append(new_item) + else: + new_items.append(item) + return tuple(new_items), changed + if isinstance(value, list): + new_items = [] + changed = False + for item in value: + if not changed: + new_item, changed = _detach_first_tensor(item) + new_items.append(new_item) + else: + new_items.append(item) + return new_items, changed + if isinstance(value, dict): + new_value = dict(value) + for key, item in value.items(): + new_item, changed = _detach_first_tensor(item) + if changed: + new_value[key] = new_item + return new_value, True + return value, False + + +@contextmanager +def _probe_baseline_output_layer_forward(args: Namespace, model): + memory_probe = _env_flag("VIME_LINEAR_LOGP_MEMORY_PROBE") + runtime_timer = baseline_linear_logp_timer_enabled(args) + event_timer = baseline_cuda_event_timer_enabled(args) + detach_hidden = _env_flag("VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN") + if getattr(args, "enable_rl_kernel", False) or not (memory_probe or runtime_timer or event_timer or detach_hidden): + yield + return + + module = model + while hasattr(module, "module"): + module = module.module + output_layer = getattr(module, "output_layer", None) + if output_layer is None: + yield + return + + state: dict[str, object] = {} + backward_handles = [] + + def register_backward_event_timer(module, input_tensor, start_event): + watched_tensor = None + # Prefer the layer input for full-gradient timing; tied/shared output + # weights may get gradient contributions outside the output layer. + for candidate in (input_tensor, getattr(module, "weight", None), getattr(module, "bias", None)): + if isinstance(candidate, torch.Tensor) and candidate.requires_grad: + watched_tensor = candidate + break + if watched_tensor is None: + return + + handle_box = {} + + def hook(grad): + end_event = torch.cuda.Event(enable_timing=True) + end_event.record() + queue_baseline_output_layer_forward_backward_cuda_event( + start_event, + end_event, + ) + handle = handle_box.get("handle") + if handle is not None: + handle.remove() + if handle in backward_handles: + backward_handles.remove(handle) + return grad + + handle_box["handle"] = watched_tensor.register_hook(hook) + backward_handles.append(handle_box["handle"]) + + def pre_hook(_module, inputs, kwargs): + input_tensor = kwargs.get("input_") if isinstance(kwargs, dict) else None + if input_tensor is None: + input_tensor = _first_tensor(inputs) + if input_tensor is None or not input_tensor.is_cuda: + state.clear() + if detach_hidden: + new_inputs, inputs_changed = _detach_first_tensor(inputs) + new_kwargs, kwargs_changed = _detach_first_tensor(kwargs) + if inputs_changed or kwargs_changed: + return new_inputs, new_kwargs + return None + new_inputs = inputs + new_kwargs = kwargs + if detach_hidden: + if isinstance(kwargs, dict) and isinstance(kwargs.get("input_"), torch.Tensor): + new_kwargs = dict(kwargs) + new_kwargs["input_"] = kwargs["input_"].detach() + input_tensor = new_kwargs["input_"] + else: + new_inputs, _changed = _detach_first_tensor(inputs) + input_tensor = _first_tensor(new_inputs) + probe_device = input_tensor.device + state.clear() + tokens = int(input_tensor.numel() // input_tensor.size(-1)) if input_tensor.dim() > 0 else 0 + if memory_probe: + torch.cuda.synchronize(probe_device) + state.update( + { + "device": probe_device, + "input_shape": tuple(input_tensor.shape), + "tokens": tokens, + } + ) + if runtime_timer: + state["timer_start_s"] = time.perf_counter() + if event_timer: + start_event = torch.cuda.Event(enable_timing=True) + start_event.record() + state["cuda_event_start"] = start_event + if torch.is_grad_enabled(): + register_backward_event_timer(_module, input_tensor, start_event) + if memory_probe: + state["alloc_before"] = torch.cuda.memory_allocated(probe_device) + state["reserved_before"] = torch.cuda.memory_reserved(probe_device) + update_peak_memory_tracker("actor_train", device=probe_device) + torch.cuda.reset_peak_memory_stats(probe_device) + if detach_hidden: + return new_inputs, new_kwargs + return None + + def post_hook(_module, _inputs, output): + if not state: + return + probe_device = state["device"] + timer_start_s = state.get("timer_start_s") + if timer_start_s is not None: + record_baseline_output_layer_runtime(int(state["tokens"]), time.perf_counter() - float(timer_start_s)) + cuda_event_start = state.get("cuda_event_start") + if cuda_event_start is not None: + end_event = torch.cuda.Event(enable_timing=True) + end_event.record() + queue_baseline_output_layer_cuda_event( + cuda_event_start, + end_event, + ) + if not memory_probe: + state.clear() + return + + output_tensor = _first_tensor(output) + torch.cuda.synchronize(probe_device) + alloc_before = int(state["alloc_before"]) + reserved_before = int(state["reserved_before"]) + alloc_after = torch.cuda.memory_allocated(probe_device) + reserved_after = torch.cuda.memory_reserved(probe_device) + peak_alloc = torch.cuda.max_memory_allocated(probe_device) + peak_reserved = torch.cuda.max_memory_reserved(probe_device) + update_peak_memory_tracker( + "actor_train", + peak_alloc=peak_alloc, + peak_reserved=peak_reserved, + device=probe_device, + ) + logger.info( + "Baseline output_layer memory_probe: op=%s input_shape=%s output_shape=%s " + "tokens=%d alloc_before_mb=%.2f peak_alloc_mb=%.2f peak_delta_mb=%.2f " + "alloc_after_mb=%.2f alloc_after_delta_mb=%.2f reserved_before_mb=%.2f reserved_after_mb=%.2f", + type(_module).__name__, + state["input_shape"], + None if output_tensor is None else tuple(output_tensor.shape), + int(state["tokens"]), + alloc_before / (1024**2), + peak_alloc / (1024**2), + (peak_alloc - alloc_before) / (1024**2), + alloc_after / (1024**2), + (alloc_after - alloc_before) / (1024**2), + reserved_before / (1024**2), + reserved_after / (1024**2), + ) + state.clear() + + pre_handle = output_layer.register_forward_pre_hook(pre_hook, with_kwargs=True) + post_handle = output_layer.register_forward_hook(post_hook) + try: + yield + finally: + pre_handle.remove() + post_handle.remove() + + +def _forward_only_should_return_hidden_for_linear_logp( + f: Callable[..., dict[str, list[torch.Tensor]]], + args: Namespace, +) -> bool: + return f is get_log_probs_and_entropy and should_use_linear_logp_model_output( + args, + with_entropy=args.use_rollout_entropy, + ) + + +def _train_should_return_hidden_for_linear_logp(args: Namespace, *, return_schedule_plan: bool) -> bool: + if not is_rl_kernel_op_enabled(args, "linear_logp"): + return False + + if args.loss_type not in {"policy_loss", "sft_loss"}: + return False + + if return_schedule_plan: + warn_linear_logp_fallback(args, "schedule-plan forward path is not supported") + return False + + if getattr(args, "enable_mtp_training", False): + warn_linear_logp_fallback(args, "MTP training path is not supported") + return False + + with_entropy = args.loss_type == "policy_loss" and getattr(args, "entropy_coef", 0.0) != 0 + return should_use_linear_logp_model_output(args, with_entropy=with_entropy) + + def _iter_critic_output_layers(model: Sequence[DDP]): for chunk_id, module in enumerate(unwrap_model(model)): output_layer = getattr(module, "output_layer", None) @@ -342,17 +607,26 @@ def forward_step( } if batch["multimodal_train_inputs"] is not None: forward_kwargs.update(batch["multimodal_train_inputs"]) - output_tensor = model(**forward_kwargs) + linear_logp_context = None + if _forward_only_should_return_hidden_for_linear_logp(f, args): + linear_logp_context = get_linear_logp_context_from_model(args, model) + + with _probe_baseline_output_layer_forward(args, model): + with return_hidden_states_for_linear_logp(args, model, linear_logp_context): + output_tensor = model(**forward_kwargs) + + callback_kwargs = { + "args": args, + "unconcat_tokens": unconcat_tokens, + "total_lengths": total_lengths, + "response_lengths": response_lengths, + "with_entropy": args.use_rollout_entropy, + "max_seq_lens": batch.get("max_seq_lens", None), + } + if f is get_log_probs_and_entropy: + callback_kwargs["rl_kernel_linear_logp_context"] = linear_logp_context - return output_tensor, partial( - f, - args=args, - unconcat_tokens=unconcat_tokens, - total_lengths=total_lengths, - response_lengths=response_lengths, - with_entropy=args.use_rollout_entropy, - max_seq_lens=batch.get("max_seq_lens", None), - ) + return output_tensor, partial(f, **callback_kwargs) # Turn on evaluation mode which disables dropout. for model_module in model: @@ -512,6 +786,10 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p old_stage = os.environ["ROUTING_REPLAY_STAGE"] os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" + linear_logp_context = None + if _train_should_return_hidden_for_linear_logp(args, return_schedule_plan=return_schedule_plan): + linear_logp_context = get_linear_logp_context_from_model(args, model) + if return_schedule_plan: assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b" position_ids = None @@ -539,12 +817,21 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p if args.enable_mtp_training: forward_kwargs["mtp_kwargs"] = {"mtp_labels": batch["tokens"]} - output_tensor = model(**forward_kwargs) + with _probe_baseline_output_layer_forward(args, model): + with return_hidden_states_for_linear_logp(args, model, linear_logp_context): + output_tensor = model(**forward_kwargs) if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": os.environ["ROUTING_REPLAY_STAGE"] = old_stage - return output_tensor, partial(loss_function, args, batch, num_microbatches, step_global_batch_size) + return output_tensor, partial( + loss_function, + args, + batch, + num_microbatches, + step_global_batch_size, + rl_kernel_linear_logp_context=linear_logp_context, + ) # Forward pass. forward_backward_func = get_forward_backward_func() @@ -795,6 +1082,109 @@ def train( # Per-step gbs — uneven step sizes are easy to miss without this. log_dict[f"train/{role_tag}global_batch_size"] = global_batch_sizes[step_id] + if role == "actor" and getattr(args, "enable_rl_kernel", False): + runtime_totals = get_rl_kernel_runtime_counters() + runtime_delta = get_rl_kernel_runtime_counter_delta() + for key, value in runtime_totals.items(): + log_dict[f"train/rl_kernel_{key}_total"] = value + for key, value in runtime_delta.items(): + log_dict[f"train/rl_kernel_{key}_delta"] = value + total_calls = runtime_totals.get("linear_logp_call_count", 0.0) + delta_calls = runtime_delta.get("linear_logp_call_count", 0.0) + log_dict["train/rl_kernel_linear_logp_tokens_per_call_total"] = ( + runtime_totals.get("linear_logp_token_count", 0.0) / total_calls if total_calls > 0 else 0.0 + ) + log_dict["train/rl_kernel_linear_logp_tokens_per_call_delta"] = ( + runtime_delta.get("linear_logp_token_count", 0.0) / delta_calls if delta_calls > 0 else 0.0 + ) + if role == "actor" and baseline_linear_logp_timer_enabled(args): + runtime_totals = get_baseline_linear_logp_runtime_counters() + runtime_delta = get_baseline_linear_logp_runtime_counter_delta() + for key, value in runtime_totals.items(): + log_dict[f"train/baseline_{key}_total"] = value + for key, value in runtime_delta.items(): + log_dict[f"train/baseline_{key}_delta"] = value + + total_calls = max( + runtime_totals.get("output_layer_call_count", 0.0), + runtime_totals.get("native_logprob_call_count", 0.0), + ) + delta_calls = max( + runtime_delta.get("output_layer_call_count", 0.0), + runtime_delta.get("native_logprob_call_count", 0.0), + ) + total_tokens = max( + runtime_totals.get("output_layer_token_count", 0.0), + runtime_totals.get("native_logprob_token_count", 0.0), + ) + delta_tokens = max( + runtime_delta.get("output_layer_token_count", 0.0), + runtime_delta.get("native_logprob_token_count", 0.0), + ) + total_elapsed_s = ( + runtime_totals.get("output_layer_dispatch_elapsed_s", 0.0) + + runtime_totals.get("native_logprob_dispatch_elapsed_s", 0.0) + ) + delta_elapsed_s = ( + runtime_delta.get("output_layer_dispatch_elapsed_s", 0.0) + + runtime_delta.get("native_logprob_dispatch_elapsed_s", 0.0) + ) + total_cuda_event_elapsed_s = ( + runtime_totals.get("output_layer_cuda_event_elapsed_s", 0.0) + + runtime_totals.get("native_logprob_cuda_event_elapsed_s", 0.0) + ) + delta_cuda_event_elapsed_s = ( + runtime_delta.get("output_layer_cuda_event_elapsed_s", 0.0) + + runtime_delta.get("native_logprob_cuda_event_elapsed_s", 0.0) + ) + total_forward_backward_cuda_event_count = runtime_totals.get( + "output_layer_forward_backward_cuda_event_count", 0.0 + ) + delta_forward_backward_cuda_event_count = runtime_delta.get( + "output_layer_forward_backward_cuda_event_count", 0.0 + ) + total_forward_backward_cuda_event_elapsed_s = runtime_totals.get( + "output_layer_forward_backward_cuda_event_elapsed_s", 0.0 + ) + delta_forward_backward_cuda_event_elapsed_s = runtime_delta.get( + "output_layer_forward_backward_cuda_event_elapsed_s", 0.0 + ) + log_dict["train/baseline_linear_logp_call_count_total"] = total_calls + log_dict["train/baseline_linear_logp_call_count_delta"] = delta_calls + log_dict["train/baseline_linear_logp_token_count_total"] = total_tokens + log_dict["train/baseline_linear_logp_token_count_delta"] = delta_tokens + log_dict["train/baseline_linear_logp_dispatch_elapsed_s_total"] = total_elapsed_s + log_dict["train/baseline_linear_logp_dispatch_elapsed_s_delta"] = delta_elapsed_s + log_dict["train/baseline_linear_logp_forward_cuda_event_elapsed_s_total"] = ( + total_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_forward_cuda_event_elapsed_s_delta"] = ( + delta_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_forward_backward_cuda_event_count_total"] = ( + total_forward_backward_cuda_event_count + ) + log_dict["train/baseline_linear_logp_forward_backward_cuda_event_count_delta"] = ( + delta_forward_backward_cuda_event_count + ) + log_dict["train/baseline_linear_logp_forward_backward_cuda_event_elapsed_s_total"] = ( + total_forward_backward_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_forward_backward_cuda_event_elapsed_s_delta"] = ( + delta_forward_backward_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_cuda_event_elapsed_s_total"] = ( + total_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_cuda_event_elapsed_s_delta"] = ( + delta_cuda_event_elapsed_s + ) + log_dict["train/baseline_linear_logp_tokens_per_call_total"] = ( + total_tokens / total_calls if total_calls > 0 else 0.0 + ) + log_dict["train/baseline_linear_logp_tokens_per_call_delta"] = ( + delta_tokens / delta_calls if delta_calls > 0 else 0.0 + ) log_dict["train/step"] = accumulated_step_id logging_utils.log(args, log_dict, step_key="train/step") diff --git a/vime/backends/megatron_utils/model_provider.py b/vime/backends/megatron_utils/model_provider.py index e46e3825..39687009 100644 --- a/vime/backends/megatron_utils/model_provider.py +++ b/vime/backends/megatron_utils/model_provider.py @@ -1,6 +1,7 @@ # Adapt from https://github.com/NVIDIA/Megatron-LM/blob/b1efb3c7126ef7615e8c333432d76e08038e17ff/pretrain_gpt.py import argparse import inspect +import logging import re from contextlib import nullcontext from typing import Literal @@ -20,6 +21,8 @@ from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config from vime.utils.misc import load_function +logger = logging.getLogger(__name__) + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): @@ -269,13 +272,39 @@ def get_model_provider_func(args, role="actor"): def freeze_model_params(model: GPTModel, args: argparse.Namespace): if getattr(args, "only_train_params_name_list", None): + matched_names = [] for name, param in model.named_parameters(): param.requires_grad = False for pattern in args.only_train_params_name_list: if re.search(pattern, name): param.requires_grad = True + matched_names.append(name) break + wants_output_layer = any( + re.search(pattern, "output_layer.weight") for pattern in args.only_train_params_name_list + ) + matched_output_layer = any("output_layer" in name for name in matched_names) + if wants_output_layer and not matched_output_layer and getattr(model, "share_embeddings_and_output_weights", False): + shared_weight = None + shared_weight_fn = getattr(model, "shared_embedding_or_output_weight", None) + if callable(shared_weight_fn): + try: + shared_weight = shared_weight_fn() + except Exception: + logger.debug("Unable to read shared embedding/output weight for output-layer-only training.", exc_info=True) + if isinstance(shared_weight, torch.Tensor): + shared_weight.requires_grad_(True) + logger.info( + "Training shared embedding/output weight for output_layer pattern because " + "this model ties embeddings and output weights." + ) + else: + logger.warning( + "only_train_params_name_list requested output_layer, but no output_layer " + "parameter or shared embedding/output weight was found." + ) + if getattr(args, "freeze_params_name_list", None): for name, param in model.named_parameters(): for pattern in args.freeze_params_name_list: diff --git a/vime/backends/megatron_utils/rl_kernel.py b/vime/backends/megatron_utils/rl_kernel.py new file mode 100644 index 00000000..5ca5d304 --- /dev/null +++ b/vime/backends/megatron_utils/rl_kernel.py @@ -0,0 +1,555 @@ +from __future__ import annotations + +import logging +import os +import time +from argparse import Namespace +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import torch +from megatron.core import mpu + +from vime.utils.memory_utils import update_peak_memory_tracker +from vime.utils.rl_kernel import is_rl_kernel_op_enabled + +from .cuda_event_timer import CudaEventTimerQueue + +logger = logging.getLogger(__name__) + +_LOGP_OP = None +_LOGP_OP_LOAD_ERROR: Exception | None = None +_LINEAR_LOGP_OP = None +_LINEAR_LOGP_OP_LOAD_ERROR: Exception | None = None +_WARNED_FALLBACK_REASONS: set[str] = set() +_FALLBACK_COUNTS: dict[str, int] = {"logp": 0, "linear_logp": 0} +_LINEAR_LOGP_SAVE_PROBS_CAST_LOGGED = False +_RUNTIME_COUNTER_KEYS = ( + "linear_logp_call_count", + "linear_logp_token_count", + "linear_logp_dispatch_elapsed_s", + "linear_logp_forward_cuda_event_count", + "linear_logp_forward_cuda_event_elapsed_s", + "linear_logp_forward_backward_cuda_event_count", + "linear_logp_forward_backward_cuda_event_elapsed_s", +) +_RUNTIME_COUNTERS: dict[str, float] = dict.fromkeys(_RUNTIME_COUNTER_KEYS, 0.0) +_RUNTIME_COUNTER_LAST_SNAPSHOT: dict[str, float] = dict.fromkeys(_RUNTIME_COUNTER_KEYS, 0.0) +_CUDA_EVENT_TIMER_QUEUE = CudaEventTimerQueue() + + +def _env_flag(name: str) -> bool: + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} + + +def _env_bool(name: str) -> bool | None: + value = os.getenv(name) + if value is None or value.strip() == "": + return None + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean flag, got {value!r}") + + +@dataclass(frozen=True) +class LinearLogpContext: + lm_head_weight: torch.Tensor + bias: torch.Tensor | None + tp_group: Any + vocab_start_index: int = 0 + global_vocab_size: int | None = None + sequence_parallel: bool = False + + +def get_rl_kernel_fallback_count(op: str | None = None) -> int: + if op is not None: + return _FALLBACK_COUNTS.get(op, 0) + return sum(_FALLBACK_COUNTS.values()) + + +def reset_rl_kernel_runtime_counters() -> None: + _CUDA_EVENT_TIMER_QUEUE.clear() + for key in _RUNTIME_COUNTER_KEYS: + _RUNTIME_COUNTERS[key] = 0.0 + _RUNTIME_COUNTER_LAST_SNAPSHOT[key] = 0.0 + + +def get_rl_kernel_runtime_counters() -> dict[str, float]: + _CUDA_EVENT_TIMER_QUEUE.flush_ready() + return dict(_RUNTIME_COUNTERS) + + +def get_rl_kernel_runtime_counter_delta() -> dict[str, float]: + current = get_rl_kernel_runtime_counters() + delta = { + key: current.get(key, 0.0) - _RUNTIME_COUNTER_LAST_SNAPSHOT.get(key, 0.0) for key in _RUNTIME_COUNTER_KEYS + } + _RUNTIME_COUNTER_LAST_SNAPSHOT.update(current) + return delta + + +def _record_linear_logp_runtime(token_count: int, elapsed_s: float) -> None: + _RUNTIME_COUNTERS["linear_logp_call_count"] += 1.0 + _RUNTIME_COUNTERS["linear_logp_token_count"] += float(token_count) + _RUNTIME_COUNTERS["linear_logp_dispatch_elapsed_s"] += float(elapsed_s) + + +def _record_linear_logp_forward_event_runtime(elapsed_s: float) -> None: + _RUNTIME_COUNTERS["linear_logp_forward_cuda_event_count"] += 1.0 + _RUNTIME_COUNTERS["linear_logp_forward_cuda_event_elapsed_s"] += float(elapsed_s) + + +def _record_linear_logp_forward_backward_event_runtime(elapsed_s: float) -> None: + _RUNTIME_COUNTERS["linear_logp_forward_backward_cuda_event_count"] += 1.0 + _RUNTIME_COUNTERS["linear_logp_forward_backward_cuda_event_elapsed_s"] += float(elapsed_s) + + +def _cuda_event_timer_enabled(tensor: torch.Tensor) -> bool: + return _env_flag("VIME_RL_KERNEL_CUDA_EVENT_TIMER") and tensor.is_cuda + + +def _should_detach_linear_logp_hidden(args: Namespace) -> bool: + override = _env_bool("VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN") + if override is not None: + return override + + patterns = tuple(getattr(args, "only_train_params_name_list", ()) or ()) + return bool(patterns) and all("output_layer" in str(pattern) for pattern in patterns) + + +def _linear_logp_needs_bf16_fast_path_cast() -> bool: + return _env_flag("RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16") or _env_flag( + "RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL" + ) + + +def _maybe_cast_hidden_for_bf16_fast_path(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + global _LINEAR_LOGP_SAVE_PROBS_CAST_LOGGED + if not _linear_logp_needs_bf16_fast_path_cast(): + return hidden_states + if not (hidden_states.is_cuda and weight.is_cuda and hidden_states.device == weight.device): + return hidden_states + if weight.dtype != torch.bfloat16 or hidden_states.dtype == torch.bfloat16: + return hidden_states + if not hidden_states.is_floating_point(): + return hidden_states + + if not _LINEAR_LOGP_SAVE_PROBS_CAST_LOGGED: + logger.info( + "Casting RL-Kernel linear_logp hidden states from %s to bf16 " + "to enable bf16 fast path; hidden_requires_grad=%s.", + hidden_states.dtype, + hidden_states.requires_grad, + ) + _LINEAR_LOGP_SAVE_PROBS_CAST_LOGGED = True + return hidden_states.to(dtype=torch.bfloat16) + + +def _register_linear_logp_backward_event_timer( + *, + start_event: torch.cuda.Event, + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> None: + watched_tensor = None + # In full-gradient runs, shared/tied output weights can receive other + # gradient contributions later in the model backward. Watch the op input + # first so this timer captures the linear_logp backward boundary. + for candidate in (hidden_states, weight, bias): + if isinstance(candidate, torch.Tensor) and candidate.requires_grad: + watched_tensor = candidate + break + if watched_tensor is None: + return + + handle_box = {} + + def _hook(grad): + end_event = torch.cuda.Event(enable_timing=True) + end_event.record() + _CUDA_EVENT_TIMER_QUEUE.enqueue( + start_event, + end_event, + _record_linear_logp_forward_backward_event_runtime, + ) + handle = handle_box.get("handle") + if handle is not None: + handle.remove() + return grad + + handle_box["handle"] = watched_tensor.register_hook(_hook) + + +def _warn_fallback(args: Namespace, op: str, reason: str) -> None: + _FALLBACK_COUNTS[op] = _FALLBACK_COUNTS.get(op, 0) + 1 + if getattr(args, "rl_kernel_strict", False): + raise RuntimeError(f"RL-Kernel {op} is enabled but unavailable: {reason}") + warning_key = f"{op}: {reason}" + if warning_key not in _WARNED_FALLBACK_REASONS: + logger.warning("Falling back to vime logprob path because RL-Kernel %s is unavailable: %s", op, reason) + _WARNED_FALLBACK_REASONS.add(warning_key) + + +def _get_logp_op(args: Namespace): + global _LOGP_OP, _LOGP_OP_LOAD_ERROR + if _LOGP_OP is not None: + return _LOGP_OP + if _LOGP_OP_LOAD_ERROR is not None: + _warn_fallback(args, "logp", str(_LOGP_OP_LOAD_ERROR)) + return None + + try: + from rl_engine.kernels.registry import kernel_registry + + _LOGP_OP = kernel_registry.get_op("logp") + logger.info("Using RL-Kernel logp op: %s", type(_LOGP_OP).__name__) + return _LOGP_OP + except Exception as exc: # pragma: no cover - exercised with missing optional package in integration envs + _LOGP_OP_LOAD_ERROR = exc + _warn_fallback(args, "logp", str(exc)) + return None + + +def _get_linear_logp_op(args: Namespace): + global _LINEAR_LOGP_OP, _LINEAR_LOGP_OP_LOAD_ERROR + if _LINEAR_LOGP_OP is not None: + return _LINEAR_LOGP_OP + if _LINEAR_LOGP_OP_LOAD_ERROR is not None: + _warn_fallback(args, "linear_logp", str(_LINEAR_LOGP_OP_LOAD_ERROR)) + return None + + try: + forced_backend = os.getenv("VIME_RL_KERNEL_LINEAR_LOGP_BACKEND", "").strip().lower() + if forced_backend in {"triton", "triton_linear_logp"}: + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + _LINEAR_LOGP_OP = TritonLinearLogpOp() + elif forced_backend in {"cuda", "sm90", "cuda_sm90", "fused_sm90"}: + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + _LINEAR_LOGP_OP = FusedLinearLogpSM90Op() + elif forced_backend in {"", "auto", "registry"}: + from rl_engine.kernels.registry import kernel_registry + + _LINEAR_LOGP_OP = kernel_registry.get_op("linear_logp") + else: + raise ValueError( + "unknown VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=" + f"{forced_backend!r}; expected triton, cuda, or registry" + ) + logger.info("Using RL-Kernel linear_logp op: %s", type(_LINEAR_LOGP_OP).__name__) + return _LINEAR_LOGP_OP + except Exception as exc: # pragma: no cover - exercised with missing optional package in integration envs + _LINEAR_LOGP_OP_LOAD_ERROR = exc + _warn_fallback(args, "linear_logp", str(exc)) + return None + + +def _unwrap_model_chunk(model): + while hasattr(model, "module"): + model = model.module + return model + + +def _is_pipeline_last_stage_for_model(model) -> bool: + module = _unwrap_model_chunk(model) + vp_stage = getattr(module, "vp_stage", None) + try: + vp_world_size = mpu.get_virtual_pipeline_model_parallel_world_size() + except Exception: + vp_world_size = None + + try: + if vp_world_size is not None and vp_stage is not None: + return bool(mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)) + return bool(mpu.is_pipeline_last_stage(ignore_virtual=True)) + except Exception: + return True + + +def _get_lm_head_weight(model, output_layer) -> torch.Tensor | None: + if getattr(model, "share_embeddings_and_output_weights", False): + shared_weight = getattr(model, "shared_embedding_or_output_weight", None) + if callable(shared_weight): + try: + weight = shared_weight() + if isinstance(weight, torch.Tensor): + return weight + except Exception: + logger.debug("Unable to read shared embedding/output weight for RL-Kernel linear_logp.", exc_info=True) + + weight = getattr(output_layer, "weight", None) + if isinstance(weight, torch.Tensor): + return weight + + shared_weight = getattr(model, "shared_embedding_or_output_weight", None) + if callable(shared_weight): + try: + weight = shared_weight() + if isinstance(weight, torch.Tensor): + return weight + except Exception: + logger.debug("Unable to read shared embedding/output weight for RL-Kernel linear_logp.", exc_info=True) + + return None + + +def get_linear_logp_context_from_model(args: Namespace, model) -> LinearLogpContext | None: + if not is_rl_kernel_op_enabled(args, "linear_logp"): + return None + + if not _is_pipeline_last_stage_for_model(model): + return None + + module = _unwrap_model_chunk(model) + output_layer = getattr(module, "output_layer", None) + if output_layer is None: + _warn_fallback(args, "linear_logp", "model output_layer is unavailable") + return None + + weight = _get_lm_head_weight(module, output_layer) + if weight is None: + _warn_fallback(args, "linear_logp", "LM-head weight is unavailable") + return None + + bias = getattr(output_layer, "bias", None) + if not isinstance(bias, torch.Tensor): + bias = None + + tp_world_size = int(mpu.get_tensor_model_parallel_world_size()) + tp_group = mpu.get_tensor_model_parallel_group() if tp_world_size > 1 else None + vocab_start_index = 0 + global_vocab_size = None + if tp_world_size > 1: + local_vocab_size = int(weight.size(0)) + vocab_start_index = int(mpu.get_tensor_model_parallel_rank()) * local_vocab_size + global_vocab_size = getattr(args, "padded_vocab_size", None) + if global_vocab_size is None: + global_vocab_size = local_vocab_size * tp_world_size + + return LinearLogpContext( + lm_head_weight=weight, + bias=bias, + tp_group=tp_group, + vocab_start_index=vocab_start_index, + global_vocab_size=None if global_vocab_size is None else int(global_vocab_size), + sequence_parallel=bool(getattr(output_layer, "sequence_parallel", getattr(args, "sequence_parallel", False))), + ) + + +def _linear_logp_runtime_blocker(args: Namespace, *, with_entropy: bool) -> str | None: + if with_entropy: + return "entropy is requested" + if getattr(args, "qkv_format", "thd") != "thd": + return "only qkv_format=thd is supported by RL-Kernel linear_logp" + if mpu.get_context_parallel_world_size() != 1 or getattr(args, "allgather_cp", False): + return "context parallel logprob redistribution is not supported by RL-Kernel linear_logp" + if getattr(args, "rollout_temperature", 1.0) <= 0: + return "rollout_temperature must be positive" + return None + + +def should_use_linear_logp_model_output(args: Namespace, *, with_entropy: bool) -> bool: + if not is_rl_kernel_op_enabled(args, "linear_logp"): + return False + reason = _linear_logp_runtime_blocker(args, with_entropy=with_entropy) + if reason is not None: + _warn_fallback(args, "linear_logp", reason) + return False + return True + + +def warn_linear_logp_fallback(args: Namespace, reason: str) -> None: + _warn_fallback(args, "linear_logp", reason) + + +@contextmanager +def return_hidden_states_for_linear_logp(args: Namespace, model, context: LinearLogpContext | None): + if context is None: + yield False + return + + module = _unwrap_model_chunk(model) + if not hasattr(module, "post_process"): + _warn_fallback(args, "linear_logp", "model post_process flag is unavailable") + yield False + return + + old_post_process = module.post_process + module.post_process = False + try: + yield True + finally: + module.post_process = old_post_process + + +def maybe_compute_logp( + logits: torch.Tensor, + tokens: torch.Tensor, + *, + args: Namespace, + with_entropy: bool, +) -> torch.Tensor | None: + """Return selected log-probs from RL-Kernel when this runtime is safe. + + The first integration deliberately limits itself to forward-only logprob + precompute paths: no autograd, no vocab tensor parallelism, no CP + redistribution, and no entropy. Unsupported cases fall back to vime's + Megatron-aware implementation. + """ + if not is_rl_kernel_op_enabled(args, "logp"): + return None + + if with_entropy: + _warn_fallback(args, "logp", "entropy is requested") + return None + + if logits.requires_grad or torch.is_grad_enabled(): + _warn_fallback(args, "logp", "autograd is enabled") + return None + + if mpu.get_tensor_model_parallel_world_size() != 1: + _warn_fallback(args, "logp", "tensor-parallel vocab shards are not supported by RL-Kernel logp") + return None + + if mpu.get_context_parallel_world_size() != 1 or getattr(args, "allgather_cp", False): + _warn_fallback(args, "logp", "context parallel logprob redistribution is not supported by RL-Kernel logp") + return None + + if logits.size(0) == 0: + return logits.new_zeros((0,), dtype=torch.float32) + + op = _get_logp_op(args) + if op is None: + return None + + try: + if hasattr(op, "apply_fp32"): + log_prob = op.apply_fp32(logits, tokens) + else: + log_prob = op(logits, tokens).float() + except Exception as exc: + _warn_fallback(args, "logp", str(exc)) + return None + + return log_prob.reshape(-1) + + +def maybe_compute_linear_logp( + hidden_states: torch.Tensor, + target_ids: torch.Tensor, + *, + context: LinearLogpContext | None, + args: Namespace, + with_entropy: bool, +) -> torch.Tensor | None: + if not is_rl_kernel_op_enabled(args, "linear_logp"): + return None + + reason = _linear_logp_runtime_blocker(args, with_entropy=with_entropy) + if reason is not None: + _warn_fallback(args, "linear_logp", reason) + return None + + if context is None: + _warn_fallback(args, "linear_logp", "hidden-state linear_logp context is unavailable") + return None + + if target_ids.numel() == 0: + return hidden_states.new_zeros((0,), dtype=torch.float32) + + op = _get_linear_logp_op(args) + if op is None: + return None + + weight = context.lm_head_weight + bias = context.bias + rollout_temperature = float(getattr(args, "rollout_temperature", 1.0)) + if rollout_temperature != 1.0: + weight = weight / rollout_temperature + if bias is not None: + bias = bias / rollout_temperature + if _should_detach_linear_logp_hidden(args): + hidden_states = hidden_states.detach() + hidden_states = _maybe_cast_hidden_for_bf16_fast_path(hidden_states, weight) + + start_s = time.perf_counter() + event_timer = _cuda_event_timer_enabled(hidden_states) + forward_start_event = None + if event_timer: + forward_start_event = torch.cuda.Event(enable_timing=True) + forward_start_event.record() + memory_probe = _env_flag("VIME_LINEAR_LOGP_MEMORY_PROBE") and hidden_states.is_cuda + if memory_probe: + probe_device = hidden_states.device + torch.cuda.synchronize(probe_device) + probe_before_alloc = torch.cuda.memory_allocated(probe_device) + probe_before_reserved = torch.cuda.memory_reserved(probe_device) + update_peak_memory_tracker("actor_train", device=probe_device) + torch.cuda.reset_peak_memory_stats(probe_device) + try: + log_prob = op( + hidden_states, + weight, + target_ids.long(), + bias, + tp_group=context.tp_group, + vocab_start_index=context.vocab_start_index, + global_vocab_size=context.global_vocab_size, + ) + except Exception as exc: + _warn_fallback(args, "linear_logp", str(exc)) + return None + + _record_linear_logp_runtime(target_ids.numel(), time.perf_counter() - start_s) + + if event_timer and forward_start_event is not None: + forward_end_event = torch.cuda.Event(enable_timing=True) + forward_end_event.record() + _CUDA_EVENT_TIMER_QUEUE.enqueue( + forward_start_event, + forward_end_event, + _record_linear_logp_forward_event_runtime, + ) + if log_prob.requires_grad: + _register_linear_logp_backward_event_timer( + start_event=forward_start_event, + hidden_states=hidden_states, + weight=weight, + bias=bias, + ) + + if memory_probe: + torch.cuda.synchronize(probe_device) + probe_after_alloc = torch.cuda.memory_allocated(probe_device) + probe_after_reserved = torch.cuda.memory_reserved(probe_device) + probe_peak_alloc = torch.cuda.max_memory_allocated(probe_device) + probe_peak_reserved = torch.cuda.max_memory_reserved(probe_device) + update_peak_memory_tracker( + "actor_train", + peak_alloc=probe_peak_alloc, + peak_reserved=probe_peak_reserved, + device=probe_device, + ) + logger.info( + "RL-Kernel linear_logp memory_probe: op=%s hidden_shape=%s weight_shape=%s " + "tokens=%d alloc_before_mb=%.2f peak_alloc_mb=%.2f peak_delta_mb=%.2f " + "alloc_after_mb=%.2f reserved_before_mb=%.2f reserved_after_mb=%.2f", + type(op).__name__, + tuple(hidden_states.shape), + tuple(weight.shape), + int(target_ids.numel()), + probe_before_alloc / (1024**2), + probe_peak_alloc / (1024**2), + (probe_peak_alloc - probe_before_alloc) / (1024**2), + probe_after_alloc / (1024**2), + probe_before_reserved / (1024**2), + probe_after_reserved / (1024**2), + ) + + return log_prob.float().reshape(-1) diff --git a/vime/backends/megatron_utils/update_weight/common.py b/vime/backends/megatron_utils/update_weight/common.py index 89555aca..8f056990 100644 --- a/vime/backends/megatron_utils/update_weight/common.py +++ b/vime/backends/megatron_utils/update_weight/common.py @@ -1,4 +1,5 @@ import inspect +import os import re from argparse import Namespace from collections.abc import Iterator, Sequence @@ -126,6 +127,9 @@ def named_params_and_buffers( else: ans = _named_params_and_buffers_vanilla(model) + if os.environ.get("VIME_SYNC_TRAINABLE_WEIGHTS_ONLY", "0") == "1": + ans = ((name, tensor) for name, tensor in ans if getattr(tensor, "requires_grad", False)) + if translate_gpu_to_cpu: ans = ((name, _maybe_get_cpu_backup(tensor)) for name, tensor in ans) @@ -141,6 +145,43 @@ def _maybe_get_cpu_backup(x: torch.Tensor): return x +def _set_tensor_parallel_attrs(tensor: torch.Tensor, source: torch.Tensor, *, partition_dim: int) -> torch.Tensor: + tensor.tensor_model_parallel = getattr(source, "tensor_model_parallel", False) + tensor.partition_dim = partition_dim + tensor.partition_stride = 1 + tensor.parallel_mode = getattr(source, "parallel_mode", None) + return tensor + + +def _iter_grouped_mlp_expert_weights( + args: Namespace, + layer_idx: int, + rest: str, + param: torch.Tensor, + *, + ep_size: int, + expert_offset: int, +) -> Iterator[tuple[str, torch.Tensor]]: + num_local_experts = args.num_experts // ep_size + prefix = f"module.module.decoder.layers.{layer_idx}.mlp.experts" + + if rest == "mlp.experts.weight1": + expert_tensors = param.view(num_local_experts, args.hidden_size, -1).transpose(-1, -2) + partition_dim = 0 + target = "linear_fc1" + elif rest == "mlp.experts.weight2": + expert_tensors = param.view(num_local_experts, -1, args.hidden_size).transpose(-1, -2) + partition_dim = 1 + target = "linear_fc2" + else: + return + + for local_expert_idx, expert_tensor in enumerate(expert_tensors): + expert_idx = expert_offset + local_expert_idx + name = f"{prefix}.{target}.weight{expert_idx}" + yield name, _set_tensor_parallel_attrs(expert_tensor, param, partition_dim=partition_dim) + + def _named_params_and_buffers_vanilla(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]: for vp_stage, model_module in enumerate(model): @@ -208,6 +249,17 @@ def _named_params_and_buffers_global( layer_idx, rest = match.groups() layer_idx = int(layer_idx) + layer_offset + if rest in {"mlp.experts.weight1", "mlp.experts.weight2"}: + yield from _iter_grouped_mlp_expert_weights( + args, + layer_idx, + rest, + param, + ep_size=ep_size, + expert_offset=expert_offset, + ) + continue + # this is hardcoded for te grouped matmul expert_pattern = r"mlp\.experts\.(.+)\.(weight|bias)(\d+)" match = re.match(expert_pattern, rest) diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index d345adde..65819396 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -1,4 +1,5 @@ import dataclasses +import time from argparse import Namespace from collections.abc import Sequence @@ -22,13 +23,41 @@ def __init__(self, *args, **kwargs): def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): rank = dist.get_rank() + total_buckets = len(self.megatron_local_param_info_buckets) - for megatron_local_param_infos in tqdm( - self.megatron_local_param_info_buckets, disable=rank != 0, desc=progress_desc + for bucket_idx, megatron_local_param_infos in enumerate( + tqdm(self.megatron_local_param_info_buckets, disable=rank != 0, desc=progress_desc), start=1 ): + if rank == 0: + bucket_bytes = _bucket_size_bytes(megatron_local_param_infos) + print( + f"[weight-sync] bucket {bucket_idx}/{total_buckets}: " + f"{len(megatron_local_param_infos)} params, {bucket_bytes / 1024**2:.1f} MiB begin", + flush=True, + ) + t0 = time.perf_counter() megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights) + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}/{total_buckets}: " + f"megatron all-gather done in {time.perf_counter() - t0:.2f}s", + flush=True, + ) + t1 = time.perf_counter() hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos) + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}/{total_buckets}: " + f"HF convert done in {time.perf_counter() - t1:.2f}s " + f"({len(hf_named_tensors)} tensors)", + flush=True, + ) yield hf_named_tensors + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}/{total_buckets}: finished", + flush=True, + ) del megatron_full_params def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tensor], param_infos: list[ParamInfo]): @@ -53,7 +82,9 @@ def _get_megatron_full_params( if dist.get_rank() == info.src_rank: params.append( torch.nn.Parameter( - megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True), + megatron_local_weights[info.name] + .to(device=torch.cuda.current_device(), non_blocking=True) + .contiguous(), requires_grad=False, ) ) @@ -103,6 +134,17 @@ def _get_megatron_full_params( return gathered_params +def _bucket_size_bytes(param_infos: Sequence[ParamInfo]) -> int: + total = 0 + for info in param_infos: + if ".experts." in info.name: + tp_size = mpu.get_expert_tensor_parallel_world_size() + else: + tp_size = mpu.get_tensor_model_parallel_world_size() + total += info.size * tp_size + return total + + def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torch.nn.Module]) -> list[list[ParamInfo]]: """ Partition params into buckets ≤ update_weight_buffer_size (with TP replication). diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 2e8343f9..4ad117ce 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -14,6 +14,7 @@ from __future__ import annotations import os +import time from argparse import Namespace from collections.abc import Callable, Iterable, Mapping, Sequence from typing import Any @@ -289,9 +290,29 @@ def update_weights(self) -> None: megatron_local_weights = self.weights_getter() - for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): + for bucket_idx, hf_named_tensors in enumerate( + self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights), start=1 + ): + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}: sending {len(hf_named_tensors)} HF tensors to rollout", + flush=True, + ) + t0 = time.perf_counter() refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}: IPC payload submitted in {time.perf_counter() - t0:.2f}s " + f"({len(refs)} refs)", + flush=True, + ) + t1 = time.perf_counter() ray.get(refs) + if rank == 0: + print( + f"[weight-sync] bucket {bucket_idx}: rollout update returned in {time.perf_counter() - t1:.2f}s", + flush=True, + ) # Free GPU tensors so the caching allocator can reuse the blocks, # then release CUDA IPC cache entries whose consumers (vLLM engines) # have already closed their IPC handles. diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 4d0b1b3f..351b13a3 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -61,6 +61,7 @@ def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]: env.setdefault("NCCL_CUMEM_ENABLE", "0") env["CUDA_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"] env.setdefault("VLLM_SERVER_DEV_MODE", "1") + env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") if getattr(args, "vllm_enable_deterministic_inference", False): env["VLLM_BATCH_INVARIANT"] = "1" if getattr(args, "colocate", False): @@ -83,6 +84,13 @@ def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]: def _run_vllm_server(kwargs: dict, env: dict) -> None: os.environ.update(env) + if os.environ.get("VIME_VLLM_FAULTHANDLER", "0") == "1": + import faulthandler + import signal + import sys + + faulthandler.enable(file=sys.stderr, all_threads=True) + faulthandler.register(signal.SIGUSR1, file=sys.stderr, all_threads=True) from vllm.entrypoints.cli.serve import ServeSubcommand from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args @@ -445,12 +453,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/actor_group.py b/vime/ray/actor_group.py index db3500dc..0e601671 100644 --- a/vime/ray/actor_group.py +++ b/vime/ray/actor_group.py @@ -63,6 +63,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): import torch_memory_saver for path in [ + "torch_memory_saver_hook_mode_preload_cu13.abi3.so", "torch_memory_saver_hook_mode_preload_cu12.abi3.so", "torch_memory_saver_hook_mode_preload.abi3.so", ]: 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: diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 147f4d2b..a186a756 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -12,6 +12,7 @@ from vime.backends.vllm_utils.arguments import vllm_parse_args from vime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list from vime.utils.logging_utils import configure_logger +from vime.utils.rl_kernel import normalize_rl_kernel_args logger = logging.getLogger(__name__) @@ -157,6 +158,24 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--enable-rl-kernel", + action="store_true", + default=False, + help="Enable optional RL-Kernel acceleration for supported vime training/forward-only paths.", + ) + parser.add_argument( + "--rl-kernel-ops", + type=str, + default="linear_logp", + help="Comma-separated RL-Kernel ops to enable. Current production integration supports: linear_logp.", + ) + parser.add_argument( + "--rl-kernel-strict", + action="store_true", + default=False, + help="Raise instead of falling back when an enabled RL-Kernel op is unavailable or unsupported.", + ) parser.add_argument( "--only-train-params-name-list", type=str, @@ -1606,6 +1625,8 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def vime_validate_args(args): + normalize_rl_kernel_args(args) + args.eval_datasets = _resolve_eval_datasets(args) if args.kl_coef != 0 or args.use_kl_loss: @@ -1849,6 +1870,17 @@ def vime_validate_args(args): if hasattr(args, k): logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") setattr(args, k, v) + # vllm launch_server_process distinguishes "user-supplied value" from + # "argparse default" via ``args._vllm_user_provided``. YAML overrides + # bypass argparse, so we register them explicitly here — without this, + # YAML values that happen to equal the vllm-side default (e.g. + # ``vllm_gpu_memory_utilization: 0.92``) would be treated as "default" + # and silently replaced by vime's preferred value. + if isinstance(k, str) and k.startswith("vllm_"): + if not hasattr(args, "_vllm_user_provided"): + args._vllm_user_provided = set() + args._vllm_user_provided.add(k) + normalize_rl_kernel_args(args) if args.eval_max_context_len is None: logger.info( diff --git a/vime/utils/memory_utils.py b/vime/utils/memory_utils.py index d4b2d893..6cf887db 100644 --- a/vime/utils/memory_utils.py +++ b/vime/utils/memory_utils.py @@ -6,6 +6,7 @@ import torch.distributed as dist logger = logging.getLogger(__name__) +_PEAK_MEMORY_CONTEXTS: dict[str, dict[str, int]] = {} def clear_memory(clear_host_memory: bool = False): @@ -48,3 +49,62 @@ def print_memory(msg, clear_before_print: bool = False): f"[Rank {dist.get_rank()}] Memory-Usage {msg}{' (cleared before print)' if clear_before_print else ''}: {memory_info}" ) return memory_info + + +def reset_peak_memory_tracker(name: str, device=None): + device = torch.cuda.current_device() if device is None else device + alloc_before = int(torch.cuda.memory_allocated(device)) + reserved_before = int(torch.cuda.memory_reserved(device)) + torch.cuda.reset_peak_memory_stats(device) + _PEAK_MEMORY_CONTEXTS[name] = { + "alloc_before": alloc_before, + "reserved_before": reserved_before, + "peak_alloc": alloc_before, + "peak_reserved": reserved_before, + } + + +def update_peak_memory_tracker(name: str, *, peak_alloc=None, peak_reserved=None, device=None): + device = torch.cuda.current_device() if device is None else device + if peak_alloc is None: + peak_alloc = torch.cuda.max_memory_allocated(device) + if peak_reserved is None: + peak_reserved = torch.cuda.max_memory_reserved(device) + state = _PEAK_MEMORY_CONTEXTS.setdefault( + name, + { + "alloc_before": int(torch.cuda.memory_allocated(device)), + "reserved_before": int(torch.cuda.memory_reserved(device)), + "peak_alloc": int(torch.cuda.memory_allocated(device)), + "peak_reserved": int(torch.cuda.memory_reserved(device)), + }, + ) + state["peak_alloc"] = max(state["peak_alloc"], int(peak_alloc)) + state["peak_reserved"] = max(state["peak_reserved"], int(peak_reserved)) + + +def get_peak_memory_tracker(name: str, device=None): + device = torch.cuda.current_device() if device is None else device + state = _PEAK_MEMORY_CONTEXTS.get(name) + if state is None: + alloc_before = int(torch.cuda.memory_allocated(device)) + reserved_before = int(torch.cuda.memory_reserved(device)) + peak_alloc = max(alloc_before, int(torch.cuda.max_memory_allocated(device))) + peak_reserved = max(reserved_before, int(torch.cuda.max_memory_reserved(device))) + else: + alloc_before = state["alloc_before"] + reserved_before = state["reserved_before"] + peak_alloc = max(state["peak_alloc"], int(torch.cuda.max_memory_allocated(device))) + peak_reserved = max(state["peak_reserved"], int(torch.cuda.max_memory_reserved(device))) + alloc_after = int(torch.cuda.memory_allocated(device)) + reserved_after = int(torch.cuda.memory_reserved(device)) + peak_alloc = max(peak_alloc, alloc_after) + peak_reserved = max(peak_reserved, reserved_after) + return { + "alloc_before": alloc_before, + "reserved_before": reserved_before, + "alloc_after": alloc_after, + "reserved_after": reserved_after, + "peak_alloc": peak_alloc, + "peak_reserved": peak_reserved, + } diff --git a/vime/utils/rl_kernel.py b/vime/utils/rl_kernel.py new file mode 100644 index 00000000..e0e4a0b7 --- /dev/null +++ b/vime/utils/rl_kernel.py @@ -0,0 +1,79 @@ +import os +from argparse import Namespace +from collections.abc import Iterable + + +RL_KERNEL_SUPPORTED_OPS = ("linear_logp",) +RL_KERNEL_INTEGRATED_OPS = ("linear_logp",) +_TRUE_VALUES = {"1", "true", "yes", "on"} +_FALSE_VALUES = {"0", "false", "no", "off"} + + +def parse_rl_kernel_ops(value: str | Iterable[str] | None) -> tuple[str, ...]: + """Parse a comma/space separated RL-Kernel op list.""" + if value is None: + return ("linear_logp",) + + if isinstance(value, str): + raw_items = value.replace(",", " ").split() + else: + raw_items = [] + for item in value: + raw_items.extend(str(item).replace(",", " ").split()) + + ops: list[str] = [] + for item in raw_items: + op = item.strip().lower() + if not op: + continue + if op not in RL_KERNEL_SUPPORTED_OPS: + supported = ", ".join(RL_KERNEL_SUPPORTED_OPS) + raise ValueError(f"Unsupported RL-Kernel op '{op}'. Supported ops: {supported}.") + if op not in ops: + ops.append(op) + + return tuple(ops) if ops else ("linear_logp",) + + +def _env_bool(name: str) -> bool | None: + value = os.getenv(name) + if value is None: + return None + normalized = value.strip().lower() + if normalized in _TRUE_VALUES: + return True + if normalized in _FALSE_VALUES: + return False + raise ValueError(f"{name} must be one of: 1/0, true/false, yes/no, on/off.") + + +def normalize_rl_kernel_args(args: Namespace) -> Namespace: + """Apply environment overrides and validate RL-Kernel integration args.""" + env_enabled = _env_bool("VIME_RL_KERNEL") + if env_enabled is not None: + args.enable_rl_kernel = env_enabled + + env_strict = _env_bool("VIME_RL_KERNEL_STRICT") + if env_strict is not None: + args.rl_kernel_strict = env_strict + + env_ops = os.getenv("VIME_RL_KERNEL_OPS") + if env_ops is not None: + args.rl_kernel_ops = env_ops + + args.rl_kernel_ops = parse_rl_kernel_ops(getattr(args, "rl_kernel_ops", None)) + + if getattr(args, "enable_rl_kernel", False): + unsupported = [op for op in args.rl_kernel_ops if op not in RL_KERNEL_INTEGRATED_OPS] + if unsupported: + integrated = ", ".join(RL_KERNEL_INTEGRATED_OPS) + raise ValueError( + "This vime RL-Kernel integration currently supports only " + f"{integrated}. Requested future ops: {', '.join(unsupported)}." + ) + + return args + + +def is_rl_kernel_op_enabled(args: Namespace, op: str) -> bool: + return getattr(args, "enable_rl_kernel", False) and op in getattr(args, "rl_kernel_ops", ()) diff --git a/vime/utils/timer.py b/vime/utils/timer.py index ec1bdf76..766a11f1 100644 --- a/vime/utils/timer.py +++ b/vime/utils/timer.py @@ -16,6 +16,7 @@ class Timer(metaclass=SingletonMeta): def __init__(self): self.timers = {} self.start_time = {} + self.metrics = {} def start(self, name): assert name not in self.start_time, f"Timer {name} already started." @@ -34,6 +35,7 @@ def end(self, name): def reset(self, name=None): if name is None: self.timers = {} + self.metrics = {} elif name in self.timers: del self.timers[name] @@ -43,6 +45,12 @@ def add(self, name, elapsed_time): def log_dict(self): return self.timers + def metric_dict(self): + return self.metrics + + def update_metrics(self, metrics): + self.metrics.update(metrics) + @contextmanager def context(self, name): self.start(name) diff --git a/vime/utils/train_metric_utils.py b/vime/utils/train_metric_utils.py index 0782cb2b..5c86b8b3 100644 --- a/vime/utils/train_metric_utils.py +++ b/vime/utils/train_metric_utils.py @@ -15,12 +15,14 @@ def log_perf_data_raw( ) -> None: timer_instance = Timer() log_dict_raw = deepcopy(timer_instance.log_dict()) + metric_dict_raw = deepcopy(timer_instance.metric_dict()) timer_instance.reset() if not is_primary_rank: return log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()} + log_dict.update({f"perf/{key}": val for key, val in metric_dict_raw.items()}) if ("perf/actor_train_time" in log_dict) and (compute_total_fwd_flops is not None): total_fwd_flops = compute_total_fwd_flops(seq_lens=timer_instance.seq_lens) diff --git a/vime_plugins/mbridge/__init__.py b/vime_plugins/mbridge/__init__.py index 9263cbe9..05583b0a 100644 --- a/vime_plugins/mbridge/__init__.py +++ b/vime_plugins/mbridge/__init__.py @@ -7,6 +7,75 @@ from .minimax_m2 import MiniMaxM2Bridge from .qwen3_5 import Qwen3_5Bridge from .qwen3_next import Qwen3NextBridge +from mbridge.models.qwen2moe import Qwen2MoEBridge +from mbridge.models.qwen3moe import Qwen3MoEBridge +import torch + + +_QWEN_MOE_LOCAL_NORM_MAPPING = { + "input_layernorm.weight": ["model.layers.{layer_number}.input_layernorm.weight"], +} + +for _bridge_cls in (Qwen2MoEBridge, Qwen3MoEBridge): + _other_mapping = dict(getattr(_bridge_cls, "_OTHER_MAPPING", {})) + _other_mapping.update(_QWEN_MOE_LOCAL_NORM_MAPPING) + _bridge_cls._OTHER_MAPPING = _other_mapping + +_ORIGINAL_QWEN_MOE_WEIGHT_NAME_MAPPING_MLP = Qwen2MoEBridge._weight_name_mapping_mlp +_ORIGINAL_QWEN_MOE_WEIGHT_TO_MCORE_FORMAT = Qwen2MoEBridge._weight_to_mcore_format + + +def _qwen_moe_weight_name_mapping_mlp(self, name: str) -> list[str]: + layer_number = name.split(".")[2] + num_experts = self.hf_config.num_experts + + if "mlp.experts.weight1" in name: + hf_names = [] + for expert_id in range(num_experts): + hf_names.extend( + [ + f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", + f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", + ] + ) + return hf_names + + if "mlp.experts.weight2" in name: + return [ + f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight" + for expert_id in range(num_experts) + ] + + return _ORIGINAL_QWEN_MOE_WEIGHT_NAME_MAPPING_MLP(self, name) + + +def _qwen_moe_weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list): + if "mlp.experts.weight1" in mcore_weights_name: + assert len(hf_weights) == self.hf_config.num_experts * 2 + gates = hf_weights[0::2] + ups = hf_weights[1::2] + return ( + torch.cat((torch.stack(gates), torch.stack(ups)), dim=-2) + .transpose(-1, -2) + .reshape(self.hf_config.hidden_size, -1) + .contiguous() + ) + + if "mlp.experts.weight2" in mcore_weights_name: + assert len(hf_weights) == self.hf_config.num_experts + return ( + torch.stack(hf_weights) + .transpose(-1, -2) + .reshape(-1, self.hf_config.hidden_size) + .contiguous() + ) + + return _ORIGINAL_QWEN_MOE_WEIGHT_TO_MCORE_FORMAT(self, mcore_weights_name, hf_weights) + + +for _bridge_cls in (Qwen2MoEBridge, Qwen3MoEBridge): + _bridge_cls._weight_name_mapping_mlp = _qwen_moe_weight_name_mapping_mlp + _bridge_cls._weight_to_mcore_format = _qwen_moe_weight_to_mcore_format __all__ = [ "GLM4Bridge", @@ -18,4 +87,6 @@ "Qwen3_5Bridge", "MimoBridge", "DeepseekV32Bridge", + "Qwen2MoEBridge", + "Qwen3MoEBridge", ]