19 Repeat the project with every weight trainable

For this chapter's commands, open a fresh terminal at the companion root, run cd examples/llm, then source .venv-llm/bin/activate. If you skipped the first small-language-model project, complete its explicit environment setup first. Do not run these relative paths from an earlier project directory.

Now keep the same checkpoint, dataset and evaluation but set --mode full. This changes the number of trainable parameters, not the task. You are doing full fine-tuning of a pretrained model. You are not starting from random weights.

CUDA_VISIBLE_DEVICES=0 python train_small_lm.py \
  --train data/style_train.jsonl \
  --eval data/style_valid.jsonl \
  --out runs/style-full-smoke \
  --mode full --steps 3

The recipe lowers the initial learning rate to 1e-5, explicitly loads FP32 weights, keeps FP32 gradients and ordinary AdamW state, and uses BF16 autocast for suitable computations. These are reviewed implementation choices, not a promise that this learning rate wins. Full-update mode refuses models above one billion unique parameters. Adapter modes refuse more than three billion total parameters.

Count memory before you spend time

Training memory is a sum of several different things:

M_peak = M_weights + M_gradients + M_optimizer + M_master

+ M_activations + M_temporary + M_runtime.

Weights are the model’s learned numbers. Gradients describe the local direction in which the loss changes. AdamW maintains two optimizer tensors per trainable parameter: moving estimates of first and second gradient moments. A master copy is an extra high-precision copy used by some mixed-precision schemes. Activations are intermediate values needed for backward computation. Temporary tensors, attention workspaces, allocator reservation and the GPU runtime add further overhead. L16 L17

“BF16 training uses two bytes per parameter” accounts only for one copy of the weights. It is not a training-memory estimate. Distinguish persistent parameter-related state from the peak that depends on the batch and implementation.

The table uses decimal GB and hypothetical exact parameter counts. It contains calculated persistent-state totals, not measured VRAM:

These rows are different implementations, not three simultaneous charges. Do not add a master copy to a scheme whose FP32 weights already serve that purpose. Conversely, do not assume an optimizer uses FP32 moments merely because the model was loaded in BF16; verify the implementation and actual state tensor dtypes.

Our full-update example uses the middle row. At approximately 0.6B parameters, roughly 9.6GB is reserved for its parameter-related training state before activations and other overhead. Its actual unique parameter count is printed at runtime. This leaves a meaningful working margin on a 24GB-class card at a small microbatch and short sequence length, but only the smoke test can establish the real peak. At 3B, the same state alone is 48GB. Reducing sequence length cannot make those 48GB disappear.

A one-billion-parameter full update is an advanced boundary, not an unconditional fit statement. The exact architecture, vocabulary size, gradient dtypes, optimizer implementation, sequence length and temporary buffers matter. An eight-bit optimizer or CPU optimizer offload can move that boundary, but changes the recipe and sometimes numerical behavior or speed. Benchmark such a configuration separately. Do not claim that every nominal “3B” model fits full AdamW training on one 24GB card.

Sequence length can dominate the surprise

Let B be the microbatch size, T the sequence length, d the hidden width and L the number of layers. Many saved activations grow roughly with B T d L, with substantial architecture-dependent constants. A straightforward attention implementation can materialize score tensors proportional to B H T², where H is the number of heads. Efficient attention kernels avoid storing the full quadratic score matrix, but do not remove every sequence-length-dependent cost.

There is another large tensor: logits. An unchunked tensor for B=1, T=2,048, V=151,936 vocabulary entries contains about 311 million values. At four bytes each, that is about 1.24GB for one tensor, before related intermediates. This is a calculated illustration, not a measurement of the supplied recipe. A model with a large multilingual vocabulary can have expensive logits even when its transformer blocks are small.

Gradient checkpointing saves fewer forward intermediates and recomputes them during backward. It exchanges compute for memory. Gradient accumulation computes several small microbatches before updating weights. On one GPU, effective_examples ≈ microbatch × accumulation; variable-length examples mean tokens per update vary. Neither technique removes full-model optimizer state. Accumulation also requires correct loss normalization when examples have different numbers of supervised tokens, which is one reason to use a reviewed trainer rather than improvising division by a batch count.

use_cache=False disables the generation KV cache during training. A KV cache stores past attention keys and values to accelerate autoregressive generation; it is not the same thing as the activations kept for backpropagation. Long advertised context windows describe architectural support, not affordable training lengths on your card.

Measure the run you actually made

The script records torch.cuda.max_memory_allocated() and max_memory_reserved() in GiB, where one GiB is 2^30 bytes. Allocated memory is live tensor memory known to the allocator. Reserved memory includes its cached blocks. Neither alone captures every driver-level allocation, so also inspect nvidia-smi and leave headroom for your display or other processes.

Record a warm-up and several complete optimizer steps. Some optimizer states appear only at the first update, so one forward pass is not a training fit test. Include validation and checkpoint saving in the check; those stages can fail even if the training step fits. Do not extrapolate throughput from one cold iteration that includes compilation, loading, or cache setup.

If you hit out-of-memory, diagnose in this order:

1. Check other GPU processes and actual free memory
2. Confirm the requested dtype and count trainable parameters
3. Confirm there is only one training model and no unintended reference copy
4. Keep microbatch at one and reduce sequence length
5. Enable checkpointing and avoid storing evaluation logits
6. Move to LoRA or QLoRA if persistent full-update state is the limiting term
7. Consider a smaller model before adding complex offload machinery

Exercise. Repeat the same tiny run at sequence limits 256, 512 and 1,024 with examples that actually approach those lengths. Record allocated/reserved memory and tokens per second. A higher maximum alone does not lengthen short examples, so do not call that a sequence-length experiment unless the data changed accordingly.

When QLoRA earns its complexity

QLoRA keeps a quantized frozen base model and trains floating-point adapters through it. The original method combines four-bit NormalFloat storage, quantized quantization constants and paged optimizers to reduce memory pressure. Its published large-model results are method evidence, not a benchmark for your RTX or this example. L18

The starter switch is:

CUDA_VISIBLE_DEVICES=0 python train_small_lm.py \
  --train data/style_train.jsonl \
  --eval data/style_valid.jsonl \
  --out runs/style-qlora-smoke \
  --mode qlora --steps 3

The code uses NF4, double quantization, BF16 compute and ordinary AdamW for the relatively small adapters. It does not need a paged optimizer to qualify as adapter training over a quantized base. prepare_model_for_kbit_training freezes/prepares the base before adapters are attached. Loading a four-bit model and asking ordinary full fine-tuning to update all its quantized weights is not this recipe. L19 L20

If Pq parameters are actually quantized to four bits, their ideal packed payload is 0.5 Pq bytes. Add quantization scales, unquantized modules, adapters, gradients, optimizer state and activations. Some preparation paths upcast unquantized modules to FP32. An embedding-heavy small model may therefore save less than the simple “four bits everywhere” calculation suggests.

For a 0.6B model, plain BF16 LoRA may already fit easily and be simpler. QLoRA becomes more attractive as you move toward the larger allowed checkpoints or longer sequences. Measure both before choosing. Quantization can change outputs and throughput; evaluate the exact deployed precision, not just the floating-point training checkpoint.