run_training.sh

#!/usr/bin/env bash
# =============================================================================
# run_training.sh — launch the (paid) RFT training job with sane defaults.
#
# This is the TRAINING counterpart to run_model.sh:
#   * run_model.sh      -> training_run.sh --localtest   (local eval, no weights change)
#   * run_training.sh   -> training_run.sh --go          (the paid RFT job)
#
# It is a thin wrapper: it only fixes the defaults you asked for, then hands off
# to training_run.sh (the ONE source of truth for the actual `create rft` command).
#
# The defaults baked in here (edit via env if you need to):
#   BASE_MODEL          = accounts/fireworks/models/gpt-oss-20b   (the frozen base to tune)
#   DATASET             = a 5-Easy + 5-Medium subset built HERE from banking_easy_medium.jsonl
#   USE_LEANMCP_GATEWAY = 0   (NO LeanMCP gateway — direct Fireworks)
#   EPOCHS              = 1        (one epoch)
#   RL_LOSS             = DAPO     (uppercase — the Fireworks SDK Literal requires it)
#   N (rollouts)        = 16, TEMPERATURE = 0.8, MAX_TOKENS = 4096, LORA_RANK = 8
#     ^ these come straight from training_run.sh's own defaults; overridable by env.
#
# GUARDED like training_run.sh: by default it DRY-RUNS (builds the subset + prints
# the command). Pass --go to actually launch the paid job.
#
#   bash FIREWORKS_TRAINING/run_training.sh          # dry run: build subset + print cmd
#   bash FIREWORKS_TRAINING/run_training.sh --go     # LAUNCH the paid RFT job
#
# Per repo convention, Claude does NOT run this — you do.
# =============================================================================
set -euo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
ROOT="$(cd "$HERE/.." && pwd)"   # repo root
cd "$ROOT"

# ----------------------------- your choices ---------------------------------
# Base model to fine-tune. NOTE: no `fireworks_ai/` prefix here — that prefix is a
# LiteLLM *inference* tag; `create rft --base-model` wants the bare account path.
BASE_MODEL="${BASE_MODEL:-accounts/fireworks/models/gpt-oss-20b}"
# Name of the LoRA output that lands in YOUR account.
OUTPUT_MODEL="${OUTPUT_MODEL:-gpt-oss-20b-banking-rft-v1}"
# NO LeanMCP gateway (direct Fireworks). Explicit, per request.
export USE_LEANMCP_GATEWAY="${USE_LEANMCP_GATEWAY:-0}"
# One epoch, DAPO loss (these are also training_run.sh's defaults; set explicitly
# so this file is self-documenting and immune to a future default change there).
export EPOCHS="${EPOCHS:-1}"
# Fireworks SDK Literal is UPPERCASE ({GRPO, DAPO, DPO, ORPO, GSPO_TOKEN});
# training_run.sh normalizes case, but keep the canonical form here.
export RL_LOSS="${RL_LOSS:-DAPO}"

# ----------------------------- build the 5+5 subset -------------------------
# For the REAL RFT job the training tasks are exactly the rows of $DATASET (that
# file is what gets uploaded to Fireworks). The env-var caps MAX_EASY/MAX_MEDIUM
# only reliably subset in the LOCAL dataset_adapter, so to guarantee 5 Easy + 5
# Medium in training we materialize a 10-row subset file and point DATASET at it.
SRC="FIREWORKS_TRAINING/data/banking_easy_medium.jsonl"
SUBSET="FIREWORKS_TRAINING/data/banking_5easy_5medium.jsonl"
[ -f "$SRC" ] || { echo "ERROR: source dataset $SRC missing."; exit 1; }
# JSONL = one task per line, so grepping the bucket field returns whole task rows.
# NOTE: `head -5` closes the pipe after 5 lines, so `grep` (21 Medium matches) gets
# SIGPIPE and exits non-zero. Under `set -o pipefail` that non-zero would trip
# `set -e` and kill the script silently — so we drop pipefail just for this build.
set +o pipefail
{
  grep -E '"bucket": *"Easy"'   "$SRC" | head -5
  grep -E '"bucket": *"Medium"' "$SRC" | head -5
} > "$SUBSET"
set -o pipefail
n_easy="$(grep -Ec '"bucket": *"Easy"'   "$SUBSET" || true)"
n_medium="$(grep -Ec '"bucket": *"Medium"' "$SUBSET" || true)"
rows="$(wc -l < "$SUBSET" | tr -d ' ')"
if [ "$n_easy" -ne 5 ] || [ "$n_medium" -ne 5 ]; then
  echo "ERROR: subset build got $n_easy Easy + $n_medium Medium (want 5 + 5). Check $SRC."
  exit 1
fi
export DATASET="$SUBSET"
# Belt-and-suspenders: if EP ever re-subsets via the adapter, keep it at 5+5 too.
export MAX_EASY="${MAX_EASY:-5}"
export MAX_MEDIUM="${MAX_MEDIUM:-5}"

# ----------------------------- summary --------------------------------------
echo "== run_training.sh =="
echo "  base model     : $BASE_MODEL"
echo "  output model   : $OUTPUT_MODEL"
echo "  dataset        : $DATASET  ($rows rows: ${n_easy} Easy + ${n_medium} Medium)"
echo "  gateway        : $([ "$USE_LEANMCP_GATEWAY" = 1 ] && echo "LeanMCP" || echo "OFF (direct Fireworks)")"
echo "  epochs         : $EPOCHS"
echo "  rl loss        : $RL_LOSS"
echo "  rollouts (N)   : ${N:-16}   temperature: ${TEMPERATURE:-0.8}   max-tokens: ${MAX_TOKENS:-4096}   lora-rank: ${LORA_RANK:-8}"
echo

# ----------------------------- hand off -------------------------------------
# Everything above is exported, so training_run.sh reads BASE_MODEL / OUTPUT_MODEL
# / DATASET / EPOCHS / RL_LOSS / USE_LEANMCP_GATEWAY straight from the env. We pass
# through the run mode: default = dry run; --go = launch the paid job.
export BASE_MODEL OUTPUT_MODEL
case "${1:-}" in
  --go)
    bash "$HERE/training_run.sh" --go
    ;;
  --skip)
    # Fast relaunch: --go but skip the ~15-min local validation, force evaluator
    # re-upload, and reuse the last uploaded dataset (see training_run.sh --skip).
    bash "$HERE/training_run.sh" --skip
    ;;
  *)
    echo "Dry run only (subset built above). Re-run with --go to LAUNCH the paid RFT job:"
    echo "    bash FIREWORKS_TRAINING/run_training.sh --go"
    echo
    bash "$HERE/training_run.sh"   # training_run.sh's own dry run: prints the create-rft command
    ;;
esac