#!/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}"