code/TaoTrain/scripts/remote/submit_ssm_improvement_sweep.sh

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

RUN_ID="${RUN_ID:-ssm-improvement-sweep-$(date +%Y%m%d-%H%M%S)}"
DATA_PATH="${DATA_PATH:-/home/student/Data/TaoData/pretrain.jsonl}"
TOKENIZER_PATH="${TOKENIZER_PATH:-/home/student/YouZheng/tokenizers/taodata_pilot_8k/tokenizer.model}"
SSM_REPO_PATH="${SSM_REPO_PATH:-/home/student/YouZheng/gamma_ssm_repo}"
PYTHON_BIN="${PYTHON_BIN:-/home/student/.venv/bin/python}"
JOB_ROOT="${JOB_ROOT:-/home/student/YouZheng/jobs/taotern}"
REMOTE_REPO="${REMOTE_REPO:-$(pwd)}"

# This is intentionally below the eventual 200M scale. It is a model-selection
# sweep that can run unattended and resume completed benchmark cases.
MAX_TOKENS="${MAX_TOKENS:-100000000}"
MAX_RECORDS="${MAX_RECORDS:-150000}"
TRAIN_STEPS="${TRAIN_STEPS:-5000}"
EVAL_BATCHES="${EVAL_BATCHES:-96}"
BATCH_SIZES="${BATCH_SIZES:-32,64}"
SEQ_LEN="${SEQ_LEN:-512}"

JOB_COMMAND="PYTHONPATH=$REMOTE_REPO/src:$SSM_REPO_PATH $PYTHON_BIN scripts/benchmark_taonet_real_tokens.py \

  --data-path $DATA_PATH \

  --text-field text \

  --tokenizer-type sentencepiece \

  --tokenizer-path $TOKENIZER_PATH \

  --max-records $MAX_RECORDS \

  --max-tokens $MAX_TOKENS \

  --eval-fraction 0.1 \

  --architectures taonet,taonet_ssm,taonet_hybrid \

  --batch-sizes $BATCH_SIZES \

  --seq-len $SEQ_LEN \

  --hidden-dim 256 \

  --num-layers 4 \

  --num-heads 4 \

  --d-latent-kv 192 \

  --hidden-dim-ff 1024 \

  --ssm-core dplr \

  --ssm-hidden-dims 16,32 \

  --ssm-mixer-dims 128,256 \

  --ssm-num-lanes-list 1,2 \

  --ssm-lane-combine channel \

  --ssm-lane-modes full,split \

  --ssm-split-mixes none,hadamard \

  --ssm-rank 1 \

  --ssm-kernel-mode conv \

  --no-ssm-finite-tail-correction \

  --ssm-gate-types channel \

  --hybrid-patterns attention_first,ssm_first,single_ssm_middle,single_ssm_late \

  --dtype bf16 \

  --device cuda \

  --warmup 2 \

  --repeats 3 \

  --backward \

  --train-steps $TRAIN_STEPS \

  --learning-rate 0.0008 \

  --weight-decay 0.01 \

  --eval-batches $EVAL_BATCHES \

  --ssm-local-shift \

  --ssm-local-shift-per-channel \

  --ssm-local-shift-init 0.1 \

  --output-dir \"\$REPOBRIDGE_OUTPUT_DIR\" \

  --resume-completed \

  --incremental-output \

  --save-case-checkpoints \

  --checkpoint-dir \"\$TAOTERN_CHECKPOINT_DIR\""

export RUN_ID JOB_ROOT JOB_COMMAND
export OUTPUT_DIR="${OUTPUT_DIR:-$JOB_ROOT/$RUN_ID/outputs}"
export CHECKPOINT_DIR="${CHECKPOINT_DIR:-$JOB_ROOT/$RUN_ID/checkpoints}"

bash scripts/remote/submit_detached_job.sh