I trained GPT-2 with full-batch gradient descent so you don’t have to

shawnghu14 Oct 2026 23:09 UTC

67 points

Note: The results of this are not as clean as they would naively appear, since a lot of training parameters are tuned for non-full-batch gradient descent. See discussion in the comments.

I’ve left this up in case it’s still useful, but be careful about drawing conclusions naively from the result, and if you’re trying to do anything mathematical and precise look at the reproduction code for confounders.

The impetus for doing this was provided in large part by Charles Renshaw-Whitman.

Learning dynamics people suggested to me that noise was essential in the training of neural networks. I thought this was kind of counterintuitive—I usually naively operate from a model in which the pretraining data basically specifies the outcome by virtue of the samples all “pulling” on the loss landscape in some relatively uniform way which incentivizes converging to a point that generalizes well, but instead it seems to be the case that the actual dynamics of the optimization process itself have predictable and measurable effects on the qualitative outcomes of training. The idea seems reasonable based on vague intuitions I have about simulated annealing.

I wanted to see directly that this is true at some nontrivial scale, I happened to have access to a pool of free compute that no one else seemed to be using, and I figured this would be an interesting object of study for compute-constrained researchers with an interest in studying learning dynamics closely, so I did some very rough calculations and realized this was feasible to do with GPT-2.

The main result is that on a fixed 800M-token subset of FineWeb, ordinary training of GPT-2 (based on the NanoGPT speedrun) with a batch size of 393K, gets to 3.30 validation loss in one epoch (training loss is the same as validation loss given no repeats on the data). Exact full-batch gradient descent on exactly the same tokens fits the training set better (3.09 training loss in 350 steps), but its validation loss converges at about 3.50 in somewhere between 260 and 310 steps (presumably the training loss hasn’t converged by this point).

This suggests the batch selection noise is in fact doing something essential. (Note that this removes the effect of noise due to batch selection, but not noise due to finite step size.)

The code is on GitHub. Everything the runs produced is on Hugging Face: metrics and dynamics statistics for every run, the full gradient and update at every step of the full-batch run, per-token losses on all 800M training tokens every 10 steps, checkpoints, and the exact training set.

Appendix: some methodological choices

These particular parameters are somewhat arbitrary, and were tuned around the size of dataset which, for some batch size, we could use to train a model to achieve approximately 3.28 validation loss (as a heuristic to save on time spent waiting). 3.28 is completely arbitrary and taken from the NanoGPT objective.

For example, I determined that it is also possible to achieve a validation loss of 3.23 in the same 350 steps with a batch size of 50M tokens (consequently requiring 18B tokens). It is also possible to reach this loss with only 329M tokens and 1200 steps using some batch size/learning rate scheduling (as the NanoGPT speedrun record does, but I didn’t want to be confounded by optimizations like this). 3.28 is also not the point at which any relevant dataset or batch size configuration actually converges, e.g, if you keep training in the speedrun record configuration for longer you’ll get closer to 3.0.

Consequently, it’s possible that these checkpoints still lack some information about what happens when we get close to convergence points, etc., and generally we haven’t “finished training GPT-2” in a sense which can be said of any realistic training configuration. Nevertheless, they likely contain information relevant to early developmental phases of training.

An info dump from Claude about learning rate scheduling, as well as datapoints about other possibly relevant training runs follows:

Setup

The model is the current speedrun record (#360, Sept 2026), with the parts that don’t make sense for full-batch training removed:

The hashed n-gram embedding table. It’s a 65B-parameter lookup table keyed on the previous one or two tokens, and it memorizes any data that repeats: validation loss explodes the moment a token is seen twice. Full batch on a fixed set repeats every token every step, so it has to go.
The sampled softmax, which makes the loss biased and stochastic.
Every schedule that changes the loss function during training (sequence length, attention window, multi-token-prediction weights, batch size). The whole run uses the record’s final-stage settings.

What remains is an 11-layer, width-768 model with the record’s FP8 and FlashAttention-3 kernels and its optimizer: a Muon-family optimizer (ANVIL) for the matrices and Adam for everything else. These changes cost about 40% more steps to reach 3.28 than the record (about 1,625 instead of 1,190, at 393K tokens per step).

The dataset is the first 800M tokens of FineWeb10B. I picked this size because small-batch training on it barely overfits: 2,372 steps at 393K tokens per step, about 1.2 passes over the data, reaches 3.248 validation loss with a train/val gap of 0.06.

For full batch, each step accumulates the exact gradient of the mean loss over all 800M tokens (16,272 microbatches of 49K tokens, with fp32 accumulators), then takes one optimizer step. I ran 350 steps. The budget comes from a sweep of large batches on fresh data, which extrapolates to roughly 300 to 330 steps to reach 3.28 at infinite batch; a fresh-data run at 50M tokens per step reaches 3.23 in 350 steps. The full-batch run took about 100 GPU-hours on 4 H200s on preemptible compute, surviving 106 preemptions by checkpointing both between and partway through steps.

Result

	Tokens per step	Steps	Training loss	Validation loss
Small batch, 800M set	393K	2,372	3.186	3.248
Full batch, 800M set	800M	350	3.094	3.521 (best 3.498)
Fresh data, no repeats	50M	350	3.228	3.230

“Training loss” here is loss on a fixed 10M-token sample of the training set. At equal training loss, full batch’s validation loss is about 0.2 nats worse than small batch’s. Full batch’s validation loss is flat from about step 260 to the end, while its training loss keeps falling. Intermediate batch sizes sampled from the same 800M set (4.7M to 200M tokens per step) fall progressively further behind fresh data at the same batch size as each token gets seen more times.

What the dynamics look like

The run sits at the edge of stability the whole way through:

Consecutive full gradients point in nearly opposite directions: the median cosine between them is −0.945. On 98% of steps, every single microbatch gradient is anti-aligned with the previous step’s full gradient.
Measuring the loss along each step’s actual update, the step overshoots the minimum in its own direction by about 2× and lands at about the same loss on the far side.
The exact loss over all 800M training tokens goes up on about a third of steps.
The top Hessian eigenvalue keeps growing: 48, 82 and 113 at steps 100, 200 and 300.

This is the picture from Cohen et al. on edge-of-stability, here with a Muon-family optimizer on a 2026 speedrun architecture.

Caveats

One run per configuration. Seed-to-seed noise in validation loss for this model is about 0.003, so the gap isn’t noise, but its exact size isn’t pinned down.
Full batch on a fixed set means 350 passes over the same tokens, so this doesn’t separate “no gradient noise” from “too many repetitions”. Full batch plus artificially injected minibatch-shaped noise would separate the two; I haven’t run that.
The learning rate was tuned on fresh data, not on the fixed set.
I also tried a schedule ten times as long on fresh data as a convergence reference. In that run the model’s FP8 training forward pass and its bf16 evaluation forward pass drifted apart (3.15 against 4.62 on the same unseen data), so I don’t trust its number. The full-batch run doesn’t have this problem: its two forward passes agree within 0.015 at the end.

All runs

Every run that contributed data to this post, grouped by purpose. “Tokens processed” counts every token in every step, including repeats. “Train” is loss on a fixed 10M-token sample of the 800M training set for runs on that set, and on the training subset for the ablation subset runs; “—” means it wasn’t measured. “Val” is the standard speedrun validation set (10.5M tokens) at the end of the run, with the best value in parentheses where it differs. Unless noted, runs use the fixed-objective model described above, on 4 H200s or 8 H100s; metrics for all of them are in the Hugging Face dataset.

Main comparison, on the fixed 800M-token set:

Run	Tokens per step	Steps	Tokens processed	Unique tokens	Train	Val	Notes
Full batch	800M	350	280B	0.8B	3.094	3.521 (3.498)	the main run; lr multiplier 1.2
Small batch (control)	393K	2,372	0.93B	0.8B	3.186	3.248	about 1.16 epochs; lr multiplier 1.73; at the end of the first epoch (step ~2,034) train ~3.26, val ~3.30

Batch-size ladder on the same 800M set (minibatches drawn without replacement within each pass), with fresh-data runs at the same batch size, learning rate and schedule. Runs marked “halted” were stopped by a rule comparing them against fresh data at matched steps:

Run	Tokens per step	Steps	Tokens processed	Unique tokens	Train	Val	Notes
Fixed set, 4.7M	4.7M	500	2.4B	0.8B	3.140	3.271	3 epochs
Fixed set, 4.7M, lr 1.73	4.7M	500	2.4B	0.8B	3.148	3.277	learning-rate check
Fresh data, 4.7M	4.7M	500	2.4B	2.4B	3.261	3.244
Fixed set, 19M	19M	277 of 400	5.2B	0.8B	3.247	3.411	halted; 6.5 epochs
Fresh data, 19M	19M	400	7.5B	7.5B	3.252	3.228
Fixed set, 50M	50M	160 of 350	8.1B	0.8B	3.503	3.621	halted; 10 epochs
Fresh data, 50M	50M	350	17.6B	17.6B	3.228	3.230	FineWeb10B then FineWeb100B
Fixed set, 100M	100M	150 of 350	15B	0.8B	3.555	3.675	halted; 19 epochs
Fixed set, 200M	200M	150 of 350	30B	0.8B	3.537	3.674	halted; 38 epochs

Large-batch sweep on fresh data, used to estimate how many steps full batch would need (about 300 to 330 at infinite batch). Steps to 3.28 is where the validation curve crosses 3.28, for runs that got there:

Tokens per step	Steps	LR multiplier	Tokens processed	Val	Steps to 3.28
393K	1,200 / 1,800 / 2,400	1.73	0.47B / 0.71B / 0.94B	3.341 / 3.274 / 3.235	— / 1,775 / 2,155
1.2M	700 / 1,100	1.73	0.83B / 1.3B	3.294 / 3.212	— / 887
2.4M	500	0.8 / 1.2 / 1.73 / 2.6	1.2B	3.322 / 3.301 / 3.303 / 3.364	—
2.4M	900	1.73	2.1B	3.180	638
4.7M	400	0.8 / 1.2 / 1.73 / 2.6	1.9B	3.320 / 3.300 / 3.314 / 3.349	—
4.7M	800	1.73	3.8B	3.152	515
9.4M	350	0.8 / 1.2 / 1.73 / 2.6	3.3B	3.320 / 3.297 / 3.322 / 3.371	—
9.4M	700	1.73	6.6B	3.139	440
19M	300 / 600	1.73	5.7B / 11.3B	3.341 / 3.134	— / 381
19M	500	1.73	9.4B	3.178	367

Full-batch pilots on smaller fixed sets (all stopped early once they clearly overfit):

Fixed set size	Steps	Tokens processed	Val (best)	Notes
25M	282	7.1B	10.73 (5.60)	memorized the set: training loss 0.04 by step 224
50M	111	5.6B	5.62 (4.96)	validation bottomed out at step 49
100M	78	7.9B	4.594 (4.579)
200M	90	18B	4.338 (4.308)

Convergence references on fresh data, at 1×, 3× and 10× the schedule length that reaches 3.28:

Run	Tokens per step	Steps	Tokens processed	Val	Notes
This model, 1×	393K	1,720	0.68B	3.282
This model, 3×	393K	5,160	2.0B	3.170
This model, 10×	393K	17,200	6.8B	not meaningful	FP8 training and bf16 eval passes drifted apart (3.15 against 4.62 on the same unseen data)
Old record #18 (Jan 2025), 1×	524K	1,390	0.73B	3.278	the original record code, bf16, no FP8
Old record #18, 3×	524K	4,170	2.2B	3.103
Old record #18, 10×	524K	13,900	7.3B	2.986

Ablation of the original record toward a full-batch-compatible configuration, using the record’s own code on 8 H100s (each change is cumulative with the ones above it; “subset” runs train on only the first N tokens, repeated):

Run	Steps	Tokens processed	Train	Val	Notes
Stock record, 2 seeds	1,194	0.33B	—	3.279 / 3.276	batch-size ramp 131K to 393K
Full softmax, 2 seeds	1,194	0.33B	—	3.276 / 3.272
Full softmax	1,072 / 1,162	0.30B / 0.32B	—	3.301 / 3.283	step-count sweep
Fixed objective, tied embeddings	1,194	0.47B	—	3.247	constant 393K batch from here on
Fixed objective	772 / 922 / 1,052 (2 seeds)	0.30B / 0.36B / 0.41B	—	3.349 / 3.304 / 3.274, 3.278	step-count sweep
Fixed objective, constant LR	1,194	0.47B	—	3.247
Fixed objective, with eval window extension	1,052	0.41B	—	3.281
Fixed objective, 262K batch	1,194	0.31B	—	3.319
Fixed objective, untied embeddings	1,194	0.47B	—	3.260
Fixed objective, subset 400M / 200M	1,052	0.41B	1.79 / 1.15	4.050 / 5.837	hash table memorizes repeats
Fixed objective, subset 100M / 50M	1,194	0.47B	0.95 / 1.08	7.613 / 8.834
No hash table	1,194 / 1,602 (2 seeds) / 1,772 / 2,372	0.47B / 0.63B / 0.70B / 0.93B	—	3.333 / 3.279, 3.287 / 3.261 / 3.217	step-count sweep
No hash table, subset 400M / 300M / 200M	1,602	0.63B	3.21 / 3.03 / 2.93	3.292 / 3.326 / 3.468
No hash table, subset 800M / 400M / 200M	2,372	0.93B	3.12 / 3.08 / 2.76	3.221 / 3.259 / 3.542
No hash table, subset 100M / 50M	1,194	0.47B	2.58 / 1.62	4.010 / 5.724

Related work

Unveiling the Basin-Like Loss Landscape in Large Language Models (Chen et al., 2025)
The Break-Even Point on Optimization Trajectories of Deep Neural Networks (Jastrzębski et al., 2020) (contains a study of roughly the same thing, but on CIFAR-10)
Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective (Wen et al., 2024)
Understanding Optimization in Deep Learning with Central Flows (Cohen et al.)

shawnghu14 Oct 2026 23:09 UTC

67 points

7 comments2 min readLW link

leogao 15 Oct 2026 1:34 UTC
20 points
0
did you retune for optimal step size of full batch GD?
- shawnghu 15 Oct 2026 21:10 UTC
  3 points
  0
  Parent
  not adequately--
  - i ran with a global learning rate multiplier of 0.8/1.2/1.73 on nonrepeating data with batch sizes 2.36m, 4.7m, 9.4m and found that 1.2 did best on val loss
  - i ran with 1.2/1.73 on the fixed 800m token corpus with batch size 4 and found that 1.2 was better on val loss.
  - i didn’t tune decay schedule whatsoever, in all cases it’s max LR for the first 20% and then linear decay to 0.3x max LR by the end of training.
  between this and what’s in the replies to the other comment, i think that this invalidates a lot of conclusions one might naively draw.
  
  i think the existing results still suggest that “just fitting to a large training set” is not sufficient for generalization performance? since if you decreased the learning rate presumably you can only overfit even harder, and if you increase the learning rate to improve on val you are admitting you need a form of noise (step size noise) to make things work.
  
  but the fact that various other aspects of the training were tuned for a different setup makes even this not completely established.
  
  i’ll leave the thing up but with a note at the top of the post, since it might still be useful for something or other.
  - leogao 15 Oct 2026 22:43 UTC
    2 points
    0
    Parent
    unfortunately it’s hard for me to talk about these things for obvious reasons. all i’ll say is it’s worth being really careful what conclusions you draw.
rotatingpaguro 15 Oct 2026 7:22 UTC
6 points
0
Why is 3.30 ⇒ 3.50 “essential”? I also wanted to ask the same question as leogao, did you retune the method for the change in general? Unlike leogao I don’t want to ask just about step size, but about any small detail in the method. If you train GPT2 with the method GPT2 was intended to be trained with, vs. another one, there’s a lot of stuff that’s probably implicitly optimized. I ask this because in the past I remember a case in which I was trying a new data analysis method to do something and it took a lot of time to make sure I was comparing the previous and new method in conditions that made them comparable in a meaningful sense.
- shawnghu 15 Oct 2026 21:18 UTC
  1 point
  0
  Parent
  Yeah, “essential” is not really the right word. I just meant that this was the main result I thought was worth reporting.
  
  Having thought about what you said, and then checked the details of training, I realize that there are a ton of ways in which the original method was tuned for fast training of GPT2 (LR schedule, a special Adam beta for the embeddings, a bunch of specific learning rates for the different components, tricks for normalization and orthogonalization in Muon), and this all significantly complicates the analysis.
  
  I did try to set the dataset size and number of optimization steps in an attempt to get as direct of a comparison as possible, but that’s pretty much it.
  
  One may still be able to surmise some general trends by viewing the relationship between this run and the various other runs whose stats are reported in the Claude-appendix, but it’s undoubtedly much messier than I made it seem.
programjames 15 Oct 2026 20:12 UTC
1 point
0
In case you’re wondering why noise is important, here is something I wrote (and researched) that I never published.

Regularization for General Intelligence

The main result is modifying the learning flow from

to

via the unbiased estimator

There are several ways to arrive at this result.

Information Theory

Shorter programs, as measured by bitlength, are known to be more generally applicable. Thus, neural networks are typically trained to minimize the number of bits a particular input would need to reconstruct the output. While a fixed neural network could be considered an interpreter, it is a little silly to do the same when we are actively training it. Indeed, according to the learning flow

we are adding

bits of information to the network. To minimize this information, we could use a regularization term

That regularization term adds information, so we really need

As

and

we can use the unbiased estimator

Measure Theory

Suppose is -continuous, in the sense that

The total error over all possible weights is

In frequency space, we find

so the naïve gradient flow will take much longer to correct high-frequency errors, the exact errors that cause a model to fail to generalize. Instead, we can use the Sobolev inner product

By definition,

so inspection of the above inner product gives

Neural networks are almost linear, so , and we recover the learning flow

Probability Theory

It is an incorrect assumption, but suppose the true posterior of weights during training can be modeled by independent Gaussians

We wish to minimize

where the prior distribution for the weights is

The asymmetric divergence of Gaussians is easy to compute,

giving the learning flow

The variance changes according to

During gradient descent, the curvature of the loss landscape oscillates just around (in)stability, so

giving

with a single stable point

Overall, we get a learning flow somewhat like

which is close, but not quite the same as the results from information and measure theory.

Modification for Stochastic Gradient Descent

(This section likely has sign errors, didn’t proofread it carefully.)

When training, the exact loss is not computed. Instead, a small batch of data substitutes for the entire dataset, and the loss within the batch is averaged (rather than summed) for stability. Additionally, data is reused in every epoch. This effectively weighs the penalty for information in the error correction stage by

so the flow should be modified to

Also, weights are typically updated with finite step sizes, rather than using an adaptive method that keeps them within some tolerance of the true flow. Let be the eigenpairs of the Hessian . If the error is

then

The second term will produce oscillations, so , but , and in expectation over time

If, after training long enough, the curvature stabilizes so , the chain rule gives

Empirically, the curvature becomes about uniform in every direction, so the eigenvectors may as well be randomly chosen and is almost parallel to . Also, neural networks are almost linear, so the magnitudes are about equal. Then

and the expected flow for gradient descent is, up to third order,

To match the infinitesimal flow, we need to multiply by a second , giving the learning flow

In code, an unbiased estimate can be computed with
```
s ~ Exp(1)
t ~ Exp(1)
Z ~ N(0, I)
λ ← EPOCHS / BATCH_SIZE
w' ← w + sqrt(2λt-2s)Z
Loss ← L(w')
```
For SGD:
```
s ~ Exp(1)
t ~ Exp(1)
Z ~ N(0, I)
η ← LEARNING_RATE
λ ← EPOCHS / BATCH_SIZE
w' ← w + sqrt(2ληt-ηs)Z
Loss ← L(w')
```
For Adam:
```
s ~ Exp(1)
t ~ Exp(1)
Z ~ N(0, I)
η ← LEARNING_RATE
m ← FIRST_MOMENT / SECOND_MOMENT # Adam's μ/√ν
λ ← EPOCHS / BATCH_SIZE
w' ← w + sqrt(2ληt-mηs)Z
Loss ← L(w')
```
References

Cohen et al., 2025
Understanding Optimization in Deep Learning with Central Flows. In International Conference on Learning Representations (ICLR), pages 39644–39758. Editors: Y. Yue, A. Garg, N. Peng, F. Sha, R. Yu. Link

Duvenaud, 2020
Bayesian Neural Networks. University of Toronto. Link

Cite
```
@techreport{camacho2025inforeg}
  author      = {James Camacho},
  title       = {Weight Regularization for Generalization},
  institution = {WedgeProduct, Inc.},
  year        = {2025}
```
ionel chiosa 15 Oct 2026 1:18 UTC
1 point
0
This looks super useful! I wanted to train something like this on toy distributions for a while, but I didn’t get the time or compute to justify it.

I trained GPT-2 with full-batch gradient descent so you don’t have to

Setup

Result

What the dynamics look like

Caveats

All runs

Related work

Regularization for General Intelligence

Information Theory

Measure Theory

Probability Theory

Modification for Stochastic Gradient Descent

References

Cite