#!/usr/bin/env bash
set -Eeuo pipefail

cd /workspace/agillm-4

export TOKENIZERS_PARALLELISM=false
export TOKENIZER_ID="${TOKENIZER_ID:-deepseek-ai/DeepSeek-V4-Pro}"
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True
export AGILLM_ATTN_BACKEND=sdpa

if [ -f /root/.cache/huggingface/token ]; then
  export HF_TOKEN="$(tr -d '\r\n' < /root/.cache/huggingface/token)"
  export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
fi

PRESET="${AGILLM4_4090_PRESET:-agillm4_floor}"
BLOCK="${AGILLM4_4090_BLOCK:-1280}"
TOKEN_PARAM_RATIO="${AGILLM4_4090_TOKEN_PARAM_RATIO:-100}"
SAVE_DIR="${AGILLM4_4090_SAVE_DIR:-/workspace/agillm4_4090_ckpts}"

WARMSTART_ARGS=()
if [ -n "${AGILLM4_4090_WARMSTART_FROM:-}" ]; then
  WARMSTART_ARGS+=(--warmstart_from "$AGILLM4_4090_WARMSTART_FROM")
fi

mkdir -p "$SAVE_DIR"

echo "START_AGILLM4_4090_LONG_BLOCK $(date -u +%Y-%m-%dT%H:%M:%SZ) host=$(hostname)"
echo "This is production AGILLM-4 training on 24GB, not a local toy test."
echo "preset=$PRESET block=$BLOCK token_param_ratio=$TOKEN_PARAM_RATIO sat_every=1 nat_every=${AGILLM4_4090_NAT_EVERY:-4} warmstart=${AGILLM4_4090_WARMSTART_FROM:-none}"

exec python -u /workspace/agillm-4/nB300_agillm4.py train \
  --preset "$PRESET" \
  "${WARMSTART_ARGS[@]}" \
  --batch_size 1 \
  --block "$BLOCK" \
  --amp \
  --attn_backend sdpa \
  --grad_checkpoint \
  --optimizer "${AGILLM4_4090_OPTIMIZER:-paged_adamw8bit}" \
  --sat_every 1 \
  --nat_every "${AGILLM4_4090_NAT_EVERY:-4}" \
  --nat_loss_weight "${AGILLM4_4090_NAT_LOSS_WEIGHT:-1.0}" \
  --nat_expand "${AGILLM4_4090_NAT_EXPAND:-2}" \
  --nat_max_tokens "${AGILLM4_4090_NAT_MAX_TOKENS:-768}" \
  --token_param_ratio "$TOKEN_PARAM_RATIO" \
  --save_dir "$SAVE_DIR" \
  --save_every_sec "${AGILLM4_4090_SAVE_EVERY_SEC:-86400}" \
  --delta_every_steps "${AGILLM4_4090_DELTA_EVERY_STEPS:-25000}" \
  --delta_max_keep "${AGILLM4_4090_DELTA_MAX_KEEP:-1}" \
  --max_ckpts "${AGILLM4_4090_MAX_CKPTS:-1}"