Supervised Fine-Tuning (SFT) - Instruction Data, Chat Templates, Packing & Loss Masking

"""SFT data plumbing: chat template -> tokens -> labels with loss masking.

Real trainers use the model's own tokenizer and its chat template (for example a
Jinja template shipped with the tokenizer). Here a whitespace "tokenizer" keeps the
mechanics visible. The key rule: the loss is computed only on tokens the model
should learn to PRODUCE (assistant turns + end-of-turn marker), never on the
system/user text it is merely conditioned on. Ignored positions get label -100,
the convention PyTorch's CrossEntropyLoss(ignore_index=-100) uses.
"""
import math, random

IGNORE = -100

def render_chatml(messages):
    """Render to a ChatML-style string and remember which spans are assistant output."""
    pieces = []  # (text, trainable)
    for m in messages:
        pieces.append((f"<|im_start|> {m['role']}", False))
        pieces.append((m["content"], m["role"] == "assistant"))
        # The end-of-turn token after an assistant reply IS trained: the model must learn to stop.
        pieces.append(("<|im_end|>", m["role"] == "assistant"))
    return pieces

def tokenize(pieces, vocab):
    ids, labels = [], []
    for text, trainable in pieces:
        for tok in text.split():
            idx = vocab.setdefault(tok, len(vocab))
            ids.append(idx)
            labels.append(idx if trainable else IGNORE)
    return ids, labels

convo = [
    {"role": "system", "content": "You are a terse SQL assistant ."},
    {"role": "user", "content": "count users created today"},
    {"role": "assistant", "content": "SELECT count(*) FROM users WHERE created_at::date = current_date ;"},
    {"role": "user", "content": "only active ones"},
    {"role": "assistant", "content": "SELECT count(*) FROM users WHERE active AND created_at::date = current_date ;"},
]
vocab = {}
ids, labels = tokenize(render_chatml(convo), vocab)
inv = {v: k for k, v in vocab.items()}
# Compact view: trained tokens are wrapped in [brackets], ignored ones are plain.
view = " ".join(f"[{inv[i]}]" if l != IGNORE else inv[i] for i, l in zip(ids, labels))
for line in view.replace("<|im_start|>", "\n<|im_start|>").strip().splitlines():
    print(" ", line.rstrip())
trained = sum(l != IGNORE for l in labels)
print(f"\n{trained}/{len(labels)} tokens contribute to the loss ({trained/len(labels):.0%})")

# Shifted next-token loss on fake logits: position t predicts token t+1.
random.seed(0)
V = len(vocab)
def xent(logits, target):
    m = max(logits); lse = m + math.log(sum(math.exp(x - m) for x in logits))
    return lse - logits[target]
logits = [[random.gauss(0, 1) for _ in range(V)] for _ in ids]
def mean_loss(mask):
    tot, n = 0.0, 0
    for t in range(len(ids) - 1):
        tgt = labels[t + 1] if mask else ids[t + 1]
        if tgt == IGNORE:
            continue
        tot += xent(logits[t], tgt); n += 1
    return tot / n, n
for mask in (True, False):
    loss, n = mean_loss(mask)
    print(f"masked={mask!s:<5}  loss over {n:2d} positions = {loss:.3f}")
print("\nFailure path: with masked=False the model is also trained to WRITE system prompts and user")
print("questions, wasting capacity and teaching it to role-play the user. A wrong template at")
print("inference (different special tokens than training) silently degrades quality.")

  <|im_start|> system You are a terse SQL assistant . <|im_end|>
  <|im_start|> user count users created today <|im_end|>
  <|im_start|> assistant [SELECT] [count(*)] [FROM] [users] [WHERE] [created_at::date] [=] [current_date] [;] [<|im_end|>]
  <|im_start|> user only active ones <|im_end|>
  <|im_start|> assistant [SELECT] [count(*)] [FROM] [users] [WHERE] [active] [AND] [created_at::date] [=] [current_date] [;] [<|im_end|>]

22/49 tokens contribute to the loss (45%)
masked=True   loss over 22 positions = 3.729
masked=False  loss over 48 positions = 3.828

Failure path: with masked=False the model is also trained to WRITE system prompts and user
questions, wasting capacity and teaching it to role-play the user. A wrong template at
inference (different special tokens than training) silently degrades quality.

// Padding vs packing for SFT batches, with the cross-contamination detail.
// Padding: every example sits alone in a row of length L; the rest is pad (wasted compute).
// Packing: concatenate several examples into one row of length L (first-fit decreasing here).
// Naive packing lets example B attend to example A in the same row; correct packing resets
// position ids and uses a block-diagonal ("document") attention mask so examples stay isolated.

function lcg(seed: number): () => number {
  let s = seed >>> 0;
  return () => ((s = (Math.imul(s, 1664525) + 1013904223) >>> 0) / 2 ** 32);
}

const L = 2048; // max sequence length per row
const rnd = lcg(7);
// SFT data is usually long-tailed: many short chats, a few long ones.
const lengths = Array.from({ length: 4000 }, () => Math.min(L, Math.floor(60 + 1400 * rnd() ** 3)));
const tokens = lengths.reduce((a, b) => a + b, 0);

// Padding: one example per row.
const padRows = lengths.length;
const padUtil = tokens / (padRows * L);

// Packing: first-fit decreasing bin packing.
const bins: number[][] = [];
const free: number[] = [];
for (const len of [...lengths].sort((a, b) => b - a)) {
  let i = free.findIndex((f) => f >= len);
  if (i === -1) { bins.push([]); free.push(L); i = bins.length - 1; }
  bins[i].push(len);
  free[i] -= len;
}
const packUtil = tokens / (bins.length * L);
console.log(`examples=${lengths.length} real tokens=${tokens}`);
console.log(`padding : rows=${padRows}  utilisation=${(padUtil * 100).toFixed(1)}%`);
console.log(`packing : rows=${bins.length}  utilisation=${(packUtil * 100).toFixed(1)}%  -> ~${(padRows / bins.length).toFixed(1)}x fewer rows`);

// Position ids and attention isolation inside one packed row.
const row = bins[bins.length - 1];
const positionIds: number[] = [];
const docIds: number[] = [];
row.forEach((len, d) => { for (let p = 0; p < len; p++) { positionIds.push(p); docIds.push(d); } });
const canAttend = (q: number, k: number, isolated: boolean) => k <= q && (!isolated || docIds[q] === docIds[k]);
const q = positionIds.length - 1; // last token of the last example in this row
let leakNaive = 0, leakIsolated = 0;
for (let k = 0; k <= q; k++) {
  if (docIds[k] !== docIds[q]) {
    if (canAttend(q, k, false)) leakNaive++;
    if (canAttend(q, k, true)) leakIsolated++;
  }
}
console.log(`\nlast packed row holds ${row.length} examples; position ids restart at 0 for each: [${positionIds.slice(0, 3).join(",")}, ... ${positionIds.slice(row[0] - 1, row[0] + 2).join(",")} ...]`);
console.log(`tokens from OTHER examples the last token can see: naive causal mask=${leakNaive}, block-diagonal mask=${leakIsolated}`);
console.log("Failure path: naive packing trains on cross-example context; losses look fine but the model learns spurious dependencies.");

examples=4000 real tokens=1620521
padding : rows=4000  utilisation=19.8%
packing : rows=793  utilisation=99.8%  -> ~5.0x fewer rows

last packed row holds 27 examples; position ids restart at 0 for each: [0,1,2, ... 59,0,1 ...]
tokens from OTHER examples the last token can see: naive causal mask=1560, block-diagonal mask=0
Failure path: naive packing trains on cross-example context; losses look fine but the model learns spurious dependencies.