"""Build the course's small instruction dataset and split it into train, validation and test. Purpose: write about two hundred short instruction examples in the course's own two-line house style, then split them with a fixed seed into the files the two training paths expect: TRL's conversational prompt-completion shape for Tracks S, X and N, and mlx-lm's completions shape for Track M. Platform: all (standard library only) Minimum memory: 8 GB Assumes: nothing beyond Python 3.10 or newer; writes into the directory given by --out-dir, which is created if it does not exist. Usage: python make-dataset.py --out-dir . --seed 0 Writes sample-instructions.jsonl (the whole dataset, one example per line), data/{train,valid,test}.jsonl (TRL conversational prompt-completion) data-mlx/{train,valid,test}.jsonl (mlx-lm completions) Every answer is two lines: Answer: Because: The content is deliberately narrow. The arithmetic items are generated from the formulas taught in Part 4 and Part 11, so every answer in this file is correct by construction rather than by an author's memory; the written items restate the course's own definitions. Nothing here is scraped, and nothing here is anyone else's copyright, so the dataset carries the course's content licence. """ from __future__ import annotations import argparse import hashlib import json import random from pathlib import Path GB = 1e9 GIB = 1024 ** 3 # --------------------------------------------------------------------------- # Written examples: the course's own definitions and rules of thumb. # --------------------------------------------------------------------------- WRITTEN: list[tuple[str, str, str]] = [ ("What is a parameter?", "A parameter is one of the numbers a model is made of.", "A model's size in billions, such as 8B, is a count of its parameters."), ("What is the difference between a weight and a parameter?", "A weight is the value of a parameter, and the weights are the checkpoint.", "Parameter names the slot; weight names the number currently in it."), ("What is the difference between a context window and a context length?", "The context window is what the model was trained to handle; the context length is what you asked the engine to allocate for this run.", "One is a property of the model, the other a setting of the run."), ("What is prefill?", "Prefill is the stage that reads the prompt.", "It processes many tokens at once, so it is compute-bound and measured in prompt tokens per second."), ("What is decode?", "Decode is the stage that writes the answer, one token at a time.", "Each token reads every weight once, so it is bandwidth-bound."), ("Is decode limited by compute or by memory bandwidth?", "By memory bandwidth.", "Every generated token reads the whole model from memory and does only one multiply-add per parameter."), ("Is prefill limited by compute or by memory bandwidth?", "By compute.", "The weights are read once and reused across all the prompt's tokens, so the arithmetic dominates."), ("What do total and active parameters mean for a mixture-of-experts model?", "Total parameters is what must sit in memory; active parameters is what is computed with per token.", "A mixture-of-experts model routes each token to a few experts, so it is large to hold and cheap to run."), ("What is quantisation?", "Quantisation is rounding the weights to fewer bits.", "It reduces bytes per parameter, which reduces both the memory the model needs and the bytes read per token."), ("What is the difference between quantisation and compression?", "Quantisation rounds the weights to fewer bits; compression removes parameters or layers.", "One keeps every parameter at lower precision, the other has fewer parameters afterwards."), ("What does open weight mean?", "Open weight means the weights are downloadable.", "It does not by itself mean the training data or the training code are available."), ("What is a fine-tune?", "A fine-tune adjusts an existing model on new data.", "Training from scratch instead makes a model from raw text."), ("What is tensor parallelism?", "Tensor parallelism splits every layer across devices.", "The devices exchange activations inside each layer, so it needs a fast link between them."), ("What is pipeline parallelism?", "Pipeline parallelism gives each device some of the layers.", "Only the boundary activations cross the link, so it tolerates a slow one."), ("What is an agent, as this course uses the word?", "An agent is a loop in which the model decides the next step.", "A workflow is a fixed sequence that the program decides instead."), ("What is a memory floor?", "A memory floor is the smallest memory that can follow a page's primary path.", "It is stated per lab so a reader knows before starting whether their machine qualifies."), ("Why does training need more memory than inference for the same model?", "Training holds gradients, optimiser states and activations as well as the weights.", "Inference holds only the weights and the key-value cache."), ("What is gradient checkpointing for?", "Gradient checkpointing trades compute for memory during training.", "Intermediate activations are recomputed in the backward pass instead of being kept from the forward pass."), ("What is gradient accumulation for?", "Gradient accumulation gives a large effective batch on a machine that cannot hold one.", "Several small batches are summed before the optimiser takes a step."), ("What does LoRA train?", "LoRA trains a pair of small low-rank matrices beside each frozen weight matrix.", "The base model's weights are not updated, so only the adapters need gradients and optimiser states."), ("What does the rank of a LoRA adapter control?", "The rank sets the size of the two update matrices and therefore the number of trainable parameters.", "A lower rank means smaller matrices and fewer parameters to train."), ("What is QLoRA?", "QLoRA trains LoRA adapters on top of a base model quantised to four bits.", "Gradients are backpropagated through the frozen quantised weights into the adapters."), ("Can a LoRA adapter be merged into the base model?", "Yes, the adapter can be merged into the base weights to produce a standalone model.", "After merging there is one set of weights again, so the adapter adds no inference cost."), ("Why does a chat template matter when fine-tuning?", "The template decides the exact control tokens the model sees around each turn.", "Training with one template and serving with another gives the model a format it was never trained on."), ("What is completion-only loss?", "Completion-only loss computes the training loss on the answer tokens and ignores the prompt tokens.", "The model is being taught what to answer, not how to restate the question."), ("What is packing in a training run?", "Packing groups several examples into one fixed-length sequence.", "It reduces the padding that would otherwise be computed and thrown away."), ("Why must a validation split never be trained on?", "Because a split that has been trained on cannot detect memorisation.", "Validation loss is only informative about unseen data if the data really was unseen."), ("What is leakage between a training set and an evaluation set?", "Leakage is any overlap that lets a model score well by recall rather than by generalisation.", "Near-duplicates leak as surely as exact duplicates do."), ("Why record a seed with every training run?", "Because without the seed the run cannot be repeated even on the same machine.", "The seed fixes the shuffling, the dropout and the initialisation."), ("Does a fixed seed guarantee identical results on two different machines?", "No, it does not.", "Results are not reproducible across releases, commits or platforms unless deterministic algorithms are also requested."), ("What should a run log record besides the final loss?", "The model, the dataset and its hash, the hyperparameters, the seed, the hardware, the versions and the date.", "A loss without its settings cannot be compared with anything later."), ("Why hash the dataset file in the run log?", "Because the hash ties a result to the exact bytes that produced it.", "A file with the same name can be edited between two runs without anyone noticing."), ("What does an evaluation loss that rises while training loss falls indicate?", "It indicates overfitting.", "The model is fitting the training examples rather than the pattern they share."), ("Which checkpoint should be kept when validation loss turns?", "The checkpoint from the epoch with the lowest validation loss.", "Later epochs fit noise, and the validation curve is what identifies the turn."), ("What is a base model?", "A base model has been pretrained but not instruction-tuned.", "It continues text rather than answering a question in a chat format."), ("Why is BF16 preferred over FP16 for training?", "Because BF16 keeps the range of FP32.", "Small gradients underflow to zero in FP16 unless the loss is scaled first."), ("How many bytes per parameter does BF16 use?", "Two bytes per parameter.", "Sixteen bits is two bytes, which is the baseline every quantised format is compared against."), ("Why is Q8_0 about 1.06 bytes per parameter rather than exactly 1.0?", "Because each block of 32 weights carries a scale as well.", "32 eight-bit weights plus a 16-bit scale is 272 bits for 32 weights, which is 8.5 bits each."), ("Which head count belongs in the key-value cache formula?", "The key-value head count, not the query head count.", "Grouped-query attention lets several query heads share one key-value head, so the two numbers differ."), ("Does the key-value cache matter during training?", "No, it is switched off during training.", "Training computes a forward and a backward pass over a fixed sequence and generates nothing."), ("What is unified memory, on the machines this course uses?", "Unified memory is one pool shared by the CPU and the GPU with no copy between them.", "It trades bandwidth for capacity compared with a discrete card's own memory."), ("Why is a model that fits exactly into memory a bad idea?", "Because a machine at its memory ceiling swaps or compresses instead of failing cleanly.", "The symptom is a model that loads and runs absurdly slowly rather than an error."), ] # --------------------------------------------------------------------------- # Generated examples: arithmetic that is correct by construction. # --------------------------------------------------------------------------- FORMATS = [("BF16", 2.0), ("FP8", 1.0), ("Q8_0", 1.06), ("Q6_K", 0.82), ("Q4_K_M", 0.61)] PARAMS_B = [0.6, 1.7, 3.0, 4.0, 7.0, 8.0, 12.0, 14.0, 20.0, 27.0, 32.0, 70.0] KV_SHAPES = [ ("28 layers, 8 key-value heads and a head dimension of 128", 28, 8, 128), ("36 layers, 8 key-value heads and a head dimension of 128", 36, 8, 128), ("40 layers, 8 key-value heads and a head dimension of 128", 40, 8, 128), ("48 layers, 4 key-value heads and a head dimension of 128", 48, 4, 128), ("64 layers, 8 key-value heads and a head dimension of 128", 64, 8, 128), ("24 layers, 8 key-value heads and a head dimension of 64", 24, 8, 64), ("36 layers, 8 key-value heads and a head dimension of 64", 36, 8, 64), ("32 layers, 8 key-value heads and a head dimension of 128", 32, 8, 128), ("16 layers, 4 key-value heads and a head dimension of 64", 16, 4, 64), ("60 layers, 8 key-value heads and a head dimension of 128", 60, 8, 128), ] SHAPES = [(1024, 2048), (1024, 3072), (2048, 6144), (2560, 9728), (4096, 12288), (4096, 4096), (2048, 2048), (3072, 1024), (6144, 2048), (5120, 13824), (8192, 28672), (1536, 4096)] FILE_GB = [1.1, 2.5, 4.3, 5.03, 8.0, 12.12, 18.56, 19.76, 32.48, 63.39] def _n(value: float, places: int = 2) -> str: """Trim trailing zeros so the answers read like a person wrote them.""" text = f"{value:.{places}f}".rstrip("0").rstrip(".") return text or "0" def generated() -> list[tuple[str, str, str]]: items: list[tuple[str, str, str]] = [] for params in PARAMS_B: for name, bpp in FORMATS: gb = params * bpp items.append(( f"How much memory do the weights of a {_n(params)} billion parameter model need at {name}?", f"About {_n(gb)} GB.", f"{_n(params)} billion parameters at {_n(bpp)} bytes each is {_n(gb)} x 10^9 bytes.", )) for name, bpp in FORMATS: items.append(( f"How many bytes per parameter is {name}?", f"About {_n(bpp)} bytes per parameter.", "The figure includes the per-block scales that are stored alongside the weights." if bpp not in (1.0, 2.0) else f"{int(bpp * 8)} bits is {_n(bpp)} bytes, with no per-block scale to add.", )) for label, layers, kv_heads, head_dim in KV_SHAPES: per_token = 2 * layers * kv_heads * head_dim * 2 items.append(( f"A model has {label}. How many bytes per token is its FP16 key-value cache?", f"{per_token:,} bytes per token.", f"2 x {layers} x {kv_heads} x {head_dim} x 2 = {per_token:,}.", )) for ctx in (8192, 32768): total = per_token * ctx / GB items.append(( f"A model has {label}. How much memory is its FP16 key-value cache at {ctx:,} tokens?", f"About {_n(total)} GB.", f"{per_token:,} bytes per token times {ctx:,} tokens is {_n(total)} x 10^9 bytes.", )) for params in PARAMS_B: items.append(( f"How much memory do the weights, gradients, master weights and Adam states of a " f"{_n(params)} billion parameter model need for a full fine-tune in mixed precision?", f"About {_n(params * 16)} GB.", "Two bytes of BF16 weights, two of gradients, four of FP32 master weights and eight of " f"Adam states is 16 bytes per parameter, so {_n(params)} x 16 = {_n(params * 16)} GB.", )) items.append(( f"How much memory do the Adam states of a {_n(params)} billion parameter model need at FP32?", f"About {_n(params * 8)} GB.", f"Two FP32 states per parameter is 8 bytes, so {_n(params)} x 8 = {_n(params * 8)} GB.", )) items.append(( f"How much memory do BF16 gradients for a {_n(params)} billion parameter model need?", f"About {_n(params * 2)} GB.", f"One gradient per parameter at two bytes is {_n(params)} x 2 = {_n(params * 2)} GB.", )) for inputs, outputs in SHAPES: for rank in (8, 16): added = rank * (inputs + outputs) items.append(( f"A LoRA adapter of rank {rank} is added to a linear layer with {inputs:,} inputs and " f"{outputs:,} outputs. How many trainable parameters does it add?", f"{added:,} trainable parameters.", f"The two matrices are {rank} x {inputs:,} and {outputs:,} x {rank}, so " f"{rank} x ({inputs:,} + {outputs:,}) = {added:,}.", )) for gb in FILE_GB: gib = gb * GB / GIB items.append(( f"A file listing gives a model file as {_n(gb)} GB. What will a file manager counting in gibibytes show?", f"About {_n(gib)} GiB.", f"{_n(gb)} x 10^9 bytes divided by 2^30 is {_n(gib)}, because a gibibyte is about seven per cent larger than a gigabyte.", )) return items def build() -> list[dict[str, str]]: rows = [{"instruction": q, "answer": f"Answer: {a}\nBecause: {b}"} for q, a, b in WRITTEN + generated()] seen: set[str] = set() unique = [] for row in rows: if row["instruction"] in seen: continue seen.add(row["instruction"]) unique.append(row) return unique def write_jsonl(path: Path, lines: list[dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for line in lines: handle.write(json.dumps(line, ensure_ascii=False) + "\n") def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--out-dir", default=".", help="directory to write the dataset and splits into") parser.add_argument("--seed", type=int, default=0, help="seed for the shuffle before splitting") parser.add_argument("--valid-frac", type=float, default=0.15) parser.add_argument("--test-frac", type=float, default=0.10) args = parser.parse_args() rows = build() random.Random(args.seed).shuffle(rows) n_valid = int(len(rows) * args.valid_frac) n_test = int(len(rows) * args.test_frac) splits = { "valid": rows[:n_valid], "test": rows[n_valid:n_valid + n_test], "train": rows[n_valid + n_test:], } out = Path(args.out_dir) write_jsonl(out / "sample-instructions.jsonl", rows) for name, part in splits.items(): # TRL: conversational prompt-completion. Loss is computed on the completion only. write_jsonl(out / "data" / f"{name}.jsonl", [ {"prompt": [{"role": "user", "content": r["instruction"]}], "completion": [{"role": "assistant", "content": r["answer"]}]} for r in part ]) # mlx-lm: the completions format, two plain strings. write_jsonl(out / "data-mlx" / f"{name}.jsonl", [ {"prompt": r["instruction"], "completion": r["answer"]} for r in part ]) digest = hashlib.sha256((out / "sample-instructions.jsonl").read_bytes()).hexdigest() print(f"examples: {len(rows)}") print(f" train {len(splits['train'])} valid {len(splits['valid'])} test {len(splits['test'])}") print(f"written: {out / 'sample-instructions.jsonl'}, {out / 'data'}/, {out / 'data-mlx'}/") print(f"sha256(sample-instructions.jsonl): {digest}") print("record that hash in the run log; it is what ties a result to this exact dataset.") if __name__ == "__main__": main()