Step 5: Train with SFT (the clean, modern way)

Section 6 of 10-~ 10 min read-Synced from Cuantum content

Now we get to the core of the project: training the model. This is where all your preparation pays off. You have a clean dataset, a proper format, and a working environment. Now you just need to point a trainer at your data and let it run.

We'll use TRL's SFTTrainer, which is purpose-built for supervised fine-tuning on instruction datasets. It's a high-level wrapper around Hugging Face's Trainer that handles a lot of the common patterns for instruction tuning: tokenization, packing, attention masking, and more. You could build this yourself with raw PyTorch, but SFTTrainer gives you a production-quality implementation with very little code.

Create scripts/train_sft.py:

from datasets import load_datasetfrom transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArgumentsfrom trl import SFTTrainer MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" def main():    train_ds = load_dataset("json", data_files="data/train.jsonl", split="train")    eval_ds = load_dataset("json", data_files="data/eval.jsonl", split="train")     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)    if tokenizer.pad_token is None:        tokenizer.pad_token = tokenizer.eos_token     model = AutoModelForCausalLM.from_pretrained(        MODEL_NAME,        device_map="auto"    )     args = TrainingArguments(        output_dir="outputs/ch1_sft_tinyllama",        num_train_epochs=3,        per_device_train_batch_size=2,        gradient_accumulation_steps=8,        learning_rate=2e-5,        warmup_ratio=0.03,        logging_steps=25,        eval_strategy="steps",        eval_steps=100,        save_steps=100,        save_total_limit=2,        fp16=True,  # If your GPU supports bf16 better, you can swap fp16->bf16        report_to="none"    )     trainer = SFTTrainer(        model=model,        args=args,        train_dataset=train_ds,        eval_dataset=eval_ds,        dataset_text_field="text",        tokenizer=tokenizer,        max_seq_length=1024,  # start lower for VRAM safety; raise later if needed        packing=True  # packs multiple samples into one sequence (often faster)    )     trainer.train()    trainer.save_model("outputs/ch1_sft_tinyllama/final")    tokenizer.save_pretrained("outputs/ch1_sft_tinyllama/final") if __name__ == "__main__":    main()

Let's break down what's happening here, because every line matters.

Loading the datasets: We use Hugging Face's load_dataset to read the JSONL files we created in the previous step. The split="train" argument tells it to load the entire file as a single split (even though one of them is technically our eval set—this is just how the API works).

Loading the tokenizer and model: We're using TinyLlama, a 1.1B parameter model that's small enough to train on a consumer GPU but large enough to show real improvement after fine-tuning. The use_fast=True flag enables the fast Rust-based tokenizer, which is significantly faster than the Python version. The device_map="auto" argument tells Hugging Face to automatically place the model on the best available device (GPU if you have one, CPU otherwise). If you have multiple GPUs, it will even split the model across them.

Setting the pad token: Some models don't have a padding token defined by default. If that's the case, we set it to the EOS (end-of-sequence) token. This is a common pattern and it works fine for causal language modeling.

Training arguments: This is where you control the training process. Let's go through the key parameters:

num_train_epochs=3 — We'll train for 3 full passes through the dataset. For a small dataset (100–500 examples), this is usually enough to see significant improvement without overfitting.
per_device_train_batch_size=2 — We'll process 2 examples at a time on each GPU. This is intentionally small to avoid running out of memory.
gradient_accumulation_steps=8 — We'll accumulate gradients over 8 steps before doing a weight update. This effectively gives us a batch size of 2 × 8 = 16, which is large enough to get stable gradients but doesn't require loading 16 examples into memory at once.
learning_rate=2e-5 — This is a standard learning rate for fine-tuning pre-trained models. It's small enough to avoid catastrophic forgetting (where the model forgets what it learned during pre-training) but large enough to make meaningful updates.
warmup_ratio=0.03 — We'll linearly increase the learning rate from 0 to 2e-5 over the first 3% of training steps. This helps stabilize training at the start.
logging_steps=25 — We'll print training metrics every 25 steps. This gives you a sense of progress without flooding your terminal.
eval_strategy="steps" and eval_steps=100 — We'll run evaluation every 100 training steps. This lets you track how well the model is generalizing to unseen data.
save_steps=100 and save_total_limit=2 — We'll save a checkpoint every 100 steps, but only keep the 2 most recent checkpoints. This prevents your disk from filling up with old checkpoints you don't need.
fp16=True — We'll use mixed precision training, which runs some operations in 16-bit floating point instead of 32-bit. This uses less memory and runs faster on modern GPUs, with minimal impact on training quality. If your GPU supports bfloat16 (like Ampere or newer), you can use bf16=True instead for even better numerical stability.
report_to="none" — We're not using any experiment tracking tools like Weights & Biases or TensorBoard for this first run. You can enable them later if you want more detailed metrics.

SFTTrainer configuration: This is where we configure the trainer itself. The key parameters are:

dataset_text_field="text" — This tells the trainer which field in the JSONL contains the training text. In our case, it's the text field we created in the previous step.
max_seq_length=1024 — This is the maximum number of tokens in a single training example. Anything longer will be truncated. We're starting with 1024, which is a safe default for most GPUs. If you have a smaller GPU or want to train faster, you can reduce this to 512 or 768. If you have a large GPU and your examples are long, you can increase it to 2048 or higher.
packing=True — This enables sequence packing, which combines multiple short examples into a single sequence to reduce padding and improve training efficiency. This is almost always a good idea for instruction tuning, where examples vary in length. However, packing can sometimes increase peak memory usage depending on how your examples are distributed, so if you run out of VRAM, try turning this off.

Training and saving: Finally, we call trainer.train() to start training. This will run for 3 epochs, logging progress and saving checkpoints along the way. When it's done, we save the final model and tokenizer to outputs/ch1_sft_tinyllama/final. This is the checkpoint you'll use for inference.

Run training:

python scripts/train_sft.py

You'll see a lot of output as training runs. The key things to watch for are:

Loss decreasing: The training loss should go down steadily over time. If it's not decreasing, something is wrong with your data or your hyperparameters.
Eval loss: The evaluation loss should also decrease, though it will usually be a bit higher than the training loss. If the eval loss starts increasing while the training loss keeps decreasing, that's a sign of overfitting—the model is memorizing the training set instead of learning general patterns.
No VRAM errors: If you run out of memory, the script will crash with a CUDA out-of-memory error. If that happens, see the troubleshooting section below.

On a modern GPU (like an RTX 3090 or 4090), training should take 10–30 minutes depending on your dataset size. On a smaller GPU (like a GTX 1080 or RTX 2060), it might take 30–60 minutes. On a CPU, it will take several hours—this is not recommended unless you have no other option.

If you run out of VRAM

Running out of memory is one of the most common issues when fine-tuning models, especially if you're working with a consumer GPU that has limited VRAM. Here's how to fix it, in order from least to most drastic:

Reduce max_seq_length (1024 → 768 → 512) — Shorter sequences use less memory. This is usually the first thing to try. Start at 768 and see if that's enough. If not, go down to 512. Most instruction-following tasks don't require sequences longer than 512 tokens anyway.
Reduce per_device_train_batch_size (2 → 1) — Processing fewer examples at once reduces peak memory usage. If you do this, consider increasing gradient_accumulation_steps to keep the effective batch size the same (e.g., if you go from batch size 2 to batch size 1, increase accumulation from 8 to 16).
Increase gradient_accumulation_steps (8 → 16) — This reduces the frequency of weight updates, which can sometimes help with memory fragmentation. It also lets you maintain a large effective batch size even if you have to reduce the actual batch size.
Turn off packing — While packing usually helps with efficiency, it can sometimes increase peak memory usage if your dataset has a lot of long examples. Try setting packing=False and see if that helps.
Enable 8-bit quantization — If you installed bitsandbytes earlier, you can load the model in 8-bit mode by adding load_in_8bit=True to the from_pretrained call. This reduces memory usage by about 50% with minimal impact on quality.

If none of these work, you might need to use a smaller base model (like a 0.5B or 0.7B parameter model instead of 1.1B), or train on a machine with more VRAM. But in most cases, the above steps will be enough to get training working on an 8GB GPU.