The structure of Practical LLM in Rust

Extend a complete model to larger workloads

A working small language model makes every operation visible. Larger workloads put pressure on several connected parts at once: more text increases tokenization and batching work; deeper models carry more weights and intermediate activations; learning retains gradients and optimizer state. Reducing numerical precision can shrink that storage while changing rounding and overflow behavior. Generation repeatedly reads an expanding prefix. A change that saves storage can also change which tokens contribute to loss or which earlier positions attention can read. The practical task is to improve how the model runs while keeping those mathematical choices explicit.

Practical LLM in Rust starts from the completed model in LLM from scratch. Its separate implementation lives in rust/crates/llm-from-scratch-practical. We extend its tokenizer, tensors, automatic differentiation, causal decoder, training, evaluation, checkpoints and cached generation as one continuing system. Chapter 0 locates those parts and the course chapters; Chapter 1 begins the experiments.

Bring the first course’s model with you

Complete LLM from scratch through Chapter 39 or bring equivalent experience with Rust, byte-pair encoding, tensors, gradients, causal attention and next-token training. The first course derives and implements those mechanisms. This course uses that foundation to explain changes in data layout, numerical execution, training and inference.

The default CPU build and CPU inference work without NVIDIA hardware. Published Chapter 11 adds explicit CPU or CUDA selection for matrix products and sum/mean reductions in the shared forward and backward paths. Selecting CUDA requires compatible NVIDIA hardware and runtime libraries; an unavailable CUDA request is refused. This boundary does not make every operation GPU-resident, and other GPU backends are not implemented.

The external corpus-preparation example separately uses NVIDIA NeMo Curator and needs a device and driver compatible with its pinned CUDA/RAPIDS environment. NeMo prepares text outside the Rust model. Another preparation tool can supply the same document format, and Rust can load already prepared documents on CPU.

Continue from the completed small model

The first course builds the mechanisms and assembles a complete causal language model. The practical course begins at that completed system and changes how its parts handle larger workloads. The first course keeps its own implementation; the practical copy evolves independently. Each extension should therefore have a clear connection to the model already present and a concrete reference for comparison.

Follow one model through learning and generation

Prepared documents become token IDs and document-aware examples. Batching keeps each example’s permitted context and real next-token targets intact. The text encoder uses embedding lookup to turn input IDs into vectors, accompanied by position and attention information. The causal transformer turns those vectors into hidden features. The text head uses the same learned embedding table to map the features to vocabulary scores, called logits.

Learning and generation use those scores differently. Learning compares scores with observed next-token targets to calculate loss. Gradients describe how changes in weights affect that loss, and the optimizer changes the same weights used by the next forward calculation. Generation chooses another token according to a decoding policy, returns its ID to the text encoder and runs a subsequent decoder step. Compatible cached attention keys and values let that step reuse earlier prefix work.

The shared forward path

Prepared text supplies token batches; encoding, transformation and the text head produce next-token scores in that order.

1. 1 Prepared text and batches

   Tokenize documents and arrange examples without joining unrelated contexts.

   Chapters 2–5

   ↓

2. 2 Text encoder

   Embedding lookup turns input token IDs into vectors with positions and attention layout.

   Chapter 10

   ↓

3. 3 Causal transformer

   Repeated decoder blocks transform vectors into hidden features using the permitted prefix.

   Chapters 6–10

   ↓

4. 4 Tied text head

   The text head reuses the input embedding lookup table to map hidden features to one score per vocabulary token.

   Chapter 10

Learning and generation use that forward path

Learning changes weights

Scores and observed targets give loss; gradients and the optimizer change the model weights.

↺ Updated weights feed the next forward calculation.

Chapters 13–14; planned extensions: Chapters 23–25

Generation extends a sequence

A decoding policy chooses a token from the scores.

↺ The chosen token ID returns to embedding lookup for the next cached decoder step.

Chapters 18–20; planned extensions: Chapters 21–22, 26–27, 29–30

Execution and state support the same model

Numerical execution

CPU/CUDA matrix and reduction primitives support forward and backward. CPU half-format matrix policies and loss-scaled updates add deliberate numerical choices; accumulation, recomputation and scheduled updates extend training.

Chapters 11–14

Save, load and resume the composition

Preserve model weights, semantic settings, tokenizer assets and the shared alias for a new CPU inference process. Additional optimizer, data, random and update state support training continuation within a narrower execution envelope.

Chapter 15

The figure follows model computation. The roadmap below follows the order in which the course studies and extends that computation. For example, tensor execution supports both learning and generation; saving a model preserves a state that later inference can use. Neither is another token-processing stage between the encoder and the text head.

Chapter 10 already implements the text encoder, source-independent transformer backbone and tied text head. The backbone consumes vectors and their layout without requiring a tokenizer or vocabulary. Observed targets enter the learning branch outside that backbone. Text remains the only worked input and output modality. A compatible encoder for images, binary data or entities would still need suitable encoding and training; the shared vector boundary alone does not give the text model those capabilities.

Preserve learned weights for later use

The starting model already has checkpoint mechanisms. Published Chapter 15 extends that work to saving and loading the encoder, backbone and head composition. It saves named weights, semantic settings and the tokenizer assets needed by the complete text model, stores the tied table once and restores its shared identity. The model can be saved, its process can exit, and a fresh local CPU process can load it for inference. The generic backbone remains independent of a text tokenizer.

Portable saved weights do not contain mandatory device addresses or CUDA handles. The fuller training checkpoint additionally preserves optimizer, update-policy, random-generator, data-position, loss-scaler and accumulation state. Exact continuation is checked at complete-update boundaries with cleared gradients, matching data and the same-binary scalar CPU f64 execution envelope. These training conditions are narrower than the requirements for CPU inference.

The separate physical portability experiment performs a CUDA training update, saves the model, exits and loads the same learned weights on CPU without retraining. That experiment remains pending; the implemented CPU save/load and continuation fixtures do not establish it.

Read the complete chapter roadmap

Chapter 0 is this orientation; Chapters 1–20 are published. Chapters 21–32 are planned and appear without links. A planned lesson may extend an existing mechanism, such as sampling or checkpointing. Its status describes the lesson’s availability. The practical course index lists the available lessons.

Establish the reference and the data path

0. Chapter 0: The structure of Practical LLM in Rust You are here. Connect the model's roles and find each step in the practical course.

1. Chapter 1: Keep a fixed reference before extending the decoder Published. Fix the small reference experiment's inputs and results before extending its decoder.

2. Chapter 2: Prepare text externally, then load it in Rust Published. Separate external text preparation from Rust loading while retaining document and split identity.

3. Chapter 3: Extend BPE with bounded local updates Published. Extend byte-pair encoding (BPE) with bounded ingestion, local pair updates and ranked application of learned merges.

4. Chapter 4: Pad variable-length batches without adding predictions Published. Batch unequal examples while excluding padding from context, targets and mean loss.

5. Chapter 5: Pack sequences without joining their contexts Published. Place examples together in storage while keeping each example's causal context and loss contribution separate.

Make the model and its execution explicit

6. Chapter 6: Initialize residual projection weights for decoder depth Published. Scale fresh residual projection weights for decoder depth and inspect activations and gradients before training.

7. Chapter 7: Configure one decoder core before construction Published. Check dimensions and resource plans before constructing one core shared by full and cached text paths.

8. Chapter 8: Make optional dropout explicit and replayable Published. Distinguish training from evaluation and preserve saved-mask gradients and local random replay.

9. Chapter 9: Keep model operations and failures explicit Published. Identify refused operations and the model, optimizer or checkpoint state preserved by each failure boundary.

10. Chapter 10: Separate input encoding from the transformer Published. Compose the implemented text encoder, source-independent backbone and tied head while preserving values and gradients.

11. Chapter 11: Run tensor primitives on a GPU Published. Select CPU or CUDA for existing matrix products and sum/mean reductions used by the shared model and reverse mode.

Extend numerical training and preserve its results

12. Chapter 12: Train with deliberate precision choices Published. Use CPU half-format matrix arithmetic and loss-scaled updates, inspect actual dtypes and preserve model and optimizer state when an overflowing update is skipped.

13. Chapter 13: Accumulate gradients and recompute activations Published. Combine smaller training batches with the correct token denominator, and trade saved activations for recomputation with matching randomness.

14. Chapter 14: Schedule AdamW updates and clip global gradients Published. Choose warmup and cosine rates by successful update while inspecting existing global clipping and AdamW state.

15. Chapter 15: Save, load and resume the composed model Published. Save the composed model, exit and load it for local CPU inference; then restore complete-update training state under its declared compatibility conditions.

16. Chapter 16: Compare training and held-out language-model loss Published. Compare training and held-out loss, a simple baseline and seed variation to bound conclusions about overfitting and quality.

17. Chapter 17: Store weights in four bits and measure the error Published. Represent weights with scales and rounding, then compare storage and numerical error using a portable dequantization path.

Extend attention, caches and token selection

18. Chapter 18: Share key/value heads and make context explicit Published. Share key/value heads among query heads and preserve the intended causal or windowed context in full and cached computation.

19. Chapter 19: Accumulate attention without keeping the full grid Published. Accumulate attention through running normalization and tiles, comparing outputs and memory use with reference attention.

20. Chapter 20: Keep token order across cache pages Published. Map logical token positions to cache pages, account for unused storage and preserve cached model scores.

21. Chapter 21: Sampling distributions Planned. Transform next-token probabilities with top-p selection, repetition penalties and allowed-token masks, keeping processor order explicit.

22. Chapter 22: Batched autoregressive inference Planned. Generate unequal sequences together while retaining each row's positions, masks, cache and token decisions.

Adapt the same model and complete a local workflow

23. Chapter 23: Response-masked SFT and LoRA Planned. Use response-only supervised fine-tuning (SFT) and low-rank adaptation (LoRA), tracing adapter gradients while base weights stay frozen.

24. Chapter 24: Direct preference optimization Planned. Use chosen and rejected response probabilities from policy and reference models to explain a preference loss and its gradients.

25. Chapter 25: LoRA over a quantized base Planned. Combine the existing quantization and adapter paths, exposing dequantization and gradients without updating the frozen base.

26. Chapter 26: Prefix key/value-cache reuse Planned. Reuse prepared-prefix state only when encoding, layout, model, adapter and position identities remain compatible.

27. Chapter 27: RoPE and context scaling Planned. Extend rotary position embeddings (RoPE) with explicit rotation and context-scaling choices, keeping position changes distinct from learned weights.

28. Chapter 28: Local training and adaptation capstone Planned. Connect training or adaptation, saving, loading, resuming, evaluation, quantization and generation in one bounded local workflow.

Study further decoding and model mathematics

29. Chapter 29: Beam search Planned. Expand and prune candidate sequences using cumulative scores while keeping each branch's cache consistent.

30. Chapter 30: Speculative verification Planned. Use fixed proposals to trace acceptance and correction probabilities with the target decoder, without inferring a speedup.

31. Chapter 31: Distributed-training mathematics Planned. Explain sharding, normalized gradient aggregation, ring collectives and pipeline occupancy through bounded local examples.

32. Chapter 32: Mixture-of-experts mathematics Planned. Reuse feed-forward blocks to study routing, expert capacity, weighted outputs and load-balancing gradients in a tiny model.

Keep the experiments at their stated scale

Small fixtures expose values that are easy to inspect: token boundaries, allowed attention positions, a loss, a gradient or a restored state. They establish the behavior shown for those inputs. A CPU fixture cannot establish a physical-GPU result; a bounded CUDA primitive comparison cannot establish whole-model training throughput. Neither by itself establishes full-corpus throughput, broad language quality or a larger model’s memory use. Full-data preparation and larger hardware experiments need their own measurements.

The course focuses on LLM mathematics, numerical choices, training and inference. A thin local command-line interface is intended to make experiments with this same model convenient. Application-specific web services, retrieval and tool frameworks are outside the roadmap.

Begin with the fixed reference experiment

The data path supplies the model’s inputs; learning changes its weights; generation uses those weights to extend a sequence. Later lessons extend these connected roles and compare the result with a declared reference. Start with Chapter 1: Keep a fixed reference before extending the decoder to establish the small experiment and the limits of what it demonstrates.