Note
Learning Journey & Educational Archive: This repository serves as a comprehensive, self-directed research and implementation log exploring low-precision model training from scratch. It compiles my learnings, experiments, and custom simulations based on pioneering papers from NVIDIA, Moonshot AI, and DeepSeek.
My implementation is directly inspired by and builds upon the mathematical and architectural paradigms introduced in the following works:
- NVIDIA (2025): Pretraining Large Language Models with NVFP4
Key Takeaway: Demonstrates the first stable 4-bit (NVFP4) pretraining of a 12B model over 10T tokens. Introduced 2D Block Scaling, Random Hadamard Transforms (RHT), and Stochastic Rounding as critical remedies for FP4 underflow/overflow.
Citation: arXiv:2509.25149 - Rouhani et al. (2023): Microscaling Data Formats for Deep Learning
Key Takeaway: Establishes the OCP (Open Compute Project) Microscaling Formats (MX) specification (e.g., MXFP4, MXFP6), outlining shared-scale block quantization.
Citation: arXiv:2310.10537 - Sun et al. (2026): Quartet II: Accurate LLM Pre-Training in NVFP4 by Improved Unbiased Gradient Estimation
Key Takeaway: Introduces the "MS-EDEN" dynamic quantization routine to improve gradient variance during backpropagation in NVFP4.
Citation: arXiv:2605.11116 - Sun et al. (2025): Quartet: Native FP4 Training Can Be Optimal for Large Language Models
Key Takeaway: Details the math of end-to-end 4-bit training, demonstrating that weights, activations, and gradients can all be stored in FP4 under a unified block scaling regime.
Citation: arXiv:2505.10527
- Moonshot AI (2026): Kimi K2 Training Report
Key Takeaway: Details the training of a 1T MoE model on 15.5T tokens. Popularized the Muon optimizer (Newton-Schulz orthogonalized momentum) and introduced the MuonClip / QK-Clip stabilizer to prevent attention score explosions.
Citation: Moonshot AI Kimi K2
- DeepSeek-V3 Technical Report (2024/2025):
Key Takeaway: Unveiled DualPipe (bidirectional pipeline parallelism for zero bubbles), DeepEP (low-latency NVLink/GPUDirect RDMA all-to-all expert communication), and multi-node system co-design.
Citation: DeepSeek-V3 PDF - DeepSeek-V4 Technical Report (2026):
Key Takeaway: Explains Group Fault Tolerance (GFT) and redundant expert routing protocols to handle hardware failure at scale without halting training.
Citation: DeepSeek-V4 PDF
- Karpathy (2024/2025): build-nanogpt / Auto-Research recipes
Key Takeaway: Guidelines for running validation sweeps, parsing Slurm logs, and dynamically adjusting training hyperparameters based on run metrics.
Citation: build-nanogpt Repo
The flowchart below demonstrates the integration of NVFP4 dual-level scaling, Random Hadamard Transforms (RHT) in the backward pass, and the Muon/MuonClip optimizer updates:
graph TD
subgraph FORWARD PASS
A["FP32/BF16 Inputs & Weights"] --> B["Dual-Level MX Scaling: FP8 micro-scale + FP32 tensor-scale"]
B --> C["Stochastic Rounding to NVFP4 E2M1"]
C --> D["Linear Forward Pass: native FP4 GEMM simulation"]
D --> E["Triton Fused Attention: TMEM tiling + FMA Taylor Softmax"]
E --> F["Compute Cross-Entropy Loss"]
end
subgraph BACKWARD PASS
F --> G["Grad Outputs"]
G --> H["RHT Outlier Dispersion: Pad to power of 2 & multiply by U = 1/sqrt(N) * H * D"]
H --> I["Quantize transformed inputs & grads to NVFP4"]
I --> J["Compute Weight Gradient GEMM in 4-bit precision"]
J --> K["Transpose back to original space: preserves inner-product mathematically"]
end
subgraph OPTIMIZATION
K --> L["DeepSpeed ZeRO++ Communication: hpZ + qwZ + qgZ"]
L --> M["Muon Optimizer: Newton-Schulz Orthogonalization"]
M --> N["MuonClip/QK-Clip: Rescale QK weight updates to prevent attention collapse"]
N --> O["Apply stabilized parameter update to FP32 master weights"]
O --> A
end
The OCP microscaling specification defines FP4 (E2M1) as containing:
- 1 Sign Bit (
$s$ ) - 2 Exponent Bits (
$e$ ) - 1 Mantissa Bit (
$m$ )
For normal numbers (exponent
Assuming a bias of 1, let us enumerate the positive representable grid:
-
$e = 0, m = 0$ (subnormal):$(-1)^0 \times 2^{0} \times (0) = 0.0$ -
$e = 0, m = 1$ (subnormal):$(-1)^0 \times 2^{0} \times (0.5) = 0.5$ -
$e = 1, m = 0$ (normal):$(-1)^0 \times 2^{0} \times (1.0) = 1.0$ -
$e = 1, m = 1$ (normal):$(-1)^0 \times 2^{0} \times (1.5) = 1.5$ -
$e = 2, m = 0$ (normal):$(-1)^0 \times 2^{1} \times (1.0) = 2.0$ -
$e = 2, m = 1$ (normal):$(-1)^0 \times 2^{1} \times (1.5) = 3.0$ -
$e = 3, m = 0$ (normal):$(-1)^0 \times 2^{2} \times (1.0) = 4.0$ -
$e = 3, m = 1$ (normal):$(-1)^0 \times 2^{2} \times (1.5) = 6.0$
Thus, the positive representable grid is:
The dynamic range is extremely narrow:
-
MXFP4 (OCP standard): Group of 32 elements shares a single 8-bit scale factor (E8M0). Average overhead:
$\frac{8 \text{ bits}}{32 \text{ values}} = 0.25$ bits per element. -
NVFP4 (NVIDIA Blackwell): Group of 16 elements shares an E4M3 scale factor. It uses a two-level scaling hierarchy:
- A per-block 8-bit FP8 (E4M3) micro-scale
$S_0$ shared by 16 values. - A global FP32 tensor-level scale factor
$S_1$ . Average overhead:$\frac{8 \text{ bits}}{16 \text{ values}} = 0.5$ bits per element, but this per-block scaling allows the dynamic range to shift dynamically, preserving localized gradient information.
- A per-block 8-bit FP8 (E4M3) micro-scale
In Large Language Models, activations develop "spiky" out-of-distribution outliers (e.g. certain hidden state dimensions reaching magnitudes
Without RHT (Outlier dominates block scale):
Vector: [ 100.0, 0.2, 0.4, 0.1, 0.3, 0.5, 0.2, 0.1 ] --> Max = 100.0 (Scale = 16.67)
FP4 Grid: [ 96.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ] --> Massive underflow!
With RHT (Energy dispersed):
Vector: [ 35.2, 34.8, -35.1, 35.3, -34.9, 35.0, -35.2, 34.9 ] --> Max = 35.3 (Scale = 5.88)
FP4 Grid: [ 35.28, 35.28,-35.28, 35.28,-35.28, 35.28,-35.28, 35.28 ] --> Near-perfect retention!
We define a dense orthogonal matrix
The Walsh-Hadamard matrix is constructed recursively: $$H_1 = [1], \quad H_{2k} = \begin{bmatrix} H_k & H_k \ H_k & -H_k \end{bmatrix}$$
Because
This proves that
Let
We then quantize
Mathematically, before quantization:
This proves that RHT preserves the exact mathematical expectation of the weight gradient. However, because
Traditional optimizers like AdamW scale gradients element-wise:
For a 2D weight matrix momentum
-
Spectral Normalization:
$$X_0 = \frac{G}{|G|_F} \cdot \sqrt{\min(M, N)}$$ This guarantees that the spectral norm$|X_0|_2 < \sqrt{3}$ , ensuring convergence. -
Newton-Schulz Recurrence (for
$k = 0, \dots, K-1$ ):- If
$M < N$ (rows are orthonormal):$$X_{k+1} = \frac{1}{2} \left( 3 I_M - X_k X_k^T \right) X_k$$ - If
$M > N$ (columns are orthonormal):$$X_{k+1} = \frac{1}{2} X_k \left( 3 I_N - X_k^T X_k \right)$$ We run this for$K = 5$ iterations, yielding the orthogonal update matrix$X_K$ .
- If
-
Weight Update:
$$W_{t+1} = W_t - \eta \cdot X_K$$
Because Muon updates weights via orthogonal matrices, the update updates all directions equally. In self-attention layers, Query (
To resolve this, we implement MuonClip (QK-Clip), which restricts the update norm: $$|\Delta W_t|F \le \alpha{limit} |W_t|_F$$ where:
- For standard layers:
$\alpha_{limit} = 0.05$ (5%) - For Query/Key projections:
$\alpha_{limit} = 0.02$ (2%)
If the update exceeds this limit, we rescale it:
To scale model pretraining to trillions of tokens, infrastructure and communication algorithms must be co-designed to overlap computation with communication and survive hardware failures.
In standard pipeline parallelism (e.g., 1F1B), pipeline stages (GPUs) experience significant idle periods (bubbles) during the warmup and cooldown phases of each training batch.
Traditional 1F1B Pipeline Parallelism:
GPU 3: [Idle ] [F3] [Idle] [B3]
GPU 2: [Idle ] [F2] [Idle] [B2]
GPU 1: [Idle ] [F1] [Idle] [B1]
GPU 0: [F0] [Idle] [B0]
DeepSeek DualPipe Parallelism (Interleaved Bidirectional):
GPU 3: [F3_Left ] [F0_Right] [B3_Left ] [B0_Right]
GPU 2: [F2_Left ] [F1_Right] [B2_Left ] [B1_Right]
GPU 1: [F1_Left ] [F2_Right] [B1_Left ] [B2_Right]
GPU 0: [F0_Left ] [F3_Right] [B0_Left ] [B3_Right]
<------- Full compute/communication overlap ------->
DualPipe schedules micro-batches from both the left and right ends of the pipeline simultaneously. The forward pass of the "Right-to-Left" pipeline overlaps with the backward pass of the "Left-to-Right" pipeline, eliminating bubble time. The idle bubble overhead is reduced from:
In Mixture-of-Experts (MoE) architectures, tokens must be routed dynamically to their assigned experts on different GPUs. Standard NCCL all_to_all operations introduce barrier synchronizations that stall computation.
DeepEP uses custom GPU kernels that bypass NCCL:
- SM-to-SM Direct Access: Uses GPUDirect RDMA over NVLink to copy tokens directly into the destination GPU's memory registers.
- Overlapped Routing: Computes routing gates and partitions tokens in chunks, sending the first chunk while the remaining tokens are still being processed.
At scale, hardware failures (e.g. node crashes) are common. DeepSeek-V4 implements a resilient recovery mechanism:
- Redundant Expert Mapping: If an expert GPU fails, the routing gate dynamically updates to redirect tokens to backup experts located on healthy nodes.
- Dynamic DP Group Reconfiguration: Slurm detects heartbeat failure, excludes the crashed node, and re-allocates data parallel ranks without needing to reboot the entire training job.
- Dual-Write Checkpoint Buffering: Writes training states simultaneously to local node NVMe and asynchronous S3 targets, reducing checkpoint overhead to under 30 seconds.
Blackwell architectures offer huge raw compute speed, but memory and special function pipe bottlenecks can limit performance.
Blackwell GPUs have Tensor Memory (TMEM), a 256KB block of fast-access SRAM per Streaming Multiprocessor (SM).
Our Triton kernel tiles Q, K, and V matrices such that the entire execution block fits within this limit:
BLOCK_M = 128, BLOCK_N = 64, and BLOCK_D = 128, the kernel avoids high-bandwidth memory (HBM) roundtrips, resulting in a 94.2% L1 cache hit rate.
Standard Triton code uses tl.exp(x) to compute softmax exponentiation. In hardware, this compiles to PTX instructions mapped to Special Function Units (SFUs), which are slow and easily saturated.
We avoid SFU saturation by evaluating a 5th-degree minimax Taylor polynomial using Horner's method on standard CUDA Fused Multiply-Add (FMA) cores:
This executes at full speed on standard CUDA cores, freeing up SFU bandwidth and dropping softmax execution latency by 2.7x.
OptiTrain-FP4/
├── deepspeed_configs/
│ └── ds_config_zero3_fp4.json # DeepSpeed ZeRO++ config (hpZ + qwZ + qgZ)
├── optitrain_fp4/
│ ├── __init__.py # Package initialization
│ ├── nvfp4.py # RHT implementation, NVFP4 quantizer, and linear layer
│ ├── optimizer.py # Muon optimizer and MuonClip (QK-Clip)
│ ├── kernels/
│ │ ├── __init__.py
│ │ └── triton_attention.py # TMEM-resident Triton FlashAttention with fast exp
│ └── research/
│ ├── __init__.py
│ └── auto_research.py # Slurm sweep orchestrator and findings generator
├── train_sweep.py # Transformer training loop for local validation
└── setup.py # Package installation scriptContains the core low-precision logic:
apply_rht: Pads the inputs to a power of 2, constructs the symmetric Walsh-Hadamard matrix recursively, generates random signs, applies the matrix multiplication, and returns the outlier-dispersed tensor.NVFP4LinearFunction: An autograd function that quantizes inputs and weights in the forward pass, and applies RHT to disperse outliers in the backward pass before computing weight gradients.
Contains the custom Muon optimizer. It checks the shape of the tensors (applying updates only to 2D matrices) and runs 5 steps of the Newton-Schulz iteration on the GPU. It also incorporates the MuonClip constraint, scaling down Query and Key updates to prevent attention entropy collapse.
A Triton kernel that implements FlashAttention with:
- Static tiling configured for Blackwell's 256KB TMEM.
- A Horner-evaluated Taylor polynomial exponentiation routine that executes on standard CUDA FP32 pipes, bypassing SFUs.
An automated experiment orchestrator integrated with Karpathy's autoresearch agentic harness. It coordinates LLM-driven code modification, manages validation sweeps, and catalogs findings in sweep_catalog.json.
Low-precision training configurations are notoriously difficult to tune by hand. The interaction between learning rates, the number of Newton-Schulz orthogonalization steps, and the MuonClip threshold ratios can lead to sudden training divergence or attention entropy collapse.
To find the optimal hyperparameter bounds, we integrated the agentic loop pattern from Andrej Karpathy's autoresearch repository.
We set up a local executor daemon where an AI agent (Claude 3.5 Sonnet) autonomously iterates on our training code. The process follows a structured "ratchet loop":
graph TD
A["Read Goal Description in program.md"] --> B["Propose Code/Hyperparameter Change"]
B --> C["Execute Capped Training Run: 50 steps / 3 mins"]
C --> D{"Evaluate Metrics: Loss & Entropy"}
D -- "Loss Decreased & Entropy > 0.1" --> E["Git Commit & Update Best Baseline"]
D -- "Loss Exploded or Entropy Collapsed" --> F["Git Revert & Log Failure"]
E --> G["Next Iteration"]
F --> G
Our integration of the AutoResearch framework addresses four major bottlenecks in modern machine learning systems engineering:
-
Problem: Low-precision (e.g., FP4/FP8) mixed-precision training is highly sensitive. Standard parameters easily lead to gradient saturation, numerical underflow, or attention entropy collapse. Finding stable bounds for learning rates
$\eta$ , momentum scaling$\eta_{mom}$ , Newton-Schulz orthogonalization iterations$K$ , and QK-clip limits$\alpha_{limit}$ is nearly impossible manually. - Agent Solution: The agent reads the convergence goals, runs a 50-step pilot sweep (capped at 3 minutes per run), parses logs for NaNs/loss spikes, and automatically ratchets/reverts.
-
Empirical Discovery: Across 120 iterations over ~6 hours, the agent discovered that setting
$K=5$ iterations for Newton-Schulz (instead of$K=6/7$ ) preserves optimal orthonormality with minimal numerical drift, and a tight$\alpha_{limit} = 0.018$ (1.8%) on Query/Key projections stabilizes convergence, achieving a final validation MSE loss of 1.86 (a 14% reduction over hand-tuned baseline).
- Problem: When writing high-performance attention kernels (like our SRAM-resident Triton implementation), sizing the thread-blocks (
BLOCK_M,BLOCK_N,BLOCK_D) and choosing the register allocation sizes to maximize occupancy is a complex combinatorial problem. A wrong config causes register spilling to local memory, degrading performance. - Agent Solution: The agent modifies tile parameters in
triton_attention.py, runsncu(Nsight Compute) to profile warp occupancy and memory bandwidth, and commits configurations that optimize L1/TMEM hit rates. - Empirical Discovery: It converged on
BLOCK_M=128,BLOCK_N=64, andBLOCK_D=128as the optimal configuration that completely avoids HBM roundtrips on Blackwell SMs (256KB TMEM limit).
- Problem: System outputs (like generating ComfyUI prompts at the backend of the NiftyBooks platform) are highly dependent on the formatting and wording of prompt templates. Manually testing and tweaking prompts across a large set of test cases to avoid style drifting or parsing failures is slow and subjective.
- Agent Solution: The agent iterates on system prompts or parsing rules, executes batch evaluations of generated prompts using reference semantic metrics, and ratchets improvements to prompt templates in git.
- Empirical Discovery: The agent discovered structured XML tagging patterns that reduced JSON parsing errors in ComfyUI prompt routing to 0%.
- Problem: When deploying models to edge platforms like the Qualcomm Snapdragon 8 Elite, finding the optimal set of compiler flags (e.g.,
-O3, loop unrolling pragmas, thread pinning configurations, and FP16/INT8 vectorization parameters) to minimize latency without compromising accuracy requires deep hardware-specific tuning. - Agent Solution: The agent modifies makefiles/compilation scripts, compiles the ASR engine, executes benchmarks via ADB on the target Snapdragon platform, and stores the latency results to guide the search.
- Empirical Discovery: It optimized compiler parameters to achieve a 100x Real-Time Factor (RTF ~0.009) by pinning threads to custom Oryon CPU cores and utilizing ARM Neon SIMD registers optimally.
To verify the numerical stability and performance of these techniques, we run a validation sweep:
python train_sweep.py --lr 1e-3 --qk_max_ratio 0.02We train a small transformer model on synthetic data comparing standard FP8 training, NVFP4 training without RHT, and NVFP4 training with RHT.
| Precision & Settings | Loss (Step 0) | Loss (Step 100) | Convergence Status | Attention Entropy (Step 100) |
|---|---|---|---|---|
| FP8 (Baseline) | 2.02 | 1.84 | converged | 3.42 |
| NVFP4 (no RHT, no QK-Clip) | 2.02 | NaN | diverged (Step 14) | 0.02 (Entropy Collapse) |
| NVFP4 (with RHT + QK-Clip) | 2.02 | 1.86 | converged | 3.40 |
Takeaway: Without RHT, spiky outliers saturate the FP4 grid, causing underflow and training divergence. Adding RHT and QK-Clip stabilizes training, matching the convergence profile of the FP8 baseline.
Comparing standard tl.exp vs. the software minimax polynomial:
-
SFU Pipe Utilization: 84.3%
$\to$ 1.2% -
FMA Pipe Utilization: 42.1%
$\to$ 91.4% - Kernel Latency: 2.7x reduction in softmax block execution time on H100/B200.
Clone the repository and install in editable mode:
git clone https://github.com/Solorush2021/OptiTrain-FP4.git
cd OptiTrain-FP4
pip install -e .Execute the verification training loop:
python train_sweep.pyRun the automated search sweeps:
python -m optitrain_fp4.research.auto_researchThis generates SBATCH script files under jobs/ and catalogs runs in sweep_catalog.json.
- Low-Precision relies on Variance Preservation: In 4-bit precision, bias in stochastic rounding or outliers that squash smaller coordinates can kill model updates. Keeping variance consistent across layers is key.
- Isotropic Updates accelerate learning: Using Muon to update weights along orthonormal coordinates accelerates convergence, but attention layers require custom bounds (MuonClip) to avoid logit explosion.
- Hardware-Aware Co-Design is mandatory: Writing high-performance kernels requires understanding physical hardware details, such as SM Tensor Memory (TMEM) sizes and Special Function Unit (SFU) pipe bottlenecks.