Run and understand the training loop

From loaded model to one update

The lab uses MLX for array computation on Apple Silicon, MLX-LM for model loading and adapter training, and the model's tokenizer for message formatting. run_experiment.py configures these libraries rather than implementing a transformer again.

model, tokenizer = load(model_path)
# Configure frozen base weights and trainable LoRA adapters.
# The trainer repeats the following conceptual steps:
batch = next_training_batch()
loss, gradients = value_and_grad(response_loss)(model, batch)
optimizer.update(model, gradients)

This is a conceptual outline, not a standalone runnable snippet. The actual trainer includes compilation, batching, validation, and saving. The following steps explain the work it performs.

1. Load and configure. Load the pinned pretrained weights and tokenizer, freeze the base, and create rank-eight adapters.
2. Make a batch. Convert a training example into token IDs. The sequence contains context and demonstrated response; target IDs are shifted one token ahead of input IDs.
3. Predict and score. Run the model and average cross-entropy over response positions. Prompt tokens provide context; prompt and padding targets do not contribute directly to the loss.
4. Calculate gradients. Backpropagation calculates how a small change to each trainable parameter would affect this loss. This computes sensitivities; it does not itself apply the update.
5. Apply an update. Adam combines the gradients with its running statistics and learning rate to change adapter weights. This run uses learning rate 0.0001.
6. Check and save. Periodically measure development loss without updates and save adapter snapshots. Continue using training examples for subsequent updates.

Batch size 1 means one example per update here. Two thousand steps visit the 2,000 training examples once. Larger batches and more passes change the experiment; they do not automatically improve held-out answers.

Run the optional local lab

Download the source and recorded results. Unzip it and open a terminal inside wikisql-sft-lab. This path requires an Apple Silicon Mac and Python 3.12. The measured job took four minutes on an M3 Max with 36 GiB RAM, with about 2.05 GiB peak whole-process memory and no new swap-outs. Setup and downloads are excluded; compilation caches were warm. An 8/16 GiB or Intel Mac has not been validated.

Create an environment outside the project. The pinned dependencies and public model download need internet access; no API key is needed. The model weights are about 1 GB, with additional disk space for the environment, data, and outputs.

export WIKISQL_HOME="$HOME/.cache/dougdoes-wikisql-sft"
python3.12 -m venv "$WIKISQL_HOME/venv"
source "$WIKISQL_HOME/venv/bin/activate"
python -m pip install -r requirements-lock.txt
python fetch.py --home "$WIKISQL_HOME"
python prepare.py --source "$WIKISQL_HOME/wikisql/data" \
  --model "$WIKISQL_HOME/model" --out "$WIKISQL_HOME/data"
python -m unittest -v

First measure the original model on development data:

python guard.py --budget-gib 2 --log "$WIKISQL_HOME/runs/base-valid.log" -- \
  python -u run_experiment.py eval --model "$WIKISQL_HOME/model" \
  --data "$WIKISQL_HOME/data" --source "$WIKISQL_HOME/wikisql/data" \
  --out "$WIKISQL_HOME/runs/base-valid" --split valid

For calibration, use the training command below with --steps 50 and a separate output directory such as calibration. Then start the full run fresh:

python guard.py --budget-gib 2.25 --log "$WIKISQL_HOME/runs/trained.log" -- \
  python -u run_experiment.py train --model "$WIKISQL_HOME/model" \
  --data "$WIKISQL_HOME/data" --source "$WIKISQL_HOME/wikisql/data" \
  --out "$WIKISQL_HOME/runs/trained" --steps 2000

The guard requires the process budget plus 1.5 GiB of free/speculative memory before launch and monitors memory pressure and swap. It stops the worker on new swapping or a budget violation. Run commands sequentially; choose new output names for repeat experiments so prior evidence is retained.

What to inspect after the process exits

Keep adapters.safetensors, adapter_config.json, metrics, and raw predictions together with the base revision and data manifest. Lower loss means the demonstrated tokens became more probable on average; it does not mean generated SQL will answer every question. Next we reload the saved adapter in a new process and check answers.

For optional mathematical depth, read the existing gradients lesson or PyTorch's backpropagation explanation. The practical distinction to retain is that gradients are calculated first and the optimizer applies the update second.