Supervised Fine-Tuning: Model Training

In the previous chapter, we built the SFT training data: 512 (prompt, response) pairs. In this chapter, we will use this data to perform supervised fine-tuning on the pre-trained model.

Supervised fine-tuning¶

The core idea of supervised fine-tuning: compute the loss only on the response part.

During pre-training, the model computes a loss at every position, because every position is learning to predict the next token. But in supervised fine-tuning, we only care whether the model has learned to generate suitable answers, and we do not need it to learn how to generate prompts.

Mask¶

The solution is to add a mask to every sample:

Prompt part: the mask is 0 and no loss is computed;
Response part: the mask is 1 and the loss is computed as usual.

This way, the prompt part produces no loss, and the model learns only from the response part.

Strategy¶

Learning rate: supervised fine-tuning usually uses a learning rate one order of magnitude smaller than pre-training. The model is already trained, and we only want to adjust the direction in which it generates text, not relearn language knowledge;
Epochs: SFT data is small and training for too long easily overfits, so the number of epochs is usually small. This chapter trains for only 1 epoch.

import pickle

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

Dataset¶

SFT dataset¶

The SFT dataset SFTDataset differs from the pre-training dataset: it processes (prompt, response) pairs:

Concatenation: join the prompt and the response into one token sequence;
Truncation: if the sequence exceeds the context size, truncate from the head of the prompt and keep the complete response;
Masking: give the prompt part a mask of 0 and the response part a mask of 1. The result is returned as (features, labels, mask).

In addition, the samples have different lengths, so the batch size is fixed at 1 and no padding is needed.

class SFTDataset(Dataset):

    def __init__(self, filename, context_size=64, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.split = split
        super().__init__(1)

    # Load function: read the SFT sample file
    def load(self):
        with open(self.filename, "rb") as f:
            examples = pickle.load(f)

        split = int(len(examples) * self.split)
        self.train_data = self._pack(examples[:split])
        self.test_data = self._pack(examples[split:])

    # Packing function: convert (prompt, response) into features, labels, and a mask
    def _pack(self, examples):
        xs, ys, masks = [], [], []
        for prompt, response in examples:
            x, y, mask = self._build_example(prompt, response)
            xs.append(x)
            ys.append(y)
            masks.append(mask)
        return xs, ys, masks

    # Build a single sample
    def _build_example(self, prompt, response):
        # Concatenate the prompt and the response
        tokens = list(prompt) + list(response)
        # Overlong sample: truncate from the head of the prompt first, keeping the response intact where possible
        if len(tokens) > self.context_size + 1:
            overflow = len(tokens) - (self.context_size + 1)
            prompt = prompt[overflow:] if overflow < len(prompt) else []
            tokens = list(prompt) + list(response)
            tokens = tokens[-(self.context_size + 1):]

        # Position of the response in the token sequence
        response_start = len(prompt)

        # Features: drop the last token
        x = np.array(tokens[:-1], dtype=np.int64)
        # Labels: drop the first token
        y = np.array(tokens[1:], dtype=np.int64)
        # Mask: only the response part holds valid tokens
        mask = np.arange(len(y)) + 1 >= response_start
        return x, y, mask

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y, mask = self.data
        return Tensor(x[s]), Tensor(y[s]), Tensor(mask[s])

Loss function¶

SFT cross-entropy loss function¶

The SFT loss function SFTLoss is a cross-entropy loss function with a mask:

compute the cross-entropy at all positions;
multiply by the mask, keeping only the loss of the response part;
divide by the length of the response part (the number of 1s in the mask) to take the average over the valid values.

During backpropagation, the gradient at the prompt positions is likewise 0.

class SFTLoss(Loss):

    def loss(self, p: Tensor, y: Tensor, mask):
        # Built-in Softmax
        exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        # Flatten
        flat_softmax = softmax.reshape(-1, softmax.shape[-1])
        flat_y = y.data.reshape(-1).astype(np.int64)
        # Mask
        flat_mask = mask.data.reshape(-1)
        n = max(np.sum(flat_mask), 1.0)
        rows = np.arange(len(flat_y))

        # Create the loss tensor
        log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
        # Apply the mask
        ce = Tensor(0 - np.sum(log * flat_mask) / n)

        # Gradient (closure) function of the loss tensor
        def gradient_fn():
            flat_grad = flat_softmax.copy()
            flat_grad[rows, flat_y] -= 1
            # Apply the mask: where the prompt mask is 0, the gradient is 0 as well
            flat_grad *= flat_mask[:, None]
            p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n

        # Set the gradient (closure) function and the parent list of the loss tensor
        return ce.attach(gradient_fn, parents=[p])

Trainer¶

SFT trainer¶

The SFT trainer SFTTrainer inherits from the GPT trainer GPTTrainer. The main differences are:

train (training function): receives the mask from the dataset and passes it to the loss function as an argument;
test (test function): no testing is needed in supervised fine-tuning, so it returns None directly.

SFT trains for only 1 epoch and saves the model parameters right after.

class SFTTrainer(GPTTrainer):

    def train(self, dataset, epochs, scheduler=None, filename=None):
        self.model.train()

        steps = 0
        for epoch in range(epochs):
            order = list(range(len(dataset)))
            np.random.shuffle(order)

            total_loss = 0.0
            for step, i in enumerate(order):
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                # Pass the mask
                feature, label, mask = dataset[i]
                prediction = self.model(feature)
                loss = self.loss_fn(prediction, label, mask)

                self.optimizer.zero_grad()
                loss.backward()
                total_loss += float(loss.data)
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

                if (step + 1) % 100 == 0:
                    lr = f" lr {self.optimizer.lr:.6f}" if scheduler is not None else ""
                    print(f"epoch {epoch + 1} step {step + 1}/{len(dataset)} loss {(total_loss / 100):.4f}{lr}")
                    total_loss = 0.0

            if filename is not None:
                self.save(filename)
                print(f"epoch {epoch + 1} saved SFT model to {filename}")

    # Test function: no testing during fine-tuning
    def test(self, dataset):
        return None

Hyperparameters¶

Learning rate¶

Supervised fine-tuning uses a smaller learning rate to avoid disturbing the pre-training result too much.

LEARNING_RATE = 0.0001

Model parameters¶

The other model parameters must stay the same as in the pre-training stage.

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

Modeling¶

Loading the model¶

The pre-training dataset CharDataset is used here only to provide the vocabulary and the conversion between text and tokens.

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = SFTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = SFTTrainer(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)

Loading the SFT data¶

The SFT samples are split 9:1 into a training set and a test set. The training set has 460 samples and the batch size is 1, so one epoch has 460 steps. The 52 samples of the test set are not used in this chapter.

SFT_SAMPLES = "../../../../data/sft-samples.pkl"

sft_dataset = SFTDataset(SFT_SAMPLES, CONTEXT_SIZE)

Training¶

Use a warmup-cosine learning rate scheduler: warm up linearly to LEARNING_RATE over the first 50 steps, then decay along a cosine curve to LEARNING_RATE / 10.

SFT_MODEL = "../../../../data/tinyshakespeare-sft.npz"

scheduler = WarmupCosineScheduler(LEARNING_RATE, len(sft_dataset), 50, LEARNING_RATE / 10)
trainer.train(sft_dataset, 1, scheduler, SFT_MODEL)

epoch 1 step 100/460 loss 1.5225 lr 0.000097
epoch 1 step 200/460 loss 1.4181 lr 0.000074
epoch 1 step 300/460 loss 1.3249 lr 0.000040
epoch 1 step 400/460 loss 1.3159 lr 0.000015
epoch 1 saved SFT model to ../../../../data/tinyshakespeare-sft.npz

Inference¶

Let the fine-tuned model generate, starting from the prompt ROMEO:.

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:
He theople, that low with many thrieve weak:
While crampudged you I know the valous in your blood,
I must my patience, and my Edward's life!

TRANIO:

ANORTEN:

POLIXENES:
I know not mine a deep a shado, your whom

PETRUCHIO:
You say you know! but know re

In this chapter, we performed supervised fine-tuning on the pre-trained model with the SFT dataset: only the loss of the response part is computed, so that the model learns the (prompt, response) conversation format. From the generated result we can see that the model keeps the “speaker: speech” turn-taking format, but the content is still incoherent.

But SFT has its own limitation: it only tells the model how to answer a question, not what a better answer is. For example, two answers may both be correct, but one is concise and clear while the other is wordy. SFT cannot tell the model which one is better.

In the next chapter, we will introduce Direct Preference Optimization (DPO), so that the model learns to pick the better one among several answers.

Exercises¶

Try raising the SFT learning rate, for example to the same value as in pre-training, and observe the training result. Think about why fine-tuning needs a smaller learning rate.