Module 9 of 14 · one week

Scale: memory, modes, sharding, and the trillion context arithmetic

Where the company's numbers come from, what the public work actually reaches, and what one 4 GB laptop GPU can measure about the constants. The quadratic cost argument does rule out vanilla attention at industrial grids. It does not rule out attention, and every physics foundation model that exists is attention-based with patching or Fourier attention.

Primary sources: earlier episode 01:09:14 to 01:11:47; company site and press; FNO cost \(O(n\log n)\), ICLR p.6; MG-TFNO arXiv 2310.00120; Incremental FNO 2211.15188; U-NO 2204.11127; Transolver++ 2502.02414; FourCastNet 3 2507.12144; PDEBench Appendix G. Lab: ../labs/03_scale/.

Why this module exists

The seed episode is titled "Trillion Token Context. No, Really". The company site claims up to one trillion parameters in training and five trillion tokens or context at inference, and press coverage repeats "5 trillion data points in a single prompt" with no benchmark attached (research_notes/company-episode.md:107-124). Those are the largest numbers in the whole course, and they are the least documented.

Anandkumar's own derivation, in the earlier episode between 01:09:14 and 01:09:54, is short and deserves a serious answer: physical data is three-dimensional plus time, industrial resolution starts near 1,000 grid points per dimension, so the context runs to "hundreds of billions to even a trillion", and therefore "forget ever having a transformer for anything of this scale" (:42). That reasoning is sound. This module carries it out with units attached, applies it to the specific claims, and then measures the constants on real hardware so the estimates are not free parameters.

The arithmetic

Points. A \(1000^3\) grid holds \(10^9\) points per time step. With \(c\) channels and \(T\) time steps that is \(10^9cT\) numbers, and at \(c=10\), \(T=100\) it is \(10^{12}\): the trillion. Stored in float32 that is 4 TB per sample, before any activations exist.

Dense attention. Self-attention over \(n\) tokens costs \(O(n^2d)\) time and \(O(n^2)\) memory just for the score matrix. At \(n=10^9\) the score matrix has \(10^{18}\) entries, which is \(4\times10^{18}\) bytes in float32, or four million one-terabyte drives. This is the arithmetic behind the sentence "all of the world's compute will not be enough", and for dense attention over points it is correct.

What that does not rule out. Attention over \(p\) patches or slices, with \(p \ll n\), costs \(O(np + p^2)\). MPP and Aurora tokenize into patches, DPOT uses Fourier attention, Transolver attends over learned slices, and all three are linear in \(n\). Module 6's Proposition 6 says attention is a neural operator layer whose kernel happens to be dense; the repair is Module 3's four answers applied to that kernel. So the quadratic argument rules out one design and not the family, and citing it as though it ruled out transformers for physics is an overreach.

FNO. One Fourier layer on \(n\) points costs \(O(n\log n)\) for the transforms and \(O(k_{\max}^dd_v^2)\) for the multipliers (ICLR p.6, research_notes/pdf-fno-jmlr.md:58-61). Time is near linear and activation memory is linear in \(n\). The parameter memory tells a different story, because it scales with the modes and the dimension: \(k_{\max}^dd_v^2\) complex numbers per layer, which at \(k_{\max}=16\), \(d_v=32\) and four layers is \[16^3 \times 32^2 \times 4 \times 8\ \text{bytes} = 4096 \times 1024 \times 32 = 128\ \text{MiB in 3D},\] against \(16^2\times32^2\times4\times8 = 8\) MiB for the same configuration in 2D. Sixteen times more, from one extra dimension. A 3D FNO at a useful \(k_{\max}\) is heavy before it has seen a single data point.

The claims, costed

The lab below measures 1,050 bytes of peak activation memory per grid point for a four-layer, 32-channel 2D FNO doing a batch-1 forward pass. That constant turns the claims into numbers, and the numbers are the point of this module. Everything here carries the assumptions with it, because an estimate whose assumptions are hidden is not an audit.

ClaimPointsActivation memory at 1,050 B/point80 GB accelerators, forward pass only
Lab 03's largest measured grid, \(1024^2\)1.05e61.02 GiB (measured)a fraction of one
Anandkumar's industrial scale, \(1000^3\) at one step1e9978 GiB13
The company's "5 trillion context at inference"5e124.66 PiBabout 65,000

Read those two right-hand numbers as an order of magnitude and not a quotation. Four assumptions go into them, and each one can move the answer. The 1,050 bytes per point comes from a 2D kernel at 32 channels, and a 3D kernel's FFT workspace is not accounted for, which pushes the estimate up. Float32 is assumed: bfloat16 halves it. Activation checkpointing routinely buys a further factor of two to four by recomputing instead of storing. And "context" in a company statement may not mean "points held in memory during one forward pass" at all; it may mean the total size of a dataset streamed through the model, in which case the memory arithmetic does not apply and a different question should be asked. Taking the most generous reading available, bfloat16 with aggressive checkpointing, the five-trillion figure still needs on the order of 8,000 accelerators for one forward pass. That is a datacentre. What does not exist in public is any benchmark showing the model runs at that scale, which is why this claim is graded C and stays C.

The honest summary is narrower than either the hype or the dismissal. The quadratic-attention argument is correct and is not the binding constraint, because nobody proposes dense attention over points. The binding constraint is linear activation memory times a large number of points, and the published mechanisms for beating it are all in the next section. Whether they compose at five trillion is exactly the claim nobody outside the company can check.

What public work reaches

WorkMechanismReported numberGrade
MG-TFNO (Kossaifi, Kovachki, Azizzadenesheli, Anandkumar, arXiv 2310.00120)Tucker-factorized FNO weights plus multigrid domain decomposition across devices"Less than half the error with over 150× compression" on turbulent Navier-Stokes. In the library as TFNO with factorization='tucker', about 90 percent fewer parameters per the docsB lineage.md:93-102; teaching.md:82
Incremental FNO (2211.15188)grow the mode count and the data resolution during training10 percent lower test error, 20 percent fewer modes, 30 percent faster trainingB lineage.md:104-113
U-NO (2204.11127)a U-Net hierarchy of operator layers, downsampling then upsampling26 percent better on Darcy, 44 percent on turbulent Navier-Stokes, 37 percent on 3D spatiotemporal Navier-StokesB lineage.md:71-80
Transolver++ (2502.02414)optimized parallelism plus a local adaptive mechanism over slicesmillion-point car and aircraft meshes on one GPU, scaling across GPUs with linear complexityB lineage.md:137-146
FourCastNet 3 (2507.12144)spherical harmonic transforms, geometric probabilistic ensemble60-day global forecast at 0.25°, six-hourly, in under 4 minutes on one GPU; 8 to 60× faster than diffusion methodsB lineage.md:262-271
FourCastNet 1 (2202.11214), as described on the podcastadaptive Fourier layers on a 0.25° grid, about 700 by a few thousand points, about 50,000 training samples, six-hour stepa week-long forecast in under 2 secondsB paper; C the 50,000 and the grid figures come from the episode at 00:28:25 and 01:01:18
PDEBench runtime study (Appendix G)ML inference against classical solvers, 64 to 1024 in 2D and 32 to 256 in 3Dinference nearly three orders of magnitude faster at matched resolution, and independent of viscosity, while the explicit solver's step obeys \(\Delta t \propto \Delta x^2/\eta\)A pdf-mgn-pdebench.md:284
PhysicsNeModomain parallelism, ShardTensor, CPU offloading in neuraloperatordocumented mechanisms; no single headline number was readB repos.md:343, :94

Line the table up against the claim and the gap is easy to state. The largest published neural operator results are a 0.25° globe, which is about \(1.4\times10^6\) points per vertical level; million-point unstructured meshes; and 3D compressible flow at \(128^3\), which is \(2\times10^6\) points, on 100 samples in PDEBench. None of them is \(10^9\) points per step with \(10^{12}\) of context, and the largest of them is three orders of magnitude short of the first of those and six short of the second. The mechanisms that would close that gap are all in the table: factorized weights, hierarchical resolution, domain sharding, and cheap attention. Nothing in the public record says they compose.

Lab 03: one forward pass, measured

LAB ../labs/03_scale/run_scale.py runs a 2D FNO with 32 channels and 4 layers, batch size 1, forward pass only, over grids from \(64^2\) to \(1024^2\) and at n_modes 16 and 64, recording peak GPU memory and the median time over five passes. It then extrapolates linearly in memory and as \(n\log n\) in time to \(1000^3\).

n_modesParamsGridPointsPeak memory (MiB)Time (ms)Bytes per point
16602,97764²4,09616.83.2—
16602,977128²16,38429.33.1—
16602,977256²65,53677.79.71,240
16602,977512²262,144272.236.21,090
16602,9771024²1,048,5761,049.7142.61,050
648,663,905128²16,384102.24.1—
648,663,905256²65,536139.211.0—
648,663,905512²262,144333.737.5—
648,663,9051024²1,048,5761,111.2143.8—

RTX 3050 Laptop GPU, batch 1, forward only, no gradients, float32, median of five passes (../labs/03_scale/results_scale.json). Run-to-run timing varies by about one percent.

Memory is linear in points, and the constant is explainable. Above \(256^2\) it settles at about 1,050 bytes per grid point. That number is not arbitrary: 32 channels at 4 bytes each is 128 bytes per point per stored copy, so 1,050 bytes means roughly eight live activation copies are held at once, which is what a four-layer network with a skip path and an FFT workspace would need. The agreement is close enough to trust the extrapolation and loose enough to remember it is an extrapolation.

Time is near linear, once you are past the launch overhead. Going from \(64^2\) to \(1024^2\) is 256 times the points and only 44 times the time, because the small grids are launch-bound at about 3 ms and the GPU is idle inside them. The honest comparison is \(256^2\) to \(1024^2\): 16 times the points for 14.7 times the time, which is linear within measurement error and consistent with \(n\log n\) at these sizes.

Modes are cheap in activations and expensive in weights. Going from 16 to 64 modes multiplies the parameter count by 14.4, from 0.60M to 8.66M, and adds 62 MiB of fixed memory. At \(1024^2\) that is a 6 percent memory increase and 1.2 ms of extra time. So mode count is close to free at inference on a large grid and dominates the model's storage. That asymmetry is why Incremental FNO's "grow the modes during training" idea is worth 30 percent of the training time.

What the 4 GB card can and cannot do. It holds this forward pass at \(1024^2\) with three quarters of its memory free. It would not hold the training pass at that grid, because storing activations for the backward pass roughly triples the requirement. Lab 02 established this the expensive way: evaluating 200 dies at \(128^2\) in a single batch alongside another training job thrashed the card. Batch size is a memory decision, not a convenience.

The extrapolation, and its direction of error. Scaling the \(1024^2\) row linearly gives \(1050 \times 10^9\) bytes, or 978 GiB, for a batch-1 forward pass on \(1000^3\) points at 32 channels, before the 128 MiB of 3D weights and before any gradient. Scaling the time as \(n\log n\) from the same row gives about 203 seconds per forward pass on this GPU. Both are lower estimates, for two reasons: they assume the 2D kernel's efficiency carries over to 3D, and they ignore the FFT workspace at that size. So one time step of a trillion-point field, at a channel width one third of the FNO paper's, needs on the order of a terabyte of activation memory, which is about thirteen 80 GB accelerators for a single forward pass and several times that to train it. That is the arithmetic behind the word "sharding" in this module's title, and MG-TFNO's domain decomposition across devices is the published mechanism for doing it.

Sources

Exit test

You can now cost an architecture in bytes rather than in adjectives, say precisely which design the quadratic-attention argument rules out and which it does not, and turn a public scale claim into a number of accelerators with its assumptions attached.

  1. Compute the score-matrix size of dense attention over a \(1000^3\) grid in bytes, then compute the same for attention over 4,096 patches, and say which of the two anybody actually proposes.
  2. Compute the parameter memory of a 3D FNO with \(k_{\max}=16\), \(d_v=32\) and four layers, and the 2D version, and explain the factor between them in one sentence.
  3. From the lab, state the measured memory per grid point, explain where the constant comes from in terms of channels and stored copies, and say whether the extrapolation to \(1000^3\) is a lower or an upper estimate and why.
  4. Take the "five trillion context" claim and produce an accelerator count. Then list the four assumptions that could move your answer, and say which direction each one moves it.
  5. Name three published mechanisms that reduce FNO memory or cost, with their reported numbers, and say which one Lab 03's mode-count measurement most directly supports.