I Rebuilt Karpathy's NanoChat in JAX. Here's What XLA Gets Right and What It Gets Dead Wrong.

DEV Community

Authentication Security Deep Dive: From Brute Force to Salted Hashing (With Java Examples) Why AI Systems Don’t Fail — They Drift Spilling beans for how i learn for exam😁"Reinforcement Learning Cheat Sheet" I Replaced Chrome with Safari for AI Browser Automation. Here's What Broke (and What Finally Worked) How Python Borrows Other People's Work The $40 Architecture: Processing 1 Billion API Requests with 99.99% Uptime Vibe Coding: A Workflow Guide (From Zero to SaaS) Most webhook security guides protect the wrong side. The scary part is delivery. Headless CMS for TanStack Start: Build a Blog with Cosmic EU Age Verification App "Hacked in 2 Minutes" — What Actually Happened Comfy Cloud’s delete function does not actually remove files Running AI Models on GPU Cloud Servers: A Beginner Guide Event-driven media intelligence with AWS Step Functions and Bedrock I scored 500 AI prompts across 8 quality dimensions — here's what broke How to Call Google Gemini API from Next.js (Free Tier, No Backend Needed) The Portal Protocol: Reclaiming Human Connection in the Age of AI How to Fix Your Team's Scattered Knowledge Problem With a Self-Hosted Forum Intro to tc Cloud Functors: A Graph-First Mental Model for the Modern Cloud Designing Multi-Tenant Backends With Both Ownership and Team Access I Built a Neumorphic CSS Library with 77+ Components — Here's What I Learned PostgreSQL Performance Optimization: Why Connection Pooling Is Critical at Scale Cómo construí un SaaS multi-rubro para gestionar expensas en Argentina con FastAPI + Vue 3 🚀 I Built an Ethical Hacking Scanner Tool – Open Source Project I Replaced /usage and /context in Claude Code With a Single Statusline A Pythonic Way to Handle Emails (IMAP/SMTP) with Auto-Discovery and AI-Ready Design I Collected 8.9 Million Polymarket Price Points — Here's What I Found About How Markets Really Move EcoTrack AI — Carbon Footprint Tracker & Dashboard Everyone's Using AI. No One Agrees How. 5 self-hosted ebook managers worth trying in 2026 Building Your First AI Agent with LangChain: From Chatbot to Autonomous Assistant Common SOC 2 Failures (Real World) Stop Vibe-Checking Your AI App: A Practical Guide to Evals How to Use SonarQube and SonarScanner Locally to Level Up Your Code Quality Your Next To-Do App Is Dead — I Replaced Mine with an OpenClaw AI Sign a Nostr event in 60 lines of Python using coincurve — no nostr-sdk, no nbxplorer, no rust toolchain ITGC Audit Explained Like You’re in Big 4 Patch Tuesday abril 2026: Microsoft parcha 163 vulnerabilidades y un zero-day en SharePoint Stop scraping everything: a better way to track competitor price changes Listing on MCPize + the Official MCP Registry while routing payments OUTSIDE the marketplace — how I kept 100% of my x402 revenue Building an AI-Powered Risk Intelligence System Using Serverless Architecture Why We Ripped Function Overloading Out of Our AI Toolchain Testing AI-Generated Code: How to Actually Know If It Works SaaS Churn Is Killing Your Business. Here Is What to Do About It (Without a Support Team) The Speed of AI Is No Longer Linear - And Self-Improving Models Are Why How to Implement RBAC for MCP Tools: A Practical Guide for Engineering Teams From Standard Quote to Persuasive Proposal: AI Automation for Arborists I built a CLI that scaffolds complete multi-tenant SaaS apps Axios CVE-2025–62718: The Silent SSRF Bug That Could Be Hiding in Your Node.js App Right Now The dashboard that ended our friendship Data Pipelines Explained Simply (and How to Build Them with Python) The Hidden Cost of AI Systems Nobody Talks About. undefined vs undeclared, and how typeof behaves Switching from file-based jobs to NATS/Kafka in Rust without changing code io_uring Adventures: Rust Servers That Love Syscalls Why Agentic AI is Killing the Traditional Database The POUR principles of web accessibility for developers and designers Quantum Neural Network 3D — A Deep Dive into Interactive WebGL Visualization How To Install Caveman In Codex On macOS And Windows Automation Pipeline Reliability: Why Your Workflow Breaks When Nobody Is Watching I Built an 'Open World' AI Coding Agent — It Works From ANY Folder From Freelancing to Product: A Tech Service Company's SaaS Transformation China's AI Giants: Adding Tencent Hunyuan & ByteDance Doubao to AI University (74 Providers) On the Vibe Coders and Their Lies clerk: Auto-Summarize Your Claude Code Sessions AI Weekly — 2026/04/10–04/17 | The Model Lockdown Is Here, but the Toolchain Is the Real Battleground AI 週報 — 2026/04/10–2026/04/17 模型封鎖潮來了，但工具鏈才是真戰場 Maybe this is how Open-Source apps are born... 🚀 Fine-Tune LLMs with LoRA and QLoRA: 2026 Guide tRPC v11 + Next.js App Router: End-to-End Type Safety Without the Boilerplate ShadCN UI in 2026: Why I Stopped Installing Component Libraries and Started Owning My Components SaaS Billing in React Server Components: Stripe + Supabase Without a Single `useEffect` Join our DEV Weekend Challenge — $1,000 in Prizes Across TEN winners! Submissions Due April 20 at 6:59 AM UTC. Implementing FSRS Spaced Repetition in Flutter + Supabase — Adding Memory Science to an AI Learning App "I Texted My Localhost From the Train — Claude Code Fixed the Bug Before I Got Home" I Built a Sales Prep AI and It Went Deeper Than Expected Design to Code #2: One JSON, Eleven Outputs Solving the 100M-Row Problem: A Summary Table Pattern for High-Volume Push Notification Logs Flutter Web With Wasm: What Actually Changes For Developers I Built 50 Royalty-Free Soundtracks for My Side Project in a Weekend Using AI Music Generation The Vibe Coding Security Checklist: 7 Things to Check Before You Ship Stop Letting Googlebot Guess Fix Your React App's SEO Right Desconstruindo o Streaming do LinkedIn: Como Criar um Engine de Extração de Vídeo de Alta Performance com HLS e FFmpeg (EDA Part-1) EDA (Exploratory Data Analysis) Explained With Real Life — Why Looking at Your Data Is the Most Important Step in Machine Learning Brand Relationship Management at Scale: Our 4-Touch Outreach System for 200+ Brands Why String.fromEnvironment() Might Return an Empty String in Dart JGuardrails 1.0.0 — Hardening Java LLM Apps Against Jailbreaks, Toxicity, and Prompt Injection Plan and Schedule a Full Week of Threads Content From One Claude Conversation Coding Cat Oran Ep3, Five Tables Changed Everything Updated: BFF Pattern I'm done watching freelancers get buried by 200 proposals. So I'm building the alternative. This is my first post BFS Algorithm in Java Step by Step Tutorial with Examples Tracking LLM Pricing Monthly: An Open Dataset for 22 AI Models How We Measure Content ROI on a Comparison Site: Revenue Attribution Without Perfect Data Introducing Nova AI Ops: The AI-Native Operating System for SRE Teams I built a free desktop video downloader for Windows — Grabbit How Talkie OCR Helps Vision-Impaired & Dyslexic Users Read the World Around Them VRCFaceTracking安装和iPhone面捕配置教程，有bug Even CrowdStrike Can't See Your Agents The Automation Gold Rush: What n8n Workflows and Claude Are Opening Up for Developers Right Now

Omotayo Aina · 2026-05-01 · via DEV Community

AI GDE TPU Sprint 2026 · Google TPU Research Cloud Quick summary: We ported Andrej Karpathy's NanoChat architecture from PyTorch to JAX and Flax NNX. The repo is about 12,400 lines across source and scripts. We trained a nano model (885K parameters) on TinyStories in under 10 minutes on a single GPU and served it through a streaming chat UI. XLA compilation eliminates Python overhead after a one-time upfront cost. The same code runs on TPU without modification. The catch: no vLLM, no Flash Attention 3, and painful debugging inside JIT-compiled functions. This post covers what worked, what did not, and when you should care. <a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2F3u0jr2uobpsqnjmzejj3.png" class="article-body-image-wrapper"><img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2F3u0jr2uobpsqnjmzejj3.png" alt="nanochat-jax transformer block. ✦ marks components absent from standard GPT and LLaMA: Value Embeddings, Smear/Backout, per-layer scalars, QK norm, logit softcap." width="800" height="723"></a> The nanochat-jax transformer block. The ✦ symbol marks nanochat-specific components not present in standard GPT or LLaMA: Value Embeddings, Smear/Backout token mixing, per-layer learnable scalars, QK L2 normalization, and logit softcap. <h2> What This Project Is </h2> Andrej Karpathy's <a href="https://github.com/karpathy/nanochat" rel="noopener noreferrer">NanoChat</a> is roughly 8,600 lines of PyTorch that tries to answer one question: what is the best ChatGPT you can train for $100? It ships with Flash Attention 3 on Hopper, the Muon optimizer (Newton-Schulz orthogonalization), FP8 mixed-precision, distributed training via <code>DistMuonAdamW</code>, SFT, RL via GRPO, and a web chat UI. The model architecture packs in a lot of modern ideas: Grouped-Query Attention, RoPE positional encoding, parameterless RMSNorm, ReLU-squared activation, Value Embeddings, Smear/Backout token mixing, per-layer learnable scalars, QK L2 normalization, and logit softcap. <a href="https://github.com/ainaomotayo/nanochat-jax" rel="noopener noreferrer">NanoChat-JAX</a> is a faithful port of that architecture to JAX and Flax NNX. Two things motivated the port. First, we wanted scaling law instrumentation that the original does not have: the ability to sweep model size, data volume, and compute budget systematically and fit Chinchilla-style power laws to the results. Second, we wanted to run the same code on GPU and TPU without a device-specific codebase. JAX's XLA backend makes that possible. This project is part of the AI GDE TPU Sprint 2026, where Google TPU Research Cloud compute is available for exactly this kind of scaling experiment. <h2> Five NanoChat Architecture Components Worth Understanding </h2> Most open-source GPT implementations stop at multi-head attention, an FFN, and RoPE. NanoChat adds five components that are not in the standard recipe. Understanding them matters because each one shows up in the JAX translation in a non-obvious way. <h3> Logit Softcap </h3> NanoChat applies <code>cap * tanh(logits / cap)</code> to attention scores before the softmax. This clamps score magnitudes to the range <code>[-cap, cap]</code> and prevents the entropy collapse that happens at depth when very sharp attention distributions starve gradients. NanoChat uses a softcap value of 15.0; our JAX port uses 30.0, which is a deliberate divergence noted in the comparison table at the end of this post. PyTorch: <div class="highlight js-code-highlight"> <pre class="highlight python"><code>scores = torch.matmul(q, k.transpose(-2, -1)) * self.attn_scale if self.logit_softcap is not None: scores = self.logit_softcap * torch.tanh(scores / self.logit_softcap) scores.masked_fill_(~mask, float('-inf')) weights = F.softmax(scores, dim=-1) </code></pre> </div> JAX: <div class="highlight js-code-highlight"> <pre class="highlight python"><code>scores = jnp.matmul( q.astype(jnp.float32), jnp.transpose(k_exp.astype(jnp.float32), (0, 1, 3, 2)) ) * self.attn_scale if self.logit_softcap is not None: cap = float(self.logit_softcap) scores = cap * jnp.tanh(scores / cap) scores = jnp.where(combined_mask, scores, jnp.float32(-1e9)) weights = jax.nn.softmax(scores, axis=-1) </code></pre> </div> Two differences stand out. <code>jnp.where</code> replaces <code>masked_fill_</code> because JAX arrays are immutable: you cannot modify them in-place, so you select between two arrays instead. We also use <code>-1e9</code> instead of <code>-inf</code>. When an entire attention row is masked (as happens with padding tokens), JAX's softmax produces <code>NaN</code> on <code>-inf</code> input rather than handling it gracefully the way PyTorch does. Switching to <code>-1e9</code> fixes the issue cleanly; we hit this bug midway through the port. <h3> The Training Step as a Pure Function </h3> In PyTorch, you write a loop: zero grads, forward, loss, backward, step. In Flax NNX, the entire training step becomes a pure function decorated with <code>@nnx.jit</code>: <div class="highlight js-code-highlight"> <pre class="highlight python"><code>@staticmethod @nnx.jit def _train_step_jit( model: TransformerLM, optimizer: nnx.Optimizer, batch: dict[str, jax.Array], ) -> dict[str, jax.Array]: def loss_fn(model: TransformerLM) -> tuple[jax.Array, dict]: logits, _ = model(batch["input_ids"], deterministic=False) loss, metrics = cross_entropy_loss( logits=logits[:, :-1, :], labels=batch["labels"][:, :-1], ) return loss, metrics (loss, metrics), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model) grad_norm = optax.global_norm(jax.tree.leaves(grads)) optimizer.update(model, grads) return {"loss": loss, "grad_norm": grad_norm, **metrics} </code></pre> </div> <code>nnx.value_and_grad</code> replaces the <code>loss.backward()</code> / <code>optimizer.step()</code> / <code>optimizer.zero_grad()</code> sequence. NNX extracts model parameters as a pytree, computes gradients as a matching pytree, and applies them through the optimizer. There is no gradient tape to manage, no <code>.detach()</code> to remember, and no <code>torch.no_grad()</code> context manager. <code>@nnx.jit</code> traces the function once and compiles it to an XLA HLO program. Every subsequent call skips Python entirely and dispatches directly to the compiled kernel. <h3> Muon Optimizer: Newton-Schulz Inside XLA </h3> Muon orthogonalizes each 2D weight gradient via Newton-Schulz iterations. The update rule is <code>X_{t+1} = 1.5X - 0.5(XX^TX)</code>, repeated for a fixed number of steps. Since the full training step runs under <code>@nnx.jit</code>, this loop must be JIT-compilable. In PyTorch, <code>torch.compile</code> unrolls a Python <code>for</code> loop at trace time. In JAX, we use <code>jax.lax.fori_loop</code>: <div class="highlight js-code-highlight"> <pre class="highlight python"><code>def newton_schulz_orthogonalize(G: jax.Array, steps: int = 10) -> jax.Array: G = G.astype(jnp.float32) G = G / (jnp.linalg.norm(G) + 1e-8) transpose = G.shape[0] > G.shape[1] if transpose: G = G.T def ns_step(_, X): A = X @ X.T return 1.5 * X - 0.5 * (A @ X) G = jax.lax.fori_loop(0, steps, ns_step, G) if transpose: G = G.T return G </code></pre> </div> <code>jax.lax.fori_loop</code> compiles to a single XLA while-loop operation. The compiled program size stays constant regardless of iteration count, and changing <code>steps</code> does not trigger recompilation. The tradeoff is that you cannot use Python control flow inside the loop body. Debugging means <code>jax.debug.print</code> rather than <code>print</code>. <h3> Value Embeddings </h3> Value Embeddings give each token a learned residual vector that is independent of context. Unlike input embeddings, which are used once at the bottom of the stack, value embeddings are added into the attention output at every layer. A single shared lookup table lives at the top-level <code>TransformerLM</code> and gets passed by reference to each block: <div class="highlight js-code-highlight"> <pre class="highlight python"><code># TransformerLM.__init__ if cfg.use_value_embeddings: self.value_embed = ValueEmbedding(cfg.vocab_size, cfg.d_model, rngs=rngs) self.layers = nnx.List([ TransformerBlock(cfg, layer_idx=i, value_embed=self.value_embed, rngs=rngs) for i in range(cfg.n_layers) ]) </code></pre> </div> The table initializes near zero (scale <code>1e-4</code>), so the model starts as if value embeddings do not exist. Their contribution grows only as gradient descent drives the table weights away from zero. This no-op initialization pattern appears throughout nanochat: new components start inert and the optimizer learns to use them. <h3> Smear and Backout Token Mixing </h3> Smear and Backout are cheap causal token-mixing operations. Smear blends each token with its immediate predecessor: <code>x[t] = (1 - alpha) * x[t] + alpha * x[t-1]</code>, where <code>alpha = sigmoid(raw_alpha)</code> is a learned per-feature vector of shape <code>(d_model,)</code>. Backout then removes the correlation introduced by Smear from the attention output, preventing double-counting when that output is added back to the residual stream. In JAX, the causal shift needs no Python loop: <div class="highlight js-code-highlight"> <pre class="highlight python"><code>x_prev = jnp.concatenate( [jnp.zeros_like(x[:, :1, :]), x[:, :-1, :]], axis=1 ) alpha = jax.nn.sigmoid(self.raw_alpha.get_value()) x_smeared = x + alpha * (x_prev - x) </code></pre> </div> Both <code>raw_alpha</code> and <code>raw_beta</code> initialize to <code>-10.0</code>. <code>sigmoid(-10)</code> is approximately <code>5e-5</code>, so at step 0 the Smear and Backout components have essentially zero effect on the forward pass. Their contribution grows only as the optimizer pushes these parameters away from their initialization. <h3> Depth-Aware Weight Initialization </h3> Residual output projections (attention <code>out_proj</code> and FFN <code>down_proj</code>) at layer index <code>l</code> are scaled by <code>1 / sqrt(2 * (l + 1))</code>. This assigns progressively smaller initialization scales to deeper layers, which controls how much each layer can perturb the residual stream as depth increases. GPT-NeoX applies <code>1 / sqrt(2 * n_layers)</code> uniformly to all layers, giving every layer the same scale regardless of its position in the stack. <div class="highlight js-code-highlight"> <pre class="highlight python"><code>def _init_weights_from_depth(self) -> None: for layer_idx, layer in enumerate(self.layers): depth_scale = 1.0 / math.sqrt(2.0 * (layer_idx + 1)) layer.attention.out_proj.kernel = nnx.Param( layer.attention.out_proj.kernel.get_value() * depth_scale ) layer.ffn.down_proj.kernel = nnx.Param( layer.ffn.down_proj.kernel.get_value() * depth_scale ) </code></pre> </div> In Flax NNX, parameter mutation works by calling <code>.get_value()</code>, applying the transform, and wrapping the result in <code>nnx.Param</code>. There is no <code>param.data.mul_()</code> equivalent. <h2> Getting It Running </h2> <div class="highlight js-code-highlight"> <pre class="highlight shell"><code># Clone and install git clone https://github.com/ainaomotayo/nanochat-jax cd nanochat-jax pip install -e ".[dev]" # Download and preprocess TinyStories (saves vocab alongside HDF5) python -m scripts.preprocess --dataset tinystories --output_dir data/ # Train the nano model (about 10 minutes on a GPU) python -m scripts.train \ --model-size nano \ --data-path data/tinystories.h5 \ --device gpu \ --steps 2000 # Start the chat UI server python -m scripts.chat_web \ --checkpoint checkpoints/tinystories_nano/latest \ --model-size nano \ --device gpu # Open http://localhost:8000 </code></pre> </div> The server loads the tokenizer from <code>data/tinystories_vocab.json</code>, restores weights from the checkpoint via <code>CheckpointManager</code>, and serves an OpenAI-compatible streaming API at <code>/v1/chat/completions</code> alongside the browser UI. <a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fbn6srp9rpd6q036y4xlq.png" class="article-body-image-wrapper"><img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fbn6srp9rpd6q036y4xlq.png" alt="nanochat-jax chat UI with the trained 885K-parameter nano model. Real checkpoint loaded, no random weights." width="800" height="338"></a> The chat UI with the trained nano model responding. Tokens stream to the browser via SSE. The orange warning banner appears when the server starts with random weights; it is absent here because a real checkpoint is loaded. <h2> The Numbers We Actually Measured </h2> <a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fsaq6i7j61s66slp5l30x.png" class="article-body-image-wrapper"><img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fsaq6i7j61s66slp5l30x.png" alt="Training and validation loss over 2,000 steps on TinyStories. val_loss = 1.295, perplexity = 3.65. Step 1 includes the ~35s XLA compile." width="800" height="396"></a> Training and validation loss for the nano model over 2,000 steps on TinyStories. The XLA compile happens on step 1 (about 35 seconds). Steady-state step time is 290ms. Final val_loss = 1.295, perplexity = 3.65. We ran one confirmed training experiment. Every number in this table came from that run: <div class="table-wrapper-paragraph"><table> <thead> <tr> <th>Metric</th> <th>Value</th> </tr> </thead> <tbody> <tr> <td>Model preset</td> <td>nano</td> </tr> <tr> <td>Parameters</td> <td>885,768</td> </tr> <tr> <td>Dataset</td> <td>TinyStories (180M tokens, 95-character vocabulary)</td> </tr> <tr> <td>Device</td> <td>GPU, bfloat16</td> </tr> <tr> <td>Steps</td> <td>2,000</td> </tr> <tr> <td>Best validation loss</td> <td>1.295</td> </tr> <tr> <td>Perplexity</td> <td>3.65</td> </tr> <tr> <td>Total training time</td> <td>9.7 minutes</td> </tr> <tr> <td>Steady-state step time</td> <td>290ms</td> </tr> <tr> <td>First step time (includes XLA compile)</td> <td>~35 seconds</td> </tr> </tbody> </table></div> Step 1 takes about 35 seconds because XLA traces and compiles the entire forward pass, backward pass, and Muon optimizer update into a single fused HLO program. Steps 2 through 2,000 average 290ms with no Python dispatch overhead. Over the full 2,000-step run, that 35-second compile adds roughly 3% to total training time. Parameter counts for model sizes other than nano are estimates based on the architecture config. We have not run training at larger scales yet. <h3> Model Scale Presets </h3> <a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fx8ysgpy8vtrhc1gazazf.png" class="article-body-image-wrapper"><img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fx8ysgpy8vtrhc1gazazf.png" alt="Five model presets. Nano (885K parameters) is confirmed from training; larger sizes are estimated from the architecture config." width="800" height="397"></a> Five model presets from 885K (nano, confirmed from training) to approximately 6.7B parameters (xlarge, estimated from config). All presets share the same nanochat architecture. Every preset uses GQA, ReLU-squared, QK normalization, logit softcap, Value Embeddings, Smear/Backout, and per-layer scalars. Only the dimensions change: <div class="table-wrapper-paragraph"><table> <thead> <tr> <th>Preset</th> <th>d_model</th> <th>n_layers</th> <th>n_heads</th> <th>n_kv_heads</th> <th>d_ff</th> <th>Vocab</th> <th>Max Seq</th> </tr> </thead> <tbody> <tr> <td>nano</td> <td>128</td> <td>4</td> <td>4</td> <td>4</td> <td>512</td> <td>256</td> <td>64</td> </tr> <tr> <td>small</td> <td>512</td> <td>6</td> <td>8</td> <td>8</td> <td>2048</td> <td>32000</td> <td>2048</td> </tr> <tr> <td>medium</td> <td>1024</td> <td>12</td> <td>16</td> <td>8</td> <td>4096</td> <td>32000</td> <td>2048</td> </tr> <tr> <td>large</td> <td>2048</td> <td>24</td> <td>32</td> <td>8</td> <td>8192</td> <td>32000</td> <td>4096</td> </tr> <tr> <td>xlarge</td> <td>4096</td> <td>32</td> <td>32</td> <td>8</td> <td>16384</td> <td>32000</td> <td>4096</td> </tr> </tbody> </table></div> <h2> Where JAX Has a Real Advantage </h2> <a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2F56j61jxulm4kgeapkto6.png" class="article-body-image-wrapper"><img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2F56j61jxulm4kgeapkto6.png" alt="Left: capability comparison for scaling research workloads. Right: XLA compilation cost vs steady-state step time, measured on the nano model." width="800" height="334"></a> <h3> Gradient bugs surface as errors, not silent failures </h3> The <code>nnx.value_and_grad</code> pattern forces training into the shape <code>(params, data) -> (loss, grads)</code>. During the port, we did not encounter forgotten <code>zero_grad()</code> calls, gradients leaking across accumulation steps, or missing <code>.detach()</code> on a loss component. Those bugs exist in PyTorch codebases and they manifest as training that runs but converges to a worse minimum, which is the hardest category of bug to diagnose because nothing crashes. JAX's explicit PRNG threading via <code>nnx.Rngs</code> also makes reproducibility the default. Two runs with the same seed produce identical step-by-step losses without any extra flags. Achieving the same in PyTorch requires <code>torch.use_deterministic_algorithms(True)</code>, which disables several optimized kernels and hurts throughput. <h3> XLA compiles the whole training step, not just parts of it </h3> When <code>@nnx.jit</code> compiles a training step, XLA sees the full computation graph from the input tensor to the updated parameters. It can fuse operations that PyTorch's operator-level dispatch handles as separate kernel launches: attention softcap, masking, and softmax can become one kernel; RMSNorm and the linear projection that follows it can become another. After compilation, each step is a single kernel launch with no Python-level overhead between operations. For the nano model at 290ms per step, that saved overhead is modest in absolute terms. The argument for XLA becomes stronger as model size grows and each kernel does proportionally more compute relative to the dispatch cost. <h3> TPU runs with a flag change, not a code change </h3> <div class="highlight js-code-highlight"> <pre class="highlight shell"><code>python -m scripts.train --device gpu # local GPU python -m scripts.train --device tpu # Google TPU Research Cloud </code></pre> </div> JAX's XLA backend compiles the same HLO program to both CUDA and TPU targets. There is no <code>torch.xla</code> bridge, no CUDA-specific memory management, and no device-specific attention kernel. For researchers with TPU Research Cloud access through the AI GDE program, this means developing on a GPU workstation and running scaling sweeps on TPU v4-8 pods without any porting step in between. <h3> The transformation system composes naturally </h3> <code>jit</code>, <code>grad</code>, <code>vmap</code>, and <code>hessian</code> are all first-class JAX transformations that compose with each other. Because <code>train_step</code> is a pure function of <code>(model_state, batch)</code>, you can apply any of them without restructuring the code. <code>jax.hessian</code> in particular lets you compute the full Hessian of the loss with respect to model parameters, which enables loss landscape geometry analysis at convergence. In PyTorch, getting there requires <code>torch.autograd.functional.hessian</code> and considerably more scaffolding. We have not yet run these analyses on nanochat-jax, but the codebase is structured to support them directly. <h2> Where PyTorch Still Wins </h2> Ecosystem. vLLM, DeepSpeed, PEFT, bitsandbytes, HuggingFace Transformers: none of these work with JAX. If your pipeline involves LoRA fine-tuning, vLLM inference, or HuggingFace benchmark evaluation, you are staying in PyTorch. No JAX-native equivalent of that integrated stack exists today. Flash Attention. NanoChat uses Flash Attention 3 on Hopper hardware via Triton, which is a hand-optimized kernel that gets close to peak memory bandwidth. The JAX equivalent requires writing Pallas custom kernels, which is significantly harder and less mature. XLA's fused attention is capable but is not equivalent to FA3 for long sequences on Hopper GPUs. Debugging. Inside a JIT-compiled JAX function, you cannot use <code>print</code> or <code>breakpoint</code>. You use <code>jax.debug.print</code>, and shape errors surface as abstract trace errors rather than concrete tensor shapes. We spent more time diagnosing shape mismatches in JAX than we would have in PyTorch for the same codebase. Distributed training. NanoChat's <code>DistMuonAdamW</code> with ZeRO-2 sharding is production-tested for multi-GPU runs. NanoChat-JAX has a <code>jax.sharding</code> stub for data parallelism that exists in the codebase but has not been tested at scale. FP8. PyTorch has <code>torch._scaled_mm</code> for FP8 matrix multiplications. JAX's FP8 support is still early. NanoChat-JAX does not implement FP8 training. <h2> Scaling Law Findings </h2> The main research goal is systematic scaling experimentation. We built three experiment types into the codebase: <code>scale_n</code> (vary model size, fix data), <code>scale_d</code> (vary data, fix model), and <code>scale_c</code> (sweep compute budget, co-optimize model and data). We fit the standard power law <code>L(N) = a * N^(-alpha)</code> to validation loss as a function of non-embedding parameter count, using 600 training steps per model size on TinyStories. Result: <code>L = 3.29 * N^(-0.027)</code> The exponent alpha = 0.027 is much flatter than published values. Kaplan et al. (2020) measured alpha around 0.076 on WebText2 (~40B tokens). Hoffmann et al. (2022, Chinchilla) measured the parametric component at around 0.34 on MassiveText (~1.4T tokens). Three things explain the gap. Training duration. 600 steps does not bring larger models to convergence. The nano model is close to its loss floor at 600 steps; the small model (~35M parameters) is still in the steep part of its loss curve. With extended training, preliminary results show alpha rising toward the 0.07 to 0.12 range, which is closer to Kaplan's values. Data scale. TinyStories is 180M tokens. The datasets used in published scaling laws are orders of magnitude larger. With limited data, larger models overfit before the power law can stabilize, which artificially flattens the scaling curve. Character-level tokenization. Character models scale differently than subword models. A single character carries less information than a BPE token, which changes both the effective data scale and the compute requirements per unit of language modeling quality. The Chinchilla analysis module (<code>nanochat/scaling/analysis.py</code>) implements the parametric loss model <code>L(N, D) = E + A/N^alpha + B/D^beta</code> with 1,000-sample bootstrap confidence intervals. For a given compute budget <code>C = 6ND</code>, it computes the optimal allocation <code>N* ~ C^(beta/(alpha+beta))</code>. Our current exponents are too noisy for reliable Chinchilla-optimal predictions. With TPU-scale compute, we plan to run full <code>scale_n</code> and <code>scale_c</code> sweeps and publish updated results. <h2> Practical Advice </h2> Use JAX if you have TPU access and care about scaling experiments. The functional programming model makes it natural to express experiments as pure functions over configuration spaces. TPU portability costs nothing once you are writing JAX. <code>jax.hessian</code> and <code>jax.vmap</code> are real, stable APIs in JAX, not future roadmap items, and the nanochat-jax codebase is structured to use them directly. Use PyTorch if you need the serving and fine-tuning stack. LoRA, vLLM, HuggingFace pipelines, Flash Attention 3: these are PyTorch-only today. The research advantages of JAX do not outweigh rebuilding that infrastructure from scratch for most applied teams. Budget time for XLA compilation. The first step of any new model configuration takes 30 to 60 seconds. Always run a warmup step before starting any timing measurement. Never report first-step timing as representative of steady-state performance. The port took about three focused weeks. Architecture translation (attention, FFN, norms, embeddings) was mechanical once we understood the Flax NNX API. The hard parts were making Muon JIT-compatible via <code>fori_loop</code>, tracing shape errors in the KV cache path, and implementing checkpoint serialization and SSE streaming without PyTorch-ecosystem conveniences. <h2> What Is Implemented vs What Is Not </h2> <div class="table-wrapper-paragraph"><table> <thead> <tr> <th>Capability</th> <th>nanochat (PyTorch)</th> <th>nanochat-jax (JAX)</th> </tr> </thead> <tbody> <tr> <td>GQA + RoPE + parameterless RMSNorm</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Logit softcap</td> <td>yes (15.0)</td> <td>yes (30.0)</td> </tr> <tr> <td>Value Embeddings</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Smear/Backout token mixing</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Per-layer learnable scalars</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Depth-aware weight initialization</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Muon optimizer</td> <td>yes, DistMuonAdamW with ZeRO-2</td> <td>yes, single-device only</td> </tr> <tr> <td>Flash Attention 3</td> <td>yes, Hopper via Triton</td> <td>no</td> </tr> <tr> <td>FP8 training</td> <td>yes</td> <td>no</td> </tr> <tr> <td>Distributed multi-GPU training</td> <td>yes</td> <td>no (stub)</td> </tr> <tr> <td>SFT and RL via GRPO</td> <td>yes</td> <td>no</td> </tr> <tr> <td>Streaming chat UI</td> <td>yes</td> <td>yes, via SSE</td> </tr> <tr> <td>OpenAI-compatible API</td> <td>yes</td> <td>yes</td> </tr> <tr> <td>Scaling law instrumentation</td> <td>no</td> <td>yes</td> </tr> <tr> <td>TPU portability</td> <td>no</td> <td>yes</td> </tr> </tbody> </table></div> <h2> Conclusion </h2> NanoChat-JAX reproduces the nanochat architecture in JAX and Flax NNX: GQA, QK L2 normalization, logit softcap, Value Embeddings, Smear/Backout, per-layer scalars, and Muon. The nano model trains to val_loss 1.295 (perplexity 3.65) on TinyStories in under 10 minutes on a single GPU. The chat UI works and streams tokens via SSE from a FastAPI server. XLA is a real tradeoff. You pay a one-time compile cost of around 35 seconds per new model configuration and get zero Python overhead on every subsequent step. TPU portability is one flag away. The ecosystem gap is equally real: no Flash Attention 3, no distributed optimizer, no PEFT or vLLM. For GPU-only practitioners who need those tools, the JAX advantages do not justify the rebuild cost. The scaling law experiments are early. The measured exponent (alpha = 0.027 at 600 steps, rising with extended training) reflects the known limitations of short runs on a small dataset. Full <code>scale_n</code> and <code>scale_c</code> sweeps on TPU are planned for the AI GDE TPU Sprint 2026. Updated results with confidence intervals will be published when those runs complete. Code: <a href="https://github.com/ainaomotayo/nanochat-jax" rel="noopener noreferrer">github.com/ainaomotayo/nanochat-jax</a> This work is part of the AI Google Developer Expert TPU Sprint 2026, with compute from the Google TPU Research Cloud program.

此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。原文来自 — 版权归原作者所有。

推荐订阅源

DEV Community