Andrej Karpathy: Reproducing GPT-2 from Scratch

Build the 124-million-parameter GPT-2, accelerate its training, and evaluate the resulting model with validation loss and HellaSwag.

AKAndrej Karpathy연사
채널
Andrej Karpathy
게시일
2024년 6월 9일
영상
4:01:26
읽기
28분
원문
영어
YouTube에서 열기
스크롤하면 영상이 오른쪽 아래로 작아져 계속 재생됩니다. 자막은 읽기 언어를 따릅니다.
읽기 언어2개 언어로 읽을 수 있음
이 글은 아직 한국어 번역이 없어 영어 원문으로 표시합니다.
요약

30초 요약

Andrej Karpathy implements the small GPT-2 architecture in PyTorch, verifies it against the released checkpoint and trains a new model. He explains batching, initialization, GPU efficiency, distributed optimization and the data and evaluation choices behind the results. The reproduction differs from the original run, so its comparisons come with important caveats.

  1. 1Reproduce the small checkpointThe exercise targets GPT-2 124M; loading the released weights checks the architecture before training from random initialization.
  2. 2Tie embeddingsThe input embedding and output projection share a tensor, reducing the number of independent parameters.
  3. 3Measure the hardwareMixed precision, compilation, flash attention and aligned dimensions improve the demonstrated GPU throughput. Their effects depend on hardware and workload.
  4. 4Build the effective batchGradient accumulation and distributed data parallel achieve a larger effective batch while respecting device memory limits.
  5. 5Interpret the results carefullyThe trained model compares favorably on the demonstrated checks, but different data, possible benchmark overlap and evaluation details limit direct comparison.

핵심 정리

  1. Loading the released checkpoint validates the implementation before training from scratch.
  2. Weight tying and initialization affect both parameter count and training behavior.
  3. Precision, kernel fusion and attention implementation can substantially change throughput on a given GPU.
  4. Gradient accumulation and data parallelism combine microbatches into the intended effective batch.
  5. Validation loss and a benchmark offer useful but incomplete comparisons, especially when training data differs.
전문

챕터별로 다듬은 전문 읽기

원본 자막만 보기
01

The target: GPT-2 124M

In this Zero to Hero lesson I will reproduce the 124-million-parameter GPT-2, the smallest model in the GPT-2 family. The larger family provides a way to see how performance on tasks changes with model scale. GPT-2 is a family rather than a single model. In the paper the variants range from the small model reproduced here to much larger ones, and larger models tend to improve downstream scores.

In this Zero to Hero lesson I will reproduce the 124-million-parameter GPT-2, the smallest model in the GPT-2 family. The larger family provides a way to see how performance on tasks changes with model scale.

GPT-2 is a family rather than a single model. In the paper the variants range from the small model reproduced here to much larger ones, and larger models tend to improve downstream scores.

The paper’s original table miscounted the parameter totals, according to the released repository. I use the corrected 124 million figure for the smallest model rather than treating the table as exact.

The model has 12 Transformer layers and a width of 768. I assume some familiarity with the architecture from the previous video; here the aim is to implement and train it.

The target is a model with low validation loss on held-out text. I estimate that reproducing this small model on the cloud hardware used in the lesson could take around an hour and ten dollars; those figures are specific to the example.

OpenAI released the GPT-2 weights, but the paper leaves some training details unclear. I will also consult the GPT-3 paper for optimization settings because its architecture is not a large departure.

02

Inspecting the released checkpoint

I begin at the end: load the released checkpoint and sample from it. The original implementation uses TensorFlow; Hugging Face has converted the weights to a PyTorch implementation that is easier for this demonstration. The released weights give us a concrete target. The training procedure is less completely specified, so I will distinguish matching the architecture from reconstructing the original run exactly.

I begin at the end: load the released checkpoint and sample from it. The original implementation uses TensorFlow; Hugging Face has converted the weights to a PyTorch implementation that is easier for this demonstration.

The released weights give us a concrete target. The training procedure is less completely specified, so I will distinguish matching the architecture from reconstructing the original run exactly.

In the Transformers interface, the model name `gpt2` refers to the 124M checkpoint; the larger 1.5B model uses `gpt2-xl`. I extract the state dictionary so we can examine the raw tensors rather than treating the library as a black box.

The model state dictionary lets us inspect parameter names, tensors and shapes. The token embedding table has 50,257 rows and 768 columns: each vocabulary token has a learned 768-dimensional representation.

Each embedding row corresponds to a token in the GPT-2 vocabulary, not necessarily to a whole word. The learned vector is the network’s representation of that token.

GPT-2 also has learned position embeddings for up to 1,024 positions. Adding position information lets the model use where a token appears as well as its identity.

A plot of the position embedding matrix puts positions on one axis and channels on the other. The patterns are learned from an initially random table, which makes their emergence during optimization interesting.

When we plot the position embeddings, some channels vary smoothly with position and resemble waves. They began as learned parameters, so the visible structure emerged during optimization; interpreting exactly what each channel does is harder.

Some positional channels respond more strongly in particular ranges of positions. The curves are imperfect and noisy, and a visualization by itself does not establish a complete interpretation.

I tentatively interpret the jaggedness in these curves as a sign that the small model is not fully trained. That is an interpretation of the plotted weights, not a direct measurement of training completeness.

This differs from the fixed sinusoidal position encoding in the original Transformer. GPT-2 learns its position representations from scratch, and some acquire sinusoidal-looking patterns.

Other weight matrices also show visible structure, although a picture alone does not tell us what the network has learned. Inspecting those patterns can be interesting for mechanistic interpretability.

The sample prompt is “Hello, I’m a language model,” and the pipeline generates five continuations. Even with a fixed seed, library or code changes can mean that these samples differ from an older example.

Sampling from the released model produces coherent continuations. That confirms we can load the checkpoint; the next step is to build our own module and match its behavior.

The checkpoint can be sampled through the Transformers library before we implement our own version. This gives us a working reference for the later port.

Being able to copy the released weights into our own implementation establishes that it belongs to the same architectural class. Training from scratch will still produce different weights and may use different data.

03

Implementing the Transformer

GPT-2 moves layer normalization relative to the original Transformer and adds a final normalization. These placement details matter when reproducing an existing checkpoint. GPT-2 is decoder-only: it omits the encoder and the cross-attention to encoder outputs shown in the original Transformer diagram. Its causal self-attention and feed-forward blocks form the main repeated stack.

GPT-2 moves layer normalization relative to the original Transformer and adds a final normalization. These placement details matter when reproducing an existing checkpoint.

GPT-2 is decoder-only: it omits the encoder and the cross-attention to encoder outputs shown in the original Transformer diagram. Its causal self-attention and feed-forward blocks form the main repeated stack.

The module names are chosen to line up with the checkpoint. Matching the structure and naming is useful because it lets us compare and transfer the individual tensors.

I organize the model into token and position embeddings, a list of Transformer blocks, a final normalization and a language-model head. The list contains twelve blocks for the small checkpoint.

The token and position embeddings are indexed tables. The language-model head projects the final 768-dimensional representation to vocabulary-sized scores and, in this implementation, has no bias.

The residual stream carries a representation through the network. Each block normalizes its input before attention and the feed-forward network, then adds each result back to the stream.

The clean residual path is useful for optimization because addition sends gradients back through the skip connection as well as through the computational branch. The branch can then contribute its own learned update.

Pre-normalization leaves a direct residual path through the block. The attention and multilayer-perceptron computations each contribute an update to that path.

Attention allows positions to exchange information, while the multilayer perceptron processes the representation at each position. This is a useful intuition for the two repeated computations, not a complete explanation of everything they learn.

The feed-forward network applies two linear projections with GELU between them. It expands the representation internally and brings it back to the residual-stream width before the addition.

The feed-forward block expands the channel dimension and projects back. To match GPT-2 closely, I use its approximate GELU activation rather than assuming every GELU implementation is numerically identical.

The approximate GELU has a historical reason: the exact error function was expensive in the earlier environment discussed in the lesson. I use the approximation because matching GPT-2 is the goal here.

In multi-head attention, projections produce queries, keys and values. Splitting them into heads gives parallel attention streams, whose results are concatenated and projected back into the residual stream.

I contrast GELU with the flat negative region of ReLU, where gradients can be zero. GPT-2’s choice is one historical activation design; newer architectures may make different choices.

For causal attention, scores for future positions are excluded before softmax. This ensures the prediction at a position cannot depend on tokens that follow it in the training sequence.

The compact implementation treats the head dimension like an extra batch dimension, allowing PyTorch to compute the heads in parallel. Reshaping and transposing make the data layout work for those operations.

The causal mask prevents a position from attending to future tokens. Softmax turns the allowed attention scores into weights, and multiplying by values forms a weighted combination of the available information.

The multiple attention heads are a rearrangement of projected channels. After attention is computed independently per head, their outputs are reassembled before the final projection.

The model configuration uses 12 attention heads as well as 12 layers and 768 channels. The GPT-2 tokenizer has 50,257 entries, including byte-level pieces, merges and an end-of-text token.

04

Loading weights and running the model

With the architecture in place, I load the Hugging Face parameters into our own module. The names and shapes must match, and certain matrices need transposition because the implementations store them differently. Some older GPT-2 layers use a convention for the orientation of the weight matrix that differs from a PyTorch linear layer. Transposing those weights preserves the intended computation.

With the architecture in place, I load the Hugging Face parameters into our own module. The names and shapes must match, and certain matrices need transposition because the implementations store them differently.

Some older GPT-2 layers use a convention for the orientation of the weight matrix that differs from a PyTorch linear layer. Transposing those weights preserves the intended computation.

Not every entry in a state dictionary is a learned parameter. The port ignores mask buffers that can be reconstructed while copying the actual trained weights and checking their shapes.

I check the keys and tensor dimensions before copying the weights. Successfully loading a state dictionary is useful, but generation provides an additional end-to-end check.

The forward pass takes token indices, adds token and position embeddings, runs the sequence through the blocks and returns logits over the vocabulary at each position.

If the input has batch and time dimensions B and T, the output has shape B by T by vocabulary size. Each vector describes scores for the token that might follow that position.

05

Sampling text

To generate, I encode a text prefix with the GPT-2 tokenizer and put multiple copies into a batch. The model and input tensors must be on the same device. A prompt is converted to token IDs before being given to the network. Generated token IDs can be decoded back to text after the sampling loop.

To generate, I encode a text prefix with the GPT-2 tokenizer and put multiple copies into a batch. The model and input tensors must be on the same device.

A prompt is converted to token IDs before being given to the network. Generated token IDs can be decoded back to text after the sampling loop.

Batching several copies of the prefix lets us obtain multiple continuations in one pass. Their contents can diverge because the next token is sampled.

The sampling loop runs the model, takes logits at the last position, converts them to probabilities and samples a next token. It appends that token and repeats until the desired length.

Top-k sampling restricts the draw to a set of high-scoring candidates. This changes the distribution of continuations and can keep samples closer to likely text, while preserving randomness.

At each iteration the model sees the growing prefix. The new column of sampled tokens is concatenated to the existing batch before the next iteration.

The output is coherent even though the exact sample does not match the earlier example. Sampling can differ because of random-number generation and other implementation details; the important check here is that loading and generation work.

I add device detection so the script can use CUDA, Apple’s MPS backend or the CPU depending on what is available. Performance differs substantially between these choices.

The model and inputs need explicit device placement. Moving only one of them to an accelerator would leave incompatible tensors in the same calculation.

06

Batches and next-token loss

Now we move from inference to training. The model needs batches of token sequences and targets shifted by one token, so every position can learn to predict its successor. For an initial debugging run I tokenize a small text dataset. This is a deliberately simple starting point, not the final training corpus.

Now we move from inference to training. The model needs batches of token sequences and targets shifted by one token, so every position can learn to predict its successor.

For an initial debugging run I tokenize a small text dataset. This is a deliberately simple starting point, not the final training corpus.

For a small batch, the input and target are neighboring slices of the same token buffer. The last token of each input row has a corresponding target because the buffer contains one extra token.

The toy data is encoded with the same GPT-2 tokenizer. For next-token prediction the targets use the same stream shifted one position to the right.

A contiguous buffer can be reshaped into B rows of T tokens. Loading one extra token gives the matching targets: the input and target arrays are offset by one.

Every row represents a training sequence. A token in that row can attend to preceding positions in its own causal context, rather than to future targets.

For example, 24 tokens can be viewed as a four-by-six array: the first six become the first row, the next six the second, and so on. The target buffer is sliced one position later before being viewed in the same shape.

The model returns B by T by C logits, where C is the vocabulary size. For loss computation, the batch and time dimensions can be flattened so each prediction lines up with its target.

Cross-entropy measures how well the predicted distribution assigns probability to the actual next token. We pass the flattened logits and targets to PyTorch’s loss function.

The expected initial loss can be estimated from the vocabulary size when probabilities are close to uniform. This gives a rough sanity check before optimization.

At initialization, a nearly uniform prediction over the vocabulary gives a characteristic baseline loss. If the initial loss is wildly different, it is a useful signal to inspect the model or data.

For a uniform distribution over 50,257 tokens, the expected cross-entropy is approximately the logarithm of the vocabulary size. The training loss should start near that scale before the model learns useful preferences.

07

Optimization and the data loader

Before attempting a full run, I try to overfit a single batch. A small training problem should show a falling loss; otherwise, there may be a bug in the forward pass, gradients or optimizer. The optimization loop clears old gradients, evaluates the loss, runs backward and calls the optimizer step. These stages should be easy to inspect when a small fixed batch fails to fit.

Before attempting a full run, I try to overfit a single batch. A small training problem should show a falling loss; otherwise, there may be a bug in the forward pass, gradients or optimizer.

The optimization loop clears old gradients, evaluates the loss, runs backward and calls the optimizer step. These stages should be easy to inspect when a small fixed batch fails to fit.

Adam maintains running statistics of gradients to scale updates. The details of its first and second moments matter for optimization, but this first experiment mainly checks that our loop can learn.

A falling loss on the same fixed batch tells us the system can fit that small target. It does not tell us that the model generalizes to unseen text.

A device mismatch exposes a common PyTorch mistake: calling `.to(device)` on a tensor returns the moved tensor, so the result must be assigned. Moving the model alone does not move an independently created input buffer.

Overfitting one fixed batch is a debugging exercise. To train a useful model we need to keep supplying fresh batches rather than repeatedly optimizing the same examples.

I write a small data loader that advances through the token buffer, returning inputs and shifted targets. Once it reaches the end, it wraps around for this toy dataset.

The loader maintains a current position and advances it by the number of input tokens in a batch. It is intentionally minimal so we can verify the basic data flow before introducing large datasets.

The input batches begin on the CPU and are copied to the selected training device. On a small dataset, the number of iterations in an epoch depends on the batch size and sequence length.

The toy loader wraps when it cannot provide another full batch. A production data pipeline needs to consider boundaries and shuffling more carefully.

08

Weight tying and initialization

The token embedding table and the output projection have compatible dimensions. GPT-2 ties these weights: the two module attributes refer to the same underlying tensor. Weight tying shares representations between reading a token and scoring possible output tokens. It also reduces the number of independent parameters compared with separate tables.

The token embedding table and the output projection have compatible dimensions. GPT-2 ties these weights: the two module attributes refer to the same underlying tensor.

Weight tying shares representations between reading a token and scoring possible output tokens. It also reduces the number of independent parameters compared with separate tables.

The shared embedding matrix contains roughly 40 million entries in this configuration. Tying it to the output projection avoids maintaining a second independent matrix of that size.

The motivation for tying input and output embeddings is that both tables relate vocabulary items to the network’s representation space. The implementation must preserve that sharing during optimization.

The shared tensors must actually refer to the same storage, not merely have equal values at initialization. A pointer check makes that distinction visible.

The same tied tensor participates at the input and at the output. During backpropagation, contributions from both uses update the shared parameters.

Tying is also an inductive bias: it assumes a useful relationship between the embedding used to read a token and the representation used to score it as an output.

Initialization needs attention too. The released GPT-2 code uses a normal distribution with a standard deviation of 0.02 for weights, while the paper leaves some details implicit.

The released code initializes ordinary weights with standard deviation 0.02 and biases at zero; it treats position embeddings somewhat differently. For this educational implementation I use 0.02 for the embedding modules.

Residual branches repeatedly add contributions to the stream. Scaling the initialization of certain residual projections helps control how the variance grows as layers accumulate.

Adding many independently initialized residual contributions can increase activation variance with depth. The adjusted initialization is intended to keep the initial residual stream better behaved.

I mark the relevant projection layers and apply the modified initialization while leaving the default scale for other linear and embedding weights. Implementation details should be checked against the architecture being reproduced.

In a simplified illustration, repeatedly adding 100 independent unit-scale random contributions grows the standard deviation to roughly 10. Scaling the contributions by the inverse square root of their count counters that growth.

The residual projections receive the depth-dependent scale during initialization. The rest of the implementation continues to use the chosen base standard deviation.

Each Transformer block has two additions to the residual stream, from attention and the MLP. That is why the scaling calculation counts twice the number of Transformer layers in this implementation.

Because the embedding and language-model head share a tensor, an initialization pass over modules can encounter that tensor twice. I point this out while checking the initialization code.

09

GPU throughput and numerical precision

We have a model, a loader and a basic training loop. The next challenge is throughput: the numerical format and GPU execution path can make a large difference to training speed. A GPU carries out many operations in parallel but also moves data between compute units and memory. A workload can be limited by arithmetic throughput, memory traffic or overhead.

We have a model, a loader and a basic training loop. The next challenge is throughput: the numerical format and GPU execution path can make a large difference to training speed.

A GPU carries out many operations in parallel but also moves data between compute units and memory. A workload can be limited by arithmetic throughput, memory traffic or overhead.

GPU memory capacity limits the size of batches and activations we can keep at once. Memory bandwidth and compute throughput are separate constraints.

Float32 uses 32 bits per number. Neural-network training can often tolerate lower precision in many operations, though numerical stability still has to be managed.

Floating-point formats trade exponent range against significand precision. The right choice depends on the values encountered and the hardware operations available.

NVIDIA Tensor Cores can accelerate matrix multiplication. TF32 uses a reduced-precision representation for the inputs to supported operations while accumulating with a wider format.

Much of the computation in this small model occurs in matrix multiplications, including the large projection from 768 channels to vocabulary-sized logits. Tensor Cores accelerate eligible matrix operations.

The Tensor Core precision modes involve both input representation and accumulator precision. Faster arithmetic comes with some loss of precision in the inputs.

I measure elapsed training time and tokens per second, not just whether the code runs. Accurate GPU timing requires accounting for asynchronous execution, for example by synchronizing before reading the clock.

In the example, the first iteration may be slower because setup and buffer initialization occur there. Later iterations provide a more representative view of steady-state throughput.

Before comparing optimizations I increase the batch to make better use of the GPU, then log timing and throughput. Benchmarking the same workload is essential for a useful comparison.

Batch size affects utilization and memory use. Once the experiment uses a larger batch, tokens per second is a more useful comparison than the duration of a step alone.

If the batch does not fit in memory, I suggest reducing it until it does. A larger batch can improve utilization, but its size must be interpreted together with the effective batch used by optimization.

Enabling TF32 can speed up eligible operations with a precision tradeoff, but it does not change every tensor in the program into a smaller storage format. The observed speedup depends on the workload.

PyTorch provides a setting for allowing reduced precision in float32 matrix multiplication. I enable it and rerun the same training workload.

The setting affects eligible matrix multiplications such as those inside linear layers. It does not mean that every activation and stored parameter has been converted to TF32.

The benefit can be smaller than the theoretical arithmetic gain if the workload is constrained elsewhere, including by moving full-precision tensors through memory.

Float16 and bfloat16 are two lower-precision formats. Bfloat16 retains a wide exponent range but fewer precision bits; that range can make it convenient for training on hardware that supports it.

The exponent controls the representable range, while the significand controls precision within that range. Bfloat16 and float16 allocate their limited bits differently.

With float16, gradient scaling can help prevent small gradients from underflowing. The exact safeguards depend on the precision format and the training setup.

PyTorch autocast selects lower precision for suitable operations while keeping others in a safer format. It is mixed precision: it does not mean that every parameter and intermediate is stored in the same format.

The autocast context applies its choices around the forward computation. We still inspect the actual tensors rather than assuming all parts of the program use a single dtype.

I inspect the dtype of weights and intermediate activations to illustrate the distinction. The optimizer can retain full-precision parameters while selected matrix multiplications use lower precision.

Autocast is placed around the forward pass and loss calculation; the backward pass and optimizer step remain outside that context. This follows the PyTorch guidance discussed in the lesson.

On the A100 used in this demonstration, autocast with bfloat16 improves the run. The numerical and speed results are hardware- and workload-specific.

The specific autocast policy matters: matrix multiplies are candidates for lower precision, while operations sensitive to precision can stay in float32. I do not assume every layer is automatically converted.

10

Compilation and kernel fusion

Next I introduce torch.compile. Compilation can reduce Python and kernel-launch overhead and fuse operations, although debugging and compatibility can affect when it is convenient to use. In this demonstration, compiling the model reduces the measured step time substantially after compilation. The number reflects the particular model, batch and A100 setup, rather than a guaranteed speedup.

Next I introduce torch.compile. Compilation can reduce Python and kernel-launch overhead and fuse operations, although debugging and compatibility can affect when it is convenient to use.

In this demonstration, compiling the model reduces the measured step time substantially after compilation. The number reflects the particular model, batch and A100 setup, rather than a guaranteed speedup.

Compilation has a startup cost, so timing should distinguish the first compiled execution from subsequent steps. The lesson’s speed comparisons concern this particular training workload.

Without fusion, intermediate results may be written to GPU memory and read back for the next operation. A fused kernel can keep more of that work close to the compute units.

A compiler can see a sequence of elementwise operations together and generate a combined kernel. In eager execution, the separate calls may each incur memory traffic and launch overhead.

Memory bandwidth matters even when arithmetic is fast. A series of simple elementwise operations can be limited by repeatedly moving intermediate arrays rather than by the number of calculations.

Fusion avoids some round trips to global memory by combining operations. The gain depends on how much work and intermediate data can stay within the fused computation.

The difference between on-chip storage and larger GPU memory explains why moving an intermediate tensor can be expensive. Reusing a piece of data while it remains close to compute avoids some of those round trips.

11

Flash attention and aligned dimensions

Flash attention reorganizes the attention computation to reduce expensive memory traffic. It may perform extra arithmetic, yet run faster because it avoids materializing and repeatedly moving a large attention matrix. The key insight behind flash attention is the arrangement of memory access. Counting floating-point operations alone can therefore give the wrong prediction about which implementation runs faster.

Flash attention reorganizes the attention computation to reduce expensive memory traffic. It may perform extra arithmetic, yet run faster because it avoids materializing and repeatedly moving a large attention matrix.

The key insight behind flash attention is the arrangement of memory access. Counting floating-point operations alone can therefore give the wrong prediction about which implementation runs faster.

PyTorch exposes optimized scaled dot-product attention. I replace the hand-written attention sequence with that operation and specify causal behavior so future tokens remain masked.

The optimized attention call must preserve the causal mask. A faster kernel is useful only if it computes the behavior needed by the language model.

Hardware kernels often work efficiently with dimensions aligned to convenient multiples. The GPT-2 vocabulary has 50,257 tokens, so I pad the output dimension to 50,304 for this run.

The tokenizer still produces only its original IDs; the extra rows are a computational padding choice. The improvement comes from a shape that is friendlier to the GPU kernels, not from adding useful new vocabulary.

Padding is a practical illustration of how apparently small dimension choices affect kernel efficiency. The best alignment is specific to the hardware and operations.

The added vocabulary entries are never emitted by the original tokenizer as input IDs. They do enter the output distribution, so the model can learn to assign them low probability.

The measured gain from padding varies with the PyTorch version and kernels in use. I report the settings in the demonstration instead of treating a single percentage as universal.

12

Optimizer and learning-rate schedule

The speed work gives us a practical training loop, but speed alone will not reproduce the result. We also need a reasonable choice of optimizer, learning-rate schedule, batch size and training data. I turn to the GPT-3 paper for more explicit hyperparameters. The reproduction uses AdamW and gradient clipping, drawing on the published settings while acknowledging that not all GPT-2 details are specified.

The speed work gives us a practical training loop, but speed alone will not reproduce the result. We also need a reasonable choice of optimizer, learning-rate schedule, batch size and training data.

I turn to the GPT-3 paper for more explicit hyperparameters. The reproduction uses AdamW and gradient clipping, drawing on the published settings while acknowledging that not all GPT-2 details are specified.

Gradient clipping measures the global norm of the gradients and rescales them when it exceeds a chosen threshold. This can limit unusually large updates that would otherwise destabilize training.

Printing gradient norms lets us see whether clipping is being invoked. The threshold controls the maximum norm used for an update rather than guaranteeing a particular loss.

A fixed learning rate is not the only choice. The schedule in this example starts with warmup and then decays according to a cosine curve toward a smaller final rate.

The clipping example uses a global gradient-norm limit of 1.0. The reported norm is also a diagnostic: unusually large values or spikes can point to instability or an unusual batch.

Warmup gradually increases the learning rate during the initial steps. Later, cosine decay smoothly reduces it over the selected training horizon; the schedule is one design choice among many.

For the small GPT-3 configuration used as a reference, the paper gives a peak learning rate of 6e-4. The example lowers the rate toward ten percent of its peak over the chosen schedule.

The scheduler computes a learning rate from the current step. Beyond the decay horizon, the implementation can hold the rate at its minimum rather than continuing the cosine indefinitely.

The schedule in this reproduction is not an exact copy of the paper: its decay horizon and total training horizon are set together here. That is one more reason to describe this as a reproduction with chosen settings.

Weight decay is commonly applied to matrix weights but not to biases and other one-dimensional parameters. I separate parameter groups so the optimizer can treat them differently.

The GPT-3 paper also describes increasing batch size during training. I skip that ramp in this implementation to keep the token and step arithmetic simpler, while explaining why smaller early batches may be efficient.

The loader iterates through data without replacement until an epoch boundary in this simple setup. Weight decay is then added as a separate regularization choice in the optimizer.

A fused AdamW implementation can combine work that would otherwise launch many small GPU kernels. I enable it when the optimizer and CUDA device support it.

Fusing optimizer operations reduces the overhead of launching many separate kernels for parameter tensors. The benefit depends on optimizer support and the device.

I put optimizer setup in a helper that takes the relevant hyperparameters and reports the parameter groups. This keeps the training loop focused on the sequence of steps.

The fused option is checked for availability before use. Older PyTorch versions or a different device may require the ordinary optimizer path.

Learning rate, batch size, decay and Adam settings interact. In this demonstration I follow published choices as a starting point rather than claiming each value is universally optimal.

13

Gradient accumulation

The target batch can be larger than fits in GPU memory. Gradient accumulation lets us divide that batch into microbatches, run several forward and backward passes, then make one optimizer update. The number of accumulation steps depends on the desired total tokens and the tokens in each microbatch. The purpose is to preserve the effective batch size without allocating it all at once.

The target batch can be larger than fits in GPU memory. Gradient accumulation lets us divide that batch into microbatches, run several forward and backward passes, then make one optimizer update.

The number of accumulation steps depends on the desired total tokens and the tokens in each microbatch. The purpose is to preserve the effective batch size without allocating it all at once.

Gradients accumulate in parameter buffers across backward calls until they are cleared. The optimizer step is performed only after the intended microbatches have contributed.

Because a batch loss is normally averaged, each microbatch loss must be divided by the number of accumulation steps before backward. Summing those gradients then corresponds to the average over the effective batch.

I demonstrate the arithmetic with a small regression example. Computing the full-batch gradient and accumulating appropriately scaled individual losses produces matching gradients apart from numerical rounding.

The regression example reveals the common error: summing gradients of unscaled per-example losses gives a different scale from the gradient of their mean.

Dividing each loss by the number of microbatches restores the intended mean. The backward calls add the resulting contributions into the same gradient buffers.

In the single-GPU configuration shown here, the calculated accumulation count is 32. That number is a consequence of this run’s target batch and microbatch dimensions, not a general rule.

14

Distributed data parallel

Distributed data parallel extends the run across multiple GPUs. Each process owns a model replica and processes a different slice of the batch; gradients are synchronized so the replicas update consistently. The torchrun launcher starts the requested number of processes and sets the environment variables used to identify them. This changes how the script is launched compared with a single-process run.

Distributed data parallel extends the run across multiple GPUs. Each process owns a model replica and processes a different slice of the batch; gradients are synchronized so the replicas update consistently.

The torchrun launcher starts the requested number of processes and sets the environment variables used to identify them. This changes how the script is launched compared with a single-process run.

The launcher supplies rank, local rank and world size. Each process selects its GPU, and rank zero can handle shared duties such as logging and checkpointing.

The master-process convention is a practical way to avoid duplicate output. Every rank still participates in the distributed computation and synchronization.

Each GPU’s microbatch size is only part of the total. World size and accumulation steps both contribute to the effective batch used for one update.

The data loader must give different ranks different samples. Otherwise several GPUs can duplicate the same work and the intended effective batch will not contain the expected diversity.

Without rank-aware logging, every process can print the same messages. Restricting shared output to the master process keeps logs readable.

I launch one process per GPU in the example. A distributed process group needs orderly cleanup when the run ends; otherwise the framework can report a shutdown warning.

The loader offsets each process’s starting position and advances by the combined amount consumed by all ranks. This gives each process its own portion of the stream.

After backward, distributed training averages gradients across ranks through an all-reduce operation. Each replica can then apply the same update to its copy of the parameters.

Synchronizing gradients across processes is communication work. The distributed wrapper handles much of the mechanism, but the training script still has to manage data and logging correctly.

With gradient accumulation, synchronizing on every microbatch wastes communication. We can accumulate locally and synchronize when the final microbatch of the update is complete.

After the final accumulated backward pass, the all-reduce result gives each process the averaged gradients. Scaling within each process and averaging across processes must both match the intended objective.

Throughput accounting should include all processes and the full effective batch. A single process’s token count understates the total work performed by the distributed run.

With eight GPUs, the example divides the accumulation work among the processes, reducing the per-process accumulation count from the single-GPU setup.

The single- and multi-GPU losses need not match exactly at every step because their data-loader ordering can differ. The comparison must account for which examples each run actually sees.

15

Preparing FineWeb-Edu

The final major ingredient is data. The public web contains valuable text but also spam and other noise, so large datasets are usually filtered and processed rather than fed in raw. I discuss FineWeb and use its educational subset for this experiment. The dataset and filtering choices differ from the original GPT-2 training data, which matters when interpreting a reproduction.

The final major ingredient is data. The public web contains valuable text but also spam and other noise, so large datasets are usually filtered and processed rather than fed in raw.

I discuss FineWeb and use its educational subset for this experiment. The dataset and filtering choices differ from the original GPT-2 training data, which matters when interpreting a reproduction.

The preparation script loads documents, tokenizes them and writes token shards. Sharding makes the large dataset easier to store and load than one enormous file.

Documents are converted into a stream of token IDs before being saved. A reserved end-of-text marker separates documents in the constructed stream.

The loader can select training or validation shards. Keeping a separate validation portion gives us data that is not used for parameter updates.

I calculate the number of optimization steps from the intended token budget and effective batch size. These choices set how much of the corpus the training run will consume.

The warmup step count is derived from a token count divided by the effective batch size. Expressing the schedule in tokens makes it easier to compare configurations with different batch sizes.

The estimated total run time comes from measured step time multiplied by the planned number of steps. It is an estimate for the demonstrated cloud setup.

16

Validation and HellaSwag

Validation uses a held-out split and runs without gradient updates. Periodic validation loss gives a view of generalization beyond the training batches, while periodic samples show what the model generates. At each validation interval the loader is reset and the model is evaluated on held-out batches without gradient tracking. The losses are averaged to make the reported value less dependent on one batch.

Validation uses a held-out split and runs without gradient updates. Periodic validation loss gives a view of generalization beyond the training batches, while periodic samples show what the model generates.

At each validation interval the loader is reset and the model is evaluated on held-out batches without gradient tracking. The losses are averaged to make the reported value less dependent on one batch.

I can also evaluate the released GPT-2 checkpoint on this same held-out data. That gives a useful comparison on this dataset, while leaving the differences in training data and procedure visible.

Evaluation and sampling are scheduled less frequently than training steps to limit overhead. I also encounter a compilation-related sampling issue and adjust the demonstration while noting it needs follow-up.

For an additional benchmark I use HellaSwag, a multiple-choice completion task intended to probe commonsense-like prediction. It is a useful signal, but it is not a complete measure of a language model.

A HellaSwag item provides one shared context and four possible endings. The benchmark asks which ending is the plausible continuation.

Benchmark difficulty and model scores change as systems improve. The historical scores in the lesson should be understood in the context of the benchmark and the models evaluated at that time.

Each example has a context and several candidate endings. I tokenize the candidates and use a mask so the scoring focuses on the continuation rather than padded positions or the shared context.

Padding and masks let candidate endings of different lengths be evaluated in a batch. Averaging over the relevant continuation tokens avoids simply favoring short candidates because they contain fewer tokens.

The script compares average continuation losses and chooses the lowest-loss candidate. I evaluate the released GPT-2 checkpoint with the same procedure to establish a reference for our run.

The model scores each candidate separately; it does not see all four options together and choose after comparing them in a single prompt. We compare the calculated scores outside the model.

The script’s result for the released 124M GPT-2 is around 29.5 percent in this evaluation, compared with a 25 percent random baseline for four choices. It is the reference for the later run.

With validation, benchmark evaluation and occasional generation in place, I start the larger run. The different measurements answer different questions: predictive loss, benchmark performance and qualitative sample behavior.

For distributed benchmark evaluation, processes evaluate portions of the examples and combine the counts of correct and total answers. The final accuracy uses the combined counts.

The training script periodically logs validation and benchmark results and saves samples. These records allow the overnight run to be inspected after it finishes.

17

Results, caveats and further work

The next morning I inspect the results. The reproduced model improves substantially over training and reaches a HellaSwag score that compares favorably with the released small GPT-2 checkpoint in this setup. On the held-out FineWeb-Edu split, the new model’s validation loss falls below the released checkpoint’s. Since that checkpoint was trained on a different distribution, I treat this as a cross-check, not a perfectly fair comparison.

The next morning I inspect the results. The reproduced model improves substantially over training and reaches a HellaSwag score that compares favorably with the released small GPT-2 checkpoint in this setup.

On the held-out FineWeb-Edu split, the new model’s validation loss falls below the released checkpoint’s. Since that checkpoint was trained on a different distribution, I treat this as a cross-check, not a perfectly fair comparison.

That comparison has caveats: the training corpus and token budget differ, and benchmark contamination or differences in evaluation code could affect the result. It should not be treated as an exact replay of the original training run.

The initial run processes approximately 10 billion tokens. It exceeds the small GPT-2 reference on the displayed HellaSwag evaluation, but the model’s dataset and potential benchmark overlap complicate comparisons of training efficiency.

The benchmark may respond to differences in the training data and evaluation procedure. Comparing one score is informative, but it does not establish that two training histories are equivalent.

I notice a periodic pattern in the training loss and suspect that the FineWeb-Edu sample’s ordering may be involved. I do not know the cause for certain and flag it for investigation in the repository.

I examine generated samples as another qualitative check. A sample can be fluent or amusing without proving factual accuracy or establishing broad capability.

The longer overnight run makes roughly four passes over the sample, processing about 40 billion tokens. Its HellaSwag score rises further and approaches, but does not quite reach, the small GPT-3 reference shown.

The loader in the demonstration repeats document order between passes. Shuffling document order could reduce the model’s exposure to the same accidental document boundaries.

One data detail is how documents are joined in a stream with end-of-text markers. Repeating the same ordering across epochs can preserve accidental adjacencies, so shuffling document order is worth considering.

The longer run shows about 33.24 percent on the displayed HellaSwag evaluation. I also inspect its continuations and find them more coherent than those from the shorter run, while recognizing that samples are only a qualitative check.

Changing model width, depth or context length moves the implementation toward other members of the GPT family. A larger configuration also changes its compute and memory requirements.

A closer comparison with GPT-3 would also change the context length from 1,024 to 2,048 tokens. Keeping the same total tokens per step would require adjusting the per-device batch dimensions accordingly.

Saving model parameters is enough for some later evaluation, but faithfully resuming an interrupted training run requires more state: optimizer statistics, the step, and the data and random-number positions can all matter.

This training produces a base language model that predicts document continuations. Turning it into a conversational assistant requires further work, such as supervised fine-tuning on conversations; this lesson stops at pretraining.

I point to llm.c as a related C and CUDA implementation that can train similar models efficiently. Its details and results are not automatically identical to this PyTorch walkthrough.

A lower-level implementation can make different engineering tradeoffs from PyTorch. I suggest consulting its own code and documentation when comparing speed or reproducing results.

The lesson brings together architecture, data, precision, efficient kernels, distributed training and evaluation. The build-nanogpt repository records the changes so readers can inspect the implementation and follow up on corrections.

The video closes by identifying unresolved issues: unusual loss periodicity, compilation interfering with generation and evaluation in this version, and the need for better shuffling. The repository is where follow-up corrections are documented.

같은 채널의 다른 글