#!/bin/bash -l
# Topic 6: full training run of the N-gram MLP on WikiText-103.
#
# Submit from this directory with:   qsub train_ngram.qsub
# Extra arguments are passed on to train_ngram.py, e.g.:   qsub train_ngram.qsub --max_steps 2000
# If the job hits its walltime, just submit it again: training resumes from <WORK_DIR>/latest.pt.

#$ -P jbrcs                # <-- change to your SCC project
#$ -N ngram_wt103          # job name
#$ -l h_rt=12:00:00        # walltime limit
#$ -pe omp 8               # 8 CPU cores (sets NSLOTS=8): 7 DataLoader workers + the main process
#$ -l gpus=1               # 1 GPU ...
#$ -l gpu_c=8.0            # ... with compute capability >= 8.0 (A100, A40, A6000, L40S, H200, ...): fast bfloat16
#$ -j y                    # merge stdout and stderr into one log file: ngram_wt103.o<jobid>

# where everything is stored (the dataset itself is read-only and shared)
WORK_DIR=/projectnb/jbrcs/intropt/runs/wt103_n5       # <-- checkpoints, metrics.jsonl, model.pt
CACHE_DIR=/projectnb/jbrcs/intropt/runs/cache         # <-- tokenized data, shared between runs

module purge
module load miniconda
module load academic-ml
conda activate fall-2026-pyt

echo "job $JOB_ID on $(hostname), $NSLOTS cores, GPU(s): $CUDA_VISIBLE_DEVICES, started $(date)"
nvidia-smi -i $CUDA_VISIBLE_DEVICES --query-gpu=name,memory.total --format=csv,noheader

python -u train_ngram.py \
    --work_dir "$WORK_DIR" \
    --cache_dir "$CACHE_DIR" \
    --n 5 \
    --vocab_size 32768 \
    --n_embd 256 \
    --n_hidden 2048 \
    --batch_size 1024 \
    --max_steps 500000 \
    --eval_interval 10000 \
    --time_limit_hours 11.5 \
    "$@"

echo "finished $(date)"