Skip to main content
The scripts are in the repository, not in the installed package. Clone it, then cd recipes/04-train/fireworks before running the commands below. Browse this recipe on GitHub.
While writes the dataset and measures the gain; Fireworks supplies the training GPUs and the endpoint. format="fireworks" on export and export_preference writes what a Fireworks managed training job reads, the firectl lines below run the job and deploy the result, and prove.py compares base against trained on the same held-out tasks through wai.Fireworks. What you learn: the two shapes Fireworks trains on (SFT rows with a per-message weight, one-turn DPO pairs), how the SDK’s loss mask and its chosen/rejected pairs map onto them, what firectl needs in which order, and why the proof runs both arms on the same tasks with the same number of rollouts. Needs: nothing for the export. A Fireworks account with firectl signed in and FIREWORKS_API_KEY in the environment for the job and the proof, and a judge key for the proof (OPENAI_API_KEYfor—judge openai:…, or WHILEAI_API_KEYfor the judge While hosts).smoke.sh` needs nothing. Takes: the export under a minute. The job is Fireworks’ queue: a few minutes for a LoRA on a 4B base. The proof is two runs of 128 rollouts. Costs: the export is free. Fireworks bills managed training per training token and the deployment per GPU hour; the cost estimator has the numbers for a base and a dataset size.

Run it

The script prints the Fireworks side with your ids filled in:

What happens

  1. wai.simulate runs the stand-in agent over the two tools, offline, two rollouts per ask. Swap in wai.Fireworks(base) and a model writer to collect from the real agent.
  2. The rule grades: the seeded fault is the failure. A model agent takes wai.Judge(rubric=RUBRIC) or a verifier here.
  3. select(mode="sft").export(..., format="fireworks") writes the passes. Each line is {"messages": [...], "tools": [...]} in the OpenAI wire shape; every assistant turn carries weight 0 or 1 from the SDK’s loss mask, so mask_mode="final" trains the last turn only and the default trains every assistant turn. Nothing else the row carries is written.
  4. select_for_preference() pairs a pass with a fail on the same ask, and export_preference(..., format="fireworks") writes Fireworks’ one-turn shape: input.messages is the prefix both sides share (the ask, and often the same opening tool call and its result), preferred_output and non_preferred_output the first assistant turn where they differ. Pairs that diverge on a tool result have no one-turn contrast and are dropped; the report counts them, and counts the pairs that lost later turns as fireworks_turns_cut. When a run has no such contrast the script builds the negative the policy names (a refund with no lookup) against each pass and says so.
  5. check_shapes holds both files to the shapes the Fireworks docs name, offline, so smoke.sh catches a drift before a job is billed.
  6. firectl uploads the datasets, trains SFT then DPO warm-started from it, and deploys the result. Trained LoRAs deploy on-demand only; the default shape live-merges the adapter, so serving has no LoRA overhead.
  7. prove.py runs the base through wai.Fireworks(base) on 32 held-out tasks at four rollouts each, then the trained model through wai.Fireworks(tuned) on the same tasks (tasks=before), grades both with the same judge, and prints wai.compare: the gain on the paired tasks with its interval, and NO DIFFERENCE when the interval spans zero. The judge is never either arm.

What this is not

The offline run proves the data path, not a training gain: the stand-in agent’s faults are seeded, so the pairs teach the recipe’s shape and nothing about a real agent. Point step 1 at your agent and step 2 at your judge before a job is worth its tokens. The shapes follow the Fireworks docs for supervised fine-tuning, DPO and deploying trained models; a live job on this exact export has not been run yet, so a field Fireworks rejects is a bug here, and the fix is in whileai/simulations/export.py. Reinforcement fine-tuning on Fireworks (eval-protocol create rft) takes a rollout evaluator rather than a file; a While judge as that evaluator is the next step, not this recipe.
Last modified on September 21, 2026