VKVlad KulikovPublic / lectures
Modern LLMs13 / 20

Module 13. Supervised Fine-Tuning and Data Work

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 2588.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Notebook 13. SFT data: from a dialogue to a correctly packed batch

Main module notebook

Download .ipynb module_13_notebook_sft.ipynb · 58.9 KB

Show notebook text

Notebook 13. SFT data: from a dialogue to a correctly packed batch

Module 13 — hands-on notebook

This is an engineering-disciplined educational SFT data pipeline. One small dialogue crosses the boundaries most often broken in a real collator:

messages → input_ids/labels → shifted causal loss → FFD assignments → packed sequence → position_ids + segment_ids + cu_seq_lens → block-diagonal causal isolation → padded batch.

We then reproduce the compute, best-of-N, and LoRA calculations and train a tiny causal model, generate held-out answers, and reload its weights. It does not train a large model or reproduce a FlashAttention kernel or distributed trainer; those layers are unnecessary for verifying the data contracts.

Colab contract: Runtime → Run all completes the mandatory path on CPU in a clean session. The notebook reads no adjacent files and needs no manual upload, API key, or paid service. A frozen Module-13 numerical payload is embedded with SHA-256 verification.

Route:

  1. a chat template and real token labels rather than a 1 sentinel;

  2. the causal-LM shift and ignore_index=-100;

  3. multi-turn and tool-response masking;

  4. FFD with original indices rather than bin loads only;

  5. packed input_ids, labels, position_ids, segment_ids, and cu_seq_lens;

  6. a negative cross-example leakage test;

  7. a padded batch and padding/loss masks;

  8. compute, best-of-N, and strict GiB/MiB units in the LoRA estimate.

Python
import importlib.util
import subprocess
import sys

if importlib.util.find_spec("numpy") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "numpy>=1.24"])

from dataclasses import dataclass
import hashlib
import json
import math
from typing import Iterable, Sequence

import numpy as np

np.set_printoptions(precision=4, suppress=True, linewidth=120)
IGNORE_INDEX = -100
print('Environment ready.')

0. Embedded Module-13 numerical contract

The reference numbers belong to the module rather than mutable lecture prose. The JSON is embedded and verified by SHA-256, so future lecture edits cannot change the computational route.

Python
REFERENCE_JSON = '{\n  "best_of_n": {\n    "coverage": {\n      "1": 0.30000000000000004,\n      "16": 0.9966767069430399,\n      "2": 0.51,\n      "4": 0.7599,\n      "8": 0.94235199\n    },\n    "pass_probability": 0.3\n  },\n  "compute": {\n    "accelerator_flops": 990000000000000.0,\n    "average_sft_tokens": 1200,\n    "mfu": 0.4,\n    "n_params": 8030000000.0,\n    "pretrain_flops": 7.227e+23,\n    "pretrain_tokens": 15000000000000.0,\n    "ratio": 20833.333333333336,\n    "sft_epochs": 3,\n    "sft_examples": 200000,\n    "sft_flops": 3.46896e+19,\n    "sft_hours": 24.333333333333332,\n    "sft_tokens": 240000000\n  },\n  "lora": {\n    "d_kv": 1024,\n    "d_model": 4096,\n    "four_bit_payload_gib": 3.7392601370811462,\n    "fraction": 0.001697570112079701,\n    "full_training_states_gib": 119.65632438659668,\n    "layers": 32,\n    "lora_adapter_bf16_mib": 26.0,\n    "lora_training_states_mib": 208.0,\n    "params": 13631488,\n    "per_layer": 425984,\n    "qlora_parameter_floor_gib": 3.9423851370811462,\n    "rank": 16,\n    "state_saving": 589.0772892878606\n  },\n  "packing": {\n    "bins": [\n      {\n        "indices": [\n          7,\n          10,\n          18\n        ],\n        "lengths": [\n          1347,\n          626,\n          72\n        ],\n        "used": 2045\n      },\n      {\n        "indices": [\n          15,\n          11,\n          1,\n          3\n        ],\n        "lengths": [\n          754,\n          556,\n          527,\n          180\n        ],\n        "used": 2017\n      },\n      {\n        "indices": [\n          23,\n          12,\n          6,\n          0,\n          8\n        ],\n        "lengths": [\n          514,\n          443,\n          425,\n          403,\n          259\n        ],\n        "used": 2044\n      },\n      {\n        "indices": [\n          14,\n          21,\n          2,\n          4,\n          17,\n          9,\n          13,\n          20\n        ],\n        "lengths": [\n          392,\n          326,\n          315,\n          267,\n          267,\n          230,\n          174,\n          76\n        ],\n        "used": 2047\n      },\n      {\n        "indices": [\n          5,\n          22,\n          19,\n          16\n        ],\n        "lengths": [\n          165,\n          128,\n          126,\n          120\n        ],\n        "used": 539\n      }\n    ],\n    "ffd_utilization": 0.848828125,\n    "forward_reduction": 4.8,\n    "lengths": [\n      403,\n      527,\n      315,\n      180,\n      267,\n      165,\n      425,\n      1347,\n      259,\n      230,\n      626,\n      556,\n      443,\n      174,\n      392,\n      754,\n      120,\n      267,\n      72,\n      126,\n      76,\n      326,\n      128,\n      514\n    ],\n    "max_length": 2048,\n    "n_ffd": 5,\n    "n_naive": 24,\n    "naive_utilization": 0.17683919270833334\n  },\n  "s1": {\n    "accelerators": 8,\n    "flops": 1.728e+18,\n    "minutes": 9.090909090909092,\n    "n_params": 32000000000.0,\n    "tokens": 9000000,\n    "tokens_per_trace": 9000,\n    "traces": 1000\n  },\n  "schema_version": 3,\n  "template": {\n    "content_lengths": {\n      "assistant": 20,\n      "system": 12,\n      "user": 8\n    },\n    "per_message_overhead": 5,\n    "special_ids": {\n      "bos": 1,\n      "end_header": 3,\n      "eot": 4,\n      "newline": 5,\n      "pad": 0,\n      "role_assistant": 8,\n      "role_system": 6,\n      "role_tool": 9,\n      "role_user": 7,\n      "start_header": 2\n    },\n    "total_tokens": 56,\n    "trainable_fraction": 0.375,\n    "trainable_tokens": 21\n  }\n}'
EXPECTED_SHA256 = "ff1e8f3e598ed6fbe3e6b807fb4d3f783cad366dc965b1f51681b3ebef6bb840"
actual_sha256 = hashlib.sha256((REFERENCE_JSON + "\n").encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256, (actual_sha256, EXPECTED_SHA256)
REF = json.loads(REFERENCE_JSON)
assert REF["schema_version"] == 3
print('Module-13 reference payload verified', actual_sha256[:12])

1. Chat template: input_ids and −100 labels

The teaching template preserves the lecture arithmetic: BOS; four header tokens per message; content; EOT. For assistant messages the content and its EOT are trainable. System, user, tool, and every header receive −100. In a real causal LM, labels contain source token IDs rather than a 1 sentinel meaning "train here."

Python
@dataclass(frozen=True)
class Message:
    """A pre-tokenized chat message used by the teaching template."""

    role: str
    token_ids: tuple[int, ...]


@dataclass(frozen=True)
class SFTExample:
    """One causal-LM example before packing; all arrays have shape [tokens]."""

    input_ids: np.ndarray
    labels: np.ndarray
    attention_mask: np.ndarray
    position_ids: np.ndarray
    message_spans: tuple[tuple[str, int, int], ...]


def make_tokens(length: int, start: int) -> tuple[int, ...]:
    """Create deterministic content token IDs for a small offline example."""

    if length < 0:
        raise ValueError("length must be non-negative")
    return tuple(range(start, start + length))


def build_sft_example(
    messages: Sequence[Message],
    special_ids: dict[str, int],
    trainable_roles: frozenset[str] = frozenset({"assistant"}),
) -> SFTExample:
    """Encode messages and mask every non-trainable role with IGNORE_INDEX.

    Contract:
    - the sequence begins with BOS;
    - every message has four header tokens and one closing EOT token;
    - labels equal token IDs on trainable content plus its EOT;
    - headers and non-trainable roles remain IGNORE_INDEX.
    """

    role_token = {
        "system": special_ids["role_system"],
        "user": special_ids["role_user"],
        "assistant": special_ids["role_assistant"],
        "tool": special_ids["role_tool"],
    }
    input_ids: list[int] = [special_ids["bos"]]
    labels: list[int] = [IGNORE_INDEX]
    spans: list[tuple[str, int, int]] = []

    for message in messages:
        if message.role not in role_token:
            raise ValueError(f"unsupported role: {message.role}")
        if any((not isinstance(token, int)) or isinstance(token, bool) or token < 0 for token in message.token_ids):
            raise ValueError("message token_ids must be non-negative integers")

        header = [
            special_ids["start_header"],
            role_token[message.role],
            special_ids["end_header"],
            special_ids["newline"],
        ]
        body = [*message.token_ids, special_ids["eot"]]
        start = len(input_ids)
        input_ids.extend(header + body)
        labels.extend([IGNORE_INDEX] * len(header))
        if message.role in trainable_roles:
            labels.extend(body)
        else:
            labels.extend([IGNORE_INDEX] * len(body))
        spans.append((message.role, start, len(input_ids)))

    ids = np.asarray(input_ids, dtype=np.int64)
    target = np.asarray(labels, dtype=np.int64)
    return SFTExample(
        input_ids=ids,
        labels=target,
        attention_mask=np.ones(ids.shape, dtype=bool),
        position_ids=np.arange(ids.size, dtype=np.int64),
        message_spans=tuple(spans),
    )



tm = REF["template"]
special_ids = {name: int(value) for name, value in tm["special_ids"].items()}
lengths = tm["content_lengths"]
messages = [
    Message("system", make_tokens(lengths["system"], 100)),
    Message("user", make_tokens(lengths["user"], 200)),
    Message("assistant", make_tokens(lengths["assistant"], 300)),
]
example = build_sft_example(messages, special_ids)
active = example.labels != IGNORE_INDEX

assert example.input_ids.size == tm["total_tokens"] == 56
assert int(active.sum()) == tm["trainable_tokens"] == 21
assert math.isclose(float(active.mean()), tm["trainable_fraction"], abs_tol=1e-15)
assistant_role, assistant_start, assistant_end = example.message_spans[-1]
assert assistant_role == "assistant"
assert example.labels[assistant_end - 1] == special_ids["eot"]  # teach the stop token
assert np.all(example.labels[:assistant_start] == IGNORE_INDEX)
print('Single-turn dialogue: 56 tokens, 21 trainable (37.5%)')
print("trainable positions:", np.flatnonzero(active)[:4], "...", np.flatnonzero(active)[-4:])

2. What the causal loss predicts

Logits at position t predict token t+1. Many library models perform this shift inside forward. Here the caller already supplies logits for the first T−1 positions; the function pairs those rows with labels[1:] and excludes −100 targets. Do not trim the logits again inside the function. Changing logits on masked rows must not change the loss.

Python
def masked_next_token_nll(logits: np.ndarray, labels: np.ndarray) -> tuple[float, np.ndarray, np.ndarray]:
    """Return mean shifted causal NLL over targets not equal to IGNORE_INDEX.

    `logits[t]` predicts `labels[t+1]`. The implementation uses a stable log-sum-exp and
    rejects invalid active target IDs rather than silently wrapping them.
    """

    scores = np.asarray(logits, dtype=np.float64)
    target = np.asarray(labels)
    if target.dtype.kind not in "iu":
        raise ValueError("labels must be integer token IDs or IGNORE_INDEX")
    if not np.isfinite(scores).all():
        raise ValueError("logits must be finite")
    if scores.ndim != 2 or target.ndim != 1 or scores.shape[0] != target.size - 1:
        raise ValueError("logits must have shape [len(labels)-1, vocab_size]")
    shifted_target = target[1:]
    active = shifted_target != IGNORE_INDEX
    if not np.any(active):
        raise ValueError("at least one next-token target must be trainable")
    active_target = shifted_target[active]
    if np.any(active_target < 0) or np.any(active_target >= scores.shape[1]):
        raise ValueError("an active target ID is outside the vocabulary")

    active_scores = scores[active]
    row_max = active_scores.max(axis=1)
    log_partition = row_max + np.log(np.exp(active_scores - row_max[:, None]).sum(axis=1))
    token_nll = log_partition - active_scores[np.arange(active_target.size), active_target]
    return float(token_nll.mean()), token_nll, active


rng = np.random.default_rng(11)
vocab_size = int(example.input_ids.max()) + 32
logits = rng.normal(0.0, 0.5, size=(example.input_ids.size - 1, vocab_size))
shifted_target = example.labels[1:]
active_rows = shifted_target != IGNORE_INDEX
for row, token_id in zip(np.flatnonzero(active_rows), shifted_target[active_rows]):
    logits[row, token_id] += 5.0

loss, token_nll, active_rows = masked_next_token_nll(logits, example.labels)
corrupted_masked_rows = logits.copy()
corrupted_masked_rows[~active_rows] = 1e6 * rng.normal(size=corrupted_masked_rows[~active_rows].shape)
loss_after_corruption, _, _ = masked_next_token_nll(corrupted_masked_rows, example.labels)
assert math.isclose(loss, loss_after_corruption, rel_tol=0.0, abs_tol=1e-12)
assert token_nll.size == tm["trainable_tokens"]
print('Shifted NLL uses assistant targets only; masked rows do not matter', f"mean NLL={loss:.4f}")

3. Multi-turn and tool use

The assistant tool call and final assistant answer are trainable; the tool result is masked. This is not a special loss case: assign roles correctly before creating labels.

Python
conversation = [
    Message("system", make_tokens(4, 100)),
    Message("user", make_tokens(3, 200)),
    Message("assistant", make_tokens(5, 300)),       # tool call
    Message("tool", make_tokens(4, 400)),            # tool result: context in this objective
    Message("assistant", make_tokens(6, 500)),       # final answer
]
multi = build_sft_example(conversation, special_ids)
active = multi.labels != IGNORE_INDEX
assistant_expected = (5 + 1) + (6 + 1)
assert int(active.sum()) == assistant_expected
for role, start, end in multi.message_spans:
    if role == "tool":
        assert np.all(multi.labels[start:end] == IGNORE_INDEX)
    elif role == "assistant":
        # Four header tokens are masked; content and EOT are active.
        assert np.all(multi.labels[start:start + 4] == IGNORE_INDEX)
        assert np.all(multi.labels[start + 4:end] == multi.input_ids[start + 4:end])
print('Two assistant segments train; the tool response is fully masked', f"active targets={active.sum()}")

4. FFD must return assignments

Bin loads alone are insufficient: the collator must know which examples occupy each bin. The implementation preserves original indices, lengths, used tokens, and utilization. Sorting is stable: equal lengths retain increasing original indices.

Python
@dataclass(frozen=True)
class BinAssignment:
    """One deterministic FFD bin with traceable source examples."""

    example_indices: tuple[int, ...]
    lengths: tuple[int, ...]
    used: int
    capacity: int

    @property
    def utilization(self) -> float:
        """Return used/capacity for the bin."""

        return self.used / self.capacity


def first_fit_decreasing(lengths: Sequence[int], max_length: int) -> list[BinAssignment]:
    """Pack lengths with stable FFD and preserve original example indices."""

    if not isinstance(max_length, int) or isinstance(max_length, bool) or max_length <= 0:
        raise ValueError("max_length must be a positive integer")
    normalized: list[int] = []
    for value in lengths:
        if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or value <= 0:
            raise ValueError("every length must be a positive integer")
        if int(value) > max_length:
            raise ValueError(f"length {value} exceeds max_length={max_length}")
        normalized.append(int(value))

    mutable: list[dict[str, object]] = []
    for index in sorted(range(len(normalized)), key=lambda i: (-normalized[i], i)):
        length = normalized[index]
        for bucket in mutable:
            if int(bucket["used"]) + length <= max_length:
                bucket["indices"].append(index)
                bucket["lengths"].append(length)
                bucket["used"] = int(bucket["used"]) + length
                break
        else:
            mutable.append({"indices": [index], "lengths": [length], "used": length})

    return [
        BinAssignment(tuple(b["indices"]), tuple(b["lengths"]), int(b["used"]), max_length)
        for b in mutable
    ]


pk = REF["packing"]
assignments = first_fit_decreasing(pk["lengths"], pk["max_length"])
assert [list(b.example_indices) for b in assignments] == [b["indices"] for b in pk["bins"]]
assert [b.used for b in assignments] == [b["used"] for b in pk["bins"]]
assert len(assignments) == pk["n_ffd"] == 5
naive_util = sum(pk["lengths"]) / (len(pk["lengths"]) * pk["max_length"])
ffd_util = sum(pk["lengths"]) / (len(assignments) * pk["max_length"])
assert math.isclose(naive_util, pk["naive_utilization"], abs_tol=1e-15)
assert math.isclose(ffd_util, pk["ffd_utilization"], abs_tol=1e-15)
print('FFD: 24 examples → 5 bins; assignments match the frozen contract')
for i, bucket in enumerate(assignments):
    print(f"bin {i}: used={bucket.used:4d}/{bucket.capacity}, indices={bucket.example_indices}")

5. The complete packed artifact

Each bin carries:

  • input_ids and token labels;
  • an attention_mask for valid tokens;
  • position_ids reset to zero at each example boundary;
  • segment_ids for explicit isolation checks;
  • cu_seq_lens and max_seqlen, the compact varlen form;
  • original example_indices for tracing the result back to the dataset.
Python
@dataclass(frozen=True)
class PackedSequence:
    """A traceable packed SFT sequence in both dense and varlen-friendly form."""

    input_ids: np.ndarray
    labels: np.ndarray
    attention_mask: np.ndarray
    position_ids: np.ndarray
    segment_ids: np.ndarray
    cu_seq_lens: np.ndarray
    max_seqlen: int
    example_indices: tuple[int, ...]
    lengths: tuple[int, ...]


def synthetic_example(length: int, example_index: int) -> SFTExample:
    """Create a deterministic variable-length example for packing tests."""

    base = 10_000 + example_index * 3_000
    ids = np.arange(base, base + length, dtype=np.int64)
    labels = ids.copy()
    labels[: max(1, length // 3)] = IGNORE_INDEX  # a prompt-like masked prefix
    return SFTExample(
        input_ids=ids,
        labels=labels,
        attention_mask=np.ones(length, dtype=bool),
        position_ids=np.arange(length, dtype=np.int64),
        message_spans=(),
    )


def pack_examples(examples: Sequence[SFTExample], assignment: BinAssignment) -> PackedSequence:
    """Concatenate one bin while preserving per-example positions and boundaries."""

    pieces_ids: list[np.ndarray] = []
    pieces_labels: list[np.ndarray] = []
    pieces_pos: list[np.ndarray] = []
    pieces_seg: list[np.ndarray] = []
    cu = [0]
    lengths: list[int] = []

    for segment, example_index in enumerate(assignment.example_indices):
        ex = examples[example_index]
        n = int(ex.input_ids.size)
        if n != assignment.lengths[segment]:
            raise ValueError("assignment length does not match the source example")
        if not (ex.labels.shape == ex.attention_mask.shape == ex.position_ids.shape == ex.input_ids.shape):
            raise ValueError("all example arrays must have the same [tokens] shape")
        if not np.all(ex.attention_mask):
            raise ValueError("source examples must be unpadded before packing")
        if not np.array_equal(ex.position_ids, np.arange(n)):
            raise ValueError("source position_ids must start at zero and be contiguous")

        pieces_ids.append(ex.input_ids)
        target = ex.labels.copy()
        target[0] = IGNORE_INDEX  # prevent a next-token target crossing the boundary
        pieces_labels.append(target)
        pieces_pos.append(ex.position_ids)
        pieces_seg.append(np.full(n, segment, dtype=np.int32))
        lengths.append(n)
        cu.append(cu[-1] + n)

    input_ids = np.concatenate(pieces_ids)
    labels = np.concatenate(pieces_labels)
    position_ids = np.concatenate(pieces_pos)
    segment_ids = np.concatenate(pieces_seg)
    if input_ids.size != assignment.used:
        raise AssertionError("packed size must equal the bin's used count")
    return PackedSequence(
        input_ids=input_ids,
        labels=labels,
        attention_mask=np.ones(input_ids.shape, dtype=bool),
        position_ids=position_ids,
        segment_ids=segment_ids,
        cu_seq_lens=np.asarray(cu, dtype=np.int32),
        max_seqlen=max(lengths),
        example_indices=assignment.example_indices,
        lengths=tuple(lengths),
    )


source_examples = [synthetic_example(length, i) for i, length in enumerate(pk["lengths"])]
packed = [pack_examples(source_examples, assignment) for assignment in assignments]
for item in packed:
    assert item.cu_seq_lens[0] == 0 and item.cu_seq_lens[-1] == item.input_ids.size
    assert np.array_equal(np.diff(item.cu_seq_lens), np.asarray(item.lengths))
    for boundary in item.cu_seq_lens[:-1]:
        assert item.position_ids[int(boundary)] == 0
    # Preserve source labels except the explicitly ignored segment start.
    for segment, original_index in enumerate(item.example_indices):
        lo, hi = item.cu_seq_lens[segment:segment + 2]
        assert item.labels[lo] == IGNORE_INDEX
        assert np.array_equal(item.labels[lo+1:hi], source_examples[original_index].labels[1:])
print('Packed artifacts preserve labels, boundaries, position resets, and cu_seq_lens')
print("first bin cu_seq_lens:", packed[0].cu_seq_lens, "max_seqlen:", packed[0].max_seqlen)

6. Packing without isolation leaks context

A global causal mask allows a later example to read an earlier one. The correct block-diagonal causal mask requires valid tokens, equal segment IDs, and key_position ≤ query_position. The full square matrix is pedagogical; a varlen kernel usually consumes cu_seq_lens and never materializes it.

Python
def block_diagonal_causal_mask(item: PackedSequence) -> np.ndarray:
    """Return [query, key] permission: valid, same segment, and key not in future."""

    valid_q = item.attention_mask[:, None]
    valid_k = item.attention_mask[None, :]
    same_segment = item.segment_ids[:, None] == item.segment_ids[None, :]
    causal = item.position_ids[None, :] <= item.position_ids[:, None]
    return valid_q & valid_k & same_segment & causal


item = packed[0]
naive = np.tril(np.ones((item.input_ids.size, item.input_ids.size), dtype=bool))
same_segment = item.segment_ids[:, None] == item.segment_ids[None, :]
block = block_diagonal_causal_mask(item)
naive_leaks = int((naive & ~same_segment).sum())
block_leaks = int((block & ~same_segment).sum())
assert naive_leaks > 0 and block_leaks == 0
# Each segment's first query can attend only to its own first key.
for boundary in item.cu_seq_lens[:-1]:
    row = block[int(boundary)]
    assert int(row.sum()) == 1
    assert row[int(boundary)]
print('The global mask leaked across examples; the block-diagonal mask did not', f"naive cross-example cells={naive_leaks:,}; block={block_leaks}")

7. A padded batch over packed bins

When a backend does not accept varlen sequences directly, bins can be collated into a regular [batch, max_length] tensor. Padding receives attention_mask=0, labels=-100, segment_ids=-1; real tokens and position IDs remain unchanged.

Python
@dataclass(frozen=True)
class PackedBatch:
    """A padded batch of packed bins with explicit masks and trace metadata."""

    input_ids: np.ndarray
    labels: np.ndarray
    attention_mask: np.ndarray
    position_ids: np.ndarray
    segment_ids: np.ndarray
    used_lengths: np.ndarray


def collate_packed_sequences(
    sequences: Sequence[PackedSequence],
    pad_to: int,
    pad_token_id: int,
) -> PackedBatch:
    """Pad packed bins to one dense shape without making padding trainable."""

    if pad_to <= 0 or any(item.input_ids.size > pad_to for item in sequences):
        raise ValueError("pad_to must fit every packed sequence")
    batch = len(sequences)
    input_ids = np.full((batch, pad_to), pad_token_id, dtype=np.int64)
    labels = np.full((batch, pad_to), IGNORE_INDEX, dtype=np.int64)
    attention = np.zeros((batch, pad_to), dtype=bool)
    positions = np.zeros((batch, pad_to), dtype=np.int64)
    segments = np.full((batch, pad_to), -1, dtype=np.int32)
    used = np.zeros(batch, dtype=np.int32)

    for row, item in enumerate(sequences):
        n = item.input_ids.size
        used[row] = n
        input_ids[row, :n] = item.input_ids
        labels[row, :n] = item.labels
        attention[row, :n] = True
        positions[row, :n] = item.position_ids
        segments[row, :n] = item.segment_ids

    return PackedBatch(input_ids, labels, attention, positions, segments, used)


batch = collate_packed_sequences(packed, pk["max_length"], special_ids["pad"])
assert batch.input_ids.shape == (pk["n_ffd"], pk["max_length"])
assert np.array_equal(batch.attention_mask.sum(axis=1), batch.used_lengths)
assert np.all(batch.labels[~batch.attention_mask] == IGNORE_INDEX)
assert np.all(batch.segment_ids[~batch.attention_mask] == -1)
assert np.array_equal(batch.labels != IGNORE_INDEX, (batch.labels != IGNORE_INDEX) & batch.attention_mask)
print('Built a [5, 2048] batch; padding and loss masks agree', "used lengths:", batch.used_lengths.tolist())

8. The compute bridge and s1

C = 6ND is a scenario approximation here. H100 throughput and MFU are explicit inputs, not a wall-clock promise for every trainer.

Python
cp = REF["compute"]
sft_flops = 6 * cp["n_params"] * cp["sft_tokens"] * cp["sft_epochs"]
pretrain_flops = 6 * cp["n_params"] * cp["pretrain_tokens"]
sft_hours = sft_flops / (cp["accelerator_flops"] * cp["mfu"]) / 3600
assert math.isclose(sft_flops, cp["sft_flops"], rel_tol=1e-15)
assert math.isclose(pretrain_flops, cp["pretrain_flops"], rel_tol=1e-15)
assert math.isclose(sft_hours, cp["sft_hours"], rel_tol=1e-15)

s1 = REF["s1"]
s1_flops = 6 * s1["n_params"] * s1["tokens"]
s1_minutes = s1_flops / (s1["accelerators"] * cp["accelerator_flops"] * cp["mfu"]) / 60
assert math.isclose(s1_flops, s1["flops"], rel_tol=1e-15)
assert math.isclose(s1_minutes, s1["minutes"], rel_tol=1e-15)
print('The SFT and s1 compute scenarios were reproduced')
print(f"SFT approximation: {sft_flops:.3e} FLOPs, ×{pretrain_flops/sft_flops:,.0f} below pretraining, {sft_hours:.1f} h")
print(f"s1 scenario: {s1_flops:.3e} FLOPs, {s1_minutes:.1f} min on {s1['accelerators']} accelerators")

9. Verifier-gated best-of-N

For independent attempts, p_pass=0.3 gives coverage 1-(1-p_pass)^N. Retained examples per generation fall as N grows. A separate quality scenario uses correctness 0.3, sensitivity 0.9, and false acceptance 0.1, giving a different pass rate of 0.34. We also enumerate a three-candidate joint law to see why pairwise correlation cannot determine coverage.

Python
best = REF["best_of_n"]
p_pass = best["pass_probability"]
coverage = {int(n): 1 - (1 - p_pass) ** int(n) for n in best["coverage"]}
for n, value in coverage.items():
    assert math.isclose(value, best["coverage"][str(n)], abs_tol=1e-15)
ordered = sorted(coverage)
assert all(coverage[a] < coverage[b] for a, b in zip(ordered, ordered[1:]))
assert all(coverage[a] / a > coverage[b] / b for a, b in zip(ordered, ordered[1:]))
print('Coverage increases while yield/N decreases')
for n in ordered:
    print(f"N={n:2d}: coverage={coverage[n]:.4f}, yield/N={coverage[n]/n:.4f}")

# A separate scenario: correctness and acceptance are different events.
from fractions import Fraction
from itertools import product
pc, sensitivity, false_accept = Fraction(3,10), Fraction(9,10), Fraction(1,10)
pass_rate = pc*sensitivity + (1-pc)*false_accept
precision = pc*sensitivity/pass_rate
assert pass_rate == Fraction(17,50) and precision == Fraction(27,34)
print("separate verifier scenario:", float(pass_rate), float(precision))

# Enumerate the symmetric law of three success indicators.
mass_by_count = [Fraction(1,10),Fraction(1,2),Fraction(1,5),Fraction(1,5)]
joint = {bits: mass_by_count[sum(bits)] / math.comb(3,sum(bits))
         for bits in product((0,1),repeat=3)}
assert sum(joint.values()) == 1
marginal = sum(prob*bits[0] for bits,prob in joint.items())
pair = sum(prob*bits[0]*bits[1] for bits,prob in joint.items())
union = sum(prob for bits,prob in joint.items() if any(bits))
independent = 1-(1-marginal)**3
assert marginal == Fraction(1,2) and pair == Fraction(4,15)
assert pair > marginal**2 and union == Fraction(9,10) > independent
perfectly_shared = marginal  # all three indicators are identical
print("coverage: joint / independent / perfectly shared:",
      float(union), float(independent), float(perfectly_shared))

10. LoRA and QLoRA: strict units

Division by 2**20 and 2**30 means MiB and GiB. The four-bit base below is only the raw parameter payload. The QLoRA parameter-state floor excludes activations, temporary dequantization buffers, allocator overhead, and sharding, so it is not total card occupancy.

Python
def lora_parameter_count(
    projections: Iterable[tuple[int, int]],
    rank: int,
    layers: int,
) -> tuple[int, int]:
    """Return (per_layer, total) for rank*(d_in+d_out) per targeted matrix."""

    if not isinstance(rank, int) or rank <= 0 or not isinstance(layers, int) or layers <= 0:
        raise ValueError("rank and layers must be positive integers")
    per_layer = 0
    for d_in, d_out in projections:
        if d_in <= 0 or d_out <= 0:
            raise ValueError("projection dimensions must be positive")
        per_layer += rank * (d_in + d_out)
    return per_layer, per_layer * layers


lo = REF["lora"]
d, d_kv, rank, layers = lo["d_model"], lo["d_kv"], lo["rank"], lo["layers"]
per_layer, params = lora_parameter_count(
    [(d, d), (d, d_kv), (d, d_kv), (d, d)],
    rank=rank,
    layers=layers,
)
assert per_layer == lo["per_layer"] and params == lo["params"]
full_gib = 16 * cp["n_params"] / 2**30
states_mib = 16 * params / 2**20
adapter_mib = 2 * params / 2**20
base_payload_gib = 0.5 * cp["n_params"] / 2**30
qlora_floor_gib = base_payload_gib + states_mib / 1024
assert math.isclose(full_gib, lo["full_training_states_gib"], abs_tol=1e-12)
assert math.isclose(states_mib, lo["lora_training_states_mib"], abs_tol=1e-12)
assert math.isclose(adapter_mib, lo["lora_adapter_bf16_mib"], abs_tol=1e-12)
assert math.isclose(qlora_floor_gib, lo["qlora_parameter_floor_gib"], abs_tol=1e-12)
print('LoRA parameters and the GiB/MiB floor estimate were reproduced')
print(f"LoRA: {params/1e6:.2f}M params ({params/cp['n_params']*100:.3f}%)")
print(f"training states: full {full_gib:.1f} GiB → LoRA {states_mib:.0f} MiB (×{full_gib*1024/states_mib:.0f})")
print(f"4-bit payload {base_payload_gib:.2f} GiB + adapter states {states_mib:.0f} MiB (includes {adapter_mib:.0f} MiB weights)")
print(f"QLoRA parameter-state floor: {qlora_floor_gib:.2f} GiB (not total card memory)")

11. From prepared data to an actual training run

A tiny causal model learns to reverse a pair of symbols. Before optimization, separate and packed dialogues must agree on logits, summed loss and gradients. We then train, generate on held-out pairs, and round-trip the weights in memory. Each answer has three targets: two symbols and an end marker. Diagonal pairs are held out, testing unseen combinations of familiar symbols rather than language generalization. The code reports accuracy without assuming perfection.

Python
"""CPU teaching experiment: serialize, pack, train, generate, and reload.

Synthetic symbol reversal exposes the data contract; it is not evidence about LLM quality.
"""
import copy
import importlib.util
import io
import subprocess
import sys

if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
torch.manual_seed(130)
VOCAB, WIDTH, BOS, USER, ASSISTANT, END = 16, 32, 0, 1, 2, 3


def dialogue(a, b):
    # Reverse two symbols. Role headers are inputs; answer and END are targets.
    ids = torch.tensor([BOS, USER, a, b, END, ASSISTANT, b, a, END])
    labels = ids.clone()
    labels[:6] = -100
    return ids, labels


def pack(records):
    ids, labels, positions, segments = [], [], [], []
    for i, (x, y) in enumerate(records):
        target = y.clone()
        target[0] = -100  # No prediction across an independent example boundary.
        ids.append(x); labels.append(target)
        positions.append(torch.arange(len(x)))
        segments.append(torch.full_like(x, i))
    return tuple(torch.cat(xs) for xs in (ids, labels, positions, segments))


class TinyCausalLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.token = nn.Embedding(VOCAB, WIDTH)
        self.position = nn.Embedding(16, WIDTH)
        self.qkv = nn.Linear(WIDTH, 3 * WIDTH)
        self.proj = nn.Linear(WIDTH, WIDTH)
        self.norm = nn.LayerNorm(WIDTH)
        self.ff = nn.Sequential(nn.Linear(WIDTH, 64), nn.GELU(), nn.Linear(64, WIDTH))
        self.head = nn.Linear(WIDTH, VOCAB)

    def forward(self, ids, positions, segments):
        h = self.token(ids) + self.position(positions)
        q, k, v = self.qkv(h).chunk(3, dim=-1)
        idx = torch.arange(len(ids))
        allowed = (segments[:, None] == segments[None, :]) & (idx[:, None] >= idx[None, :])
        scores = q @ k.T / WIDTH**.5
        h = self.norm(h + self.proj(scores.masked_fill(~allowed, -torch.inf).softmax(-1) @ v))
        return self.head(h + self.ff(h))


def loss_sum(model, batch):
    x, y, pos, seg = batch
    logits = model(x, pos, seg)
    return F.cross_entropy(logits[:-1], y[1:], ignore_index=-100, reduction="sum")


train = [dialogue(a, b) for a in range(4, 12) for b in range(4, 12) if a != b]
heldout = [dialogue(a, a) for a in range(4, 12)]
model = TinyCausalLM()

# Packing preserves logits, summed loss and gradients when positions and boundaries match.
records = train[:3]
packed_batch = pack(records)
with torch.no_grad():
    packed_logits = model(packed_batch[0], packed_batch[2], packed_batch[3])
    individual = torch.cat([model(b[0], b[2], b[3]) for b in map(lambda r: pack([r]), records)])
torch.testing.assert_close(packed_logits, individual, rtol=1e-5, atol=1e-6)
packed_loss = loss_sum(model, packed_batch)
separate_loss = sum(loss_sum(model, pack([r])) for r in records)
torch.testing.assert_close(packed_loss, separate_loss, rtol=1e-6, atol=1e-6)
gp = torch.autograd.grad(packed_loss, tuple(model.parameters()))
gs = torch.autograd.grad(separate_loss, tuple(model.parameters()))
for left, right in zip(gp, gs):
    torch.testing.assert_close(left, right, rtol=1e-4, atol=3e-6)

def score(model, records):
    with torch.no_grad():
        return sum(float(loss_sum(model, pack([r]))) for r in records) / (3 * len(records))

before = {"train": score(model, train), "heldout": score(model, heldout)}
optimizer = torch.optim.AdamW(model.parameters(), lr=.006)
history = []
for step in range(240):
    # Eight independent examples per packed row; all 56 rotate through training.
    subset = [train[(step * 8 + j) % len(train)] for j in range(8)]
    batch = pack(subset)
    loss = loss_sum(model, batch) / int((batch[1][1:] != -100).sum())
    optimizer.zero_grad(); loss.backward(); optimizer.step()
    history.append(float(loss.detach()))
after = {"train": score(model, train), "heldout": score(model, heldout)}
assert after["train"] < before["train"] / 5

@torch.no_grad()
def generate(model, prompt, max_new=4):
    sequence = prompt.clone()
    output = []
    for _ in range(max_new):
        logits = model(sequence, torch.arange(len(sequence)), torch.zeros_like(sequence))
        token = int(logits[-1].argmax())
        output.append(token)
        sequence = torch.cat([sequence, torch.tensor([token])])
        if token == END:
            break
    return output

# Generate once per held-out example, then derive every metric from stored outputs.
generated = [generate(model, ids[:6]) for ids, _ in heldout]
exact = sum(pred == ids[6:].tolist() for pred, (ids, _) in zip(generated, heldout)) / len(heldout)
stop_rate = sum(bool(pred) and pred[-1] == END for pred in generated) / len(generated)
buffer = io.BytesIO()
torch.save(model.state_dict(), buffer); buffer.seek(0)
restored = TinyCausalLM()
restored.load_state_dict(torch.load(buffer, weights_only=True))
with torch.no_grad():
    probe = pack(heldout[:2])
    torch.testing.assert_close(model(probe[0], probe[2], probe[3]), restored(probe[0], probe[2], probe[3]), rtol=0, atol=0)
print("token NLL before / after:", before, after)
print("held-out exact / stop rate:", exact, stop_rate, "stored outputs:", generated)
print("packed logits, summed loss, gradients and checkpoint reload: PASS")

PyTorch bridge · B20 — Assistant-only labels

This cell embeds the canonical NumPy→PyTorch bridge and runs its parity check. It is self-contained and does not import repository files.

Use nonnegative np.int64 / torch.long IDs and a boolean mask of the same shape; Torch inputs share a device. The result stores labels at their own positions. Align logits[..., :-1, :] with labels[..., 1:] exactly once before cross-entropy. A model that shifts internally must receive unshifted arrays. The loss function checks that active targets exist and belong to its vocabulary.

Python
import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_assistant_labels(input_ids: np.ndarray, assistant_mask: np.ndarray, *, ignore_index: int = -100) -> np.ndarray:
    ids, mask = np.asarray(input_ids), np.asarray(assistant_mask)
    if ids.dtype != np.int64 or mask.dtype != np.bool_:
        raise TypeError("use int64 token IDs and a boolean assistant mask")
    if np.any(ids < 0):
        raise ValueError("input token IDs must be nonnegative")
    if ids.shape != mask.shape:
        raise ValueError("input_ids and assistant_mask must have identical shapes")
    return np.where(mask, ids, ignore_index).astype(np.int64)

def torch_assistant_labels(input_ids: torch.Tensor, assistant_mask: torch.Tensor, *, ignore_index: int = -100) -> torch.Tensor:
    if input_ids.dtype != torch.long or assistant_mask.dtype != torch.bool:
        raise TypeError("use torch.long token IDs and a boolean assistant mask")
    if input_ids.device != assistant_mask.device:
        raise ValueError("input IDs and mask must share a device")
    if torch.any(input_ids < 0):
        raise ValueError("input token IDs must be nonnegative")
    if input_ids.shape != assistant_mask.shape:
        raise ValueError("input_ids and assistant_mask must have identical shapes")
    labels = input_ids.clone()
    return labels.masked_fill(~assistant_mask.to(torch.bool), ignore_index)

ids_np = np.array([10,11,12,13,14,15]); mask_np = np.array([0,0,0,1,1,1], dtype=bool)
np_labels = numpy_assistant_labels(ids_np, mask_np); t_labels = torch_assistant_labels(torch.tensor(ids_np), torch.tensor(mask_np))
np.testing.assert_array_equal(np_labels, t_labels.numpy()); assert np_labels.tolist() == [-100,-100,-100,13,14,15]
# Labels are stored at their own positions. Shift them once for explicit CE.
scores = torch.zeros((len(ids_np), 20), dtype=torch.float64, requires_grad=True)
loss = torch.nn.functional.cross_entropy(scores[:-1], t_labels[1:], ignore_index=-100)
torch.testing.assert_close(loss, torch.log(torch.tensor(20.,dtype=torch.float64)))
loss.backward()
assert torch.count_nonzero(scores.grad[-1]) == 0
assert torch.count_nonzero(scores.grad[:2]) == 0
for ids_bad, mask_bad in [(ids_np.astype(float),mask_np),
                           (ids_np,mask_np.astype(int)),
                           (np.array([-1,1]),np.array([False,True]))]:
    for run in (lambda: numpy_assistant_labels(ids_bad,mask_bad),
                lambda: torch_assistant_labels(torch.tensor(ids_bad),torch.tensor(mask_bad))):
        try:
            run()
        except (TypeError,ValueError):
            pass
        else:
            raise AssertionError("invalid label input was accepted")
print("B20 assistant labels: PASS")

PyTorch bridge · B21 — Packed-segment isolation

This cell embeds the canonical NumPy→PyTorch bridge and runs its parity check. It is self-contained and does not import repository files.

Python
import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_segment_position_ids(segment_ids: np.ndarray, *, padding_id: int = -1) -> np.ndarray:
    """Reset positions at every contiguous segment boundary; padding receives position zero."""
    segments = np.asarray(segment_ids)
    if segments.ndim not in (1, 2):
        raise ValueError("segment_ids must be [T] or [B,T]")
    batch = segments[None, :] if segments.ndim == 1 else segments
    positions = np.zeros_like(batch, dtype=np.int64)
    for row in range(batch.shape[0]):
        for col in range(1, batch.shape[1]):
            current, previous = batch[row, col], batch[row, col - 1]
            if current != padding_id and current == previous:
                positions[row, col] = positions[row, col - 1] + 1
    return positions[0] if segments.ndim == 1 else positions

def torch_segment_position_ids(segment_ids: torch.Tensor, *, padding_id: int = -1) -> torch.Tensor:
    """Pure-tensor counterpart of :func:`numpy_segment_position_ids`."""
    if segment_ids.ndim not in (1, 2):
        raise ValueError("segment_ids must be [T] or [B,T]")
    batch = segment_ids.unsqueeze(0) if segment_ids.ndim == 1 else segment_ids
    positions = torch.zeros_like(batch, dtype=torch.long)
    for col in range(1, batch.shape[-1]):
        continues = (batch[:, col] != padding_id) & (batch[:, col] == batch[:, col - 1])
        positions[:, col] = torch.where(continues, positions[:, col - 1] + 1, 0)
    return positions[0] if segment_ids.ndim == 1 else positions

def numpy_block_causal_mask(segment_ids: np.ndarray, *, padding_id: int = -1) -> np.ndarray:
    """Bool allowed-mask for causal attention inside each contiguous non-padding segment."""
    segments = np.asarray(segment_ids)
    if segments.ndim not in (1, 2):
        raise ValueError("segment_ids must be [T] or [B,T]")
    batch = segments[None, :] if segments.ndim == 1 else segments
    t = batch.shape[-1]
    boundaries = np.ones_like(batch, dtype=bool)
    if t > 1:
        boundaries[:, 1:] = batch[:, 1:] != batch[:, :-1]
    run_ids = np.cumsum(boundaries, axis=-1)
    same = run_ids[:, :, None] == run_ids[:, None, :]
    valid = (batch[:, :, None] != padding_id) & (batch[:, None, :] != padding_id)
    causal = np.arange(t)[None, None, :] <= np.arange(t)[None, :, None]
    out = same & valid & causal
    return out[0] if segments.ndim == 1 else out

def torch_block_causal_mask(segment_ids: torch.Tensor, *, padding_id: int = -1) -> torch.Tensor:
    if segment_ids.ndim not in (1, 2):
        raise ValueError("segment_ids must be [T] or [B,T]")
    batch = segment_ids.unsqueeze(0) if segment_ids.ndim == 1 else segment_ids
    t = batch.shape[-1]
    boundaries = torch.ones_like(batch, dtype=torch.bool)
    if t > 1:
        boundaries[:, 1:] = batch[:, 1:] != batch[:, :-1]
    run_ids = boundaries.to(torch.long).cumsum(dim=-1)
    same = run_ids[:, :, None] == run_ids[:, None, :]
    valid = (batch[:, :, None] != padding_id) & (batch[:, None, :] != padding_id)
    idx = torch.arange(t, device=batch.device)
    causal = idx[None, None, :] <= idx[None, :, None]
    out = same & valid & causal
    return out[0] if segment_ids.ndim == 1 else out

segments_np = np.array([0,0,0,1,1,1,-1]); np_pos = numpy_segment_position_ids(segments_np); np_mask = numpy_block_causal_mask(segments_np)
t_segments = torch.tensor(segments_np); t_pos = torch_segment_position_ids(t_segments); t_mask = torch_block_causal_mask(t_segments)
np.testing.assert_array_equal(np_pos, t_pos.numpy()); np.testing.assert_array_equal(np_mask, t_mask.numpy())
rng = np.random.default_rng(21); x = torch.tensor(rng.standard_normal((1,1,7,4)), dtype=torch.float64); mask = t_mask[None,None]
out_a = torch.nn.functional.scaled_dot_product_attention(x, x, x, attn_mask=mask)
x_changed = x.clone(); x_changed[..., :3, :] += 1000
out_b = torch.nn.functional.scaled_dot_product_attention(x_changed, x_changed, x_changed, attn_mask=mask)
torch.testing.assert_close(out_a[..., 3:6, :], out_b[..., 3:6, :], rtol=0, atol=0)
repeated = np.array([0,0,1,1,0,0])
np.testing.assert_array_equal(numpy_segment_position_ids(repeated), [0,1,0,1,0,1])
repeat_mask = numpy_block_causal_mask(repeated)
assert not repeat_mask[4,0] and repeat_mask[5,4]
print("B21 packed isolation: PASS")

Exercises

  1. [auto] Add an analysis channel, make it trainable only behind an explicit flag, and verify the active-target count.

  2. [auto] Replace FFD with best-fit decreasing and compare bin counts over 100 log-normal length sets.

  3. [auto] Corrupt one position_id after a segment boundary and write an assertion that localizes the error to the original example_index.

  4. [auto] Implement a mode without a materialized block mask: verify isolation from cu_seq_lens and a random sample of query/key pairs.

  5. [project] Connect a real tokenizer through apply_chat_template, compare its labels with this teaching contract, and feed the packed batch to a small causal LM.

The mandatory path is CPU-safe and repository-independent. It is a complete teaching path through SFT tensor preparation, not a production training system.

Module 13 project-exercise demo — From dialogues to training; full updates and LoRA

Project demonstration

Download .ipynb module_13_project_demo.ipynb · 14.3 KB

Show notebook text

Module 13 project-exercise demo — From dialogues to training; full updates and LoRA

Implementation demo, not a completed full-scale project. Demo 1 serializes, packs, compares gradients, trains, generates and reloads a small CPU model. Demo 2 isolates adapter rank with a shared starting matrix. All data is synthetic; no network or keys are required. All inputs are small, deterministic, offline, and CPU-friendly. The goal is to validate interfaces, metrics, invariants, and report structure before spending GPU or API budget.

Python
from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

Demo 1. From dialogue labels to learning

The first experiment follows labels all the way into learning. Each answer has three targets: the reversed symbol pair and END. Before training, separate and packed examples must agree on logits, summed loss, and gradients. Fifty-six unequal pairs train the network from scratch; eight equal pairs are held out. This tests the data contract and transfer to unseen combinations, rather than adaptation of a pretrained language model.

Python
"""CPU teaching experiment: serialize, pack, train, generate, and reload.

Synthetic symbol reversal exposes the data contract; it is not evidence about LLM quality.
"""
import copy
import importlib.util
import io
import subprocess
import sys

if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
torch.manual_seed(130)
VOCAB, WIDTH, BOS, USER, ASSISTANT, END = 16, 32, 0, 1, 2, 3


def dialogue(a, b):
    # Reverse two symbols. Role headers are inputs; answer and END are targets.
    ids = torch.tensor([BOS, USER, a, b, END, ASSISTANT, b, a, END])
    labels = ids.clone()
    labels[:6] = -100
    return ids, labels


def pack(records):
    ids, labels, positions, segments = [], [], [], []
    for i, (x, y) in enumerate(records):
        target = y.clone()
        target[0] = -100  # No prediction across an independent example boundary.
        ids.append(x); labels.append(target)
        positions.append(torch.arange(len(x)))
        segments.append(torch.full_like(x, i))
    return tuple(torch.cat(xs) for xs in (ids, labels, positions, segments))


class TinyCausalLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.token = nn.Embedding(VOCAB, WIDTH)
        self.position = nn.Embedding(16, WIDTH)
        self.qkv = nn.Linear(WIDTH, 3 * WIDTH)
        self.proj = nn.Linear(WIDTH, WIDTH)
        self.norm = nn.LayerNorm(WIDTH)
        self.ff = nn.Sequential(nn.Linear(WIDTH, 64), nn.GELU(), nn.Linear(64, WIDTH))
        self.head = nn.Linear(WIDTH, VOCAB)

    def forward(self, ids, positions, segments):
        h = self.token(ids) + self.position(positions)
        q, k, v = self.qkv(h).chunk(3, dim=-1)
        idx = torch.arange(len(ids))
        allowed = (segments[:, None] == segments[None, :]) & (idx[:, None] >= idx[None, :])
        scores = q @ k.T / WIDTH**.5
        h = self.norm(h + self.proj(scores.masked_fill(~allowed, -torch.inf).softmax(-1) @ v))
        return self.head(h + self.ff(h))


def loss_sum(model, batch):
    x, y, pos, seg = batch
    logits = model(x, pos, seg)
    return F.cross_entropy(logits[:-1], y[1:], ignore_index=-100, reduction="sum")


train = [dialogue(a, b) for a in range(4, 12) for b in range(4, 12) if a != b]
heldout = [dialogue(a, a) for a in range(4, 12)]
model = TinyCausalLM()

# Packing preserves logits, summed loss and gradients when positions and boundaries match.
records = train[:3]
packed_batch = pack(records)
with torch.no_grad():
    packed_logits = model(packed_batch[0], packed_batch[2], packed_batch[3])
    individual = torch.cat([model(b[0], b[2], b[3]) for b in map(lambda r: pack([r]), records)])
torch.testing.assert_close(packed_logits, individual, rtol=1e-5, atol=1e-6)
packed_loss = loss_sum(model, packed_batch)
separate_loss = sum(loss_sum(model, pack([r])) for r in records)
torch.testing.assert_close(packed_loss, separate_loss, rtol=1e-6, atol=1e-6)
gp = torch.autograd.grad(packed_loss, tuple(model.parameters()))
gs = torch.autograd.grad(separate_loss, tuple(model.parameters()))
for left, right in zip(gp, gs):
    torch.testing.assert_close(left, right, rtol=1e-4, atol=3e-6)

def score(model, records):
    with torch.no_grad():
        return sum(float(loss_sum(model, pack([r]))) for r in records) / (3 * len(records))

before = {"train": score(model, train), "heldout": score(model, heldout)}
optimizer = torch.optim.AdamW(model.parameters(), lr=.006)
history = []
for step in range(240):
    # Eight independent examples per packed row; all 56 rotate through training.
    subset = [train[(step * 8 + j) % len(train)] for j in range(8)]
    batch = pack(subset)
    loss = loss_sum(model, batch) / int((batch[1][1:] != -100).sum())
    optimizer.zero_grad(); loss.backward(); optimizer.step()
    history.append(float(loss.detach()))
after = {"train": score(model, train), "heldout": score(model, heldout)}
assert after["train"] < before["train"] / 5

@torch.no_grad()
def generate(model, prompt, max_new=4):
    sequence = prompt.clone()
    output = []
    for _ in range(max_new):
        logits = model(sequence, torch.arange(len(sequence)), torch.zeros_like(sequence))
        token = int(logits[-1].argmax())
        output.append(token)
        sequence = torch.cat([sequence, torch.tensor([token])])
        if token == END:
            break
    return output

# Generate once per held-out example, then derive every metric from stored outputs.
generated = [generate(model, ids[:6]) for ids, _ in heldout]
exact = sum(pred == ids[6:].tolist() for pred, (ids, _) in zip(generated, heldout)) / len(heldout)
stop_rate = sum(bool(pred) and pred[-1] == END for pred in generated) / len(generated)
buffer = io.BytesIO()
torch.save(model.state_dict(), buffer); buffer.seek(0)
restored = TinyCausalLM()
restored.load_state_dict(torch.load(buffer, weights_only=True))
with torch.no_grad():
    probe = pack(heldout[:2])
    torch.testing.assert_close(model(probe[0], probe[2], probe[3]), restored(probe[0], probe[2], probe[3]), rtol=0, atol=0)
print("token NLL before / after:", before, after)
print("held-out exact / stop rate:", exact, stop_rate, "stored outputs:", generated)
print("packed logits, summed loss, gradients and checkpoint reload: PASS")

Agreement before training checks more than attention geometry: it also tests target alignment and gradient weighting. Afterward, distinguish answer exact match from stopping rate; a well-terminated answer may still be wrong. Every generation metric uses the same stored outputs. Reloading weights preserves predictions on a common probe, but none of these checks establishes natural-language generalization.

Demo 2. Adapter rank and held-out error

The second task has a known rank-four change to one matrix. Full updates and adapters use the same starting matrix and data. The code multiplies row vectors, so its factor order and shapes are transposed relative to the lecture’s column-vector notation. Set alpha/r to one. Randomly initializing one factor and zeroing the other preserves the initial prediction while leaving a route for learning.

Python
# Same base and data for full updates and LoRA; the target change has rank four.
torch.manual_seed(131)
x_train, x_test = torch.randn(256, 32), torch.randn(128, 32)
base = torch.randn(32, 32) / 32**.5
delta = torch.randn(32, 4) @ torch.randn(4, 32) / 8
y_train, y_test = x_train @ (base + delta), x_test @ (base + delta)
results = []
for rank in (None, 2, 4, 8):
    torch.manual_seed(132)
    frozen = base.clone()
    if rank is None:
        update = torch.zeros_like(base, requires_grad=True)
        params = [update]
        predict = lambda x: x @ (frozen + update)
    else:
        a = (torch.randn(32, rank) * .03).requires_grad_()
        b = torch.zeros(rank, 32, requires_grad=True)
        params = [a, b]
        predict = lambda x: x @ frozen + (x @ a) @ b  # alpha/r = 1
    optimizer = torch.optim.Adam(params, lr=.05)
    for _ in range(240):
        loss = (predict(x_train) - y_train).square().mean()
        optimizer.zero_grad(); loss.backward(); optimizer.step()
    with torch.no_grad():
        train_mse = float((predict(x_train)-y_train).square().mean())
        test_predictions = predict(x_test)
        test_mse = float((test_predictions-y_test).square().mean())
    assert torch.equal(frozen, base)
    assert test_predictions.shape == y_test.shape
    assert torch.isfinite(test_predictions).all()
    assert torch.isfinite(torch.tensor([train_mse, test_mse])).all()
    results.append(dict(rank=rank, trainable=sum(p.numel() for p in params),
                        train_mse=train_mse, test_mse=test_mse))
for result in results: print(result)
comparison = results[2]['test_mse'] < results[1]['test_mse']
label = ('Ошибка ранга 4 меньше ошибки ранга 2 в этом запуске:' if DEMO_LANGUAGE == 'ru'
         else 'Rank 4 has lower held-out error than rank 2 in this run:')
print(label, comparison)
print('This is a matrix adaptation experiment, not a measured GPU-memory or LLM-quality comparison.')

The full update has 1,024 parameters; ranks 2, 4, and 8 use 128, 256, and 512. Rank two cannot exactly express the specified rank-four change. The larger adapters can represent it, although a finite optimization run need not reach the exact solution. Read held-out MSE as the result of this run. Vary the training-set size and the true update rank, then compare the errors. With few examples, the higher-rank adapter may generalize less well. That is a valid observation; the checks enforce shapes, finite computations, and an unchanged base matrix. The example measures neither runtime nor memory, and its deliberately low-rank task cannot establish that adapters always match full fine-tuning.

Scaling to the full project

For the full project, pin the official tokenizer and template, split task families before synthesis, and preserve token weighting during gradient accumulation. Compare full tuning and LoRA from the same base using held-out tasks and measured runtime and memory.

Minimum report

  1. Data, model, and code revisions, plus the random seed.

  2. The primary metric, plus quality, cost, latency, and memory where applicable.

  3. A baseline, a matched budget, and at least one negative control.

  4. Failed runs and the limits of the supported conclusions.