125 lines
5.5 KiB
Bash
125 lines
5.5 KiB
Bash
#!/bin/bash
|
||
# POLARIS-1.7B PLAIN-SGD rerun.
|
||
#
|
||
# Upstream (ChenxinAn-fdu/POLARIS) trains Qwen3-1.7B-Base with AdamW lr=1e-6,
|
||
# GRPO+ (use_kl_loss=False, entropy_coeff=0), strict on-policy
|
||
# (mini_batch=train_batch=128, 1 grad step per rollout batch).
|
||
# This script swaps ONLY the optimizer: AdamW → PLAIN SGD (momentum=0).
|
||
# Everything else kept faithful within a 4×B200 budget (response length capped at 8192).
|
||
#
|
||
# Scientific non-negotiable: momentum=0.0, nesterov=false, dampening=0.0,
|
||
# weight_decay=0.0. If the run diverges, lower the LR — never turn on momentum.
|
||
|
||
set -euo pipefail
|
||
|
||
export PYTHONUNBUFFERED=1
|
||
|
||
# Sources venv + vLLM-critical env: local TMPDIR for AF_UNIX sockets (NFS refuses
|
||
# AF_UNIX bind), LD_LIBRARY_PATH for system libstdc++ GLIBCXX_3.4.32, CCCL headers
|
||
# for flashinfer, VLLM_USE_V1=1, stale-PYTHONPATH cleanup, CC/CXX/TRITON_CC on
|
||
# /usr/bin/gcc (the /home/utils/gcc-11.2.0 + binutils-2.37 pair can't link
|
||
# against the system glibc's .relr.dyn section).
|
||
# shellcheck disable=SC1091
|
||
source /home/scratch.pjayasinha_gpu/rs/vllm_env.sh
|
||
|
||
# Force HF offline after vllm_env.sh to avoid 4 simultaneous vLLM EngineCores
|
||
# hitting the HF hub → 429 rate-limit.
|
||
export HF_HUB_OFFLINE=1
|
||
export TRANSFORMERS_OFFLINE=1
|
||
|
||
export CUDA_VISIBLE_DEVICES=0,1,2,3 # GPUs 4-7 are held by the SVD sweep
|
||
# Force Ray to respect CUDA_VISIBLE_DEVICES for num_gpus=0 actors (e.g. the
|
||
# vLLM http server coordinator). Without this, Ray's legacy default overrides
|
||
# our CUDA_VISIBLE_DEVICES for coordinator actors → vLLM's EngineCore picks up
|
||
# physical GPU 4 and starts rolling out on a GPU we promised to leave alone.
|
||
export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0
|
||
|
||
# NOTE: do NOT prepend /usr/bin to PATH — vllm_env.sh already sources the venv,
|
||
# and /usr/bin first would shadow the venv's python3 (no verl installed there).
|
||
# Triton uses CC=/usr/bin/gcc (set in vllm_env.sh); that pulls in /usr/bin/ld
|
||
# automatically via gcc's default linker path. No explicit LD override needed.
|
||
|
||
export WANDB_API_KEY=$(cat /home/scratch.pjayasinha_gpu/rs/wandb_api_key.txt)
|
||
|
||
# SMOKE=1 runs a tiny sanity check: single epoch, short response length, small val slice.
|
||
SMOKE="${SMOKE:-0}"
|
||
LR="${LR:-1e-1}" # lr=1 and lr=2e-1 both collapsed (max-length babble, rewards pinned at -1). lr=1e-1 was the stable-but-slow regime where val-core actually started rising from 0. Locking in here.
|
||
|
||
if [[ "$SMOKE" == "1" ]]; then
|
||
EXP_NAME="polaris_1p7b_sgd_lr${LR}_SMOKE"
|
||
TOTAL_EPOCHS=1
|
||
MAX_RESP=${MAX_RESP:-8192} # default now matches the real run; override to 2048 for fast debug
|
||
SAVE_FREQ=10
|
||
TEST_FREQ=999999
|
||
else
|
||
EXP_NAME="polaris_1p7b_sgd_lr${LR}_bs128_gpu4_ep5_n4"
|
||
TOTAL_EPOCHS=5
|
||
MAX_RESP=${MAX_RESP:-8192}
|
||
SAVE_FREQ=25
|
||
TEST_FREQ=50
|
||
fi
|
||
|
||
CKPT_DIR="/home/scratch.pjayasinha_gpu/rs/checkpoints/${EXP_NAME}"
|
||
LOG_DIR="/home/scratch.pjayasinha_gpu/rs/logs"
|
||
mkdir -p "$CKPT_DIR" "$LOG_DIR"
|
||
LOG_PATH="${LOG_DIR}/${EXP_NAME}.log"
|
||
|
||
echo "=============================================="
|
||
echo "EXP_NAME: $EXP_NAME"
|
||
echo "MODEL: Qwen/Qwen3-1.7B-Base"
|
||
echo "OPTIMIZER: PLAIN SGD (momentum=0, nesterov=false, dampening=0, wd=0)"
|
||
echo "LR: $LR"
|
||
echo "MAX_RESP: $MAX_RESP"
|
||
echo "TOTAL_EPOCHS: $TOTAL_EPOCHS"
|
||
echo "CKPT_DIR: $CKPT_DIR"
|
||
echo "LOG_PATH: $LOG_PATH"
|
||
echo "GPUs: $CUDA_VISIBLE_DEVICES"
|
||
echo "=============================================="
|
||
|
||
python3 -m verl.trainer.main_ppo \
|
||
algorithm.adv_estimator=grpo \
|
||
algorithm.use_kl_in_reward=False \
|
||
data.train_files=/home/scratch.pjayasinha_gpu/rs/data_polaris/train.parquet \
|
||
data.val_files="['/home/scratch.pjayasinha_gpu/rs/data_polaris/val.parquet']" \
|
||
data.train_batch_size=128 \
|
||
data.max_prompt_length=1024 \
|
||
data.max_response_length=${MAX_RESP} \
|
||
data.filter_overlong_prompts=True \
|
||
data.truncation='error' \
|
||
actor_rollout_ref.model.path=Qwen/Qwen3-1.7B-Base \
|
||
actor_rollout_ref.model.use_remove_padding=True \
|
||
actor_rollout_ref.model.enable_gradient_checkpointing=False \
|
||
actor_rollout_ref.model.use_fused_kernels=True \
|
||
actor_rollout_ref.actor.optim.optimizer=SGD \
|
||
actor_rollout_ref.actor.optim.optimizer_impl=torch.optim \
|
||
actor_rollout_ref.actor.optim.lr=${LR} \
|
||
actor_rollout_ref.actor.optim.weight_decay=0.0 \
|
||
actor_rollout_ref.actor.optim.override_optimizer_config='{momentum: 0.0, nesterov: false, dampening: 0.0}' \
|
||
actor_rollout_ref.actor.ppo_mini_batch_size=128 \
|
||
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
|
||
actor_rollout_ref.actor.use_kl_loss=False \
|
||
actor_rollout_ref.actor.kl_loss_coef=0.0 \
|
||
actor_rollout_ref.actor.entropy_coeff=0.0 \
|
||
actor_rollout_ref.actor.fsdp_config.param_offload=False \
|
||
actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
|
||
actor_rollout_ref.rollout.name=vllm \
|
||
actor_rollout_ref.rollout.n=4 \
|
||
actor_rollout_ref.rollout.temperature=1.0 \
|
||
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
|
||
actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
|
||
actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \
|
||
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
|
||
actor_rollout_ref.ref.fsdp_config.param_offload=False \
|
||
trainer.critic_warmup=0 \
|
||
trainer.logger='["console","wandb"]' \
|
||
trainer.project_name='sgd_rerun_polaris_1p7b' \
|
||
trainer.experiment_name="${EXP_NAME}" \
|
||
+trainer.wandb.entity='pavanjayasinha-university-of-illinois-urbana-champaign' \
|
||
trainer.n_gpus_per_node=4 \
|
||
trainer.nnodes=1 \
|
||
trainer.save_freq=${SAVE_FREQ} \
|
||
trainer.test_freq=${TEST_FREQ} \
|
||
trainer.total_epochs=${TOTAL_EPOCHS} \
|
||
trainer.default_local_dir="${CKPT_DIR}" \
|
||
2>&1 | tee "${LOG_PATH}"
|