| #!/usr/bin/env bash
|
| set -euo pipefail
|
|
|
| ROOT_OUTPUT_DIR="${REPOBRIDGE_OUTPUT_DIR:-$(pwd)/results/200m-until-selection}"
|
| ROOT_CHECKPOINT_DIR="${TAOTERN_CHECKPOINT_DIR:-$ROOT_OUTPUT_DIR/checkpoints}"
|
| SEQ_LEN="${SEQ_LEN:-512}"
|
| BATCH_SIZE="${BATCH_SIZE:-8}"
|
| PILOT_TOKENS="${PILOT_TOKENS:-300000000}"
|
| SERIOUS_TOKENS="${SERIOUS_TOKENS:-1000000000}"
|
| PILOT_EVAL_BATCHES="${PILOT_EVAL_BATCHES:-64}"
|
| SERIOUS_EVAL_BATCHES="${SERIOUS_EVAL_BATCHES:-128}"
|
| LEARNING_RATE="${LEARNING_RATE:-0.0003}"
|
| WEIGHT_DECAY="${WEIGHT_DECAY:-0.01}"
|
|
|
| ceil_div() {
|
| local numerator="$1"
|
| local denominator="$2"
|
| echo $(( (numerator + denominator - 1) / denominator ))
|
| }
|
|
|
| run_phase() {
|
| local phase="$1"
|
| local target_tokens="$2"
|
| local eval_batches="$3"
|
| local tokens_per_step=$((BATCH_SIZE * SEQ_LEN))
|
| local train_steps
|
| train_steps="$(ceil_div "$target_tokens" "$tokens_per_step")"
|
|
|
| printf '\n============================================================\n'
|
| printf '200M until-selection phase: %s\n' "$phase"
|
| printf 'target_tokens=%s batch=%s seq_len=%s train_steps=%s eval_batches=%s\n' \
|
| "$target_tokens" "$BATCH_SIZE" "$SEQ_LEN" "$train_steps" "$eval_batches"
|
| printf '============================================================\n'
|
|
|
| REPOBRIDGE_OUTPUT_DIR="$ROOT_OUTPUT_DIR/$phase" \
|
| TAOTERN_CHECKPOINT_DIR="$ROOT_CHECKPOINT_DIR/$phase" \
|
| BATCH_SIZES="$BATCH_SIZE" \
|
| SEQ_LEN="$SEQ_LEN" \
|
| TRAIN_STEPS="$train_steps" \
|
| EVAL_BATCHES="$eval_batches" \
|
| LEARNING_RATE="$LEARNING_RATE" \
|
| WEIGHT_DECAY="$WEIGHT_DECAY" \
|
| bash scripts/remote/run_200m_base_suite.sh
|
| }
|
|
|
| mkdir -p "$ROOT_OUTPUT_DIR" "$ROOT_CHECKPOINT_DIR"
|
|
|
| cat > "$ROOT_OUTPUT_DIR/run_plan.json" <<JSON
|
| {
|
| "stopping_point": "selection_after_1b_all_four_variants",
|
| "batch_size": $BATCH_SIZE,
|
| "seq_len": $SEQ_LEN,
|
| "learning_rate": $LEARNING_RATE,
|
| "weight_decay": $WEIGHT_DECAY,
|
| "phases": [
|
| {
|
| "name": "pilot_300m",
|
| "target_tokens_per_variant": $PILOT_TOKENS,
|
| "train_steps": $(ceil_div "$PILOT_TOKENS" $((BATCH_SIZE * SEQ_LEN))),
|
| "eval_batches": $PILOT_EVAL_BATCHES
|
| },
|
| {
|
| "name": "serious_1b",
|
| "target_tokens_per_variant": $SERIOUS_TOKENS,
|
| "train_steps": $(ceil_div "$SERIOUS_TOKENS" $((BATCH_SIZE * SEQ_LEN))),
|
| "eval_batches": $SERIOUS_EVAL_BATCHES
|
| }
|
| ],
|
| "variants": [
|
| "attention_196m",
|
| "pure_ssm_196m_hadamard",
|
| "pure_ssm_196m_nomix",
|
| "hybrid_ssm_first_199m"
|
| ]
|
| }
|
| JSON
|
|
|
| run_phase "pilot_300m" "$PILOT_TOKENS" "$PILOT_EVAL_BATCHES"
|
| run_phase "serious_1b" "$SERIOUS_TOKENS" "$SERIOUS_EVAL_BATCHES"
|
|
|
| echo "Selection gate reached after pilot_300m and serious_1b completed for all four variants."
|
|
|
