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

cd /workspace/model-trainer

RUN_NAME="convgpt_v2_b200_full_synth_20h_$(date +%Y%m%d_%H%M%S)"
LOG="logs/${RUN_NAME}.log"
mkdir -p logs

export PYTHONUNBUFFERED=1
export TOKENIZERS_PARALLELISM=false
export WANDB_MODE=${WANDB_MODE:-online}
export HF_HOME=${HF_HOME:-/workspace/hf_cache}
export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-/workspace/hf_cache/datasets}
export TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE:-/workspace/hf_cache/transformers}

HF_PUSH_ARGS=()
if [[ -n "${HF_PUSH_REPO_ID:-}" ]]; then
  HF_PUSH_ARGS+=(--hf-push-repo-id "${HF_PUSH_REPO_ID}" --hf-push-private)
fi

# Conservative single-B200 baseline. Tune batch/grad-accum after first 200 steps
# if memory headroom is clearly large.
source .venv/bin/activate
accelerate launch train_convgpt_v2_2d_pleias_long.py \
  --tokenizer /workspace/model-trainer/results/convgpt_v2_long_2d/convgpt_v2_rope_nope_hybrid_clean_h512_l16_seq512_bs8_fullloss_20260523_161605/checkpoint-46000 \
  --output-root /workspace/model-trainer/results/convgpt_v2_long_2d \
  --run-name "${RUN_NAME}" \
  --cache-dir /workspace/hf_cache \
  --pleias-files all \
  --train-limit 0 \
  --eval-limit 1000 \
  --max-length 2048 \
  --max-position-embeddings 65536 \
  --per-device-train-batch-size 128 \
  --per-device-eval-batch-size 1 \
  --gradient-accumulation-steps 2 \
  --learning-rate 0.001 \
  --warmup-ratio 0.02 \
  --max-steps 10000000 \
  --max-train-seconds 72000 \
  --save-steps 200 \
  --eval-steps 1000 \
  --logging-steps 20 \
  --save-total-limit 5 \
  --hidden-size 512 \
  --layers 32 \
  --grid-size 256 \
  --packing hilbert \
  --conv2d-backend triton_gather \
  --conv2d-chunk-size 1024 \
  --use-1d-branch \
  --use-2d-branch \
  --two-d-every 1 \
  --two-d-start-layer 0 \
  --position-embedding-type rope_nope \
  --rope-theta 10000 \
  --router-rope-fraction 1.0 \
  --no-use-row-col-embeddings \
  --conv1d-gate-init -6.0 \
  --conv2d-gate-init 2.0 \
  --router-type chunk_token_memory \
  --retrieval-every 2 \
  --retrieval-top-k 4 \
  --retrieval-num-slots 64 \
  --chunk-memory-size 64 \
  --chunk-memory-top-k 2 \
  --chunk-memory-token-top-k 64 \
  --chunk-memory-gate-init -4.0 \
  --branch-dropout 0.01 \
  --dropout 0.01 \
  --gradient-checkpointing \
  --bf16 \
  --torch-compile-model \
  --torch-compile-mode default \
  --report-to wandb \
  --wandb-mode "${WANDB_MODE}" \
  --dataloader-num-workers 2 \
  --early-stopping-patience 0 \
  --gen-prompts-json /workspace/model-trainer/convgpt_v2_gen_eval_prompts.json \
  --init-from-checkpoint /workspace/model-trainer/results/convgpt_v2_long_2d/convgpt_v2_rope_nope_hybrid_clean_h512_l16_seq512_bs8_fullloss_20260523_161605/checkpoint-46000 \
  --gen-max-new-tokens 2048 \
  --shuffle-buffer-size 10000 \
  --curriculum-lengths 2048 \
  --curriculum-stage-steps 0 \
  --tst-bag-sizes 1 \
  "${HF_PUSH_ARGS[@]}" \
  2>&1 | tee -a "${LOG}"