gilesthomas-com

The author builds and trains a GPT-2 small model from scratch in JAX, starting from a basic bigram-style model and incrementally adding components like LayerNorm and Transformer blocks. Achieved a final loss of 3.418, beating their PyTorch version (3.538) and original GPT-2 small (3.499) on the same test dataset.

Building a JAX training loop for an LLM from scratch using Flax NNX and Optax. First validated the harness with a minimal "A-to-A" model (embedding then projection) before adding Transformer layers. Key challenges included fixing JAX's GPU memory default and slow data iteration by committing data to CPU memory.

A paper finds LLMs rely on text tone rather than unfakeable role tags to determine speaker identity, explaining many jailbreaks. Text mimicking a model's reasoning trace can trick it even when tagged as user input.

The author debugged a Flax NNX training loop where the loss was stuck at 10.82, indicating random guessing. By hashing the model parameters and comparing hashes across steps, they discovered the parameters weren't changing. The root cause was using @jax.jit instead of @nnx.jit, which is needed for proper in-place state propagation of parameter updates in NNX.

An overheating 10GBASE-T SFP+ module in a home switch was replaced with a Broadcom-based model, fixing link flapping and lowering CPU temperature by about 5°C, though the new module does not report temperature and impersonates a fibre-optic Intel module in its EEPROM.

A JAX developer found that array lookups on CPU-hosted data took over a second each, because arrays created with the `default_device` context manager are uncommitted to that device. Using `jax.device_put` to explicitly commit the array to the CPU reduced lookup times from ~1.2s to under 0.0002s, fixing a severe bottleneck in LLM training data loading.

A developer porting PyTorch code to JAX hit a GPU OOM error loading a 19 GiB dataset, because JAX defaults to the GPU backend. They solved it by using `jax.default_device(jax.devices("cpu")[0])` as a context manager to load data into RAM instead.

Using Safetensors with Flax requires a flat dictionary of string keys to arrays, not nested dicts. The author uses nnx.to_flat_state to flatten the model state, then converts it to dot-separated keys for save_file, and reverses the process for loading.

An experienced PyTorch user explores JAX, highlighting its mathematical purity and functional approach. Unlike PyTorch's procedural, piecewise-optimized style, JAX uses JIT compilation and transformations like jax.grad to compute derivatives automatically. The post argues PyTorch is engineering-focused while JAX is minimalist and math-oriented.

A user attached mini-heatsinks to a 10GBASE-T SFP+ module in their home network switch, achieving a temperature reduction of about 3.5°C. The discussion also highlights that older 30-meter-rated SFP+ modules (like the MikroTik S+RJ10 with a Marvell chip) run hotter than newer 100-meter-rated Broadcom-based modules, which could be a future upgrade.

The author upgraded his home network from 2.5Gb/s to 10Gb/s Ethernet using MikroTik switches, DAC cables, and a Protectli router. Tests confirmed the existing cabling could handle 10Gb/s, achieving ~8–9 Gb/s to the ISP. Some SFP+ modules run very hot (~93°C), requiring monitoring.

The author describes upgrading home networking to 10Gb/s after their ISP began offering the speed, covering the history of Ethernet from 10BASE2 to modern switches. They explain key challenges with 10GBASE-T—including heat and cable requirements—and recommend cooler alternatives like SFP+ Direct Attach Copper (DAC) cables or Active Optical Cables (AOCs) for short-range, in-room connections.

The author reviews the appendices of "Build a Large Language Model (from Scratch)" and found useful material on PyTorch basics, DistributedDataParallel training, and LoRA implementation. While these sections could have saved time during their explorations, they believe working through concepts independently provided deeper learning than simply reading explanations.

The author completed training a GPT-2-like model in 44 hours on a local machine, achieving performance close to GPT-2 small. Through systematic testing of various interventions, they identified learning rate adjustments and dropout removal as most effective for improving model loss. The author plans to next implement an LLM from scratch using JAX without reference to their book.

Updated instruction fine-tuning tests on GPT-2-style models show OpenAI's models performed best. Some custom models with similar test loss scores showed unexpected variations in instruction-following ability, with no clear pattern emerging across all tested models.

A researcher trained a GPT-2-style language model on 3.2 billion tokens and tracked its progress through 57 checkpoints. The model evolved from generating incoherent text to producing coherent, motivational content by processing about one-third of the training data.

The author explains gradient accumulation techniques to match cloud training batch sizes locally. By accumulating gradients over multiple forward-backward passes before optimizer updates, they achieve the stabilization benefits of larger batches without requiring more GPU memory. This allows local training with effective batch sizes comparable to cloud setups.

The author tested multiple training interventions on a 163M-parameter GPT-2-style model to improve its performance. The best result came from combining gradient clipping, removing dropout, using a higher scheduled learning rate, and changing weight decay to 0.01, achieving a test loss of 3.577761. This was better than the baseline loss of 3.691526 but still above the original GPT-2's 3.500.

The author tested various training interventions on a GPT-2-style model, finding learning rate scheduling provided the best improvement. Random seed experiments showed weight initialization significantly affects results, with losses ranging from 3.653 to 3.692, suggesting some intervention effects may be within noise levels.

The author tested training a GPT-2 small model without PyTorch's AMP and lower-precision matrix multiplication optimizations. The full float32 training took over 8 hours and cost $135, more than double the time and triple the cost of baseline runs. The resulting model showed only a tiny test loss improvement of 0.013, suggesting AMP provides significant speed benefits with minimal quality impact.

The author created a tool called lambda-manager to automate launching Lambda Labs instances. It monitors availability of specific instance types and launches them when they become available, then sends Telegram notifications. The tool has been running for six hours without finding the desired 8x A100 instance.

The article examines weight tying in LLMs, a technique that reduces parameters by sharing weights between input and output layers. The author tests this approach on a GPT-2 style model to see if it improves performance, despite research suggesting it typically worsens model quality.

The article examines weight decay as a regularization technique in training a GPT-2 small model from scratch. It explains that weight decay adds a penalty based on the squared L2 norm of model weights to the loss function to prevent overfitting. The author explores the mathematical formulation and its implementation in the AdamW optimizer.

The author explores learning rate scheduling for training an LLM from scratch, examining why fixed learning rates can fail and discussing various decay methods including step, exponential, and cosine decay. The post focuses on implementing a cosine learning rate scheduler with warmup, following recommendations from the Chinchilla paper.

The author experiments with adding bias to attention weight matrices in a GPT-2 small model trained from scratch. Surprisingly, this intervention improved test loss by 0.023 compared to the baseline, contradicting conventional wisdom that such bias doesn't help modern LLMs. The model showed slightly better stability during training despite adding only minimal extra parameters.