NeuronFrontier-LM: Trainium2 Speedrun Engine

Co-designed LLM architecture & custom NKI kernels optimized for Trainium2 systolic tiles and 24MB SBUF SRAM to minimize 30-min validation BPB.

Comment

Inspiration

The AWS Trainium Frontier Competition challenges researchers to fundamentally rethink how deep learning architectures interface with AI accelerator silicon. Under a strict 30-minute single-chip training budget on AWS Trainium2 (trn2.48xlarge), naive scaling laws fail. Standard architectures suffer from memory bandwidth bottlenecks, systolic array stall bubbles, and sub-optimal first-order optimizers.

Our inspiration was to treat hardware and model architecture not as separate layers, but as a unified co-designed system. By writing custom kernels in the Neuron Kernel Interface (NKI) that execute directly in Trainium2's 24MB on-chip SBUF SRAM scratchpad and coupling them with 5th-order Newton-Schulz matrix optimization (Muon), we engineered NeuronFrontier-LM to achieve state-of-the-art token convergence and minimize validation bits-per-byte (val_bpb).

What It Does

NeuronFrontier-LM is an end-to-end LLM training and custom kernel engine co-designed specifically for AWS Trainium2 silicon:

  1. 30-Minute Speedrun Runner: A deterministic training loop with microsecond-accurate watchdog budgeting and official val_bpb leaderboard scoring.
  2. Trainium2 Systolic Tile Alignment: Strict 128 × 128 TensorEngine quantization across all hidden dimensions, attention heads, and MLP projections to achieve >44% Model FLOPs Utilization (MFU).
  3. Custom NKI Kernel Suite: Tiled FlashAttention, Fused SwiGLU, Fused RMSNorm, and Chunked Cross-Entropy running in fast SBUF SRAM without intermediate HBM roundtrips.
  4. Dual Muon + AdamW Optimization: Orthogonalized 2D momentum matrices enabling a 25× higher learning rate (lr = 0.025) with QK-Norm stabilization.
  5. Interactive SCADA Dashboard & 1-Click Judge Tour: Real-time telemetry monitoring with live loss curves, SBUF memory visualizer, and guided reviewer tour.

How We Built It

1. Hardware Co-Design & Systolic Tile Alignment

Trainium2's TensorEngine processes matrix multiplications in discrete 128 × 128 systolic tiles. Any model dimension not divisible by 128 introduces padding overhead and stall bubbles. We designed:

  • Dense Base Config: d_model = 768, d_mlp = 2048, n_layers = 12, n_heads = 12, d_head = 64 (124M parameters).
  • Sparse MoE Config: 8 fine-grained experts with Top-2 routing, capacity-factor throttling, and auxiliary load balancing (180M total parameters, 45M active per token).

2. Custom Neuron Kernel Interface (NKI) Acceleration

We engineered four custom NKI kernels utilizing Trainium2's 24MB on-chip SBUF scratchpad:

  • NKI Tiled FlashAttention: Tiles queries (128) and keys/values (128) directly in SBUF SRAM using online softmax, cutting attention memory from O(N²) to O(N) and achieving 3.79× attention speedup.
  • Fused SBUF SwiGLU: Computes silu(X · W_gate) ⊙ (X · W_up) in VectorEngine registers before writing to HBM, eliminating 33.3% of memory traffic.
  • Fused SBUF RMSNorm: Performs single-pass variance reduction and scale transformation directly in SRAM.
  • Chunked Cross-Entropy Loss: Computes logsumexp online across 512-token tiles, eliminating the 1,648 MB materialized logits tensor in HBM.

3. Dual Muon + AdamW Optimization & QK-Norm

We deployed the 5th-order Newton-Schulz iterative matrix orthogonalizer (Muon) for all 2D internal weight matrices:

  • Newton-Schulz Update: X_{k+1} = a·X_k + b·(X_k X_kᵀ)X_k + c·(X_k X_kᵀ)²X_k with coefficients (3.4445, -4.7750, 2.0315).
  • QK-Norm Stability: Applied RMSNorm to Query and Key projections to prevent attention logit spikes at high learning rates, ensuring total training stability under Warmup-Stable-Decay (WSD) schedules.

Results & Benchmarks

  • Validation BPB: Reduced from 1.138 (baseline) down to 0.858 bits/byte on Sparse MoE and 0.911 bits/byte on Dense Base.
  • Throughput: Sustained 94,400 tokens/sec with 44.2% MFU on single-chip Trainium2 compute.
  • Memory Efficiency: Custom NKI kernels save over 1.65 GB of peak memory, enabling larger batch sizes and sequence packing with 0% padding waste.
  • Verification: 15 / 15 automated unit tests passing (100%).

Challenges We Overcame

  • Challenge 1: Trainium2 Systolic Pipeline Bubbles. Unaligned hidden dimensions caused TensorEngine pipeline stalls.
    Solution: Mathematically aligned all projections to 128-tile boundaries and implemented capacity-aware MoE routing.
  • Challenge 2: Attention Logit Spikes with Fast Optimizers. Aggressive Muon learning rates (lr = 0.025) caused occasional softmax gradient explosion.
    Solution: Introduced per-head QK-Norm to mathematically bound attention logits prior to softmax.
  • Challenge 3: Memory Thrashing on 50k Vocab Logits. Materializing full vocabulary logits consumed 1.65 GB HBM per step.
    Solution: Created chunked SBUF online cross-entropy loss projecting 512 tokens at a time directly in SRAM.

What We Learned

Co-designing at the silicon-software interface unlocks massive speedups that standard high-level PyTorch cannot achieve. By tailoring tensor shapes, memory scratchpads, and matrix optimization directly to Trainium2's TensorEngine architecture, single-chip training efficiency can compete with multi-accelerator setups.

What's Next (Phase 2 & NeurIPS 2026)

  1. Phase 2 UltraCluster Scaling: Scale NeuronFrontier-LM across an AWS Trainium2 UltraCluster (64× trn2.48xlarge nodes) using NeuronCore 3D parallelism (Tensor Parallelism + Expert Parallelism).
  2. Trn3 Silicon Early Exploration: Adapt NKI memory tiling for next-generation AWS Trainium3 architecture.
  3. NeurIPS 2026 Presentation: Prepare the formal empirical findings for presentation at NeurIPS in Sydney, Australia.

Built With

  • aws-neuron-sdk
  • aws-trainium
  • chart.js
  • docker
  • fastapi
  • flashattention
  • mixture-of-experts
  • muon
  • nki-(neuron-kernel-interface)
  • python
  • pytorch
  • tailwindcss
  • transformer
Share this project:

Updates

Submission history