Skip to content

Add RL-Kernel 8GPU runbook, metrics, and config support #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
inaniloquentee wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Add RL-Kernel 8GPU runbook, metrics, and config support #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ac79a46
Add RL-Kernel linear_logp integration with TP2 benchmark config
inaniloquentee Jun 27, 2026
1aef1fe
Document 2xH100 RL-Kernel dev validation
inaniloquentee Jul 1, 2026
47980cf
Fix 2xH100 RL-Kernel benchmark defaults
inaniloquentee Jul 1, 2026
7447813
Add RL-Kernel runtime counters for TP2 benchmark
inaniloquentee Jul 1, 2026
8f1ccf6
Add RL-Kernel TP2 benchmark instrumentation
inaniloquentee Jul 6, 2026
9d4f9df
Add 8-GPU RL-Kernel linear_logp benchmark runner
inaniloquentee Jul 6, 2026
82904ec
Add RL-Kernel 8GPU runbook metrics and config support
codex Jul 7, 2026
88ba1ff
Document RL framework workflow for long-run report
codex Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
349 changes: 349 additions & 0 deletions scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,349 @@
#!/usr/bin/env bash
set -euo pipefail

usage() {
cat >&2 <<'EOF'
usage: scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh CONFIG MODE

CONFIG: T01 | T02 | T03 | T04 | T05 | T06 | T07 | T08
MODE: baseline | cuda

Important environment overrides:
TRAIN_SCOPE=full|output_layer default: full
TRACE_MODE=none|train|rollout|all default: all
TRACE_ROLLOUTS=3 nsys capture rollout ids
RUN_ROOT=/workspace log/trace root
NSYS=/path/to/nsys nsys binary
EOF
}

CONFIG_ID="${1:-}"
MODE="${2:-}"
if [[ -z "${CONFIG_ID}" || -z "${MODE}" ]]; then
usage
exit 2
fi
if [[ "${MODE}" != "baseline" && "${MODE}" != "cuda" ]]; then
usage
exit 2
fi

WORKSPACE_ROOT="${WORKSPACE_ROOT:-/workspace}"
RUN_ROOT="${RUN_ROOT:-${WORKSPACE_ROOT}}"
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
VIME_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)"
VENV="${VIME_PYTHON_ENV:-${WORKSPACE_ROOT}/vime-rlk-env}"
NSYS="${NSYS:-/opt/nvidia/nsight-compute/2024.1.1/host/target-linux-x64/nsys}"

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

NSYS default points at the Nsight Compute tree, not Nsight Systems.

The script runs nsys profile, but the nsys binary ships with Nsight Systems (e.g. /opt/nvidia/nsight-systems/<ver>/bin/nsys or /usr/local/cuda/bin/nsys); the nsight-compute directory ships ncu. With TRACE_MODE != none, this default fails the -x check at Line 293 unless NSYS is overridden. Point the default at a Nsight Systems path.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh` at line 36, The
NSYS default in the benchmark script points to the Nsight Compute installation
instead of the Nsight Systems binary used by nsys profile. Update the NSYS
fallback in the script’s environment setup to a valid Nsight Systems path (for
example, a bin/nsys location), so the existing checks in the profiling flow
succeed without requiring an override.


export HF_HOME="${HF_HOME:-${WORKSPACE_ROOT}/.cache/huggingface}"
export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}"
export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}"
export PIP_CACHE_DIR="${PIP_CACHE_DIR:-${WORKSPACE_ROOT}/.cache/pip}"
export TORCH_HOME="${TORCH_HOME:-${WORKSPACE_ROOT}/.cache/torch}"
export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${WORKSPACE_ROOT}/.cache}"
export TMPDIR="${TMPDIR:-${WORKSPACE_ROOT}/.cache/tmp}"
export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${WORKSPACE_ROOT}}"
export FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST:-9.0}"
mkdir -p "${HF_HOME}" "${HUGGINGFACE_HUB_CACHE}" "${TRANSFORMERS_CACHE}" \
"${PIP_CACHE_DIR}" "${TORCH_HOME}" "${XDG_CACHE_HOME}" "${TMPDIR}" \
"${FLASHINFER_WORKSPACE_BASE}/.cache/flashinfer"

case "${CONFIG_ID}" in
T01)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-2048}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-512}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.40}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-2048}"
;;
T02)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-3072}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-1024}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.42}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-3072}"
;;
T03)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-4096}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-1536}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
T04)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-4096}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-2048}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
T05)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-2}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-6144}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3072}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
T06)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3072}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
T07)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3968}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.46}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
T08)
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-4}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-2}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}"
ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-3584}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.45}"
VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-4096}"
;;
*)
usage
exit 2
;;
esac

TRAIN_SCOPE="${TRAIN_SCOPE:-full}"
if [[ "${TRAIN_SCOPE}" != "full" && "${TRAIN_SCOPE}" != "output_layer" ]]; then
echo "TRAIN_SCOPE must be full or output_layer, got ${TRAIN_SCOPE}" >&2
exit 2
fi

TRACE_MODE="${TRACE_MODE:-all}"
case "${TRACE_MODE}" in
none|train|rollout|all) ;;
*)
echo "TRACE_MODE must be none, train, rollout, or all; got ${TRACE_MODE}" >&2
exit 2
;;
esac

RUN_NAME="${RUN_NAME:-8gpu_${CONFIG_ID}_${TRAIN_SCOPE}_${MODE}_$(date +%Y%m%d_%H%M%S)}"
LOG_DIR="${LOG_DIR:-${RUN_ROOT}/logs/${RUN_NAME}}"
TRACE_DIR="${TRACE_DIR:-${RUN_ROOT}/nsys_traces/${RUN_NAME}}"
mkdir -p "${LOG_DIR}" "${TRACE_DIR}"

export NUM_GPUS="${NUM_GPUS:-8}"
export MEGATRON_TP="${MEGATRON_TP:-2}"
export MEGATRON_EP="${MEGATRON_EP:-8}"
export MEGATRON_CP="${MEGATRON_CP:-1}"
export ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-8}"
export NUM_ROLLOUT="${NUM_ROLLOUT:-12}"
export VIME_DISABLE_SAVE="${VIME_DISABLE_SAVE:-1}"
export VIME_SKIP_EVAL_BEFORE_TRAIN="${VIME_SKIP_EVAL_BEFORE_TRAIN:-1}"
export VIME_USE_KL_LOSS="${VIME_USE_KL_LOSS:-0}"
export VIME_SKIP_ZERO_ENTROPY_METRIC="${VIME_SKIP_ZERO_ENTROPY_METRIC:-1}"
export VIME_VLLM_ENFORCE_EAGER="${VIME_VLLM_ENFORCE_EAGER:-1}"
export VLLM_WORKER_MULTIPROC_METHOD="${VLLM_WORKER_MULTIPROC_METHOD:-spawn}"
export VIME_RL_KERNEL_OPS="${VIME_RL_KERNEL_OPS:-linear_logp}"
export VIME_RL_KERNEL_LINEAR_LOGP_BACKEND="${VIME_RL_KERNEL_LINEAR_LOGP_BACKEND:-cuda}"
export VIME_RL_KERNEL_CUDA_EVENT_TIMER="${VIME_RL_KERNEL_CUDA_EVENT_TIMER:-1}"
export VIME_LINEAR_LOGP_MEMORY_PROBE="${VIME_LINEAR_LOGP_MEMORY_PROBE:-1}"
export VIME_TRAIN_MEMORY_MARGIN_BYTES="${VIME_TRAIN_MEMORY_MARGIN_BYTES:-1073741824}"
export VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK="${VIME_SKIP_MOE_GROUPED_GEMM_CAPABILITY_CHECK:-1}"
export MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ="${MEGATRON_LOCAL_ATTENTION_SINGLE_PACKED_SEQ:-1}"
export VIME_USE_DISTRIBUTED_OPTIMIZER="${VIME_USE_DISTRIBUTED_OPTIMIZER:-1}"
export VIME_USE_PRECISION_AWARE_OPTIMIZER="${VIME_USE_PRECISION_AWARE_OPTIMIZER:-1}"
export VIME_OPTIMIZER_CPU_OFFLOAD="${VIME_OPTIMIZER_CPU_OFFLOAD:-1}"
export VIME_USE_FP32_GRAD_BUFFER="${VIME_USE_FP32_GRAD_BUFFER:-0}"
export VIME_GRAD_REDUCE_IN_BF16="${VIME_GRAD_REDUCE_IN_BF16:-1}"
export VIME_SEQUENCE_PARALLEL="${VIME_SEQUENCE_PARALLEL:-0}"
export MEGATRON_ALLOW_MOE_TP_WITHOUT_SP="${MEGATRON_ALLOW_MOE_TP_WITHOUT_SP:-1}"
export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-0}"
export NCCL_CUMEM_ENABLE="${NCCL_CUMEM_ENABLE:-0}"
export CUDA_MODULE_LOADING="${CUDA_MODULE_LOADING:-EAGER}"
export VIME_SKIP_PROCESS_CLEANUP="${VIME_SKIP_PROCESS_CLEANUP:-0}"

export ROLLOUT_BATCH_SIZE
export N_SAMPLES_PER_PROMPT
export GLOBAL_BATCH_SIZE
export MAX_TOKENS_PER_GPU
export ROLLOUT_MAX_RESPONSE_LEN
export VLLM_GPU_MEMORY_UTILIZATION
export VLLM_MAX_MODEL_LEN

if [[ "${MODE}" == "cuda" ]]; then
export VIME_RL_KERNEL=1
export VIME_BASELINE_LINEAR_LOGP_TIMER=0
export VIME_BASELINE_CUDA_EVENT_TIMER=0
if [[ "${TRAIN_SCOPE}" == "output_layer" ]]; then
export VIME_ONLY_TRAIN_PARAMS_NAME_LIST="${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-output_layer}"
export VIME_SYNC_TRAINABLE_WEIGHTS_ONLY="${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY:-1}"
export VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN="${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-1}"
export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16="${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-1}"
export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL="${RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL:-0}"
else
export VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN="${VIME_RL_KERNEL_LINEAR_LOGP_DETACH_HIDDEN:-0}"
export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16="${RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16:-0}"
export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL="${RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL:-1}"
fi
else
export VIME_RL_KERNEL=0
export VIME_BASELINE_LINEAR_LOGP_TIMER=1
export VIME_BASELINE_CUDA_EVENT_TIMER=1
export RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16=0
export RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL=0
if [[ "${TRAIN_SCOPE}" == "output_layer" ]]; then
export VIME_ONLY_TRAIN_PARAMS_NAME_LIST="${VIME_ONLY_TRAIN_PARAMS_NAME_LIST:-output_layer}"
export VIME_SYNC_TRAINABLE_WEIGHTS_ONLY="${VIME_SYNC_TRAINABLE_WEIGHTS_ONLY:-1}"
export VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN="${VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN:-1}"
else
export VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN="${VIME_BASELINE_OUTPUT_LAYER_DETACH_HIDDEN:-0}"
fi
fi

case "${TRACE_MODE}" in
none)
export VIME_NSYS_CAPTURE_ROLLOUTS=""
export VIME_NSYS_CAPTURE_ROLE="${VIME_NSYS_CAPTURE_ROLE:-actor}"
;;
train)
export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}"
export VIME_NSYS_CAPTURE_ROLE="actor"
;;
rollout)
export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}"
export VIME_NSYS_CAPTURE_ROLE="rollout"
;;
all)
export VIME_NSYS_CAPTURE_ROLLOUTS="${TRACE_ROLLOUTS:-${VIME_NSYS_CAPTURE_ROLLOUTS:-3}}"
export VIME_NSYS_CAPTURE_ROLE="all"
;;
esac

if [[ "${TRACE_MODE}" == "rollout" || "${TRACE_MODE}" == "all" ]]; then
if [[ -z "${VLLM_PROFILER_CONFIG:-}" ]]; then
export VLLM_PROFILER_CONFIG='{"profiler":"cuda"}'
else
export VLLM_PROFILER_CONFIG
fi
fi

TOTAL_SAMPLES_PER_ROLLOUT=$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT))
TRAIN_DP_SIZE=$((NUM_GPUS / MEGATRON_TP / MEGATRON_CP))
if (( GLOBAL_BATCH_SIZE > TOTAL_SAMPLES_PER_ROLLOUT )); then
echo "GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE} exceeds ROLLOUT_BATCH_SIZE*N_SAMPLES_PER_PROMPT=${TOTAL_SAMPLES_PER_ROLLOUT}" >&2
exit 2
fi

if (( TOTAL_SAMPLES_PER_ROLLOUT < TRAIN_DP_SIZE )); then
echo "RBS*NSP=${TOTAL_SAMPLES_PER_ROLLOUT} is smaller than training DP size ${TRAIN_DP_SIZE}; each DP rank needs at least one sample" >&2
exit 2
fi

export CONFIG_ID MODE TRAIN_SCOPE TRACE_MODE RUN_NAME LOG_DIR TRACE_DIR

echo "run_name=${RUN_NAME}" | tee "${LOG_DIR}/run_env.log"
env | sort | grep -E '^(CONFIG_ID|MODE|TRAIN_SCOPE|TRACE_MODE|NUM_GPUS|MEGATRON_|ROLLOUT_|N_SAMPLES|GLOBAL_BATCH|MAX_TOKENS|VLLM_|VIME_|RL_KERNEL_|NCCL_|CUDA_MODULE_LOADING|LOG_DIR|TRACE_DIR|RUN_NAME)=' \
| tee -a "${LOG_DIR}/run_env.log"

run_train() {
bash "${VIME_ROOT}/scripts/run-qwen3-30B-A3B.sh"
}

NSYS_RAY_PID=""
cleanup() {
set +e
if [[ -n "${NSYS_RAY_PID}" ]] && kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then
"${VENV}/bin/ray" stop --force >/dev/null 2>&1 || true
for _ in $(seq 1 120); do
if ! kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then
break
fi
sleep 1
done
if kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then
kill -TERM "${NSYS_RAY_PID}" >/dev/null 2>&1 || true
fi
wait "${NSYS_RAY_PID}" >/dev/null 2>&1 || true
fi
pkill -TERM -f "[v]llm serve" >/dev/null 2>&1 || true
pkill -TERM -f "[V]LLM::" >/dev/null 2>&1 || true
}
trap cleanup EXIT

cd "${VIME_ROOT}"
if [[ "${TRACE_MODE}" == "none" ]]; then
run_train 2>&1 | tee "${LOG_DIR}/ray_job_${MODE}.log"
else
if [[ ! -x "${NSYS}" ]]; then
echo "nsys not found or not executable: ${NSYS}" >&2
exit 1
fi

"${VENV}/bin/ray" stop --force >/dev/null 2>&1 || true
pkill -9 -f "[v]llm serve" >/dev/null 2>&1 || true
pkill -9 -f "[V]LLM::" >/dev/null 2>&1 || true
sleep 2

"${NSYS}" profile \
--trace=cuda,nvtx,osrt,cublas,cudnn \
--sample=none \
--cpuctxsw=none \
--capture-range=cudaProfilerApi \
--capture-range-end=repeat:64 \
--flush-on-cudaprofilerstop=true \
--wait=all \
--force-overwrite=true \
--export=sqlite \
--output="${TRACE_DIR}/${RUN_NAME}" \
"${VENV}/bin/ray" start \
--head \
--node-ip-address 127.0.0.1 \
--num-gpus "${NUM_GPUS}" \
--disable-usage-stats \
--dashboard-host=0.0.0.0 \
--dashboard-port=8265 \
--block \
>"${LOG_DIR}/nsys_ray_start.log" 2>&1 &
NSYS_RAY_PID=$!

for _ in $(seq 1 120); do
if curl -fsS "http://127.0.0.1:8265/api/version" >/dev/null 2>&1; then
break
fi
if ! kill -0 "${NSYS_RAY_PID}" >/dev/null 2>&1; then
echo "nsys/ray process exited early; see ${LOG_DIR}/nsys_ray_start.log" >&2
exit 1
fi
sleep 1
done
if ! curl -fsS "http://127.0.0.1:8265/api/version" >/dev/null 2>&1; then
echo "ray dashboard did not become ready; see ${LOG_DIR}/nsys_ray_start.log" >&2
exit 1
fi

VIME_SKIP_PROCESS_CLEANUP=1 VIME_SKIP_RAY_START=1 run_train 2>&1 | tee "${LOG_DIR}/ray_job_${MODE}.log"
fi

cleanup
trap - EXIT

echo "trace files:"
find "${TRACE_DIR}" -maxdepth 1 -type f -print | sort || true
echo "logs:"
find "${LOG_DIR}" -maxdepth 1 -type f -print | sort || true
7 changes: 5 additions & 2 deletions scripts/models/qwen3-30B-A3B.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
)

if [[ "${VIME_NO_MOE_PERMUTE_FUSION:-0}" != "1" ]]; then
MODEL_ARGS+=(--moe-permute-fusion)
fi
Loading