Engineering Notes from the Open-Source Frontier: PSGD, Levanter, TPU Orchestration, and the Inference Economics of LLMs

1. Introduction: Systems Challenges in Modern Pre-Training

In large foundation model pre-training, raw parameter counts tell only part of the story. The primary bottlenecks in production clusters emerge at the intersection of non-convex optimization geometry, distributed scheduling on transient hardware, cloud storage I/O latency, and the autoregressive inference memory wall.

Recent engineering discussions in open-source pre-training initiatives (such as the Marin Community) surface four foundational architectural trade-offs:

2. The Curvature Breakthrough: PSGD vs. First-Order Optimizers

The Canyon Dilemma

Loss surfaces in multi-billion parameter spaces routinely form ill-conditioned ravines: transverse cliff walls exhibit massive curvature and gradient steepness, while the longitudinal floor slopes gently toward the minimum with near-zero curvature.

First-order methods (SGD, AdamW) evaluate only the gradient $\nabla L$. Dominated by the steep transverse walls, the optimizer bounces violently between cliffs while making negligible forward headway. While Adam scales coordinates via moving variance ($\sqrt{v_t}$), it remains fundamentally blind to cross-parameter coupling.

Warping the Canyon into a Symmetric Bowl

Second-order optimization incorporates the Hessian matrix $H_{ij} = \frac{\partial^2 L}{\partial \theta_i \partial \theta_j}$, which tracks gradient curvature. Preconditioned updates take the form:

Mathematically, preconditioning acts as a coordinate whitening transform. Dampening high-curvature directions and scaling up flat directions geometrically warps the distorted canyon into an isotropic bowl, directing the gradient straight toward the global minimum.

The Dimension Catastrophe of Explicit Hessians

For a standard 70B parameter model ($N = 7 \times 10^{10}$), an explicit Hessian contains $N^2 = 4.9 \times 10^{21}$ elements. In float32 (4 bytes/element):

Materializing or inverting a 19.6 ZB matrix is physically impossible.

The PSGD-Kron Solution

PSGD (Preconditioned Stochastic Gradient Descent) eliminates explicit Hessian computation:

1. Lie Group & Manifold Optimization: PSGD updates a preconditioning matrix online via stochastic gradient fitting and whitening transforms on matrix Lie groups.

2. Kronecker Factorization (Kron): For 2D linear weight matrices $W \in \mathbb{R}^{d_{out} \times d_{in}}$, PSGD decomposes the curvature matrix into two compact factors:

   \[P \approx P_{left} \otimes P_{right}\]

   where $P_{left} \in \mathbb{R}^{d_{out} \times d_{out}}$ and $P_{right} \in \mathbb{R}^{d_{in} \times d_{in}}$.

This reduces memory and compute overhead from quadratic in total layer weights to quadratic in layer dimensions, delivering second-order curvature acceleration with memory overhead competitive with AdamW.

3. Distributed Fault-Tolerance: Levanter on Preemptible Compute

Spot Infrastructure Economics

Cloud providers sell surplus datacenter capacity as Preemptible / Spot VMs at 60% to 90% discounts, but reclaim them with only 30 to 120 seconds of notice. Under standard All-Reduce collectives (FSDP / Tensor Parallelism), if a single worker node is preempted, the entire cluster job crashes.

The Levanter Stack

To achieve fault-tolerant 70B parameter training on preemptible clusters, modern pipelines deploy Levanter over JAX:

• Functional Purity (JAX + Equinox): Model weights and states are encapsulated as immutable Python dataclasses, enabling seamless serialization.
• Named Tensors (Haliax): Replaces error-prone positional slicing (.view(), .transpose()) with explicit semantic axis labels (batch, seq, hidden).
• Bitwise Reproducibility: Explicit PRNG key management (jax.random.PRNGKey) guarantees bit-for-bit exact resumption upon restarting failed workers from object storage checkpoints.
• Asynchronous Checkpointing: Checkpoints stream to cloud buckets in the background, making high-frequency snapshots cheap.

Why Home PCs Cannot Train 70B Models

A common misconception is whether fault-tolerant frameworks enable volunteer crowd-sourcing across consumer PCs (like Folding@home):

Consumer broadband latency stalls accelerator matrix units by orders of magnitude. Spot computing succeeds because instances share the same ultra-high-bandwidth datacenter fabric; only the reservation is transient.

4. Cluster Architecture & Topology: Head Nodes, Ray, and TPU v6e

Coordination Hierarchy

Efficient distributed orchestration enforces strict separation between control and compute:

1. Worker Nodes (TPUs/GPUs): Solely execute compiled XLA compute graphs and collective communications.
2. Head Node (CPU VM): A low-cost VM executing scheduling, logging, and health checks.
3. Ray Head: Runs Ray’s Global Control Store (GCS), job submission endpoints, and distributed actor lifecycles across the cluster.

TPU v6e (Trillium): Pod vs. Rack Architecture

Google’s TPU v6e delivers ~4.7x compute FLOPs and 2x HBM capacity/bandwidth over v5e, mapping one functional core per physical chip:

• Physical Rack: A standard 42U datacenter cabinet housing server chassis, power units, and Top-of-Rack Ethernet switches.
• TPU Pod (Slice): A logical supercomputer interconnected via Google’s proprietary Inter-Chip Interconnect (ICI) in a 2D/3D Torus topology. A large Pod spans dozens of physical racks stitched together with thousands of dedicated optical fibers.

The AlgoPerf Benchmark Standard

To evaluate optimizers without academic bias, practitioners rely on MLCommons’ AlgoPerf:

• The Problem: Papers frequently over-tune novel optimizers via massive grid searches while comparing against poorly tuned AdamW defaults.
• The Standard: AlgoPerf enforces a strict hyperparameter tuning budget and ranks optimizers on true Time-to-Target wall-clock execution across diverse workloads.

5. Storage Architecture Matrix: GCS FUSE vs. Hyperdisk ML

The Golden Rule: Decoupling Read and Write Planes

Distributed machine learning requires two conflicting I/O patterns:

GCS FUSE vs. Hyperdisk ML Trade-offs

Cloud Storage FUSE (gcsfuse)

Mounts GCS buckets as local POSIX directories.

• The Safetensors Bottleneck: Modern .safetensors files rely on mmap and non-sequential byte seeking. Under FUSE, unbuffered random seeks trigger cascading HTTP round-trips and I/O timeouts unless heavily cached (--file-cache, --stat-cache-capacity).

Hyperdisk ML

Google Cloud’s specialized block storage for AI workloads.

• Multi-Reader: Allows up to thousands of TPU/GPU instances to mount the exact same volume in Read-Only mode with aggregate throughput in hundreds of GB/s.
• No Direct Writes: Hyperdisk ML strictly prohibits direct writes while mounted. Data must be written to an intermediary disk, snapshotted, and instantiated as a read-only volume.
• Multi-Writer Warning: While normal Hyperdisk (Balanced/Extreme) supports multi-writer attachments, it requires an external clustered filesystem (e.g., Lustre, GPFS) with distributed locking; mounting with standard ext4/xfs corrupts disk metadata instantly.

The Idle Provisioned Billing Trap

Hyperdisk bills primarily for Provisioned Throughput (MB/s):

• Provisioning 1,200 MB/s ensures ultra-fast model loading.
• However, Google Cloud bills for that 1,200 MB/s pipe continuously (24/7), even if instances sit completely idle. An unmonitored development disk can easily spike daily spend from $5 to over $100.

Serving with vLLM TPU

• Approach A (FUSE): Inexpensive, but requires careful tuning to prevent safetensors seek latency.
• Approach B (Hyperdisk ML): Instantaneous mmap loading across hundreds of nodes, but requires pre-baked disk images and incurs provisioned throughput costs.
• Approach C (Boot Download): Parallel CLI download (gcloud storage cp) directly to local VM scratch space. Zero persistent disk costs, but incurs cold-start latency.

6. The Inference Memory Wall: Linear Attention Economics

The KV Cache Explosion

During autoregressive generation, storing Key ($K$) and Value ($V$) vectors in accelerator HBM avoids $\mathcal{O}(N^2)$ recomputation:

At context lengths of 128k to 1M tokens, the KV Cache consumes hundreds of gigabytes per concurrent request, quickly dwarfing static model weights and hitting an impenetrable Inference Memory Wall.

Linear Attention: The 3-Line Mathematical Derivation

Standard Softmax attention couples queries and keys non-linearly:

Linear attention replaces Softmax with kernel feature maps $\phi(\cdot)$:

By the associative property of matrix multiplication, we alter the execution order:

Rather than computing the $N \times N$ matrix $(\phi(Q) \phi(K)^T)$, we evaluate the compact state matrix $S = \phi(K)^T V \in \mathbb{R}^{d_k \times d_v}$ first!

In recurrent decoding, state updates incrementally:

Because $S$ has a fixed dimension $(d_k \times d_v)$, it never grows with sequence length. Decoding token 1 and token 1,000,000 consumes the exact same constant memory and execution time ($\mathcal{O}(1)$).

Frontier Implementations

• BASED: Combines Taylor-expansion linear attention with localized 1D convolutions to solve associative recall while maintaining linear throughput.
• MiniMax-01: An open-weights foundation model utilizing Lightning Attention (linear attention variant) paired with sparse Mixture-of-Experts (MoE), serving multi-million-token contexts at a fraction of standard Transformer serving costs.

7. Knowledge Check: Interactive Engineering Challenge

Verify your architectural intuition on second-order curvature, spot-compute fault tolerance, storage trade-offs, and linear attention economics. Select your answers below and click Submit Answers for an instant diagnostic score and detailed post-mortem.

8. Summary & Engineering Takeaways

1. Curvature Over Gradient Brute Force: Parameter landscapes are anisotropic ravines. Second-order Lie-group preconditioning (PSGD-Kron) transforms narrow canyons into isotropic bowls, converging faster than AdamW without explicit Hessian storage.
2. Embrace Hardware Transience: Spot compute slashes infrastructure bills by 60%–90%. Achieving stability requires pairing purely functional, bitwise-reproducible stacks (JAX, Equinox, Haliax, Levanter) with sub-minute checkpoint streaming.
3. Decouple Storage Planes: Never mix read and write paths. Stream dynamic checkpoints asynchronously to object stores (GCS Buckets), and restrict high-throughput block volumes (Hyperdisk ML) to immutable dataset loading and fast inference cold-starts—while strictly monitoring provisioned throughput billing meters.
4. Architect for Long-Context Inference: Pre-training costs are paid once; inference costs scale with every generated token. To break the quadratic KV cache barrier at multi-million token contexts, production architectures must shift toward constant-state linear attention mechanisms (BASED, Lightning Attention / MiniMax-01).