#!/bin/bash
# GRPO on the GridSim v1 bundle through AgentEnv's generate function, as the run the slime page plots ran it on four L40S
# GPUs in slimerl/slime:latest (main 540ea91): the Qwen3-4B SFT warm start as a full fine-tune on Megatron, colocated
# with SGLang, lr 1e-6 constant, 32 rollouts a step (4 days x 8) at temperature 1.1, thinking off, at most 4 model turns
# a rollout, and evals on the 200 eval days at steps 0, 20 and 40. Run it in slime's container from /root/slime, with
# agentenv-framework[slime] installed, the rows python -m agent_env.training.slime wrote in /work/data,
# /work/qwen3-4b-gridsim-sft and its torch_dist conversion beside it; with a server that asks for a token, export the
# variable the rows' --token-var names first. Env: STEPS (40).
set -ex
pkill -9 sglang || true; ray stop --force || true; sleep 3
export PYTHONUNBUFFERED=1
source scripts/models/qwen3-4B.sh
ray start --head --node-ip-address 127.0.0.1 --num-gpus 4 --disable-usage-stats --dashboard-host=127.0.0.1 --dashboard-port=8265
ARGS=(
  --actor-num-nodes 1 --actor-num-gpus-per-node 4 --colocate
  "${MODEL_ARGS[@]}"
  --hf-checkpoint /work/qwen3-4b-gridsim-sft --ref-load /work/qwen3-4b-gridsim-sft_torch_dist
  --prompt-data /work/data/train.jsonl --input-key messages --metadata-key metadata --rollout-shuffle
  --apply-chat-template-kwargs '{"enable_thinking": false}'
  --num-rollout "${STEPS:-40}" --rollout-batch-size 4 --n-samples-per-prompt 8 --global-batch-size 32
  --rollout-temperature 1.1 --rollout-max-response-len 1024 --rollout-max-context-len 8192
  --eval-interval 20 --eval-prompt-data eval /work/data/eval.jsonl --n-samples-per-eval-prompt 1
  --eval-temperature 1.1 --eval-max-response-len 1024
  --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 --context-parallel-size 1
  --recompute-granularity full --recompute-method uniform --recompute-num-layers 1
  --use-dynamic-batch-size --max-tokens-per-gpu 4096
  --advantage-estimator grpo --kl-loss-coef 0.0 --entropy-coef 0.0 --eps-clip 0.2 --eps-clip-high 0.28
  --optimizer adam --lr 1e-6 --lr-decay-style constant --weight-decay 0.0 --adam-beta1 0.9 --adam-beta2 0.98
  --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer
  --rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.5 --sglang-tool-call-parser qwen25
  --custom-generate-function-path agent_env.training.slime.generate.generate
  --dynamic-sampling-filter-path agent_env.training.slime.generate.keep_scored_groups
  --attention-dropout 0.0 --hidden-dropout 0.0 --attention-softmax-in-fp32
  --attention-backend unfused
)
# Ray joins an entrypoint's words for a shell, so the command goes in a script the job runs.
printf '%q ' cd /root/slime > /work/train-cmd.sh; printf '\n' >> /work/train-cmd.sh
printf '%q ' exec python3 train.py "${ARGS[@]}" >> /work/train-cmd.sh; printf '\n' >> /work/train-cmd.sh
ENV='{"env_vars": {"PYTHONPATH": "/root/Megatron-LM/", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1"}}'
if ray job submit --address http://127.0.0.1:8265 --runtime-env-json "$ENV" -- bash /work/train-cmd.sh; then echo TRAIN-DONE; else echo TRAIN-FAILED; fi