A5 — Training Infrastructure

Files: llmc/schedulers.h, llmc/logger.h, llmc/outlier_detector.h, llmc/global_norm.cuh, train_gpt2.cu
Needs: 1 NVIDIA GPU (schedulers/logger/detector need none)
Builds on: A0–A3
Check with: ./grade A5

---

Overview

A training run that cannot be resumed is not a training run — it’s a gamble that your machine stays up for weeks. A5 implements the infrastructure that makes a production training loop: learning rate schedulers, gradient clipping, structured logging, checkpoint save/restore, and MFU reporting.

These components are individually small, but together they are what separates a “it trains once” script from a system you would actually use to train GPT-3.

After completing A5, your training binary will:

• Ramp the learning rate up linearly, then decay it (cosine, linear, or WSD)
• Clip gradients to prevent loss spikes
• Write checkpoints that resume bit-exactly
• Write a structured log file and report MFU
• Detect outlier loss and gradient norms and skip those updates

To list your work:

grep -rn "TODO (A5)" llmc/ train_gpt2.cu

When these notes and the in-file HINT disagree, the code is authoritative — please report the discrepancy so it can be fixed.

The scheduler, logger, and outlier detector are CPU-only and testable without a GPU. global_norm.cuh and the train_gpt2.cu parts need one.

---

Where this fits

flowchart LR
    A0["A0
tokens, batches, PRNG"] --> A1["A1
forward pass (CPU)"]
    A1 --> A2["A2
backward + AdamW (CPU)"]
    A2 --> A3["A3
port to CUDA (1 GPU)"]
    A3 --> A4["A4
scale to N GPUs"]
    A3 --> A5["A5
schedules, clipping, checkpoints"]
    style A5 fill:#2563eb,stroke:#1e40af,color:#fff

Everything up to here trains a model once, in one process, until you stop it. A4 and A5 make it a system you would actually leave running for a week: a learning rate that changes on a schedule, gradient clipping so one bad batch cannot destroy the run, structured logs, and checkpoints you can resume from without changing the loss curve.

flowchart TB
    start(["step"]) --> fwd["forward + backward
over grad_accum_steps micro-batches"]
    fwd --> norm["gpt2_calculate_grad_norm
global L2 over all gradients — Part 2"]
    norm --> det["outlier detector
z-score vs last 128 steps — Part 3"]
    det -->|"z > threshold"| skip["skip this update"]
    det -->|"normal"| lr["scheduler
warmup → decay — Part 1"]
    lr --> clip["grad_scale = min(1, clip / norm)"]
    clip --> upd["AdamW, scaled gradient"]
    upd --> log["logger — Part 4"]
    log --> ck{"checkpoint
step?"}
    ck -->|yes| save["save_state — Part 5
weights + m,v,master
+ RNG + dataloader position"]
    ck -->|no| start
    save --> start
    skip --> log

Why a warmup at all? Adam’s v estimate starts at zero and is meaningless for the first few dozen steps, so a full-size learning rate applied immediately takes a large, badly-scaled step and can put the model somewhere it never recovers from. Warmup is a cheap insurance policy against losing a multi-day run in its first minute.

Why clip on the global norm and not per tensor. Clipping each tensor separately changes the direction of the update, not just its length. Scaling every gradient by one shared factor preserves the direction and only shortens the step — which is why Part 2 has to reduce a single scalar across every parameter before anything is scaled.

What “resumable bit-exactly” actually requires. It is not enough to save the weights. A resumed run must reproduce the original loss curve, so you also need the optimizer moments, the RNG state at the last update (to replay stochastic rounding identically), the RNG state now, and the DataLoader’s position and its shuffle permutation. Miss any one and the run continues plausibly but diverges — the failure mode that is hardest to notice and most expensive to discover a week later.

flowchart LR
    subgraph ck["a checkpoint at step N"]
        direction TB
        m1["model_0000N.bin
weights, rank 0 only"]
        m2["state_0000N_00000.bin
m, v, master, RNG,
dataloader pos + permutation
one per rank"]
        m3["DONE_0000N
written only after a barrier"]
    end
    m3 -.->|"the sentinel is what makes
a half-written checkpoint safe"| m1

---

Background Reading

Source	Read	Why it matters here
The training loop in train_gpt2.cu	the for (int step ...) body	Everything you implement is called from here, in this order. Reading it first makes the five parts below fall into place.
Loshchilov & Hutter 2017, SGDR	§3	Cosine annealing, the source of the 0.5 * (1 + cos(pi * t)) in Part 1.
Hägele et al. 2024, Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations	§2 and Fig. 2	The WSD schedule exactly as implemented here — this is the paper cited in schedulers.h. Explains why a constant plateau followed by a short decay beats cosine when you do not know the step budget in advance.
Pascanu et al. 2013, On the difficulty of training RNNs	§3.2	Gradient clipping, and the argument for scaling by the global norm rather than clipping elementwise.
Kaplan et al. 2020, Scaling Laws for Neural Language Models	§2.1	Where 6 * N flops per token comes from — the formula you implement in Part 6. The gpt2_estimate_mfu comment cites this directly.
Chowdhery et al. 2022, PaLM	§5.1	The definition of MFU as a hardware-efficiency metric, and why it is reported instead of raw FLOP/s.
Hoffmann et al. 2022, Chinchilla	§3	Optional. How to choose a step budget in the first place — the input to every schedule in Part 1.

---

Part 1 — Learning Rate Schedulers (llmc/schedulers.h)

typedef struct {
    const char* type;               // "cosine" | "linear" | "constant" | "wsd"
    float learning_rate;            // peak LR
    int warmup_iterations;
    int train_num_batches;          // total steps
    float final_learning_rate_frac; // min_lr = learning_rate * this
} LearningRateScheduler;

lr_scheduler_init(...) — plain field assignment, no computation.

Warmup is shared by all three decaying schedules, and it is 1-based:

if step < warmup_iterations:
    return learning_rate * (step + 1) / warmup_iterations

Note the step + 1. At step == 0 the LR is lr / warmup_iterations, not zero — a zero first step wastes an iteration. Getting this off by one is the most common way to fail the 6-significant-figure comparison, and it affects cosine, linear, and WSD.

Throughout, min_lr = learning_rate * final_learning_rate_frac.

get_learning_rate_cosine

decay_ratio = (step - warmup_iterations) / (train_num_batches - warmup_iterations)   // in [0, 1]
coeff       = 0.5 * (1 + cos(pi * decay_ratio))                                      // 1 -> 0
return min_lr + coeff * (learning_rate - min_lr)

get_learning_rate_linear

decay_ratio = (step - warmup_iterations) / (train_num_batches - warmup_iterations)
return learning_rate - decay_ratio * (learning_rate - min_lr)

get_learning_rate_constant

return learning_rate

get_learning_rate_wsd — Warmup, Stable, Decay. The decay occupies the last 20% of training and uses a 1 - sqrt curve, not a linear one:

decay_point = (int)(0.8 * train_num_batches)

if step < warmup_iterations:
    return learning_rate * (step + 1) / warmup_iterations      // warmup
elif step < decay_point:
    return learning_rate                                       // stable plateau
else:
    decay_ratio = (step - decay_point) / (train_num_batches - decay_point)
    return min_lr + (1 - sqrt(decay_ratio)) * (learning_rate - min_lr)

decay_point is an integer truncation of 0.8f * train_num_batches — do not keep it in floating point, or your boundary step will differ from the reference. WSD is normally run with final_learning_rate_frac = 0.0.

Assert 0 <= decay_ratio <= 1 in the decay branches; it catches step >= train_num_batches immediately.

get_learning_rate — strcmp dispatch over the four names. An unrecognised type is a fatal error, not a silent fallback to constant.

---

Part 2 — Gradient Norm (llmc/global_norm.cuh)

Gradient clipping needs the global L2 norm over all parameter gradients, reduced across many CUDA blocks.

global_norm_squared_kernel

Each block computes a partial sum of squares over a chunk of the gradients, using x128 loads and blockReduce:

float partial = 0.0f;
for each x128 chunk assigned to this thread:
    x = load128cs(grads + idx)
    for k in 0..x128::size: partial += (float)x[k] * (float)x[k]
// blockReduceSum across all threads, then write to out[blockIdx.x]

global_norm_squared (launcher)

Launches over a 2D grid: one row per slice, where non-layer tensors use num_slices = 1 and per-layer tensors use num_slices = num_layers. The reset argument zeroes the output buffer before accumulating, so the first tensor clears it and the rest add into it.

get_max_num_block_sums(num_slices_arr, len)

Returns the largest total block count over all parameter tensors — this sizes the intermediate buffer. Computes CEIL_DIV(shard_size, block_size) per tensor and takes the max across them.

global_sum_deterministic (provided in cuda_utils.cuh)

An in-order, reproducible sum of the block-sums buffer down to a single float. Use this rather than atomicAdd — A4’s cross-GPU reproducibility depends on it.

gpt2_calculate_grad_norm (train_gpt2.cu)

float* grad_norm_sq = (float*)model->acts.output;   // reuse the output activation as scratch

for each parameter tensor i:
    ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i);
    ShardInfo shard  = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1);
    bool reset = (i == 0);
    global_norm_squared(grad_norm_sq, grads + tensor.offset + shard.offset,
                        shard.size, tensor.size, num_layers, max_num_block_sums, reset, main_stream);

global_sum_deterministic(grad_norm_sq, grad_norm_sq, max_num_block_sums, main_stream);

// if MULTI_GPU and zero_stage == 1: ncclAllReduce(grad_norm_sq, ncclSum) to combine shards

cudaMemcpy(&norm_sq_cpu, grad_norm_sq, sizeof(float), cudaMemcpyDeviceToHost);
return sqrtf(norm_sq_cpu);

Where clipping actually happens. gpt2_calculate_grad_norm only measures. The training loop then computes

float grad_clip  = 1.0f;
float grad_scale = (grad_norm > grad_clip) ? grad_clip / grad_norm : 1.0f;
gpt2_update(&model, ..., grad_scale, step+1, &multi_gpu_config);

and grad_scale is applied inside adamw_update (which you wrote in A3). The gradient buffer is never modified. grad_clip is currently a hard-coded 1.0f in the driver — there is no command-line flag for it.

---

Part 3 — Outlier Detector (llmc/outlier_detector.h)

A sliding window over the last OUTLIER_DETECTOR_WINDOW_SIZE = 128 observations, used to detect training instability. The driver feeds it both the loss and the gradient norm, and skips the update when the z-score exceeds the -sl / -sg thresholds.

typedef struct {
    double buffer[OUTLIER_DETECTOR_WINDOW_SIZE];
    int count;      // observations so far, saturating at WINDOW_SIZE
    int index;      // circular write cursor
    double sum;
    double sum_sq;
} OutlierDetector;

init_detector(detector) — zero the buffer, count, index, sum, sum_sq.

update_detector(detector, new_value) — two regimes:

Warming up (count < WINDOW_SIZE): append at buffer[count], add into sum and sum_sq, increment count, and return nan(""). There is not yet enough data for a meaningful z-score.

Steady state (count == WINDOW_SIZE):

1. Evict buffer[index]: subtract it from sum and its square from sum_sq.
2. Write new_value at buffer[index] and add it into sum / sum_sq.
3. index = (index + 1) % WINDOW_SIZE.
4. Compute statistics dividing by WINDOW_SIZE, not by count:

   mean     = sum / WINDOW_SIZE
   variance = sum_sq / WINDOW_SIZE - mean * mean
   std_dev  = sqrt(variance)
   

5. Return 0.0 if std_dev == 0.0, else (new_value - mean) / std_dev.

The NaN return is load-bearing, not a nicety: the driver guards on isfinite(zloss) before deciding whether to skip an update. Returning 0.0 during warmup silently disables outlier detection for the first 128 steps; returning some large number would skip every early update.

Use double throughout — the running +=/-= on sum_sq drifts badly in float.

---

Part 4 — Logger (llmc/logger.h)

The logger is stateless between calls: it stores a path, and each log call opens the file in append mode, writes one line, and closes it. This is deliberate — a crash mid-run cannot lose buffered log lines.

typedef struct {
    int active;
    char output_log_file[512];
} Logger;

logger_init(logger, log_dir, process_rank, resume)

Set active = 0 first. Set active = 1 only when log_dir != NULL and process_rank == 0 — in a multi-GPU run only rank 0 writes logs. Build the path as "{log_dir}/main.log" with snprintf. If resume == 0, open the file in "w" mode and immediately close it to truncate any previous run’s log; if resuming, leave it alone so the new lines append.

Log calls — each is guarded by if (logger->active == 1), opens "a", fprintfs, and closes:

Function	Line written
logger_log_eval(logger, step, val)	s:%d eval:%.4f\n
logger_log_val(logger, step, val_loss)	s:%d tel:%.4f\n
logger_log_train(logger, step, train_loss, learning_rate, grad_norm)	s:%d trl:%.4f lr:%.6f norm:%.2f\n

The exact prefixes (trl:, tel:, eval:) and precisions matter — dev/ plotting scripts parse them.

There is no logger_free: nothing stays open between calls.

---

Part 5 — Checkpoint Save/Load (train_gpt2.cu)

Checkpointing must be resumable bit-exactly: given a checkpoint, continuing must produce the same loss curve as an uninterrupted run.

That requires saving: model weights, optimizer state (m, v, master weights — sharded under ZeRO-1), the RNG state, the DataLoader position and shuffle permutation, and the step counter.

gpt2_write_to_checkpoint(model, checkpoint_path)

Write a 256-int header — magic 20240326, version 3 for fp32 or 5 for bf16, then the config fields in the layout A3’s loader reads — followed by device_to_file of params_memory.

save_state(filename, step, model, loader)

header[0]  = 20240527  // magic
header[1]  = 1         // version
header[2]  = multi_gpu_config.num_processes
header[3]  = multi_gpu_config.process_rank
header[4]  = model->use_master_weights
header[5]  = loader->should_shuffle
header[10] = step
*(unsigned long long*)&header[20] = model->rng_state
*(unsigned long long*)&header[22] = model->rng_state_last_update
*(size_t*)&header[30] = loader->current_shard_idx
*(size_t*)&header[32] = loader->current_sample_idx

memset the header to zero first — the gaps are reserved padding.

Then, with device_to_file and a size of multi_gpu_config.shard_num_parameters * sizeof(float) each:

• m_memory
• v_memory
• master_weights (only if use_master_weights)

If loader->should_shuffle, also fwrite:

• loader->glob_result.gl_pathc (one size_t)
• loader->shard_indices (gl_pathc ints)
• loader->shard_num_samples (one size_t)
• loader->intra_shard_indices (shard_num_samples ints)
• loader->shuffle_rng (one mt19937_state)

load_state(step, model, loader, filename)

Read and assert the header: magic, version, and that num_processes and process_rank match the current run — a state file is per-rank and cannot be replayed under a different topology. Restore the fields, file_to_device the optimizer tensors, and restore the shuffle state.

Then, in this order:

1. Call gpt2_update(..., init_from_master_only=true) so the working-precision weights are re-derived from the master weights using model->rng_state_last_update — reproducing the exact stochastic rounding of the interrupted run.
2. Only afterwards, set model->rng_state from header[20].
3. Call dataloader_resume(loader, current_shard_idx, current_sample_idx) — not dataloader_reset, which would reshuffle and discard the restored permutation.

write_checkpoint(output_log_dir, step, model, train_loader, multi_gpu_config)

if rank == 0: gpt2_write_to_checkpoint -> "{dir}/model_{step:08d}.bin"
all ranks:    save_state               -> "{dir}/state_{step:08d}_{rank:05d}.bin"
multi_gpu_barrier()
if rank == 0: create empty "{dir}/DONE_{step:08d}"

The barrier before the DONE file is what makes the sentinel meaningful: DONE exists only once every rank has finished writing its state.

delete_checkpoint(output_log_dir, step, multi_gpu_config)

Mirror image: rank 0 removes the model and DONE files, every rank removes its own state file.

---

Part 6 — MFU Reporting (train_gpt2.cu)

gpt2_estimate_mfu(model, num_tokens, dt)

N = model->num_parameters
L = model->config.num_layers, C = model->config.channels
T = model->seq_len            // the runtime sequence length, not config.max_seq_len

flops_per_token = 6*N + 6*L*C*T
// 6*N   : all weight matmuls — each weight is used in 1 forward (2 flops) and 2 backward (4 flops)
// 6LCT  : the attention matmuls, usually a small contribution

flops_per_step  = flops_per_token * num_tokens
flops_achieved  = flops_per_step / dt                                  // FLOP/s
flops_promised  = get_flops_promised(deviceProp.name, PRECISION_MODE) * 1e12f
mfu             = flops_achieved / flops_promised

Do the flops_per_token arithmetic in size_t; 6 * L * C * T overflows int at realistic sizes.

Return -1.0f when flops_promised < 0 (unknown GPU) — the driver then prints -1 rather than a fabricated number. This is why MFU shows as -1 in A3.

N deliberately over-counts slightly: it includes LayerNorm parameters, biases, and position embeddings, which do not participate in matmuls. Token embeddings are correctly included, because GPT-2 weight-shares them with the classifier head. See the comment in the function for the full reasoning.

---

Testing

Scheduler

The schedulers are pure functions; test them directly rather than through a training run.

cat > /tmp/sched.c <<'EOF'
#include <stdio.h>
#include "llmc/schedulers.h"
int main(void) {
    LearningRateScheduler s;
    lr_scheduler_init(&s, "wsd", 1e-3f, 100, 1000, 0.0f);
    int steps[] = {0, 50, 99, 100, 500, 799, 800, 900, 999};
    for (int i = 0; i < 9; i++)
        printf("step %4d: lr=%.8f\n", steps[i], get_learning_rate(&s, steps[i]));
    return 0;
}
EOF
cc -O2 -I. -o /tmp/sched /tmp/sched.c -lm && /tmp/sched

Check the boundaries by hand: step 0 → 1e-5 (not 0), step 99 → 1e-3 (warmup ends), steps 100–799 → flat 1e-3, step 800 → 1e-3 (decay_ratio 0), step 999 → 1e-3 * (1 - sqrt(199/200)) ≈ 2.5e-6. Repeat for "cosine" and "linear"; your values must match the reference to 6 significant figures.

Outlier detector

cd dev/test && make test_outlier_detector && ./test_outlier_detector

This feeds 256 values through the detector and checks that the first 128 return NaN and the rest return sane z-scores. It exits non-zero with a specific message on failure.

Gradient norm and checkpoint resume

make train_gpt2cu

# run 1: train 50 steps, checkpointing every 20
./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
               -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin \
               -o /tmp/ckpt -n 20 -x 50 -v 10 -k cosine
# note the loss at step 50 and in /tmp/ckpt/main.log

# run 2: resume from the newest checkpoint in the same dir and finish
./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
               -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin \
               -o /tmp/ckpt -n 20 -x 50 -v 10 -k cosine -y 1
# the step-50 loss must be identical to run 1

Note how resume actually works: -y 1 scans -o’s directory for the highest-numbered DONE_* file and resumes from it. There is no flag that takes a checkpoint path. -n sets the checkpoint interval, -k selects the scheduler.

The relevant flags — again, all single-character, no GNU long options:

Flag	Meaning
-o	output log/checkpoint directory
-n	write a checkpoint every N steps
-nk / -nm	how many checkpoints to keep / major checkpoint interval
-y	1 = resume from the newest checkpoint in -o
-k	scheduler: cosine, linear, constant, wsd
-u	warmup iterations
-q	final LR fraction
-l	peak learning rate
-x	total steps
-sl / -sg	skip update above this loss / grad-norm z-score

Gradient clipping

grad_clip is hard-coded to 1.0f in the training loop, so there is no flag to toggle it. To see clipping work, watch the norm: field in main.log while pushing the learning rate up:

./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
               -o /tmp/ckpt -x 20 -l 0.1
grep norm: /tmp/ckpt/main.log

With a correct gpt2_calculate_grad_norm the norms will be well above 1.0 and the run stays finite because grad_scale shrinks each update. If you return a constant from gpt2_calculate_grad_norm, grad_scale is always 1.0 and the loss diverges to NaN within a few steps — that contrast is the test.

---

Common Mistakes

• The step + 1 in warmup. Applies to all three decaying schedules. See Part 1.
• WSD written as a linear decay over the last 10%. It is 1 - sqrt over the last 20%, from an integer decay_point.
• Dividing the outlier statistics by count instead of WINDOW_SIZE, or returning 0.0 instead of NaN while warming up.
• rng_state vs rng_state_last_update. Two RNG states are saved: the state at the current step, and the state at the last weight update (needed to reproduce stochastic rounding when re-deriving weights from master). Restoring them in the wrong order silently desynchronises a resumed run.
• Calling dataloader_reset on resume. It reshuffles and throws away the restored permutation. Use dataloader_resume.
• device_to_file vs cudaMemcpy + fwrite. GPU tensors must go through device_to_file/file_to_device, which do chunked DMA through a pinned buffer. Rolling your own is markedly slower on multi-GB optimizer state.
• Byte order. The state file is native-endian. Do not add swapping.
• Modifying the gradient buffer to clip. Clipping is a scalar passed into adamw_update; the gradients themselves are left alone.

---

Checking your work

./grade A5

Checks the schedulers against golden values to 6 significant figures, runs dev/test/test_outlier_detector, and — when a GPU is present — does a real checkpoint/resume round trip and compares the step-20 losses. The scheduler and outlier-detector checks need no GPU.

---

Scoring

What ./grade A5 checks, and what each part is worth.

Component	Points
Cosine and linear LR schedules match reference to 6 sig figs (incl. warmup off-by-one)	10
WSD LR schedule matches reference (80% decay point, 1 - sqrt)	10
global_norm_squared kernel and get_max_num_block_sums correctness	10
gpt2_calculate_grad_norm correct, deterministic, ZeRO-1-aware	10
outlier_detector passes dev/test/test_outlier_detector (incl. NaN warmup)	10
save_state writes a valid file; load_state reads it back identically	15
Checkpoint resume: step-50 loss identical after interrupting at step 20	15
gpt2_estimate_mfu within 5% of an ncu-measured value	10
Logger writes the exact documented format; rank-0-only; resume truncation correct	5
write_checkpoint / delete_checkpoint naming and DONE sentinel correct	5
Total	100