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..03a0f5ce 100644 --- a/scripts/run-qwen3-30B-A3B.sh +++ b/scripts/run-qwen3-30B-A3B.sh @@ -1,20 +1,53 @@ #!/bin/bash +WORKSPACE_ROOT=${WORKSPACE_ROOT:-/workspace} +VIME_PYTHON_ENV=${VIME_PYTHON_ENV:-${WORKSPACE_ROOT}/vime-rlk-env} +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 +55,218 @@ 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:-} +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_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 "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_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 +275,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 +313,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 +328,111 @@ 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 [[ "${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_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}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 # 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}\", + \"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_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_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_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}\", + \"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 +440,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 +448,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/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..f2d57628 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 @@ -146,7 +147,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/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..553c8b4b 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) 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/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", ]