// =============================================================================
// bluscript.cpp — Llama-style decoder LM training (fp32 model, bf16 fused attn,
// ZeRO-1 mixed precision optimizer state sharding, WSD LR schedule)
//
// Architecture (per the production spec):
// token embedding
// -> N x TransformerBlock:
// RMSNorm -> { QK-norm + RoPE + grouped-query causal attention }(fused
// kernel: gqa_fused_flash_attn_forward/backward) -> Wo -> +residual
// RMSNorm -> SwiGLU FFN (gate_up -> swiglu -> down) -> +residual
// -> final RMSNorm
// -> lm_head (weight-tied with the token embedding)
// -> sparse cross-entropy loss
// -> backward -> AdamW, or Muon+AdamW (BLU_OPTIMIZER=muon), with a WSD
// (warmup / stable / cosine-decay) LR schedule
//
// The QK-norm scales (q_gamma / k_gamma) and the RoPE rotation are FUSED in one
// CUDA kernel (autograd::fused_rms_rope), and grouped-query causal attention is
// a second fused flash-attention kernel reached through
// autograd::scaled_dot_product_attention. Those ops run bf16 internally and
// bridge back to the fp32 model. Matmuls go through the library BLAS
// (nn::Linear / autograd::matmul).
//
// Optimizer: ZeRO-1 (owt::zero::ShardedOptimizer). FP32 master weights + Adam
// moments live only in each rank's optimizer shard, not the model. Designed for
// ZeRO-1 over 8 GPUs (one rank per GPU): each rank holds the full model params
// but only 1/world_size of the optimizer state; per step, grads reduce-scatter
// (avg) across ranks, the fused Adam update runs on each rank's FP32 shard, and
// updated params all-gather back — no separate DDP grad all-reduce.
//
// BLU_OPTIMIZER=muon switches the transformer body's 2-D weight matrices to
// Muon, keeping AdamW for the token embedding and every 1-D scale. ZeRO-1 STILL
// APPLIES to the Muon set, but it cannot be the same ZeRO-1: AdamW's update is
// elementwise, so any partition of the flat parameter buffer is valid, while
// Muon's Newton-Schulz orthogonalization is a whole-MATRIX operator that a
// contiguous flat slice cuts through the middle of. The Muon set therefore goes
// to owt::zero::ShardedMuon, which shards by WHOLE PARAMETER — one owner rank
// per matrix, holding that matrix's FP32 master + momentum, reducing its
// gradient to the owner and broadcasting the updated weight back. Same ~1/W
// state reduction, and cheaper than sharded AdamW (no second moment, no flat
// staging buffer). See dist/zero/include/ShardedMuon.h.t
//
// Launch: make run-mpi FILE=bluscript.cpp NP=8 WITH_BLUBLAS=1
// BLU_OPTIMIZER=muon mpirun -np 8 ./_runner_
// =============================================================================
#include <cstdint>
#include <cstdlib>
#include <cassert>
#include <iostream>
#include <iomanip>
#include <chrono>
#include <cmath>
#include <vector>
#include <string>
#include <fstream>
#include <random>
#include <map>
#include <set>
#include <functional>
#include <algorithm>
#include <mpi.h>
#include <cuda_runtime.h>
// Tensor library
#include "TensorLib.h"
#include "autograd/AutogradOps.h" // matmul, add, reshape, transpose, swiglu, ...
#include "autograd/operations/AttentionOps.h" // fused GQA scaled_dot_product_attention
#include "autograd/operations/RoPEOps.h" // build_rope_cache
#include "autograd/operations/LossOps.h" // sparse_cross_entropy_loss
#include "autograd/operations/ReshapeOps.h"
#include "autograd/operations/EmbeddingOps.h"
#include "nn/NN.h"
#include "nn/optimizer/Optim.h"
#include "checkpointing/GradMode.h"
#include "DistributedCheckpointing.h" // DistCheckpoint::save / load / is_valid
#include "ZeroDcpBridge.h" // ZeroDcpModule, ZeroDcpOptimizer
#include "device/CachingCudaAllocator.h"
#include "device/DeviceCore.h" // OwnTensor::cuda::getCurrentStream()
// Tracker + Memory
#include "device/SizeClass.h"
#include "profiler/AllocationTracker.h"
// Data + distributed
#include "DataLoader.h"
#include "dist/distributed.h"
#include "dist/zero/include/ShardedOptimizer.h" // owt::zero::ShardedOptimizer (ZeRO-1)
#include "dist/zero/include/SymmetricShardedOptimizer.h" // owt::zero::SymmetricShardedOptimizer (ZeRO-1 v2)
#include "dist/zero/include/ShardedMuon.h" // owt::zero::ShardedMuon (ZeRO-1 over whole tensors)
#include "dist/zero/include/ZeroOptimizerBase.h" // common surface so v1/v2/muon are interchangeable
// Compile-time flag to toggle between fused RMSNorm+RoPE kernel (1) and sequential RMSNorm + RoPE (0)
#ifndef USE_FUSED_RMS_ROPE
#define USE_FUSED_RMS_ROPE 1
#endif
using namespace OwnTensor;
using owt::zero::ShardedOptimizer;
using owt::zero::SymmetricShardedOptimizer;
using owt::zero::ShardedMuon;
using owt::zero::ZeroConfig;
using owt::zero::MuonConfig;
using owt::zero::AdamVariant;
// =============================================================================
// Configuration — resolved from the bluscript.cpp template (integer-division
// placeholders evaluated to their intended values).
// =============================================================================
struct ModelConfig {
// ---- model ----
int64_t d_model = 3072;
int64_t n_layers = 36;
int64_t q_heads = 24;
int64_t kv_heads = 8; // grouped-query: G = q/kv = 3
int64_t head_dim = 128; // = 64 (fused kernel needs 64 or 128)
int64_t ffn_hidden = 8192; // 8/3 * d_model = 2048 (SwiGLU inner dim)
int64_t context_length = 4096; // = T (multiple of 32, RoPE cache length)
int rope_theta = 500000;
int64_t vocab_size = 50304; // 50257 (tiktoken) padded for GPU efficiency
bool weight_tying = false;
float rms_eps = 1e-5f; // block RMSNorm epsilon
float qk_eps = 1e-6f; // QK-norm (fused) epsilon
bool rope_interleaved = false; // false = NeoX half-split (Llama style)
// ---- activation / parameter precision ----------------------------------
// All false (the default) is the shipped mixed-precision layout and is
// bit-identical to the pre-flag code: bf16 Linears + attention, fp32
// residual stream / norms / embedding / logits, fp32 optimizer masters.
// Each flag moves exactly ONE stage to bf16, so a regression bisects to a
// single stage. Wired from BLU_ALL_BF16 (+ per-stage overrides) in main().
//
// Two things are deliberately NOT flags, because neither is reachable by
// changing a dtype here:
//
// * Optimizer masters and moments. shard_master_/shard_m_/shard_v_ are
// hardcoded Float32 (ShardedOptimizer.cpp allocate_buffers) and
// multi_tensor_adam_cuda takes float* only. They stay fp32 under every
// flag below, which is what keeps the update path safe: AdamW's
// decoupled decay is a (1 - lr*wd) = 1 - 6e-5 multiply, and bf16's
// relative grid is 2^-8 = 3.9e-3, so a bf16 master would round the
// entire weight-decay term away at any weight magnitude.
//
// * QK-norm + RoPE. fused_rms_rope_forward reads q, k, q_gamma and
// k_gamma through data<float>() with no dtype dispatch, and
// autograd::rope rejects anything but fp32 outright. data<T>() is a
// bare reinterpret_cast, so handing that kernel a bf16 tensor would
// read garbage SILENTLY rather than throw. q_gamma/k_gamma therefore
// stay fp32 and the island is bracketed by explicit casts in
// CausalGQA::forward. See qk_dtype().
// * Logits. bf16_logits is wired end to end and the sparse-CE kernels now
// accept it: LossKernels.cu's vectorized paths key every chunk off
// 16/sizeof(T) elements instead of a hard-coded 4, in the forward and in
// both backward variants (see VEC_ELEMS / Chunk16 there). Before that they
// staged 16-byte cp.async loads at a stride of 4 ELEMENTS — 8-byte aligned
// on bf16, a hard "misaligned address" fault — and
// `reinterpret_cast<float4*>(row_grad)` wrote 16 bytes per 8-byte logical
// span, overrunning the gradient buffer SILENTLY.
//
// What is NOT fixed, and is not fixable in those kernels: LossOps keys the
// LOSS SCALAR's dtype off the logits dtype, so bf16_logits also makes the
// reported loss bf16. The per-row losses are summed in fp32 but stored back
// through a bf16 tensor before the divide by batch_size, and bf16 cannot
// represent a loss near 10.8 finer than one ULP = 0.0625 either way — so the
// logged/allreduced loss quantizes to ~0.06 steps. GRADIENTS ARE UNAFFECTED
// (backward() is seeded at 1.0, which is exact in bf16, and the softmax
// stats are fp32), so this costs the metric, not the run. Decoupling the
// loss dtype from the logits dtype needs a LossT template parameter through
// LossKernels/LossOps plus a grad_output dtype check in LossBackward, which
// reads it via data<LogitT>() with no dtype guard.
bool bf16_resid = false; // residual stream + block RMSNorm scales
bool bf16_logits = false; // logits + loss scalar — see the note above
bool bf16_embed = false; // wte (under weight tying, also the lm_head weight)
Dtype resid_dtype() const { return bf16_resid ? Dtype::Bfloat16 : Dtype::Float32; }
Dtype logit_dtype() const { return bf16_logits ? Dtype::Bfloat16 : Dtype::Float32; }
Dtype embed_dtype() const { return bf16_embed ? Dtype::Bfloat16 : Dtype::Float32; }
// Not a flag — a kernel constraint. See the note above.
static Dtype qk_dtype() { return Dtype::Float32; }
// ---- optimization / data ----
int64_t B = 2; // micro-batch per rank
int64_t T = 4096; // sequence length (== context_length)
int64_t global_batch = 2097152; // tokens per optimizer step (all ranks)
float grad_clip = 1.0f;
float weight_decay = 0.1f;
// WSD schedule: linear warmup -> constant at max_lr -> cosine decay to min_lr.
// Requires warmup_steps + stable_steps < max_steps (the remainder is the decay tail).
float max_lr = 6e-4f;
float min_lr = 6e-5f; // 10% of max_lr
int warmup_steps = 200; //715;
int max_steps = 2000; //19073; // ~10B tokens at 0.5M tokens/step
int stable_steps = 1400; // steps held at max_lr after warmup
// ---- Muon (BLU_OPTIMIZER=muon) ----
// Mirrors train_muon.py's config so the C++ and PyTorch arms are the same
// experiment. Muon's group rides the SAME WSD curve as AdamW's — same
// warmup_steps, same stable window, same cosine tail — via a second get_lr()
// call with its own (max, min) pair, exactly as train_muon.py does at its
// `muon_lr = get_lr(step, cfg, cfg.muon_lr, cfg.muon_min_lr)` line. There is
// deliberately no separate muon_warmup_steps: both groups train the same
// model on the same steps, so only the SCALE should differ.
//
// Under the default adjust_lr=match_rms_adamw the Newton-Schulz update's
// entry RMS is lifted to 0.2 — what a converged AdamW update looks like —
// which puts muon_lr on the SAME SCALE as max_lr. Setting them equal makes
// the two arms differ only in update DIRECTION, which is the controlled
// experiment worth running first. 8e-4 here is 4x max_lr, the post-sweep
// value train_muon.py ships (Muon typically tolerates 2-4x more LR once it
// is working).
float muon_lr = 8e-4f;
float muon_min_lr = 8e-5f; // 10% of muon_lr
float muon_momentum = 0.95f;
// Pinned to weight_decay by default so the decay is identical across both
// groups and both arms — only meaningful because match_rms_adamw puts the
// two LRs on one scale. train_muon.py pins it the same way.
float muon_weight_decay = 0.1f;
int muon_ns_steps = 5;
int val_freq = 250;
int val_loss_steps = 20;
// ---- checkpointing ----
// OFF by default so a throughput A/B against the PyTorch arms is apples to
// apples: train_ds_zero1.py and train_zero1.py do not port checkpointing, so
// at ckpt_freq=900 a 2000-step run would spend three DCP saves (steps
// 899/1799/1999) writing 681 M params plus the sharded optimizer state that
// the PyTorch side never pays for — the Muon master+momentum alone is
// 2304 MB/GPU. BLU_CKPT=1 turns it back on; that is the setting for testing
// DCP itself, and the only one under which resume is exercised.
bool checkpointing = false;
int ckpt_freq = 900; // save every 1800 optimizer steps
int ckpt_keep = 2; // retain the most-recent N checkpoints
std::string ckpt_dir = "checkpoints_bluscript";
std::string ckpt_prefix = "blumodel";
};
// =============================================================================
// MFU — Model FLOPs Utilization
// =============================================================================
// The number frontier labs report next to throughput (PaLM 46.2%, Megatron-LM,
// MosaicML, Llama). It answers "what fraction of the GPUs' peak matmul rate is
// this run actually converting into MODEL math?" and is the only throughput
// metric comparable across model sizes, cluster sizes and codebases:
//
// MFU = model_FLOPs_per_step / (step_wall_time * world_size * peak_FLOPS)
//
// Two conventions matter, and both are implemented here.
//
// * MODEL FLOPs (this file's numerator) counts the arithmetic the MODEL
// DEFINITION requires — the forward+backward matmuls of the architecture,
// once. It is deliberately implementation-blind: activation recomputation,
// padded GEMM tiles, redundant casts and communication do NOT count. That is
// the point. A run that recomputes every layer does ~33% more hardware work
// for the SAME model FLOPs, so its MFU correctly drops. HARDWARE FLOPs
// utilization (HFU) counts what the silicon actually executed and is always
// >= MFU. THIS LOOP USES NO ACTIVATION CHECKPOINTING (the micro loop calls
// forward once and backward replays the stored graph), so here HFU == MFU
// and only one number is reported.
//
// * The ATTENTION-CORE count differs between papers. PaLM (Appendix B) and
// Megatron-LM charge the full T^2 score matrix; DeepMind/Chinchilla and
// MosaicML's llm-foundry halve it for causal masking. This kernel is
// FlashAttention with is_causal=true, which genuinely SKIPS the masked
// blocks — charging for them would credit work the GPU never did and inflate
// MFU. So the causal-aware count is the headline number, and the PaLM-
// convention count is printed alongside it so results stay comparable to
// published figures (it is always the larger, i.e. more flattering, of the
// two).
//
// Everything is derived from the architecture dims rather than from 6*N, so the
// count stays correct when weight_tying is turned off (an untied lm_head adds
// vocab*d_model PARAMETERS, but the embedding lookup is still a gather and not
// a matmul, so 6*N would over-count by 6*vocab*d_model). With tying ON the two
// agree exactly — the banner prints both as a self-check.
//
// Per convention, non-matmul work (RMSNorm, QK-norm, RoPE, SwiGLU's swish,
// softmax, and the optimizer's Adam / Newton-Schulz update) is NOT counted: it
// is a low-single-digit share of the FLOPs, it is memory-bandwidth bound rather
// than tensor-core bound, and every published MFU omits it. Counting it would
// make this run look better than the papers it is being compared against.
// =============================================================================
struct MFUMeter {
// ---- per-token FLOPs of one forward pass (matmuls only) ----
double fwd_per_token_causal = 0.0; // FlashAttention is_causal (headline)
double fwd_per_token_dense = 0.0; // PaLM / Megatron convention
// ---- fwd + bwd = 3x forward (dX = dY@W^T and dW = X^T@dY each cost a fwd)
double flops_per_token_causal = 0.0;
double flops_per_token_dense = 0.0;
double peak_flops_per_gpu = 0.0; // dense BF16 tensor core, no sparsity
std::string device_name = "unknown";
bool peak_is_guess = false; // no table entry matched
int world_size = 1;
int warmup_steps = 5; // steps excluded from the steady-state mean
// Steady-state accumulators. Aggregate (total model FLOPs / total wall
// time), not a mean of per-step ratios — one stalled step then costs
// exactly its own wall time instead of dragging an average of ratios.
double ss_tokens = 0.0;
double ss_secs = 0.0;
int ss_steps = 0;
int seen = 0; // observe() calls, INCLUDING the warmup
double best_mfu = 0.0;
// -------------------------------------------------------------------------
// Peak dense BF16 tensor-core throughput per GPU, WITHOUT the 2:1
// structured-sparsity doubling that NVIDIA's headline numbers carry. Dense
// is the correct denominator: nothing in this step uses sparsity, so the
// sparse peak would halve every MFU reported here.
//
// H200 SXM is H100 SXM's compute die paired with HBM3e, so its BF16 rate is
// identical: 989.4 TFLOP/s dense (the spec sheet's "1,979 TFLOPS" is the
// sparse figure). The PCIe / NVL parts clock lower and cannot be told apart
// by name, so use BLU_PEAK_TFLOPS there.
// -------------------------------------------------------------------------
static double peak_bf16_tflops(const std::string& name, bool& matched) {
struct Entry { const char* key; double tflops; };
static const Entry table[] = {
{"GB200", 2500.0}, // Blackwell superchip, per-GPU dense BF16
{"B200", 2250.0}, // Blackwell SXM
{"B100", 1800.0},
{"H200", 989.4}, // SXM (NVL/PCIe ~835 — set BLU_PEAK_TFLOPS)
{"H100", 989.4}, // SXM (PCIe ~756)
{"H800", 989.4},
{"L40S", 362.1},
{"A100", 312.0},
{"A800", 312.0},
};
for (const auto& e : table)
if (name.find(e.key) != std::string::npos) { matched = true; return e.tflops; }
matched = false;
return 989.4; // assume H200-class rather than divide by zero
}
static MFUMeter build(const ModelConfig& c, int world_size, int device_ordinal) {
MFUMeter m;
m.world_size = world_size;
const double d = static_cast<double>(c.d_model);
const double L = static_cast<double>(c.n_layers);
const double hd = static_cast<double>(c.head_dim);
const double Nq = static_cast<double>(c.q_heads);
const double Nkv = static_cast<double>(c.kv_heads);
const double F = static_cast<double>(c.ffn_hidden);
const double T = static_cast<double>(c.T);
const double V = static_cast<double>(c.vocab_size);
// A GEMM [m,k] x [k,n] costs 2*m*k*n FLOPs (one multiply + one add per
// MAC). Per TOKEN every weight matmul is a [1,k] x [k,n] row, so each
// costs exactly 2 * (its parameter count).
const double qkv = 2.0 * d * (Nq * hd) // wQ
+ 2.0 * d * (Nkv * hd) // wK (GQA: Nkv < Nq)
+ 2.0 * d * (Nkv * hd); // wV
const double proj = 2.0 * (Nq * hd) * d; // wO
const double ffn = 2.0 * d * (2.0 * F) // gate_up, emits [gate|up]
+ 2.0 * F * d; // down
// Attention core — the two matmuls with NO parameters. They scale with
// T, so a 6*N count is blind to them entirely:
// scores = Q @ K^T : per token, per q-head, [1,hd] x [hd,T] = 2*T*hd
// out = P @ V : per token, per q-head, [1,T] x [T,hd] = 2*T*hd
// GQA does not reduce this — the KV heads are broadcast across all Nq
// queries, so the score matrix is still Nq deep.
const double attn_dense = 4.0 * T * hd * Nq;
const double attn_causal = 0.5 * attn_dense; // is_causal skips the mask
// lm_head sits OUTSIDE the layer stack. The token embedding is a gather
// (~0 FLOPs) and is correctly uncounted even though it holds the tied
// weights.
const double head = 2.0 * d * V;
m.fwd_per_token_causal = L * (qkv + proj + ffn + attn_causal) + head;
m.fwd_per_token_dense = L * (qkv + proj + ffn + attn_dense) + head;
m.flops_per_token_causal = 3.0 * m.fwd_per_token_causal;
m.flops_per_token_dense = 3.0 * m.fwd_per_token_dense;
cudaDeviceProp prop{};
if (cudaGetDeviceProperties(&prop, device_ordinal) == cudaSuccess)
m.device_name = prop.name;
bool matched = false;
double tflops = peak_bf16_tflops(m.device_name, matched);
m.peak_is_guess = !matched;
// BLU_PEAK_TFLOPS overrides the table: PCIe/NVL parts, a clock-capped
// cluster, or an FP8 run (whose peak is 2x the BF16 figure).
if (const char* v = std::getenv("BLU_PEAK_TFLOPS")) {
const double ov = std::atof(v);
if (ov > 0.0) { tflops = ov; m.peak_is_guess = false; }
}
m.peak_flops_per_gpu = tflops * 1e12;
m.warmup_steps = getenv_int("BLU_MFU_WARMUP", 5);
return m;
}
double cluster_peak() const {
return peak_flops_per_gpu * static_cast<double>(world_size);
}
// Instantaneous MFU for one optimizer step. `tokens` is the GLOBAL token
// count for the step (all ranks, all micro-batches), which is why the
// denominator carries world_size: both sides scale together, so MFU is
// invariant to cluster size and the MFU decay of a scaling run IS its
// parallel-efficiency loss.
double mfu(double dt, int64_t tokens) const {
if (dt <= 0.0) return 0.0;
return (flops_per_token_causal * static_cast<double>(tokens)) / (dt * cluster_peak());
}
double mfu_palm(double dt, int64_t tokens) const {
if (dt <= 0.0) return 0.0;
return (flops_per_token_dense * static_cast<double>(tokens)) / (dt * cluster_peak());
}
// Achieved model TFLOP/s on ONE GPU — the absolute number to hold against
// the 989.4 ceiling when a ratio alone is ambiguous.
double tflops_per_gpu(double dt, int64_t tokens) const {
if (dt <= 0.0) return 0.0;
return (flops_per_token_causal * static_cast<double>(tokens))
/ (dt * static_cast<double>(world_size) * 1e12);
}
// Fold one step into the steady-state aggregate. The first `warmup_steps`
// are dropped: their dt is dominated by pipeline fill, the one-time ZeRO
// shard / autograd-cache allocations and cuBLAS autotuning, and they drag
// the reported figure well below what the run actually sustains.
//
// The window counts observe() CALLS, not the absolute step index: a run
// resumed from a checkpoint starts the loop at start_step > 0, and an
// absolute `step < warmup_steps` test would then skip the exclusion
// entirely while its pipeline is filling exactly like a fresh run's.
void observe(double dt, int64_t tokens) {
if (seen++ < warmup_steps || dt <= 0.0) return;
ss_tokens += static_cast<double>(tokens);
ss_secs += dt;
ss_steps += 1;
best_mfu = std::max(best_mfu, mfu(dt, tokens));
}
double steady_mfu() const {
if (ss_secs <= 0.0) return 0.0;
return (flops_per_token_causal * ss_tokens) / (ss_secs * cluster_peak());
}
double steady_tok_per_sec() const { return ss_secs > 0.0 ? ss_tokens / ss_secs : 0.0; }
void print_banner(int64_t num_params, int64_t tokens_per_step) const {
const double G = 1e9;
std::cout << "==================== MFU accounting ====================\n"
<< std::fixed << std::setprecision(3)
<< " device : " << device_name << " x " << world_size << "\n"
<< " peak BF16 (dense) : " << (peak_flops_per_gpu / 1e12) << " TFLOP/s per GPU"
<< (peak_is_guess ? " [!] name not in table - set BLU_PEAK_TFLOPS" : "") << "\n"
<< " cluster peak : " << (cluster_peak() / 1e15) << " PFLOP/s\n"
<< " fwd FLOPs/token : " << (fwd_per_token_causal / G) << " GFLOP (causal)\n"
<< " fwd+bwd /token : " << (flops_per_token_causal / G) << " GFLOP (causal, x3) <-- headline\n"
<< " fwd+bwd /token : " << (flops_per_token_dense / G) << " GFLOP (PaLM convention)\n"
<< " 6*N cross-check : " << (6.0 * static_cast<double>(num_params) / G)
<< " GFLOP (weight matmuls only; excludes the T-dependent attention core)\n"
<< " model FLOPs/step : "
<< (flops_per_token_causal * static_cast<double>(tokens_per_step) / 1e15)
<< " PFLOP over " << tokens_per_step << " tokens\n"
<< " 100% MFU would be : "
<< (flops_per_token_causal * static_cast<double>(tokens_per_step)
/ cluster_peak() * 1000.0) << " ms/step\n"
<< " note : no activation checkpointing in this loop, so HFU == MFU\n"
<< "========================================================" << std::endl;
}
};
// =============================================================================
// Shared advancing RNG for custom weight init (mirrors gpt2_fmha_ddp.cpp).
// A single std::mt19937 is advanced across every parameter draw so the init is
// reproducible from one seed. Values are drawn on CPU and copy_'d into the
// (device or host) parameter so the registered tensor identity is preserved.
// =============================================================================
struct InitRng {
std::mt19937 gen;
explicit InitRng(uint64_t seed) : gen(static_cast<uint32_t>(seed)) {}
Tensor normal(const Shape& shape, float std) {
Tensor t(shape, TensorOptions().with_dtype(Dtype::Float32));
std::normal_distribution<float> dist(0.0f, std);
float* d = t.data<float>();
const size_t n = t.numel();
for (size_t i = 0; i < n; ++i) d[i] = dist(gen);
return t;
}
};
// copy_ requires the destination's requires_grad OFF while copying a fresh leaf,
// so toggle it around the copy (matches gpt2_fmha_ddp.cpp::init_linear_gpt2).
static void copy_into_param(Tensor& dst, const Tensor& src) {
const bool was = dst.requires_grad();
if (was) dst.set_requires_grad(false);
dst.copy_(src.dtype() == dst.dtype() ? src : src.as_type(dst.dtype())); // cast fp32 draw -> bf16
if (was) dst.set_requires_grad(true);
}
static void init_linear(nn::Linear& layer, float std, InitRng& rng) {
copy_into_param(layer.weight, rng.normal(layer.weight.shape(), std));
if (layer.bias.is_valid()) { // our Llama-style Linears use bias=false, so this is a no-op
copy_into_param(layer.bias,
Tensor::zeros(layer.bias.shape(), TensorOptions().with_dtype(Dtype::Float32)));
}
}
static void init_embedding(nn::Embedding& emb, float std, InitRng& rng) {
copy_into_param(emb.weight, rng.normal(emb.weight.shape(), std));
}
// =============================================================================
// Transformer block: RMSNorm -> fused (QK-norm + RoPE + causal GQA) -> Wo -> +x
// =============================================================================
class CausalGQA : public nn::Module {
public:
nn::RMSNorm norm; // pre-attention RMSNorm over d_model
nn::Linear wQ, wK, wV; // projections (no bias, Llama style)
nn::Linear wO; // output projection
Tensor q_gamma, k_gamma;// [head_dim] fp32 QK-norm scales (learnable)
Tensor cos_sin_cache; // shared [context_length, head_dim] fp32 RoPE cache (no grad)
ModelConfig cfg;
CausalGQA(const ModelConfig& c, DeviceIndex dev, const Tensor& cache, InitRng& rng)
: norm(c.d_model, c.rms_eps, c.resid_dtype()),
wQ(c.d_model, c.q_heads * c.head_dim, /*bias=*/false, Dtype::Bfloat16),
wK(c.d_model, c.kv_heads * c.head_dim, /*bias=*/false, Dtype::Bfloat16),
wV(c.d_model, c.kv_heads * c.head_dim, /*bias=*/false, Dtype::Bfloat16),
wO(c.q_heads * c.head_dim, c.d_model, /*bias=*/false, Dtype::Bfloat16),
cos_sin_cache(cache),
cfg(c)
{
// GPT-2/Llama-style init: projections std=0.02, residual proj (Wo) scaled
// by (2*n_layers)^-0.5 to keep residual-stream variance ~constant.
const float resid_scale = 1.0f / std::sqrt(2.0f * static_cast<float>(c.n_layers));
init_linear(wQ, 0.02f, rng);
init_linear(wK, 0.02f, rng);
init_linear(wV, 0.02f, rng);
init_linear(wO, 0.02f * resid_scale, rng);
// QK-norm scales initialized to 1.0 (identity RMS scale). Pinned to fp32
// by the fused kernel, NOT by bf16_resid — see ModelConfig::qk_dtype().
auto gopts = TensorOptions().with_dtype(ModelConfig::qk_dtype()).with_device(dev);
q_gamma = Tensor::full(Shape{{c.head_dim}}, gopts, 1.0f);
k_gamma = Tensor::full(Shape{{c.head_dim}}, gopts, 1.0f);
q_gamma.set_requires_grad(true);
k_gamma.set_requires_grad(true);
norm.to(dev); wQ.to(dev); wK.to(dev); wV.to(dev); wO.to(dev);
register_module(norm);
register_module(wQ); register_module(wK); register_module(wV);
register_module(wO);
register_parameter(q_gamma);
register_parameter(k_gamma);
}
Tensor forward(const Tensor& x) override {
const int64_t B = x.shape().dims[0];
const int64_t T = x.shape().dims[1];
const int64_t hd = cfg.head_dim;
const int64_t Nq = cfg.q_heads;
const int64_t Nkv = cfg.kv_heads;
Tensor h1 = norm.forward(x); // resid dtype [B,T,d_model]
Tensor h = autograd::to_dtype(h1, Dtype::Bfloat16); // bf16 input for bf16 Linears
// (no-op when bf16_resid)
// bf16 projections. Only q and k enter the fp32 QK-norm + RoPE island, and
// that island is a kernel constraint, not a bf16_resid knob. v is NOT part
// of it: it goes straight to scaled_dot_product_attention, which consumes
// bf16, so widening it here and narrowing it again below was a pure
// round-trip (~201 MiB/layer/micro-step fwd+bwd). See F2 in CKPT_PERF_AUDIT.md.
//
// q/k are left bf16 here too and widened further down, AFTER the head-major
// layout change, so that copy moves 2 bytes/element instead of 4 (F3).
const Dtype qk = ModelConfig::qk_dtype();
Tensor q = wQ.forward(h); // bf16 [B,T,Nq*hd]
Tensor k = wK.forward(h); // bf16 [B,T,Nkv*hd]
Tensor v = wV.forward(h); // bf16 [B,T,Nkv*hd]
Tensor q_bf, k_bf;
#if USE_FUSED_RMS_ROPE
// Fused QK-norm + RoPE kernel (expects 4D head-major layout [B, H, T, hd], fp32).
//
// The kernel needs fp32; the reshape/transpose/contiguous copy that feeds it
// does not. Doing the copy in bf16 and widening once at the end hands the
// kernel bit-identical values (bf16 -> fp32 is exact) for half the bytes.
//
// Keep to_dtype AFTER contiguous(). Tensor::as_type materializes a strided
// source itself (the stride contract at the top of core/AsTypeTensor.cpp), so
// the reverse order is correct -- but it pays for a bf16 copy AND an fp32 one
// instead of just the bf16 copy, which is the opposite of the point here.
Tensor q_4d = autograd::to_dtype(autograd::contiguous(
autograd::transpose(autograd::reshape(q, Shape{{B, T, Nq, hd}}), 1, 2)), qk);
Tensor k_4d = autograd::to_dtype(autograd::contiguous(
autograd::transpose(autograd::reshape(k, Shape{{B, T, Nkv, hd}}), 1, 2)), qk);
auto [q_norm_rope, k_norm_rope] = autograd::fused_rms_rope(
q_4d, k_4d, q_gamma, k_gamma, cos_sin_cache,
/*pos_offset=*/0, cfg.qk_eps, /*interleaved=*/false);
q_bf = autograd::to_dtype(q_norm_rope, Dtype::Bfloat16);
k_bf = autograd::to_dtype(k_norm_rope, Dtype::Bfloat16);
#else
// Sequential QK-norm (RMSNorm) FIRST, then RoPE (both fp32). The widen has
// to happen here: rms_norm/rope are fp32-only. reshape is free either way,
// so this is bandwidth-neutral against casting before the reshape.
q = autograd::to_dtype(autograd::reshape(q, Shape{{B, T, Nq, hd}}), qk);
k = autograd::to_dtype(autograd::reshape(k, Shape{{B, T, Nkv, hd}}), qk);
q = autograd::rms_norm(q, q_gamma, static_cast<int>(hd), cfg.qk_eps);
k = autograd::rms_norm(k, k_gamma, static_cast<int>(hd), cfg.qk_eps);
q = autograd::rope(q, cos_sin_cache);
k = autograd::rope(k, cos_sin_cache);
q_bf = autograd::transpose(autograd::to_dtype(q, Dtype::Bfloat16), 1, 2);
k_bf = autograd::transpose(autograd::to_dtype(k, Dtype::Bfloat16), 1, 2);
#endif
// v never leaves bf16 (F2): wV already produced bf16 and SDPA consumes bf16.
Tensor v_bf = autograd::transpose(autograd::reshape(v, Shape{{B, T, Nkv, hd}}), 1, 2);
Tensor attn = autograd::scaled_dot_product_attention(
q_bf, k_bf, v_bf, static_cast<int>(Nq), static_cast<int>(Nkv),
/*is_causal=*/true, /*dropout_p=*/0.0f, autograd::SDPBackend::FlashAttention); // [B,Nq,T,hd] bf16
Tensor merged = autograd::reshape(autograd::transpose(attn, 1, 2),
Shape{{B, T, Nq * hd}}); // bf16
Tensor proj = autograd::to_dtype(wO.forward(merged), cfg.resid_dtype()); // bf16 matmul -> resid
return autograd::add(x, proj); // residual
}
};
// =============================================================================
// SwiGLU feed-forward: RMSNorm -> gate_up -> swiglu -> down -> +x
// gate_up: d_model -> 2*ffn_hidden, output laid out [gate | up];
// swiglu(z) = swish(gate) * up -> ffn_hidden.
// =============================================================================
class SwiGLUMLP : public nn::Module {
public:
nn::RMSNorm norm; // resid dtype (scale dtype must match its input)
nn::Linear gate_up; // d_model -> 2*ffn_hidden (bf16)
nn::Linear down; // ffn_hidden -> d_model (bf16)
Dtype resid_dt; // cached: the block does not otherwise keep cfg
SwiGLUMLP(const ModelConfig& c, DeviceIndex dev, InitRng& rng)
: norm(c.d_model, c.rms_eps, c.resid_dtype()),
gate_up(c.d_model, 2 * c.ffn_hidden, /*bias=*/false, Dtype::Bfloat16),
down(c.ffn_hidden, c.d_model, /*bias=*/false, Dtype::Bfloat16),
resid_dt(c.resid_dtype())
{
const float resid_scale = 1.0f / std::sqrt(2.0f * static_cast<float>(c.n_layers));
init_linear(gate_up, 0.02f, rng);
init_linear(down, 0.02f * resid_scale, rng);
norm.to(dev); gate_up.to(dev); down.to(dev);
register_module(norm);
register_module(gate_up);
register_module(down);
}
Tensor forward(const Tensor& x) override {
Tensor h = norm.forward(x); // resid dtype [B,T,d_model]
Tensor h_bf = autograd::to_dtype(h, Dtype::Bfloat16); // bf16 for gate_up
Tensor gu = gate_up.forward(h_bf); // bf16 [B,T,2*ffn]
Tensor act = autograd::swiglu(gu); // bf16 activation [B,T,ffn]
Tensor out = autograd::to_dtype(down.forward(act), resid_dt); // bf16 matmul -> resid
return autograd::add(x, out); // residual
}
};
// =============================================================================
// Full model
// =============================================================================
class BLU : public nn::Module {
public:
ModelConfig cfg;
private:
InitRng init_rng_; // declared before members it seeds; advanced across all draws
public:
nn::Embedding wte; // token embedding [vocab, d_model]
Tensor cos_sin_cache; // [context_length, head_dim] fp32 (shared, no grad)
std::vector<std::shared_ptr<CausalGQA>> attn;
std::vector<std::shared_ptr<SwiGLUMLP>> mlp;
nn::RMSNorm norm_f; // final RMSNorm
std::shared_ptr<nn::Linear> lm_head; // only allocated when weight_tying == false
BLU(const ModelConfig& c, DeviceIndex dev, uint64_t seed = 1234)
: cfg(c),
init_rng_(seed),
wte(static_cast<int>(c.vocab_size), static_cast<int>(c.d_model),
/*padding_idx=*/-1, c.embed_dtype()), // embed dtype
cos_sin_cache(autograd::build_rope_cache(c.context_length, c.head_dim,
static_cast<float>(c.rope_theta), dev)),
norm_f(c.d_model, c.rms_eps, c.resid_dtype()) // resid dtype
{
init_embedding(wte, 0.02f, init_rng_);
wte.to(dev);
norm_f.to(dev);
for (int i = 0; i < c.n_layers; ++i) {
auto a = std::make_shared<CausalGQA>(c, dev, cos_sin_cache, init_rng_);
auto m = std::make_shared<SwiGLUMLP>(c, dev, init_rng_);
attn.push_back(a);
mlp.push_back(m);
register_module(a.get());
register_module(m.get());
}
if (!c.weight_tying) {
lm_head = std::make_shared<nn::Linear>(static_cast<int>(c.d_model),
static_cast<int>(c.vocab_size),
/*bias=*/false, Dtype::Bfloat16); // bf16
init_linear(*lm_head, 0.02f, init_rng_);
lm_head->to(dev);
}
register_module(wte);
register_module(norm_f);
if (!c.weight_tying && lm_head) register_module(lm_head.get());
}
// indices [B, T] (UInt16) -> logits [B, T, vocab] (logit dtype)
Tensor forward(const Tensor& idx) override {
// Embedding is a gather, so its output dtype is wte's. Cast to the
// residual dtype: the two are independent flags, so bf16_embed without
// bf16_resid has to widen here (and vice versa).
Tensor x = autograd::to_dtype(wte.forward(idx), cfg.resid_dtype()); // [B,T,d_model]
// Blocks are invoked through Module::operator(), not forward(), so any
// registered module hook actually fires — a direct forward() call bypasses
// hook dispatch entirely. The ZeRO-1 staged parameter refresh registers a
// forward pre-hook on the block that begins each refresh stage (see
// register_forward_pre_hook in main). With no hooks registered operator()
// takes a fast path straight to forward(), so this costs nothing.
for (int i = 0; i < cfg.n_layers; ++i) {
x = (*attn[i])(x);
x = (*mlp[i])(x);
}
x = norm_f.forward(x); // resid dtype
// Output projection is a matmul -> always do it in bf16. Only the dtype the
// logits are HANDED BACK in is a flag: fp32 costs a [B,T,vocab] fp32 tensor
// (3.30 GB at B*T=16384, vocab=50304) purely to widen a bf16 product.
Tensor x_bf = autograd::to_dtype(x, Dtype::Bfloat16);
if (cfg.weight_tying) {
// logits = x @ wte.weight^T ; keep the autograd edge back to wte.weight.
//
// Cast BEFORE the transpose (F4(a) in CKPT_PERF_AUDIT.md): the strided
// copy is the expensive half and it does not need to run at 4
// bytes/element. wte is [vocab, d_model] = 294.8 MiB fp32, rebuilt every
// micro-step, so this drops ~30% of ~16.5 GiB/step. No-op when bf16_embed
// is on, in which case the weight is already bf16.
//
// contiguous() stays. Dropping it (the bare-transpose variant that used to
// sit here) hands matmul a strided B operand, and the forward GEMM wrapper
// has no trans_b entry point to absorb it -- that needs F4(c). It also used
// to feed to_dtype a transpose view, which as_type read in storage order
// and scrambled; that hole is closed at the source now (the stride contract
// in core/AsTypeTensor.cpp), but the GEMM side of it is not.
//
// Cast BEFORE transpose is safe for the BACKWARD too, but only because of
// that same fix: TransposeBackward returns a strided view, and with the
// cast now below the transpose in the graph, ToDtypeBackward is what
// receives that view. Before the stride contract this ordering would have
// silently permuted wte.weight's gradient under weight tying.
Tensor w_bf = autograd::to_dtype(wte.weight, Dtype::Bfloat16); // [vocab, d_model]
Tensor w_t_bf = autograd::contiguous(autograd::transpose(w_bf, 0, 1)); // [d_model, vocab]
Tensor logits = autograd::matmul(x_bf, w_t_bf); // bf16 [B,T,vocab]
return autograd::to_dtype(logits, cfg.logit_dtype());
}
return autograd::to_dtype(lm_head->forward(x_bf), cfg.logit_dtype());
}
};
// =============================================================================
// WSD learning-rate schedule: linear Warmup -> Stable (constant max_lr) ->
// cosine Decay to min_lr over the remaining steps.
// =============================================================================
static float get_lr(int step, float max_lr, float min_lr, int warmup, int max_steps, int stable_steps) {
if (step < warmup) {
return max_lr * static_cast<float>(step + 1) / static_cast<float>(warmup);
}
if ((step >= warmup) && (step < (stable_steps + warmup))) {
return max_lr;
}
float decay_ratio = static_cast<float>(step - (stable_steps + warmup)) /
static_cast<float>((max_steps - (stable_steps + warmup)) - 1);
assert(decay_ratio >= 0.0f && decay_ratio <= 1.0f);
// cosine decay
const float coeff = 0.5f * (1.0f + std::cos(M_PI * decay_ratio)); // 1 -> 0
return min_lr + coeff * (max_lr - min_lr);
// linear decay
//return min_lr + (max_lr-min_lr) * (1.0f - decay_ratio);
}
// =============================================================================
// Training
// =============================================================================
int main(int argc, char** argv) {
MPI_Init(&argc, &argv);
// Allocator size-class policy. MUST be set before the first allocation, so
// it goes here rather than with the other env parsing further down.
//
// BLU_SIZE_CLASS=exact turns on generic tight rounding. The "default"
// policy inflates every request in [1 MiB, 10 MiB) to a dedicated 20 MiB
// block, which on this config wastes 3,365 MiB on parameters plus FP32
// gradients and a further 2,016 MiB on ShardedMuon's 72 per-tensor FP32
// state buffers. Measured, see docs/ZERO1_MEMORY_AUDIT.md.
//
// Left OFF by default: tight classes cut internal fragmentation but give up
// block reuse across sizes, so judge the A/B on reserved_peak rather than on
// active_current. Any mode name from SizeClassConfigDictionary works here.
if (const char* sc = std::getenv("BLU_SIZE_CLASS"); sc && *sc) {
OwnTensor::SizeClass::set_mode(sc);
}
int rank = 0, world_size = 1;
MPI_Comm_rank(MPI_COMM_WORLD, &rank);
MPI_Comm_size(MPI_COMM_WORLD, &world_size);
const bool is_master = (rank == 0);
// One GPU per rank.
cudaSetDevice(rank);
DeviceIndex device(Device::CUDA, rank);
ModelConfig cfg;
// OwnTensor::AllocationTracker::instance().init("4b_tracker.csv");
// ---- activation / parameter precision ----------------------------------
// BLU_ALL_BF16=1 turns on all three stages at once. Each stage then has its
// own knob that overrides the master, so a regression bisects without
// editing anything:
//
// BLU_ALL_BF16=1 all three (residual, logits, embed)
// BLU_ALL_BF16=1 BLU_BF16_EMBED=0 residual + logits only
// BLU_BF16_LOGITS=1 logits only
// (unset) shipped mixed precision, unchanged
//
// The default is bit-identical to the pre-flag code: every to_dtype() this
// touches resolves to Float32, and autograd::to_dtype returns its input
// unchanged when the dtypes already match, so no cast node is even built.
//
// Optimizer masters/moments stay fp32 under all of these — that is a kernel
// constraint, not an omission. See the ModelConfig precision block.
{
const int all = getenv_int("BLU_ALL_BF16", 0);
cfg.bf16_resid = getenv_int("BLU_BF16_RESID", all) != 0;
cfg.bf16_logits = getenv_int("BLU_BF16_LOGITS", all) != 0;
cfg.bf16_embed = getenv_int("BLU_BF16_EMBED", all) != 0;
// bf16 logits are now consumable: LossKernels.cu's vectorized forward and
// both backward variants size every 16-byte chunk as 16/sizeof(T) elements
// (VEC_ELEMS/Chunk16), so the cp.async source is 16-byte aligned on a 2-byte
// dtype and the 128-bit gradient store no longer overruns the row. This is
// the single largest memory item in the model: [B,T,vocab] fp32 is 3.07 GiB
// retained for the CE backward plus a 3.07 GiB fp32 gradient, and the matmul
// already produces bf16 — the widen existed only to feed the loss kernels.
//
// It is NOT free, and the cost is not numerical in the gradients: the loss
// SCALAR comes back bf16 too (LossOps keys it off the logits dtype), which
// quantizes the reported loss to ~0.06 at this batch size. See the
// ModelConfig precision block. So warn every time rather than staying quiet,
// because the affected number is the one an A/B is read from.
if (cfg.bf16_logits && is_master) {
std::cout << "[precision] bf16 logits ON — frees ~3 GiB ([B,T,vocab] fp32 "
"activation + gradient). Gradients are unaffected (fp32 softmax "
"stats, exact 1.0 backward seed), but the REPORTED loss is bf16 "
"and quantizes to ~0.06 steps: compare tokens/s and grad-norm "
"across an A/B, not the loss digits." << std::endl;
}
}
// BLU_MAX_STEPS caps the run for benchmarking (the A/B harness in
// run_zero_ab_8xh200.sh uses it). The WSD schedule divides by
// (max_steps - warmup - stable - 1) and asserts the range below, so shrink
// warmup/stable proportionally rather than leaving them at their full-run
// values — otherwise a short benchmark aborts on "no decay tail".
if (const int ms = getenv_int("BLU_MAX_STEPS", 0); ms > 0) {
cfg.max_steps = ms;
cfg.warmup_steps = std::max(1, ms / 10);
cfg.stable_steps = std::max(1, ms / 2);
if (cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1) {
cfg.stable_steps = std::max(1, cfg.max_steps - cfg.warmup_steps - 2);
}
if (is_master) {
std::cout << "[bench] BLU_MAX_STEPS=" << ms
<< " (warmup=" << cfg.warmup_steps
<< " stable=" << cfg.stable_steps << ")" << std::endl;
}
}
// Tokens consumed per optimizer step = B * T * grad_accum * world_size.
const int64_t tokens_per_micro = cfg.B * cfg.T * static_cast<int64_t>(world_size);
if (cfg.global_batch % tokens_per_micro != 0) {
if (is_master) {
std::cerr << "ERROR: global_batch (" << cfg.global_batch
<< ") not divisible by B*T*world_size (" << tokens_per_micro << ")\n";
}
MPI_Abort(MPI_COMM_WORLD, 1);
}
const int GRAD_ACCUM = static_cast<int>(cfg.global_batch / tokens_per_micro);
if (is_master) {
std::cout << "================ Configuration ================\n"
<< " d_model : " << cfg.d_model << "\n"
<< " n_layers : " << cfg.n_layers << "\n"
<< " q_heads/kv_heads: " << cfg.q_heads << " / " << cfg.kv_heads
<< " (G=" << (cfg.q_heads / cfg.kv_heads) << ")\n"
<< " head_dim : " << cfg.head_dim << "\n"
<< " ffn_hidden : " << cfg.ffn_hidden << " (SwiGLU)\n"
<< " context_length : " << cfg.context_length << "\n"
<< " rope_theta : " << cfg.rope_theta
<< " (interleaved=" << (cfg.rope_interleaved ? "GPT-J" : "NeoX") << ")\n"
<< " vocab_size : " << cfg.vocab_size << "\n"
<< " weight_tying : " << (cfg.weight_tying ? "ON" : "OFF") << "\n"
<< " precision : resid "
<< OwnTensor::get_dtype_name(cfg.resid_dtype())
<< " | logits " << OwnTensor::get_dtype_name(cfg.logit_dtype())
<< " | embed " << OwnTensor::get_dtype_name(cfg.embed_dtype())
<< "\n"
<< " (matmuls bf16, QK-norm+RoPE fp32, "
"optimizer masters fp32 — all fixed)\n"
<< " B x T : " << cfg.B << " x " << cfg.T << "\n"
<< " world_size : " << world_size << "\n"
<< " global_batch : " << cfg.global_batch
<< " (grad_accum=" << GRAD_ACCUM << ")\n"
<< " max_lr / min_lr : " << cfg.max_lr << " / " << cfg.min_lr << "\n"
<< " warmup / steps : " << cfg.warmup_steps << " / " << cfg.max_steps << "\n"
<< " WSD stable_steps: " << cfg.stable_steps
<< " (decay tail=" << (cfg.max_steps - cfg.warmup_steps - cfg.stable_steps) << ")\n"
<< "===============================================" << std::endl;
}
// The WSD schedule needs a non-empty decay tail; get_lr()'s decay_ratio
// divides by (max_steps - warmup - stable - 1) and asserts on the range.
if (cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1) {
if (is_master) {
std::cerr << "ERROR: warmup_steps + stable_steps ("
<< (cfg.warmup_steps + cfg.stable_steps)
<< ") leaves no decay tail before max_steps (" << cfg.max_steps << ")\n";
}
MPI_Abort(MPI_COMM_WORLD, 1);
}
// ---- process group — only when multi-GPU. Used by ZeRO-1's
// reduce-scatter/all-gather (there is no separate DDP process group). ----
// BLU_ZERO_V2=1 selects the reduce-scatter/all-gather optimizer
// (owt::zero::SymmetricShardedOptimizer) instead of v1's reduce/broadcast.
// BLU_SYMMETRIC=1 additionally backs its wire buffers with NCCL symmetric
// memory so the collectives can dispatch to the NVLS (NVLink SHARP) kernels.
//
// The memory mode is fixed at process-group creation, so it must be decided
// HERE, before init_process_group.
//
// NOTE: symmetric memory is only useful with v2. NCCL 2.28 ships symmetric
// device kernels for AllGather / AllReduce / ReduceScatter ONLY — there is no
// symmetric ncclReduce and no symmetric ncclBroadcast, which are exactly the
// two collectives v1 uses. Requesting it for v1 registers windows that
// nothing can dispatch against, so we say so rather than let a benchmark
// silently measure nothing.
const bool use_zero_v2 = getenv_int("BLU_ZERO_V2", 0) != 0;
const bool want_symmetric = getenv_int("BLU_SYMMETRIC", 0) != 0;
// Read here rather than at the optimizer block below: the memory mode is
// fixed at process-group creation, so whether the Muon group can use
// symmetric memory has to be known BEFORE init_process_group.
const bool use_muon = [] {
const char* v = std::getenv("BLU_OPTIMIZER");
return v && (std::string(v) == "muon" || std::string(v) == "MUON");
}();
// ShardedMuon's default reduce-scatter/all-gather path dispatches to the same
// symmetric kernels v2 does, so Muon benefits from symmetric memory
// independently of BLU_ZERO_V2 (which only selects the AdamW remainder's
// optimizer). Its BLU_MUON_RB=1 fallback does not — that path is
// Reduce/Broadcast, which has no symmetric kernel.
const bool muon_rb = getenv_int("BLU_MUON_RB", 0) != 0;
const bool symmetric_useful = (use_zero_v2 || (use_muon && !muon_rb));
if (is_master && want_symmetric && !symmetric_useful) {
std::cout << "[ZeRO-1] WARNING: BLU_SYMMETRIC=1 has NO effect in this "
"configuration. NCCL has no symmetric kernel for ncclReduce "
"or ncclBroadcast (the collectives used by ZeRO v1 and by "
"BLU_MUON_RB=1); only AllGather/AllReduce/ReduceScatter have "
"one. Set BLU_ZERO_V2=1, or use Muon without BLU_MUON_RB, to "
"exercise NVLS." << std::endl;
}
std::shared_ptr<ProcessGroupNCCL> pg;
if (world_size > 1) {
pg = init_process_group(world_size, rank, /*stream=*/0, /*group_name=*/"",
(want_symmetric && symmetric_useful)
? dist::comm::MemoryMode::Symmetric
: dist::comm::MemoryMode::Standard);
}
if (is_master) std::cout << "Initializing model on CUDA device " << rank << "..." << std::endl;
BLU model(cfg, device, /*seed=*/1234 + rank * 0 /*same seed all ranks; DDP also broadcasts*/);
// Parameter accounting.
auto params = model.parameters();
int64_t num_params = 0;
for (auto& p : params) num_params += p.numel();
if (is_master) std::cout << "Number of parameters: " << num_params << std::endl;
// ---- MFU meter ----
// Built here rather than at first use so its banner lands next to the
// parameter count it cross-checks itself against (6*N), and so a bogus
// BLU_PEAK_TFLOPS or an unrecognized GPU is reported BEFORE the run instead
// of after 2000 steps of unusable numbers. `device_ordinal=rank` because
// this process set cudaSetDevice(rank) above — one GPU per rank.
MFUMeter mfu_meter = MFUMeter::build(cfg, world_size, rank);
if (is_master) mfu_meter.print_banner(num_params, cfg.global_batch);
// ---- optimizer: ZeRO-1 (optimizer-state sharding), mixed precision.
// No DDP wrapper: ZeRO-1's reduce-scatter *replaces* DDP's gradient
// all-reduce (running both would double-reduce the gradients), so
// model.forward() is called directly below with no DataParallel wrapping.
//
// ShardedOptimizer requires every parameter it manages to share one dtype
// (the fused Adam kernel + the reduce-scatter/all-gather wire format are
// single-dtype). This model is deliberately mixed: the Linear projections
// are BF16 (the matmuls run in bf16), while the token embedding, the
// RMSNorm scales and the QK-norm scales (q_gamma/k_gamma) stay FP32 for
// numerical stability. Group params by dtype and give each group its own
// ShardedOptimizer instance; in practice that's exactly two groups.
// Calling their step()s back-to-back (never concurrently) is safe: each
// instance's local_adam_step() stream-syncs after every fused-Adam launch,
// so the multi-tensor kernel's shared launch-metadata buffer is never
// clobbered by the next instance's launch.
//
// BLU_OPTIMIZER=muon additionally splits off the Muon set BEFORE the dtype
// grouping — see the ShardedMuon block below for why that set needs a
// different sharding scheme entirely. (use_muon is resolved up at the
// process-group block: the memory mode depends on it and is fixed there.)
// ---- which parameters Muon may orthogonalize -----------------------------
// The recipe is: Muon for the transformer body's 2-D weight matrices, AdamW
// for everything else. "Everything else" is every 1-D parameter (the RMSNorm
// scales and the QK-norm q_gamma/k_gamma — a vector has no second axis to
// orthogonalize against) plus the two vocabulary matrices, the token
// embedding and the output head. Those two are excluded because their rows
// are near-independent per-token directions rather than a map between two
// dense feature spaces, which is the structure that makes spectral
// normalization the right prior in the first place.
//
// Under weight tying lm_head IS wte, so the single exclusion covers both.
std::set<const void*> muon_excluded;
for (auto& p : model.wte.parameters()) muon_excluded.insert(p.data());
if (model.lm_head)
for (auto& p : model.lm_head->parameters()) muon_excluded.insert(p.data());
auto is_muon_param = [&](const Tensor& p) {
return use_muon && p.ndim() >= 2 && !muon_excluded.count(p.data());
};
// Muon-eligible params grouped by dtype, and the AdamW remainder grouped by
// dtype. Both optimizers are single-dtype (their fused kernels and wire
// formats are), so the dtype split happens inside each. At the default
// precision the partition is exactly: {bf16 body matrices} -> Muon,
// {fp32 wte + all the 1-D scales} -> AdamW.
//
// The precision flags move parameters BETWEEN these dtype maps, they do not
// change Muon eligibility (muon_excluded is keyed on pointer identity, so
// wte stays out of Muon whatever dtype it is). Under BLU_ALL_BF16=1 the
// AdamW side splits into a bf16 group (wte + the block RMSNorm scales) and a
// residual fp32 group holding only q_gamma/k_gamma — 2*head_dim*n_layers =
// 6144 elements, which the fused QK-norm+RoPE kernel pins to fp32. So "all
// bf16" still leaves two AdamW dtype groups, and the one-group /
// one-all-gather / one-clip-norm simplification does NOT materialize.
std::map<Dtype, std::vector<Tensor>> muon_by_dtype;
std::map<Dtype, std::vector<Tensor>> params_by_dtype;
for (auto& p : params) {
if (is_muon_param(p)) muon_by_dtype[p.dtype()].push_back(p);
else params_by_dtype[p.dtype()].push_back(p);
}
// Keeping the embedding group on ELEMENTWISE ZeRO-1 is deliberate, not an
// oversight. ShardedMuon owns whole tensors, and a group whose mass is one
// 103M-element embedding plus ~145 tensors of a few thousand elements cannot
// be balanced that way: one rank would hold 99.8% of the state. Elementwise
// sharding splits that embedding evenly, and AdamW — being elementwise — is
// free to be sharded that way. Muon is not. That asymmetry is the whole
// reason the two sets take different optimizers.
ZeroConfig zcfg;
zcfg.variant = AdamVariant::ADAMW;
zcfg.lr = cfg.max_lr;
zcfg.beta1 = 0.9f;
zcfg.beta2 = 0.95f;
zcfg.eps = 1e-8f;
zcfg.weight_decay = cfg.weight_decay;
zcfg.exclude_1d_from_weight_decay = true; // GPT/Llama recipe: no wd on 1-D params
zcfg.max_grad_norm = cfg.grad_clip;
// All ranks init weights from the same mt19937 seed (see GPT ctor, seed=1234
// on every rank) via identical CPU draws, so the model is already replica-
// consistent before the optimizer is built — the init broadcast would only
// copy data that's already equal. Disable it (verify_consistency stays on as
// the cheap startup layout check).
zcfg.init_broadcast = false;
// Overlap the gradient reduce with the final micro-step's backward: grads are
// bucketed and reduced-to-owner from autograd hooks as each bucket fills,
// instead of one serial reduce-scatter after backward. Requires arming the
// optimizer before the last backward (see arm_grad_sync below) and the
// SEQUENTIAL autograd engine (the default). No-op at world_size==1.
//
// BLU_NO_OVERLAP=1 forces BOTH overlap features off, dropping the optimizer
// onto its reference path: one monolithic gather + reduce_scatter after
// backward, then a blocking all-gather + full-model scatter at the end of
// step(). Slower, but it is the path the ZeRO-1 design is validated against
// (bit-identical to single-process AdamW at W=1), so it is the bisection that
// separates "the overlap machinery corrupts gradients or publishes bad
// weights" from "the Adam/gradient math is wrong independently of overlap".
// BISECTION (docs/ZERO1_OVERLAP_AUDIT.md). BLU_NO_OVERLAP conflated two
// independent features, so a bad loss curve could not be attributed to either.
// They are now separately switchable, and the STAGED REFRESH IS OFF BY DEFAULT:
// this build keeps the hook-driven gradient reduce and publishes parameters
// synchronously.
//
// BLU_NO_OVERLAP=1 both off — the validated reference path
// BLU_NO_OVERLAP_BWD=1 suppress (A) only
// BLU_OVERLAP_REFRESH=1 restore (B) — the previous default
//
// Why (B) is the one dropped first: (A) is the untested half. Every test in
// dist/zero/tests sets overlap_backward_comm=false and none of them calls
// arm_grad_sync(), whereas the staged refresh has direct coverage in
// test_zero_v2_staged.cpp and test_muon_rsag.cpp. Isolating the covered half
// first is backwards; this leaves the UNCOVERED half running so a surviving
// bad curve indicts it directly.
const bool no_overlap = getenv_int("BLU_NO_OVERLAP", 0) != 0;
zcfg.overlap_backward_comm = (world_size > 1) && !no_overlap
&& getenv_int("BLU_NO_OVERLAP_BWD", 0) == 0;
// (B) Overlap the post-Adam parameter refresh with the NEXT forward pass:
// instead of one blocking all-gather + full-model copy at the end of step(),
// the updated parameters are broadcast in consumption-ordered stages and
// installed just before the layers that read them (set_param_refresh_stages
// below). OFF unless explicitly requested — see the bisection note above.
//
// CAVEAT for reading the result: turning this off does NOT merely remove the
// staging. The synchronous publish it falls back to ends in a full
// cudaDeviceSynchronize() of its own (ShardedMuon::all_gather_params_sync,
// SymmetricShardedOptimizer::all_gather_params), so the host/device coupling
// that the overlapped path removes comes back with it. A curve that recovers
// here therefore means "(B) OR a race elsewhere that the sync masks", not
// "(B)" alone. The line-1916 sync below becomes redundant, not load-bearing.
zcfg.overlap_param_refresh = (world_size > 1) && !no_overlap
&& getenv_int("BLU_OVERLAP_REFRESH", 0) != 0;
// BLU_COMM_BUCKET_MB sets the target gradient-bucket size for the AdamW groups
// (ZeroConfig::comm_bucket_bytes; Muon uses BLU_MUON_BUCKETS instead, which is a
// bucket COUNT rather than a byte budget).
//
// This is the knob that trades overlap granularity against SYNCHRONISATION COUNT,
// and the nsys trace says the second term dominates here. A bucket's
// reduce-scatter is launched as soon as THIS rank has staged the bucket's
// gradients, but the NCCL device kernel cannot retire until every rank has
// launched its matching collective — so it sits on the GPU spinning, holding SMs,
// until the slowest peer arrives. Measured on 8xH100 with the default 50 MB
// (62 buckets on the 1.62 B bf16 group): the ReduceScatter kernels occupy 20x
// more GPU time with overlap on than off, for identical bytes, and 41% of that
// is exposed rather than hidden behind compute. The AllGather, whose launch
// order is rank-independent, is unaffected (1.08x) — which is what identifies
// rank skew rather than bandwidth as the cost.
//
// Larger buckets -> fewer peer-arrival waits per step (200 MB gives ~16 buckets
// instead of 62), but a bucket cannot fire until ALL of its
// parameters have landed, so the overlap window shrinks.
// Smaller buckets -> finer overlap, more spin points.
//
// 0 or unset keeps ZeroConfig's 50 MB default.
if (const int mb = getenv_int("BLU_COMM_BUCKET_MB", 0); mb > 0) {
zcfg.comm_bucket_bytes = static_cast<size_t>(mb) * 1024ull * 1024ull;
if (is_master) {
std::cout << "[ZeRO-1] BLU_COMM_BUCKET_MB=" << mb
<< ": gradient bucket target " << mb << " MB (default 50)" << std::endl;
}
}
// BLU_FP32_REDUCE=1 reduces gradients in FP32 instead of the model's compute
// dtype. The bf16 group's gradients are otherwise reduce-scattered in bf16 on
// BOTH the overlap and the reference path, so disabling overlap did not test
// it. bf16 carries 8 mantissa bits and its rounding error scales with the
// largest partial sum rather than the result, so a near-cancelling reduction
// can emerge with a plausible norm and a badly wrong direction.
//
// BLU_BF16_REDUCE=1 is the mirror knob, for A/B testing the other direction:
// it forces BF16 on the wire even for a group whose parameters are FP32 —
// on this model that is the AdamW group (wte + all 1-D scales), which
// SAME_AS_PARAM otherwise reduces in FP32 no matter how BLU_FP32_REDUCE is
// set. It exists so "both optimizers reduce in the same precision" is
// reachable from either end, not because it is a good idea:
//
// - It does NOT save memory. flat_grad_ stops aliasing flat_stage_ and
// becomes its own [padded] BF16 buffer, so on this model's FP32 group it
// costs +147.6 MiB/GPU net (+196.8 flat_grad_, -49.2 shard_grad_) — the
// v2 banner prints the exact figure at startup. For scale, the same alias
// break in the other direction (BLU_FP32_REDUCE on the BF16 Muon group)
// costs +7020 MiB/GPU, which is what the nsys buffer fingerprints show.
// - It saves almost no time. That group's whole reduce is ~1.2 ms of a
// ~3.7 s step, and its buckets already sit near the latency floor
// (137 us min vs 154 us mean), so halving the bytes does not halve the
// time. Measured headroom is ~0.3 ms/step.
// - It costs precision exactly where the model chose not to spend it. wte
// is one of the two tensors deliberately kept in FP32, and under weight
// tying its gradient carries the lm_head projection as well; the RMSNorm
// scale gradients are accumulated over B*T in FP32 by reduce_dgamma
// before being handed over. Both would be truncated to 8 mantissa bits
// on the wire, and op=SUM means that happens at 4x the average's
// magnitude.
//
// Keep it for bisection, not for production. The two are mutually exclusive
// because they name the same field.
const bool want_fp32_reduce = getenv_int("BLU_FP32_REDUCE", 0) != 0;
const bool want_bf16_reduce = getenv_int("BLU_BF16_REDUCE", 0) != 0;
if (want_fp32_reduce && want_bf16_reduce) {
if (is_master)
std::cerr << "[ZeRO-1] BLU_FP32_REDUCE=1 and BLU_BF16_REDUCE=1 both set; "
"they set the same field to opposite values." << std::endl;
MPI_Abort(MPI_COMM_WORLD, 1);
}
if (want_fp32_reduce) {
zcfg.reduce_dtype = owt::zero::ReduceDtype::FP32;
if (is_master)
std::cout << "[ZeRO-1] BLU_FP32_REDUCE=1: reducing gradients in FP32" << std::endl;
} else if (want_bf16_reduce) {
zcfg.reduce_dtype = owt::zero::ReduceDtype::BF16;
if (is_master)
std::cout << "[ZeRO-1] BLU_BF16_REDUCE=1: reducing gradients in BF16 "
"(diagnostic — costs memory on the FP32 group, saves ~0.3 ms/step)"
<< std::endl;
}
if (is_master && getenv_int("BLU_GLOBAL_CLIP", 0) != 0)
std::cout << "[ZeRO-1] BLU_GLOBAL_CLIP=1: clipping on the norm over ALL "
"parameter groups (DeepSpeed semantics) instead of per group"
<< std::endl;
zcfg.verbose = is_master;
// Always printed, not only on the reference path: which halves of the overlap
// are live is now a per-run choice, and every log has to record it or the A/B
// is unreadable after the fact.
if (is_master) {
std::cout << "[ZeRO-1] overlap: reduce="
<< (zcfg.overlap_backward_comm ? "ON (hook-driven, during backward)"
: "off (monolithic, after backward)")
<< " | refresh="
<< (zcfg.overlap_param_refresh ? "ON (staged into next forward)"
: "off (synchronous publish in step)");
if (no_overlap) std::cout << " [BLU_NO_OVERLAP=1]";
std::cout << std::endl;
}
if (is_master) {
std::cout << "[ZeRO-1] variant: " << (use_zero_v2 ? "v2 reduce-scatter/all-gather"
: "v1 reduce/broadcast")
<< " | symmetric memory: " << ((want_symmetric && use_zero_v2) ? "requested"
: "off")
<< " (BLU_ZERO_V2=" << (use_zero_v2 ? 1 : 0)
<< " BLU_SYMMETRIC=" << (want_symmetric ? 1 : 0) << ")" << std::endl;
}
// ---- ZeRO-1 for Muon --------------------------------------------------
// Muon's update is NOT elementwise, so it cannot ride the flat contiguous
// shard the AdamW groups use. Its core step is
//
// U = NewtonSchulz(M)
//
// an iterated matrix polynomial (X <- aX + (bXX' + c(XX')^2)X) in which every
// output entry depends on the WHOLE matrix. ShardedOptimizer gives rank r the
// flat range [r*S,(r+1)*S) of every parameter concatenated end to end, which
// lands mid-matrix: a rank ends up holding the last 300 rows of wQ, all of wK
// and the first 900 rows of wV. Orthogonalizing those fragments separately is
// a DIFFERENT operator, not an approximation of the right one — so reusing
// the elementwise shard here would not be a slower or less accurate Muon, it
// would not be Muon at all.
//
// ShardedMuon therefore shards by WHOLE PARAMETER: each matrix is assigned in
// full to one owner rank, which holds its FP32 master + momentum and runs the
// complete Newton-Schulz on it. Gradients reduce-to-owner, the owner updates,
// and the result broadcasts back — so the state memory still drops by ~1/W,
// which is what ZeRO-1 actually promises, while the math stays identical to
// single-process Muon.
//
// It comes out CHEAPER than sharded AdamW on the same parameters: Muon keeps
// no second moment for an orthogonalized param (8 B/elem of persistent state
// against AdamW's 12), and whole-tensor ownership needs no flat staging
// buffer at all, because a broadcast can write straight into the destination
// tensor. That removes the unsharded [padded_numel] compute-dtype buffer that
// dominates ShardedOptimizer's footprint on this model's BF16 group.
// BLU_MUON_ADJUST_LR=original reproduces published modded-nanogpt numbers.
// Note there is NO single lr that converts between the conventions: the ratio
// is 0.2*sqrt(cols), which on this model spans 4.5x (wK/wV) to 21.2x
// (gate_up). Switching the convention means re-sweeping muon_lr, so it is
// reported at startup rather than left implicit.
owt::zero::MuonAdjustLr muon_adjust = owt::zero::MuonAdjustLr::MATCH_RMS_ADAMW;
if (const char* a = std::getenv("BLU_MUON_ADJUST_LR");
a && std::string(a) == "original") {
muon_adjust = owt::zero::MuonAdjustLr::ORIGINAL;
}
// Which collectives carry the Muon group. The default is reduce-scatter/
// all-gather over a bin-packed flat buffer: whole tensors are packed into
// per-rank contiguous sub-chunks so the shard boundary lands on tensor
// boundaries, which keeps Newton-Schulz seeing whole matrices while dropping
// ~2 collectives/parameter to ~2 collectives/bucket AND making the traffic
// NVLS-eligible (NCCL ships symmetric kernels for ReduceScatter/AllGather but
// none for Reduce or Broadcast, so the old path could never reach NVLink
// SHARP however BLU_SYMMETRIC was set).
//
// BLU_MUON_RB=1 restores the original one-ncclReduce + one-ncclBroadcast per
// parameter. Keep it for bisection: it is the formulation this optimizer was
// first validated against, so a disagreement between the two isolates a bug
// in the bin-packed layout from a fault in the Muon math itself. It also
// needs no flat buffer, which matters if the ~3.3 GB the default costs on
// this model's BF16 group is the difference between fitting and not.
MuonConfig mcfg;
mcfg.collective = muon_rb
? owt::zero::MuonCollective::REDUCE_BROADCAST
: owt::zero::MuonCollective::REDUCE_SCATTER_ALL_GATHER;
mcfg.target_buckets = getenv_int("BLU_MUON_BUCKETS", 6);
mcfg.lr = cfg.muon_lr;
mcfg.momentum = cfg.muon_momentum;
mcfg.weight_decay = cfg.muon_weight_decay;
mcfg.ns_steps = cfg.muon_ns_steps;
mcfg.nesterov = true;
mcfg.adjust_lr = muon_adjust;
// The aux AdamW path inside ShardedMuon. Setting adamw_lr to the AdamW peak
// and lr to the Muon peak fixes the ratio set_lr() preserves, so feeding the
// Muon-scaled schedule value below drives BOTH paths correctly from one
// schedule. (On this model the Muon group is 100% matrices, so the aux path
// is unused — it exists so an untied lm_head or a stray 1-D parameter landing
// in a Muon dtype group is handled correctly rather than orthogonalized.)
mcfg.adamw_lr = cfg.max_lr;
mcfg.adamw_beta1 = zcfg.beta1;
mcfg.adamw_beta2 = zcfg.beta2;
mcfg.adamw_eps = zcfg.eps;
mcfg.adamw_weight_decay = cfg.weight_decay;
mcfg.exclude_1d_from_weight_decay = true;
mcfg.max_grad_norm = cfg.grad_clip;
mcfg.reduce_dtype = zcfg.reduce_dtype;
mcfg.init_broadcast = zcfg.init_broadcast;
mcfg.verify_consistency = zcfg.verify_consistency;
mcfg.overlap_backward_comm = zcfg.overlap_backward_comm;
mcfg.overlap_param_refresh = zcfg.overlap_param_refresh;
mcfg.verbose = is_master;
// Held as the shared base so v1, v2 and Muon are interchangeable. Every call
// site below (arm_grad_sync, set_param_refresh_stages, refresh_param_stage,
// drain_param_refresh, last_grad_norm, zero_grad) is virtual on
// ZeroOptimizerBase, and the base still derives from OwnTensor::nn::Optimizer
// so CheckpointManager::save/load_latest keep working unchanged.
std::vector<std::unique_ptr<owt::zero::ZeroOptimizerBase>> optimizers;
std::vector<std::string> group_label;
std::vector<std::vector<Tensor>> group_params; // parallel to `optimizers`
// Each group's own (peak, floor) LR pair, fed through the ONE shared WSD
// curve. This mirrors train_muon.py, which calls get_lr twice — once with
// (max_lr, min_lr) for the AdamW group and once with (muon_lr, muon_min_lr)
// for the Muon group — rather than scaling one curve by a constant. The two
// are equivalent only while both floors are the same FRACTION of their peak;
// carrying the pair keeps Muon's floor independently settable.
std::vector<std::pair<float,float>> group_lr;
// Muon groups first, then the AdamW remainder. Every rank builds this list in
// the same order from the same deterministic partition, which is what keeps
// the per-group NCCL collectives paired across the communicator.
for (auto& kv : muon_by_dtype) {
optimizers.push_back(std::make_unique<ShardedMuon>(kv.second, pg, mcfg));
group_params.push_back(kv.second);
// Same curve, Muon's own peak and floor: same warmup length, same stable
// window, same cosine tail.
group_lr.emplace_back(cfg.muon_lr, cfg.muon_min_lr);
int64_t n = 0;
for (auto& p : kv.second) n += p.numel();
group_label.push_back("muon:" + OwnTensor::get_dtype_name(kv.first) + "/" +
std::to_string(n / 1000000) + "M");
}
for (auto& kv : params_by_dtype) {
if (use_zero_v2) {
optimizers.push_back(std::make_unique<SymmetricShardedOptimizer>(kv.second, pg, zcfg));
} else {
optimizers.push_back(std::make_unique<ShardedOptimizer>(kv.second, pg, zcfg));
}
group_params.push_back(kv.second);
group_lr.emplace_back(cfg.max_lr, cfg.min_lr);
// Label each group for the per-step gradient-norm report. The combined
// sqrt(sum of squares) that used to be the only number reported cannot
// distinguish "the 103M-param fp32 group (tied wte + all the RMSNorm/
// QK-norm scales) blew up" from "the 1.62B-param bf16 group (every
// attention/MLP weight) blew up" — and those point at completely different
// subsystems. G0/G1 here are indices into `optimizers`, NOT step_order.
int64_t n = 0;
for (auto& p : kv.second) n += p.numel();
group_label.push_back("adamw:" + OwnTensor::get_dtype_name(kv.first) + "/" +
std::to_string(n / 1000000) + "M");
}
if (is_master) {
std::cout << "[optimizer] " << (use_muon ? "MUON + AdamW" : "AdamW")
<< " (BLU_OPTIMIZER=" << (use_muon ? "muon" : "adamw") << ")";
if (use_muon) {
std::cout << " | muon_lr " << cfg.muon_lr
<< " / " << cfg.muon_min_lr
<< " momentum " << cfg.muon_momentum
<< " wd " << cfg.muon_weight_decay
<< " ns_steps " << cfg.muon_ns_steps
<< " adjust_lr "
<< (muon_adjust == owt::zero::MuonAdjustLr::MATCH_RMS_ADAMW
? "match_rms_adamw" : "original");
}
std::cout << std::endl;
for (size_t gi = 0; gi < optimizers.size(); ++gi) {
std::cout << " G" << gi << " " << group_label[gi]
<< " : " << optimizers[gi]->variant_name()
<< " | lr " << group_lr[gi].first
<< " -> " << group_lr[gi].second << std::endl;
}
}
// ---- what the gradient clip actually covers ------------------------------
// The clip coefficient is computed from the norm over EVERY parameter, but it
// is not applied to every parameter: ShardedMuon withholds the scale from the
// orthogonalized tensors on purpose (Newton-Schulz renormalizes the momentum,
// so a step-varying coefficient would only rotate the EMA). In the usual
// partition that is every tensor it owns, which means `grad_clip` covers only
// the embedding + 1-D group — a single-digit percentage of the model.
//
// Printed because the per-step `globalclip 0.0457` line otherwise reads as
// "the model's gradients were scaled to 4.6%", which would be false. If this
// says 6% coverage, `grad_clip` is not protecting the body and something else
// has to.
if (cfg.grad_clip > 0.0f) {
int64_t covered = 0, total = 0;
for (size_t gi = 0; gi < optimizers.size(); ++gi) {
covered += static_cast<int64_t>(optimizers[gi]->clip_covered_numel());
total += optimizers[gi]->total_numel();
}
const double pct = total > 0 ? 100.0 * static_cast<double>(covered) / total : 0.0;
if (is_master) {
std::cout << "[ZeRO-1] grad_clip=" << cfg.grad_clip << " APPLIES TO "
<< covered << " / " << total << " params ("
<< std::fixed << std::setprecision(2) << pct << "%)"
<< std::defaultfloat << std::endl;
if (covered == 0) {
std::cout << "[ZeRO-1] *** WARNING: grad_clip covers NO parameters. "
"Every group withholds the coefficient (Muon path). The "
"reported norm and globalclip are DIAGNOSTIC ONLY — no "
"gradient is being scaled. ***" << std::endl;
} else if (pct < 25.0) {
std::cout << "[ZeRO-1] *** WARNING: grad_clip leaves "
<< std::fixed << std::setprecision(2) << (100.0 - pct)
<< "% of parameters unclipped (the Muon-path body). ***"
<< std::defaultfloat << std::endl;
}
}
}
// Order in which the dtype groups are STEPPED (distinct from their order in
// `optimizers`, which the checkpointing code indexes as [0] == the group the
// CheckpointManager API covers). The group carrying the bulk of the
// parameters must step LAST: its staged refresh broadcasts are issued at the
// very end of its step(), and any *other* group stepping afterwards would
// enqueue a blocking norm all-reduce behind them on the same in-order comm
// stream and then wait on it — head-of-line blocking that drains the
// broadcasts and silently destroys the overlap. Ascending parameter count
// puts the biggest group last.
// Give every group (Muon whole-tensor groups AND AdamW flat-shard groups
// alike) a distinct DCP scope. DCP names a state tensor
// "optimizer.<state_type>.<param_fqn>" and param_fqn is derived from this
// scope, so without it two groups would write the same names into the same
// checkpoint and clobber each other.
for (size_t g = 0; g < optimizers.size(); ++g) {
optimizers[g]->set_dcp_scope("zero_g" + std::to_string(g));
}
std::vector<size_t> step_order(optimizers.size());
for (size_t i = 0; i < step_order.size(); ++i) step_order[i] = i;
{
std::vector<int64_t> group_numel(optimizers.size(), 0);
for (size_t gi = 0; gi < group_params.size(); ++gi)
for (auto& p : group_params[gi]) group_numel[gi] += p.numel();
std::stable_sort(step_order.begin(), step_order.end(),
[&](size_t a, size_t b) { return group_numel[a] < group_numel[b]; });
}
// ---- refresh stage granularity ------------------------------------------
// Stage 0 holds ONLY wte. The embedding lookup is the very first op of the
// forward, so wte's transfer has nothing to hide behind and is structurally
// exposed however the stages are cut — but everything ELSE can be binned by
// the layer that reads it, so layer i waits only on layer i's own pieces
// while the later layers keep transferring underneath its compute.
//
// LAYERS_PER_STAGE trades pipelining granularity against collective count:
// 1 = one stage per layer — finest. Each layer's attention compute covers
// the following layers' transfer many times over.
// >1 = coarser: fewer, larger broadcasts. Raise it when per-collective
// latency on the interconnect starts to dominate the bytes moved.
const int LAYERS_PER_STAGE = 1;
if (zcfg.overlap_param_refresh) {
const int n_layers = static_cast<int>(cfg.n_layers);
const int layer_stages = (n_layers + LAYERS_PER_STAGE - 1) / LAYERS_PER_STAGE;
const int REFRESH_STAGES = 1 + layer_stages; // stage 0 == wte
const int last_stage = REFRESH_STAGES - 1;
// Parameter -> stage, keyed by data pointer. set_param_refresh_stages
// then maps it through input_index, so the caller never sees the flat
// ordering.
std::map<const void*, int> stage_by_ptr;
for (auto& p : model.wte.parameters()) stage_by_ptr[p.data()] = 0;
for (int i = 0; i < n_layers; ++i) {
const int s = 1 + i / LAYERS_PER_STAGE;
for (auto& p : model.attn[i]->parameters()) stage_by_ptr[p.data()] = s;
for (auto& p : model.mlp[i]->parameters()) stage_by_ptr[p.data()] = s;
}
// norm_f and an untied lm_head are read AFTER every layer, so they ride
// along with the final layer stage — installed by that layer's pre-hook,
// which fires before either is reached. (Under weight tying the logits
// matmul re-reads wte.weight, already installed as stage 0.)
for (auto& p : model.norm_f.parameters()) stage_by_ptr[p.data()] = last_stage;
if (model.lm_head)
for (auto& p : model.lm_head->parameters()) stage_by_ptr[p.data()] = last_stage;
int unmapped = 0;
for (size_t oi = 0; oi < group_params.size(); ++oi) {
std::vector<int> stage_of;
stage_of.reserve(group_params[oi].size());
for (auto& p : group_params[oi]) {
auto it = stage_by_ptr.find(p.data());
if (it != stage_by_ptr.end()) {
stage_of.push_back(it->second);
} else {
// Unrecognised parameter (a module this mapping does not know
// about). Pin it to stage 0: always CORRECT, since stage 0 is
// installed before the forward starts — it only forgoes the
// overlap for that tensor instead of risking a read-before-
// install further down the forward.
stage_of.push_back(0);
++unmapped;
}
}
optimizers[oi]->set_param_refresh_stages(stage_of, REFRESH_STAGES);
}
if (is_master) {
std::cout << "ZeRO-1 staged refresh: " << REFRESH_STAGES << " stages ("
<< LAYERS_PER_STAGE << " layer(s)/stage)";
if (unmapped > 0)
std::cout << " | WARNING: " << unmapped << " parameter(s) unmatched, "
"pinned to stage 0 (correct, not overlapped)";
std::cout << std::endl;
}
// Install each layer stage immediately before the first block that reads
// it, via a forward pre-hook on that block. Stage 0 is already in place by
// then (zero_grad installs it before the forward starts). Hooks fire in
// ascending layer order, which is also ascending stage order — the order
// refresh_param_stage() requires.
//
// Hooks fire ONLY through Module::operator(); GPT::forward invokes blocks
// as (*attn[i])(x) precisely so this dispatch happens, and a direct
// attn[i]->forward(x) call would bypass it entirely. refresh_param_stage()
// early-outs once the step's stages are consumed, so the repeat calls from
// the other GRAD_ACCUM-1 micro-steps and from the validation forward are
// free.
for (int i = 0; i < n_layers; i += LAYERS_PER_STAGE) {
const int s = 1 + i / LAYERS_PER_STAGE;
model.attn[i]->register_forward_pre_hook(
[&optimizers, s](nn::Module&, const Tensor&) -> std::optional<Tensor> {
for (auto& opt : optimizers) opt->refresh_param_stage(s);
return std::nullopt; // input unchanged
});
}
}
if (is_master) {
int64_t n_decay = 0, n_nodecay = 0;
for (const auto& p : params) {
if (p.shape().dims.size() >= 2) n_decay += p.numel();
else n_nodecay += p.numel();
}
std::cout << "decayed params: " << n_decay << ", non-decayed: " << n_nodecay
<< " | ZeRO-1 optimizer groups: " << optimizers.size()
<< " | world_size: " << world_size << std::endl;
}
// ---- data ----
const std::string data_root = "/root/.cache/huggingface/hub/datasets--navingv--flux/snapshots/5232b3d133d2eb71905725f288ff6b56dc0df773";
DataLoaderLite train_loader(static_cast<int>(cfg.B), static_cast<int>(cfg.T),
rank, world_size, "train", data_root, is_master, 100000000, rank);
DataLoaderLite val_loader(static_cast<int>(cfg.B), static_cast<int>(cfg.T),
rank, world_size, "val", data_root, is_master, 100000000, rank);
// ---- CSV log (master only) ----
std::ofstream log_file;
if (is_master) {
log_file.open(use_muon ? "muon_4b_memfix2.csv"
: "adamw_noovl_contig.csv");
// MFU columns are APPENDED, never inserted: run_zero_ab_8xh200.sh parses
// this file positionally ($1..$8 in its awk block), so adding fields at
// the end keeps the existing A/B reporting working unchanged.
log_file << "step,train_loss,val_loss,tok_per_sec,grad_norm,dt_ms,lr,elapsed_min,"
"mfu,mfu_palm,tflops_per_gpu\n";
log_file << std::fixed << std::setprecision(6);
log_file.flush();
}
// Reused GPU scalars (avoid per-step allocation).
Tensor grad_scale = Tensor::full(Shape{{1}}, TensorOptions().with_device(device),
1.0f / static_cast<float>(GRAD_ACCUM));
Tensor loss_accum_gpu = Tensor::zeros(Shape{{1}}, TensorOptions().with_device(device));
float last_val_loss = -1.0f; // carried across steps; refreshed every val_freq
auto training_start = std::chrono::high_resolution_clock::now();
// Spike diagnostics (BLU_DIAG=1). See the micro loop and the per-step report.
const bool diag = getenv_int("BLU_DIAG", 0) != 0;
std::vector<float> micro_loss;
std::vector<std::pair<int, size_t>> micro_where; // (shard_index, token_offset)
// Depth probe: the accumulated LOCAL gradient norm of one weight per probe
// layer. The per-group norms show the spike is ~45x in the 1.62B bf16 body but
// only ~6x in the embedding group, i.e. localized in the stack — this says at
// WHICH DEPTH it starts. Compare the profile on a spiking step against a quiet
// one: a blow-up that begins mid-stack and grows toward the loss points at the
// forward, while a flat multiplicative lift across all depths points at the
// optimizer/update. Read before opt->step(), so these are the pre-reduce
// per-rank gradients (step() does not modify param .grad, and zero_grad only
// runs at the top of the next step).
const std::vector<int> probe_layers = {0, 8, 17, 26, 35};
std::vector<std::pair<int, float>> layer_gnorm; // (layer, ||grad(wQ.weight)||)
std::vector<std::pair<int, float>> gamma_absmax; // (layer, max|q_gamma|)
std::vector<std::pair<int, float>> w_absmax; // (layer, max|wQ.weight|)
auto grad_l2 = [](const Tensor& p) -> float {
if (!p.has_grad()) return -1.0f;
Tensor h = p.grad_view().as_type(Dtype::Float32).to_cpu();
const float* d = h.data<float>();
const int64_t n = h.numel();
double acc = 0.0;
for (int64_t i = 0; i < n; ++i) acc += static_cast<double>(d[i]) * d[i];
return static_cast<float>(std::sqrt(acc));
};
// Replica-consistency probe. After the staged refresh has been consumed, every
// rank must hold BIT-IDENTICAL parameters: each byte of each parameter is
// produced by exactly one owner's broadcast, so there is no averaging or
// rounding path that could legitimately make two ranks differ. If they DO
// differ, the ranks are training divergent models and averaging their
// gradients is meaningless — which would show up exactly as this failure mode
// (normal-looking gradient norm, destructive update, all micro-batches worse
// the following step). Bit-exact comparison via min/max across ranks, so any
// difference at all is reported rather than hidden under a tolerance.
auto checksum = [](const Tensor& p) -> double {
Tensor h = p.as_type(Dtype::Float32).to_cpu();
const float* d = h.data<float>();
const int64_t n = h.numel();
double acc = 0.0;
// Position-weighted so a permutation or a shifted piece is caught too, not
// just a change in the multiset of values.
for (int64_t i = 0; i < n; ++i) acc += static_cast<double>(d[i]) * (i % 8191 + 1);
return acc;
};
auto abs_max = [](const Tensor& p) -> float {
Tensor h = p.as_type(Dtype::Float32).to_cpu();
const float* d = h.data<float>();
const int64_t n = h.numel();
float m = 0.0f;
for (int64_t i = 0; i < n; ++i) m = std::max(m, std::fabs(d[i]));
return m;
};
// ---- checkpointing: manager + auto-resume ----
// save() writes one shared (non-sharded) path, so saves are master-only to
// avoid an 8-rank write race; load_latest runs on EVERY rank (concurrent
// reads are safe) so each rank restores its own optimizer m/v state and
// fast-forwards its own data shard. Saves are synchronous (use_async=false).
// NOTE: CheckpointManager's save/load API takes a single nn::Optimizer&,
// so optimizers[0] rides the normal path together with the model's full
// (unsharded) parameter tensors. Any further dtype group isn't covered by
// that API, so its shard state round-trips through a small sidecar file
// keyed off the same step number.
// Muon and AdamW runs write to DIFFERENT directories. Their optimizer state
// is not interchangeable — whole-tensor master+momentum against a flat
// elementwise master/m/v shard — and optimizers[0] is a different class in
// each mode, so a cross-mode resume would (correctly) abort on the layout
// descriptor. Separating the directories means flipping BLU_OPTIMIZER starts
// a clean run instead of failing against the other mode's checkpoints.
// if (use_muon) cfg.ckpt_dir += "_muon";
// ---- checkpointing: Distributed Checkpointing (DCP) + auto-resume ----
// Every rank writes its OWN shard of every group's optimizer state (plus its
// copy of the replicated parameters) into one step directory; DCP merges the
// per-rank metadata and, on load, intersects each saved flat range against
// the resuming rank's range. That is what makes the save correct at W>1 for
// BOTH kinds of group here: ShardedOptimizer/SymmetricShardedOptimizer's flat
// elementwise shard, and ShardedMuon's whole-tensor ownership (see
// ShardedMuon::dcp_state_shard_specs) — a master-only save would persist only
// one rank's slice of either. It also lets a run resume at a DIFFERENT world
// size, including for the Muon groups: the new run derives its own owner
// assignment and each record's checkpoint range is self-contained.
//
// Two adapters bridge this trainer to DCP's API (see ZeroDcpBridge.h): the
// model is an nn::Module rather than a dnn::DModule, and there is one ZeRO
// group (Muon or AdamW) per dtype where DCP takes a single optimizer pointer.
//
// Muon and AdamW runs still write to DIFFERENT directories. Their group
// scopes are identical either way ("zero_g0", "zero_g1", ...), but which
// CONCRETE class sits behind scope g (ShardedMuon vs. ShardedOptimizer /
// SymmetricShardedOptimizer) depends on BLU_OPTIMIZER, so the same scope
// name means a different on-disk layout in each mode; loading one under the
// other surfaces as a loud DCP coverage failure rather than silent garbage,
// but there is no reason to court it. Separating the directories means
// flipping BLU_OPTIMIZER starts a clean run instead of failing against the
// other mode's checkpoints.
if (use_muon) cfg.ckpt_dir += "_muon";
std::map<std::string, Tensor*> dcp_named;
for (size_t i = 0; i < params.size(); ++i) {
char nm[32];
std::snprintf(nm, sizeof nm, "param_%06zu", i);
dcp_named[nm] = ¶ms[i];
}
owt::zero::ZeroDcpModule dcp_model(dcp_named);
std::vector<owt::zero::ZeroOptimizerBase*> dcp_groups;
for (auto& o : optimizers) dcp_groups.push_back(o.get());
owt::zero::ZeroDcpOptimizer dcp_opt(dcp_groups);
// DistCheckpoint::save/load dereference `pg` for rank/world_size, but the
// process-group setup above leaves it NULL at world_size == 1 (ZeRO-1 needs
// no collectives there). Give the checkpointer its own single-rank group so
// one-GPU runs can still save and resume instead of faulting in save().
std::shared_ptr<ProcessGroupNCCL> ckpt_pg = pg;
bool ckpt_enabled = getenv_int("BLU_CKPT", cfg.checkpointing ? 1 : 0) != 0;
if (ckpt_enabled && !ckpt_pg) {
try {
ckpt_pg = init_process_group(1, 0, /*stream=*/0, /*group_name=*/"dcp",
dist::comm::MemoryMode::Standard);
} catch (const std::exception& e) {
ckpt_enabled = false;
std::cerr << "[WARN] checkpointing disabled: could not create a "
"single-rank process group for DCP: " << e.what() << std::endl;
}
if (!ckpt_pg) ckpt_enabled = false;
}
auto ckpt_path_for = [&](int s) {
return (fs::path(cfg.ckpt_dir) /
(cfg.ckpt_prefix + "_step_" + std::to_string(s))).string();
};
// All "<prefix>_step_<N>" directories under ckpt_dir, as (step, path).
auto scan_checkpoints = [&]() {
std::vector<std::pair<int, std::string>> found;
std::error_code ec;
if (!fs::exists(cfg.ckpt_dir, ec)) return found;
const std::string pre = cfg.ckpt_prefix + "_step_";
for (const auto& e : fs::directory_iterator(cfg.ckpt_dir, ec)) {
if (!e.is_directory()) continue;
const std::string d = e.path().filename().string();
if (d.size() <= pre.size() || d.compare(0, pre.size(), pre) != 0) continue;
const std::string digits = d.substr(pre.size());
if (digits.find_first_not_of("0123456789") != std::string::npos) continue;
found.push_back({std::stoi(digits), e.path().string()});
}
std::sort(found.begin(), found.end());
return found;
};
auto prune_old_checkpoints = [&]() {
if (cfg.ckpt_keep <= 0) return;
auto found = scan_checkpoints();
if (found.size() <= static_cast<size_t>(cfg.ckpt_keep)) return;
const size_t to_delete = found.size() - static_cast<size_t>(cfg.ckpt_keep);
for (size_t i = 0; i < to_delete; ++i) {
std::error_code rm_ec;
fs::remove_all(found[i].second, rm_ec);
}
};
int start_step = 0;
if (ckpt_enabled) {
// Newest checkpoint that actually completed (is_valid == metadata.bin
// present), so a run killed mid-save falls back to the previous one.
int latest_step = -1;
std::string latest_path;
const auto all_ckpts = scan_checkpoints();
for (auto it = all_ckpts.rbegin(); it != all_ckpts.rend(); ++it) {
if (OwnTensor::DistCheckpoint::is_valid(it->second)) {
latest_step = it->first;
latest_path = it->second;
break;
}
}
if (latest_step >= 0) {
float resume_loss = 0.0f;
int resume_epoch = 0;
// Retire in-flight refreshes first: the load overwrites the very
// shard buffers they publish from.
dcp_opt.dcp_before_load();
OwnTensor::DistCheckpoint::load(dcp_model, latest_path, ckpt_pg, &dcp_opt,
&resume_epoch, &start_step, &resume_loss);
// Cover the dtype-change path, where DCP rebinds its staging tensor
// instead of writing through to the model's storage.
dcp_model.sync_back_to_model();
// Republish the restored FP32 shards into the live model tensors --
// the step the stream-based load_state() used to do internally.
dcp_opt.dcp_after_load();
// The validation loss is only recomputed every val_freq steps, so a
// resume that does not land on one would print the -1 sentinel until
// the next validation. Restore the value the checkpoint carried.
{
float v = -1.0f;
if (dcp_opt.get_extra("last_val_loss", v)) last_val_loss = v;
}
if (is_master) {
std::cout << "[resume] loaded DCP checkpoint at step " << start_step
<< " (loss " << std::fixed << std::setprecision(6) << resume_loss
<< ") from " << latest_path << std::endl;
}
// Skip the batches already consumed by steps [0, start_step).
const size_t batches_to_skip =
static_cast<size_t>(start_step) * static_cast<size_t>(GRAD_ACCUM);
if (batches_to_skip > 0) train_loader.skip_batches(batches_to_skip);
}
}
if (is_master) {
std::cout << "\nStarting training"
<< (start_step > 0 ? (" (resumed at step " + std::to_string(start_step) + ")") : "")
<< "...\n" << std::endl;
}
for (int step = start_step; step < cfg.max_steps; ++step) {
// ---------- periodic validation ----------
if (step % cfg.val_freq == 0 || step == cfg.max_steps - 1) {
autograd::NoGradGuard no_grad;
// The validation forward reads the model parameters and runs BEFORE
// zero_grad(), so the refresh still in flight from the previous step
// has to be installed. Install ONLY stage 0 here and let the
// validation forward's own pre-hooks consume stages 1..N as it walks
// the layers — exactly as the training forward does. (drain_param_
// refresh() here would pull all N stages into the prologue and run
// the broadcasts against an idle compute stream, silently costing the
// overlap on every validation step.)
//
// Safe under NoGradGuard: Module::operator() dispatches forward
// pre-hooks unconditionally; only the backward sentinel nodes are
// gated on GradMode (NN.cpp). And model.forward() reaches the blocks
// as (*attn[i])(x), so the hooks do fire. If val_loss_steps == 0 the
// remaining stages simply fall through to zero_grad() + the training
// forward, which is still correct.
for (auto& opt : optimizers) opt->refresh_param_stage(0);
val_loader.reset();
float val_accum = 0.0f;
for (int v = 0; v < cfg.val_loss_steps; ++v) {
Batch b = val_loader.next_batch();
Tensor logits = model.forward(b.input);
Tensor loss = autograd::sparse_cross_entropy_loss(logits, b.target);
// Same fp32 bridge as the training loop — data<float>() on a bf16
// scalar reads garbage without it. Under NoGradGuard this is a
// plain cast; on the default path it is a no-op.
val_accum += autograd::to_dtype(loss, Dtype::Float32)
.to_cpu().data<float>()[0] / static_cast<float>(cfg.val_loss_steps);
}
float global_val = val_accum;
if (world_size > 1) {
MPI_Allreduce(&val_accum, &global_val, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD);
global_val /= world_size;
}
last_val_loss = global_val;
if (is_master) {
std::cout << " [val] step " << step << " | val_loss "
<< std::fixed << std::setprecision(6) << global_val << std::endl;
}
}
// ---------- training step ----------
auto t0 = std::chrono::high_resolution_clock::now();
for (auto& opt : optimizers) opt->zero_grad();
loss_accum_gpu *= 0.0f;
micro_loss.clear();
micro_where.clear();
for (int micro = 0; micro < GRAD_ACCUM; ++micro) {
Batch batch = train_loader.next_batch();
Tensor logits = model.forward(batch.input);
Tensor loss_raw = autograd::sparse_cross_entropy_loss(logits, batch.target);
// LossOps keys the loss scalar's dtype off the logits, so with
// bf16_logits it comes back bf16. Two things then break unless it is
// bridged to fp32 here:
//
// * Every host read below goes through data<float>(), which is a
// bare reinterpret_cast with no dtype check — reading a bf16
// scalar as float yields garbage SILENTLY, and loss_accum_gpu is
// an fp32 tensor besides.
// * backward() must be seeded at the dtype LossBackward reads
// grad_output at, which is the LOGITS dtype. Seeding the fp32
// view is what makes that correct: ToDtypeBackward recasts the
// incoming fp32 grad back to bf16 before the CE node sees it.
//
// On the default path to_dtype returns its input unchanged, so
// loss32 IS loss_raw — same tensor, same graph, no extra node.
Tensor loss32 = autograd::to_dtype(loss_raw, Dtype::Float32);
loss_accum_gpu += loss32.detach(); // running sum for logging
// Per-micro-batch loss (BLU_DIAG=1). The reported train loss is the mean
// over GRAD_ACCUM micro-steps, so a single pathological micro-batch is
// diluted 16x while its gradient dominates the accumulated norm almost
// undiluted — which is exactly how a modest loss bump can sit next to a
// 40x norm spike. Printing the 16 values separates "one bad micro-batch"
// from "the whole step is worse", and `micro_where` names the shard and
// token offset so the responsible bytes can be decoded offline.
//
// Costs a device sync per micro-step, so it is off by default. It does
// NOT perturb the result: the loss is already fully determined by this
// point and the run is data-deterministic.
if (diag) {
micro_loss.push_back(loss32.to_cpu().data<float>()[0]);
micro_where.push_back({batch.shard_index, batch.token_offset});
}
// No DDP all-reduce here: gradients just accumulate locally across
// micro-steps; ZeRO-1's reduce-scatter inside optimizer->step()
// (below) does the cross-rank averaging once, after accumulation.
//
// Overlap: arm the bucketed reduce on the LAST micro-step only, so the
// hooks fire on the fully-accumulated gradient and reduce it to owner
// ranks concurrently with the rest of this backward. Earlier micro-
// steps just accumulate (the hooks are inert). No-op when overlap is
// disabled (world_size==1).
if (micro == GRAD_ACCUM - 1) {
for (auto& opt : optimizers) opt->arm_grad_sync();
}
loss32.backward(&grad_scale); // scaled by 1/GRAD_ACCUM
}
// Average loss for logging (across micro-steps, then across ranks).
float loss_accum;
{
float local = loss_accum_gpu.to_cpu().data<float>()[0] / static_cast<float>(GRAD_ACCUM);
if (world_size > 1) {
MPI_Allreduce(&local, &loss_accum, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD);
loss_accum /= world_size;
} else {
loss_accum = local;
}
}
// Depth probe (see declarations above). Before opt->step().
layer_gnorm.clear();
gamma_absmax.clear();
w_absmax.clear();
if (diag) {
for (int L : probe_layers) {
if (L >= static_cast<int>(cfg.n_layers)) continue;
layer_gnorm.emplace_back(L, grad_l2(model.attn[L]->wQ.weight));
gamma_absmax.emplace_back(L, abs_max(model.attn[L]->q_gamma));
// Weight magnitude. On a spiking step the model's loss decouples
// from the input (see the [micro] correlation), which is what
// saturated activations look like. A weight that has run away is
// the cheapest cause to rule in or out, and this reuses a copy the
// checksum already needs.
w_absmax.emplace_back(L, abs_max(model.attn[L]->wQ.weight));
}
// Replica consistency (see `checksum` above). Checked HERE, after the
// micro loop: the staged refresh is consumed lazily by the forward's
// per-layer pre-hooks, so this is the first point in the step at which
// every stage is guaranteed installed. wQ.weight of layer L lives in
// refresh stage 1+L, so the probe layers sample 5 different stages.
if (world_size > 1) {
double local_ck = 0.0;
for (int L : probe_layers) {
if (L >= static_cast<int>(cfg.n_layers)) continue;
local_ck += checksum(model.attn[L]->wQ.weight);
}
double ck_min = 0.0, ck_max = 0.0;
MPI_Allreduce(&local_ck, &ck_min, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD);
MPI_Allreduce(&local_ck, &ck_max, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
if (is_master && ck_min != ck_max) {
const double denom = std::max(std::fabs(ck_max), 1e-30);
std::cout << " [replica] *** RANKS DIVERGED *** min " << std::scientific
<< std::setprecision(9) << ck_min << " max " << ck_max
<< " reldiff " << ((ck_max - ck_min) / denom) << std::endl;
} else if (is_master) {
std::cout << " [replica] ranks bit-identical (ck " << std::scientific
<< std::setprecision(9) << ck_max << ")" << std::endl;
}
}
}
if (std::isnan(loss_accum) || std::isinf(loss_accum)) {
if (is_master) std::cerr << "ERROR: NaN/Inf loss at step " << step << std::endl;
if (is_master) log_file.close();
MPI_Abort(MPI_COMM_WORLD, 2);
}
// WSD LR schedule, then step each ZeRO-1 dtype group. Each
// ShardedOptimizer reduce-scatters its own grad shard (bandwidth-optimal
// cross-rank averaging — this is what replaced the DDP all-reduce),
// applies decoupled weight decay + clips to cfg.grad_clip internally
// using the exact global norm of its shard (all-reduced across ranks),
// runs the fused Adam step on its FP32 master shard, then all-gathers the
// updated params back into the live model tensors.
//
// Combined grad norm across groups is exact (not approximate): the
// per-group norms are each already the true global L2 norm of that
// group's gradients, and the groups are disjoint parameter sets, so
// sqrt(sum of squares) reconstructs the true norm over all params.
// (Each group clips independently against its own norm rather than this
// combined one.)
// ONE WSD curve, evaluated once per group with that group's own (peak,
// floor). Identical in form to train_muon.py's two get_lr calls: the
// warmup length, the stable window and the cosine tail are shared, so the
// groups stay phase-aligned and only their scale differs. `lr` is the
// AdamW-arm value, kept for the CSV/console log.
const float lr = get_lr(step, cfg.max_lr, cfg.min_lr, cfg.warmup_steps,
cfg.max_steps, cfg.stable_steps);
// ---- cross-group global gradient clipping (BLU_GLOBAL_CLIP=1) ----
// Each ZeRO instance owns a DISJOINT subset of the model and computes the
// exact global norm of that subset, so clipping inside step() clips each
// group INDEPENDENTLY. DeepSpeed / Megatron / torch.clip_grad_norm_ clip on
// the norm over ALL parameters instead.
//
// CORRECTION (measured, 2 x 2000-step runs — see docs/ZERO1_CLIPPING_AUDIT.md).
// This comment used to claim the FP32 group's norm runs ~29x the body's and
// that global clipping therefore changes the body's effective step size by
// ~2.6x. Both halves are wrong:
// * The ratio runs the OTHER way. G0 (body) / G1 (fp32) measured 4.4 at
// steps 0-199, 2.4 at 800-999, 16.2 at 1400-1599, 17.2 at 1800-1999.
// The body's norm is always the larger one.
// * There is no 2.6x step-size difference. AdamW is scale-invariant to a
// CONSTANT gradient factor; replaying the real coefficient series through
// an Adam update puts the distortion from its step-to-step VARIATION at
// ~5% RMS, not 160%.
// And in this partition the body is not clipped at all (ShardedMuon skips
// every use_muon record), so per-group vs global is a distinction without a
// difference on 94% of the model: an A/B over 104 matched steps differs by
// 0.005 mean loss, inside the 0.019 run-to-run noise floor.
//
// Phase 1 reduces every group's gradients and collects its ||g||^2; the
// groups are disjoint so the squares simply add. Phase 2 hands the single
// coefficient down and steps. step() detects that the reduction already
// happened and does not repeat it.
//
// This also removes the head-of-line blocking `step_order` works around:
// every norm all-reduce now happens in phase 1, before any group issues its
// staged parameter transfers.
//
// Off by default — it changes training semantics, so it is opt-in and
// A/B-able rather than a silent change.
const bool global_clip = getenv_int("BLU_GLOBAL_CLIP", 0) != 0;
float external_coef = -1.0f;
if (global_clip && cfg.grad_clip > 0.0f) {
float nsq_total = 0.0f; // every group: the reported diagnostic norm
float nsq_clip = 0.0f; // groups that actually APPLY the coefficient
bool skip = false;
for (size_t oi : step_order) {
const float nsq = optimizers[oi]->reduce_and_measure();
if (nsq >= 0.0f && optimizers[oi]->clip_covered_numel() > 0)
nsq_clip += nsq;
// NOTE the test is `!(nsq >= 0)`, NOT `nsq < 0`. They differ on NaN:
// `NaN < 0` is FALSE, so a NaN group norm used to slip through, make
// nsq_total NaN, and then hit std::min(1.0f, NaN) — which returns
// 1.0f (std::min is `b < a ? b : a`, and `NaN < 1.0f` is false).
// The result was a NaN gradient with NO clip applied, written
// straight into the parameters. reduce_and_measure() now returns -1
// for any non-finite norm (scaler or not); this test catches both
// that and a NaN that arrives some other way.
if (!(nsq >= 0.0f)) skip = true;
else nsq_total += nsq;
}
// WHICH norm the coefficient is derived from.
//
// "Global" clipping means one coefficient over all parameters — correct
// when the coefficient is then APPLIED to all parameters. Here it is
// not: ShardedMuon withholds the scale from every orthogonalized tensor,
// which in this partition is the whole body. Deriving the coefficient
// from the full norm and applying it only to the embedding group means
// the embedding is rescaled by a factor set by gradients it has no
// relationship to. Measured over both 2000-step runs, the AdamW group
// is only 7-11% of the total ||g||^2, so the coefficient tracked the
// body: it drifted 0.52 -> 0.18 across the stable phase purely because
// the body's norm grew from 2.0 to 4.5, damping the embedding harder and
// harder for reasons internal to the body.
//
// Default: derive from the norm of the COVERED subset, so the clip is a
// real clip of the thing it scales. The full norm is still what gets
// reported. BLU_GLOBAL_CLIP_STRICT=1 restores the previous
// full-norm-derived coefficient for A/B (and is the correct setting if
// the Muon exemption is ever removed, since then coverage is 100%).
const bool strict = getenv_int("BLU_GLOBAL_CLIP_STRICT", 0) != 0;
const float total = std::sqrt(nsq_total);
const float clip_norm = strict ? total : std::sqrt(nsq_clip);
if (!skip && std::isfinite(clip_norm) && std::isfinite(total)) {
external_coef = std::min(1.0f, cfg.grad_clip / (clip_norm + 1e-6f));
for (auto& opt : optimizers) opt->set_external_clip_coef(external_coef);
} else if (is_master) {
std::cerr << "WARNING: non-finite gradient norm at step " << step
<< " — skipping the global clip; groups fall back to their "
"own per-group coefficient this step" << std::endl;
}
}
float grad_norm_sq = 0.0f;
std::vector<float> group_norm(optimizers.size(), 0.0f);
for (size_t oi : step_order) { // biggest group last — see step_order
auto& opt = optimizers[oi];
opt->set_lr(get_lr(step, group_lr[oi].first, group_lr[oi].second,
cfg.warmup_steps, cfg.max_steps, cfg.stable_steps));
opt->step();
const float g = opt->last_grad_norm();
group_norm[oi] = g;
grad_norm_sq += g * g;
}
const float grad_norm = std::sqrt(grad_norm_sq);
// Timing / pipeline boundary.
//
// With the staged parameter refresh ON we deliberately do NOT synchronize
// here. A cudaDeviceSynchronize() drains the compute stream and makes the
// host wait for every kernel of this step before it may enqueue anything
// else — so the next forward is not launched until the GPU has gone
// completely idle. The refresh broadcasts, issued on the comm stream at
// the end of step(), would then run to completion inside that idle window
// and overlap nothing. Not a correctness bug, but it makes the staged
// refresh pointless, because there is never any compute in flight for it
// to hide behind.
//
// Dropping the sync lets the host run ahead and enqueue the next forward
// while this step's Adam/cast tail is still executing, so the broadcasts
// for stages 1..N transfer underneath layer 0's attention compute. Stage 0
// (wte) stays exposed by construction: the embedding reads it first.
//
// dt then measures HOST wall-clock per iteration rather than drained-GPU
// time. In steady state these agree: the GPU is near-saturated, so the
// launch queue back-pressures cudaLaunchKernel and the host is throttled
// to the GPU's rate. Only the first couple of steps read low while the
// pipeline fills. The non-overlap path keeps its full device sync.
//
// With the staged refresh off (the current default) this is REDUNDANT, not
// load-bearing: the synchronous publish inside step() already ends in a
// cudaDeviceSynchronize(). It is kept because it also covers a step that
// returned early (gradient overflow) and published nothing, and because it
// is what restores drained-GPU timing for dt.
if (!zcfg.overlap_param_refresh) cudaDeviceSynchronize();
auto t1 = std::chrono::high_resolution_clock::now();
const double dt = std::chrono::duration<double>(t1 - t0).count();
const int64_t tokens = cfg.B * cfg.T * GRAD_ACCUM * static_cast<int64_t>(world_size);
const double tok_per_sec = static_cast<double>(tokens) / dt;
const double elapsed_min =
std::chrono::duration<double>(t1 - training_start).count() / 60.0;
// ---- MFU for this step ----
// Uses the SAME dt as tok/s, so the caveat above applies here too: with
// overlap_param_refresh on there is no device sync, and dt is host
// wall-clock. In steady state the launch queue back-pressures the host
// to the GPU's rate so the two agree, but the first few steps read fast
// while the pipeline fills — which is exactly why observe() drops them
// from the steady-state figure (BLU_MFU_WARMUP, default 5).
const double step_mfu = mfu_meter.mfu(dt, tokens);
const double step_mfu_p = mfu_meter.mfu_palm(dt, tokens);
const double step_tflops = mfu_meter.tflops_per_gpu(dt, tokens);
mfu_meter.observe(dt, tokens);
// ---------- per-step report (every step) ----------
if (is_master) {
std::cout << "step " << std::setw(6) << step
<< " | train " << std::fixed << std::setprecision(6) << loss_accum
<< " | val " << std::fixed << std::setprecision(6) << last_val_loss
<< " | tok/s " << std::fixed << std::setprecision(2) << tok_per_sec
<< " | norm " << std::fixed << std::setprecision(4) << grad_norm
<< " | dt " << std::fixed << std::setprecision(2) << (dt * 1000.0) << " ms"
<< " | mfu " << std::fixed << std::setprecision(2) << (step_mfu * 100.0) << "%"
<< " | TF/gpu " << std::fixed << std::setprecision(1) << step_tflops
<< " | lr " << std::scientific << std::setprecision(4) << lr;
// Per-group norms: which dtype group actually spiked.
for (size_t gi = 0; gi < group_norm.size(); ++gi)
std::cout << " | G" << gi << "(" << group_label[gi] << ") "
<< std::fixed << std::setprecision(4) << group_norm[gi];
if (external_coef >= 0.0f)
std::cout << " | globalclip " << std::fixed << std::setprecision(4)
<< external_coef;
// DIAGNOSTIC ONLY: Newton-Schulz health on the Muon path. `nsrej` is
// how many owned tensors the divergence guard rejected THIS step,
// `nsratio` the worst ||update||_F / sqrt(min(rows,cols)) among the
// updates that passed it (healthy ~1.0). The guard only trips above
// element magnitude 1000, so an update can be far too large and
// still be applied — that shows up here and nowhere else.
for (auto& opt : optimizers) {
if (auto* sm = dynamic_cast<ShardedMuon*>(opt.get())) {
std::cout << " | nsrej " << sm->ns_rejected_last_step()
<< "/" << sm->rejected_ns_updates()
<< " | nsratio " << std::fixed << std::setprecision(3)
<< sm->ns_max_update_ratio();
break;
}
}
std::cout << std::endl;
if (diag && !micro_loss.empty()) {
// Flag micro-batches that stand out against this step's own median,
// so the line is readable at a glance instead of 16 bare floats.
std::vector<float> sorted = micro_loss;
std::sort(sorted.begin(), sorted.end());
const float med = sorted[sorted.size() / 2];
std::cout << " [micro] med " << std::fixed << std::setprecision(4) << med << " |";
for (size_t m = 0; m < micro_loss.size(); ++m) {
const bool hot = micro_loss[m] > med + 0.5f;
std::cout << " " << (hot ? "*" : "") << std::setprecision(3) << micro_loss[m];
}
std::cout << std::endl;
for (size_t m = 0; m < micro_loss.size(); ++m) {
if (micro_loss[m] > med + 0.5f)
std::cout << " [micro] OUTLIER m=" << m
<< " loss " << std::setprecision(4) << micro_loss[m]
<< " (median " << med << ") from shard "
<< micro_where[m].first << " token offset "
<< micro_where[m].second << std::endl;
}
}
if (diag && !layer_gnorm.empty()) {
std::cout << " [depth] ||g(wQ)|| by layer:";
for (auto& e : layer_gnorm)
std::cout << " L" << e.first << " " << std::scientific
<< std::setprecision(3) << e.second;
std::cout << " | max|q_gamma|:";
for (auto& e : gamma_absmax)
std::cout << " L" << e.first << " " << std::fixed
<< std::setprecision(4) << e.second;
std::cout << " | max|wQ|:";
for (auto& e : w_absmax)
std::cout << " L" << e.first << " " << std::fixed
<< std::setprecision(4) << e.second;
std::cout << std::endl;
}
log_file << step << ","
<< loss_accum << ","
<< last_val_loss << ","
<< tok_per_sec << ","
<< grad_norm << ","
<< (dt * 1000.0) << ","
<< lr << ","
<< elapsed_min << ","
<< step_mfu << ","
<< step_mfu_p << ","
<< step_tflops << "\n";
log_file.flush();
}
// ---------- periodic checkpoint (Distributed Checkpointing) ----------
// Labeled step+1 (the next iteration to run) so the resume path picks up
// exactly where this save left off — no duplicated or skipped iteration.
//
// COLLECTIVE, and deliberately NOT master-gated. Under ZeRO-1 each rank
// owns a DIFFERENT slice of the optimizer state — a flat elementwise
// slice for the AdamW groups, a disjoint set of whole tensors for the
// Muon groups — so a master-only save would persist only one rank's
// piece of either and every rank would then resume from rank 0's slice.
// Here every rank writes its own shard into the step directory and DCP
// merges the metadata.
if (ckpt_enabled &&
((step + 1) % cfg.ckpt_freq == 0 || step == cfg.max_steps - 1)) {
try {
// Install any refresh still in flight so the checkpoint holds
// this step's weights rather than the previous step's.
for (auto& opt : optimizers) opt->drain_param_refresh();
// Carry the most recent validation loss into the checkpoint so a
// resume reports it instead of the -1 sentinel.
dcp_opt.set_extra("last_val_loss", last_val_loss);
OwnTensor::DistCheckpoint::save(dcp_model, ckpt_path_for(step + 1),
ckpt_pg, &dcp_opt, /*epoch=*/0,
step + 1, loss_accum);
if (is_master) prune_old_checkpoints();
} catch (const std::exception& e) {
std::cerr << "[WARN] checkpoint save at step " << (step + 1)
<< " failed: " << e.what() << std::endl;
}
}
// Reclaim the extra allocator high-water mark from the first couple of
// steps' one-time warmup allocations (ZeRO shard buffers, autograd
// graph caches, etc.).
if (step == 0 || step == 5) {
CachingCUDAAllocator::instance().empty_cache();
}
if (step == 5) {
AllocationTracker::instance().shutdown();
}
}
// The final step leaves its parameter refresh in flight (nothing consumed it,
// since there is no next forward). Install it before teardown so the process
// group is not destroyed with collectives outstanding.
for (auto& opt : optimizers) opt->drain_param_refresh();
cudaDeviceSynchronize();
// Flush any in-flight async checkpoint before teardown (no-op for sync saves).
// if (cfg.checkpointing) {
// try { ckpt.wait_for_completion(); }
// catch (const std::exception& e) {
// if (is_master) std::cerr << "[WARN] checkpoint flush at shutdown: " << e.what() << std::endl;
// }
// }
if (is_master) {
log_file.close();
// ---- end-of-run MFU summary ----
// The steady-state figure is the one to quote: total model FLOPs over
// total wall time across every step past the warmup window. It is an
// aggregate rather than a mean of per-step ratios, so a stalled step
// costs exactly its own wall clock instead of skewing an average, and
// it is the same quantity a paper means by "we sustained X% MFU".
// `best` is the ceiling this configuration touched at least once — the
// gap between best and steady is jitter (checkpoints, validation,
// straggler ranks), not a kernel problem.
const double ss = mfu_meter.steady_mfu();
std::cout << "\n==================== MFU summary ====================\n"
<< std::fixed << std::setprecision(2)
<< " steady-state MFU : " << (ss * 100.0) << "%"
<< " (" << mfu_meter.ss_steps << " steps, first "
<< mfu_meter.warmup_steps << " excluded)\n"
<< " best-step MFU : " << (mfu_meter.best_mfu * 100.0) << "%\n"
<< " model TFLOP/s/gpu : "
<< (ss * mfu_meter.peak_flops_per_gpu / 1e12)
<< " of " << (mfu_meter.peak_flops_per_gpu / 1e12) << " peak\n"
<< " cluster PFLOP/s : "
<< (ss * mfu_meter.cluster_peak() / 1e15) << "\n"
<< std::setprecision(0)
<< " steady tok/s : " << mfu_meter.steady_tok_per_sec() << "\n"
<< std::setprecision(3)
<< " total model FLOPs : "
<< (mfu_meter.flops_per_token_causal * mfu_meter.ss_tokens / 1e18)
<< " EFLOP over " << (mfu_meter.ss_secs / 60.0) << " min\n"
<< "=====================================================" << std::endl;
std::cout << "\n=== Training complete ===" << std::endl;
}
MPI_Finalize();
return 0;
}
// =============================================================================
// bluscript.cpp — Llama-style decoder LM training (fp32 model, bf16 fused attn,
// ZeRO-1 mixed precision optimizer state sharding, WSD LR schedule)
//
// Architecture (per the production spec):
// token embedding
// -> N x TransformerBlock:
// RMSNorm -> { QK-norm + RoPE + grouped-query causal attention }(fused
// kernel: gqa_fused_flash_attn_forward/backward) -> Wo -> +residual
// RMSNorm -> SwiGLU FFN (gate_up -> swiglu -> down) -> +residual
// -> final RMSNorm
// -> lm_head (weight-tied with the token embedding)
// -> sparse cross-entropy loss
// -> backward -> AdamW, or Muon+AdamW (BLU_OPTIMIZER=muon), with a WSD
// (warmup / stable / cosine-decay) LR schedule
//
// The QK-norm scales (q_gamma / k_gamma) and the RoPE rotation are FUSED in one
// CUDA kernel (autograd::fused_rms_rope), and grouped-query causal attention is
// a second fused flash-attention kernel reached through
// autograd::scaled_dot_product_attention. Those ops run bf16 internally and
// bridge back to the fp32 model. Matmuls go through the library BLAS
// (nn::Linear / autograd::matmul).
//
// Optimizer: ZeRO-1 (owt::zero::ShardedOptimizer). FP32 master weights + Adam
// moments live only in each rank's optimizer shard, not the model. Designed for
// ZeRO-1 over 8 GPUs (one rank per GPU): each rank holds the full model params
// but only 1/world_size of the optimizer state; per step, grads reduce-scatter
// (avg) across ranks, the fused Adam update runs on each rank's FP32 shard, and
// updated params all-gather back — no separate DDP grad all-reduce.
//
// BLU_OPTIMIZER=muon switches the transformer body's 2-D weight matrices to
// Muon, keeping AdamW for the token embedding and every 1-D scale. ZeRO-1 STILL
// APPLIES to the Muon set, but it cannot be the same ZeRO-1: AdamW's update is
// elementwise, so any partition of the flat parameter buffer is valid, while
// Muon's Newton-Schulz orthogonalization is a whole-MATRIX operator that a
// contiguous flat slice cuts through the middle of. The Muon set therefore goes
// to owt::zero::ShardedMuon, which shards by WHOLE PARAMETER — one owner rank
// per matrix, holding that matrix's FP32 master + momentum, reducing its
// gradient to the owner and broadcasting the updated weight back. Same ~1/W
// state reduction, and cheaper than sharded AdamW (no second moment, no flat
// staging buffer). See dist/zero/include/ShardedMuon.h.t
//
// Launch: make run-mpi FILE=bluscript.cpp NP=8 WITH_BLUBLAS=1
// BLU_OPTIMIZER=muon mpirun -np 8 ./_runner_
// =============================================================================
// Tensor library
// Tracker + Memory
// Data + distributed
// Compile-time flag to toggle between fused RMSNorm+RoPE kernel (1) and sequential RMSNorm + RoPE (0)
using namespace OwnTensor;
using owt::zero::ShardedOptimizer;