modded-nanogpt

This record achieved a training time of 7.23 minutes. It implements two primary modifications to the standard architecture:

1. U-net-like skip connections: Added into the transformer architecture.
2. Doubled learning rate: The learning rate was increased to accelerate training.

This record was benchmarked on 8x SXM H100 GPUs.

This project tracks three iterations of a MUDD-gated and Lightweight DC (Direct Correction) architecture for NanoGPT.

• v1: Combines a lightweight layer-10 DC correction with a MUDD-style gate stack that generates per-token coefficients for XSA gates, attention gates, x0/bigram injections, and layer skips.
• v2: Makes the DC tensors directly generated by the MUDD gate instead of using a standalone parameter path.
• v3 (Current Version): Rebases the architecture onto an upstream base featuring the bigram sign trick and FP8 MLP changes. It optimizes the bigram dimension and vocabulary size to ensure the MUDD gate can effectively utilize the bigram embedding path.

The AsyncDataLoadAttnFinalWindow implementation introduces several optimizations to reduce training bottlenecks and improve validation accuracy:

• Asynchronous Data Prefetching: Immediately starts prefetching and indexing the next shard on the CPU during GPU-heavy workloads to prevent CPU data indexing from bottlenecking the GPU.
• Partial Shard Indexing: Partially indexes the first shard before training begins, using a parallel thread to finish indexing by the 5th training step.
• Extended Final Layer Attention: Extends the final layer attention window to 20 for validation purposes. This does not require applying YaRN for this specific layer and has been shown to drop loss by approximately 0.001.

This directory contains experimental results from training runs extended to 50B FineWeb tokens. The experiments evaluate the generalization of speedrun configurations to long durations, specifically comparing the performance of the Muon optimizer against Adam for the transformer body (while the head and embedding remain optimized by Adam).

Key comparisons include:

1. Optimizer Performance: Muon vs. Adam at 50B tokens.
2. Training Duration: 5 epochs of 10B tokens vs. a single run of 50B tokens.

Results are documented in individual text files within this folder.

This specific optimization track combines three mechanisms to achieve low loss in a reduced number of steps (2995 steps). The combination includes:

• Muown: The base architecture.
• NorMuon: An optimization mechanism.
• Contra-Muon: An additional mechanism.
• Power-law cooldown LR: A learning rate mechanism.

The performance is measured by loss across different seeds, with the goal of passing a specific margin threshold after 2995 steps.

The Modded-NanoGPT Optimization Benchmark is a competitive effort to find efficient neural network optimizers. Unlike speedruns focused on wallclock time, this benchmark aims to minimize the step count required to reach a specific performance target.

Benchmark Goal: Achieve a validation loss below 3.28 in the fewest number of steps possible.

Validation Rules:

• Do not modify the dataset, batch size, or architecture used by the baseline.
• Do not perform more than one forward-backward pass per step.
• You are free to modify the optimization algorithm and hyperparameters arbitrarily.

The NorMuonAndAdam optimizer is a combined optimizer designed for per-parameter configuration rather than per-optimizer groups. It utilizes a 'parameter bank' strategy to consolidate weight matrices (like Attention and MLP) into single large parameters. This eliminates the need for time-consuming memcpys associated with stacking and unstacking gradients during the optimizer window.

Key benefits include:

• Zero-Copy: Uses parameter banks to allow reshaping without memory movement.
• Balanced Workloads: Distributes attention and MLP matrices evenly across GPUs.
• Simplified Configuration: Replaces complex optimizer groups with a direct mapping of individual parameters to specific optimizer settings and communication strategies.