diff --git "a/source/files/bluscriptCP.cpp" "b/source/files/bluscriptCP.cpp" new file mode 100644--- /dev/null +++ "b/source/files/bluscriptCP.cpp" @@ -0,0 +1,3153 @@ +// ============================================================================= +// bluscriptCP.cpp — CONTEXT-PARALLEL Llama-style decoder LM training. +// +// Context-parallel port of bluscript.cpp: the SEQUENCE is split across GPUs +// (not the batch), so we can train at long context. Attention runs through the +// CP layer's FUSED paths (RoPE + QK-norm + GQA happen INSIDE ContextParallel): +// CP_ATTN_MODE=ring (default) -> enable_rope (ring; new fused-RoPE CP +// kernels; HeadTail load-balance; needs +// building with CP_FUSED_ROPE=1) +// CP_ATTN_MODE=ulysses -> enable_ulysses_fused (all-to-all; GQA; +// contiguous shard; default build) +// +// Topology: CP_SIZE (env, default = world_size => CP-only). dp_size = +// world_size/CP_SIZE. A 2-D DeviceMesh({dp_size, cp_size}) gives CP-only +// (dp_size==1) or 2-D DP x CP. cp_pg (mesh axis 1) drives the attention ring; +// dp_rank (axis 0) selects which batch each rank loads. DataParallel is KEPT +// over the GLOBAL pg -> init-weight broadcast + param-grad all-reduce over all +// dp*cp ranks (the correct group, since params replicate globally). The ring +// only adds the attention-activation (dQ/dK/dV) grad sum over cp_pg (disjoint +// from DDP's leaf-param reduction — no double-count). +// +// Launch (ring, 2 GPUs, CP-only): +// make CP_FUSED_ROPE=1 bluscript-cp && make CP_FUSED_ROPE=1 run-bluscript-cp NP=2 +// ============================================================================= + +#include +#include +#include // memcpy: CP_FWD_PROBE raw-loss bit dump +#include +#include +#include +#include +#include +#include +#include +#include +#include // Ms-PoE weight list parsing +#include +#include +#include + +#include +#include + +// Tensor library +#include "TensorLib.h" +#include "autograd/AutogradOps.h" // matmul, add, reshape, transpose, swiglu, ... +#include "autograd/operations/RoPEOps.h" // build_rope_cache +#include "autograd/operations/LossOps.h" // sparse_cross_entropy_loss +#include "autograd/operations/EmbeddingOps.h" +#include "nn/NN.h" +#include "nn/optimizer/Optim.h" +#include "checkpointing/GradMode.h" +#include "checkpointing/Checkpointing.h" // CheckpointManager (save / load_latest) +#include "device/CachingCudaAllocator.h" + +// Data + distributed +#include "DataLoader.h" +// NOTE: include DataParallel.hpp DIRECTLY (not the dist/distributed.h umbrella, +// which pulls the TP stack incl. a different `class DeviceMesh` that collides +// with the CP one). init_process_group comes from the canonical ProcessGroupNCCL.h +// (pulled in via device_mesh.h / ContextParallel.h). +#include "DataParallel.hpp" // DataParallel, DDP_Options +#include "dist/zero/include/ShardedOptimizer.h" // owt::zero::ShardedOptimizer (ZeRO-1) +#include "dist/zero/include/SymmetricShardedOptimizer.h" // ZeRO-1 v2 (reduce-scatter/all-gather) +#include "dist/zero/include/ZeroOptimizerBase.h" // common surface: v1/v2 interchangeable + +// Context parallelism +#include "process_group/device_mesh.h" // DeviceMesh (canonical PG, ncclCommSplit) +#include "context_parallel/ContextParallel.h" // ContextParallel, shard_sequence_pre_embed, RotatorType +#include "context_parallel/CreamPositions.h" // CREAM/PoSE/RandPos position-label generators +#include "context_parallel/LongRoPEOps.h" // build_rope_cache_longrope (searched cache) +#include "context_parallel/EvalNLLKernel.h" // on-GPU per-position NLL (eval only) +#include "context_parallel/RopeDeltas.h" // cp::local_to_global_pos (CP eval position map) +#include "ops/IndexingOps.h" // gather (cache-row select for CREAM) + +using namespace OwnTensor; + +// ============================================================================= +// CudaTimer — event-based per-phase timer (ported from gpt2_cp_test.cpp so the +// per-step [TIMING] breakdown matches that script's terminal output). +// ============================================================================= +struct CudaTimer { + cudaEvent_t start, stop; + CudaTimer() { cudaEventCreate(&start); cudaEventCreate(&stop); } + ~CudaTimer() { cudaEventDestroy(start); cudaEventDestroy(stop); } + void start_timer() { cudaEventRecord(start); } + float get_elapsed_ms() { + cudaEventRecord(stop); + cudaEventSynchronize(stop); + float ms = 0; cudaEventElapsedTime(&ms, start, stop); + return ms; + } + double get_elapsed_seconds() { return get_elapsed_ms() / 1000.0; } +}; + +// Discards all output — used to temporarily silence the DataLoader's per-shard +// debug prints during the resume fast-forward (skip_batches), which otherwise +// spam one line per shard crossing, per rank. +struct NullBuf : std::streambuf { + int overflow(int c) override { return c; } +}; + +// ============================================================================= +// Configuration +// ============================================================================= +struct ModelConfig { + // ---- model ---- + int64_t d_model = 384; + int64_t n_layers = 6; + int64_t q_heads = 6; + int64_t kv_heads = 2; // grouped-query: G = q/kv = 3 + int64_t head_dim = 384 / 12; // = 64 (fused kernel needs 64 or 128) + int64_t ffn_hidden = (8 * 384) / 3; // 8/3 * d_model = 2048 (SwiGLU inner dim) + int64_t context_length = 4096; // = T (multiple of 2*CP_SIZE, 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) + + // ---- optimization / data ---- + int64_t B = 2; // micro-batch per rank + int64_t T = 4096; // sequence length (== context_length) + int64_t global_batch = 524288; // tokens per optimizer step (all ranks) + float grad_clip = 1.0f; + float weight_decay = 0.1f; + + float max_lr = 6e-4f; + float min_lr = 6e-5f; // 10% of max_lr + int warmup_steps = 91; + // CP_BF16=1 makes the Linears (wQ/wK/wV/wO, gate_up/down, lm_head matmul) + // bf16 with fp32 norms, fp32 residual stream and fp32 logits -- the scheme + // bluscript_zero.cpp uses. Halves parameter/activation memory and moves the + // dense matmuls onto bf16 tensor cores instead of TF32. + // + // fp32 master weights come FREE: nn::Optimizer's ctor (Optim.cpp:28) builds + // an fp32 master for every param whose dtype != Float32, and AdamW::step + // (Optim.cpp:779) updates through it. No ZeRO needed for correct mixed + // precision, which is what made porting this without ShardedOptimizer safe. + // + // The CP ATTENTION path is already bf16 regardless of this flag: it casts + // q/k/v with as_type(Bfloat16) at FusedRoPESDPA.h:110. This flag only moves + // the LINEARS, so it is a smaller numerical change than it looks. + // + // DEFAULTS OFF so the fp32 trajectory -- and comparability with run100 and + // the recorded step-500 probe figures -- is preserved unless asked for. + // CP_ZERO=1 replaces DataParallel + nn::AdamW with ZeRO-1 optimizer-state + // sharding (owt::zero), the way bluscript_zero.cpp does it. REQUIRED for bf16 + // at world_size>1: DataParallel forces gradient buckets to fp32 regardless of + // parameter dtype (Data-Parallel/src/DataParallel.cpp:300), so bf16 grads and + // fp32 buckets are the same mixed-dtype hazard that was MEASURED to corrupt + // the init_sync parameter broadcast. Symptom was a correct step-0 forward + // (bf16 within 4e-3 of fp32) followed by training 2 nats WORSE by step 175 + // and a blow-up at step 700. ZeRO-1 owns its own reduction, which is why the + // team's bluscript_zero trains in bf16 without this failure. + // + // DataParallel and ZeRO must NEVER both be active: each reduces the + // gradients, so running both double-reduces them. + bool zero = false; + bool zero_v2 = false; // CP_ZERO_V2: SymmetricShardedOptimizer + bool zero_fp32_reduce = false; // CP_ZERO_FP32_REDUCE: reduce grads in fp32 + bool bf16 = false; + // CP_BF16_MASK selects WHICH sites go bf16, for bisecting a numerical failure. + // Every site casts its own output back to fp32, so the sites are independent + // and can be enabled one at a time. Default 15 = all. + // 1 = QKV projections (wQ/wK/wV) 2 = wO + // 4 = MLP (gate_up/down) 8 = lm_head / tied output matmul + int bf16_mask = 15; + // WSD stable phase: steps held at max_lr between warmup and the cosine tail + // (CP_STABLE). 0 = the pure warmup+cosine schedule this script has always + // used, so the default changes nothing. + int stable_steps = 0; + int max_steps = 917; // ~10B tokens at 0.5M tokens/step + + int val_freq = 250; + int val_loss_steps = 20; + + // ---- context parallel (env-resolved in main) ---- + int cp_size = 2; // sequence-parallel group size (CP_SIZE) + std::string attn_mode = "ulysses"; // "ring" | "ulysses" | "hybrid" (CP_ATTN_MODE) + RotatorType rotator = RotatorType::P2P; // CP_ROTATOR + bool ring_mode = true; // (attn_mode == "ring") -> HeadTail LB + // "fused" (default) -> the existing single-kernel WMMA ring path + // (GQA_fused_*_sm103_cp.cu). "split" -> Stage-2 standalone QK-norm+RoPE + // once per rank + Stage-3 plain Hopper causal-attention per ring step + // (SplitAttnSDPA_cp.h). Only meaningful for attn_mode=="ring" -- Ulysses + // still calls the separate gqa_fused_flash_attn_forward kernel entirely + // (out of scope for this flag; see .mdfiles/HOPPER_MIGRATION_PLAN.md + // section 4). CP_ATTN_FUSION env var. + std::string attn_fusion = "fused"; + // CP_ULYSSES_CPLOCAL=1 (additive): route forward_ulysses_fused at the + // CP-LOCAL fused kernel instead of libtensor gqa_fused_flash_attn_forward, + // which is Blackwell-only after the TI upgrade and fails on Hopper with + // "no kernel image is available for execution on the device". Unlike the + // Stage-3 (CP_ATTN_FUSION=split) kernels this supports hd 64 AND 128. + bool ulysses_cplocal = false; + // ---- hybrid 3-D CP (Ulysses-inner + Ring-outer); active when attn_mode=="hybrid" ---- + int ulysses_size = 1; // CP_ULYSSES_SIZE (inner all-to-all degree) + bool hybrid_mode = false; // (attn_mode == "hybrid") + + // ---- CREAM context-extension position relabeling (additive; env-resolved) ---- + // "off" (default) => byte-identical to prior behavior. "cream"|"pose"|"randpos" + // feed manipulated RoPE position LABELS via a gathered cos/sin cache. Single-GPU + // (CP_SIZE=1) only for now. factor(=YARN_SCALE) and scaled_max(=factor*T) derived. + std::string cream_mode = "off"; // CP_CREAM_MODE + double cream_sigma = 3.0; // CP_CREAM_SIGMA (truncated-Gaussian) + uint64_t cream_seed = 0; // CP_CREAM_SEED base (0 => step-derived) + int cream_factor = 0; // = YARN_SCALE (resolved in main) + int64_t cream_scaled_max = 0; // = cream_factor * T (resolved in main) + + // ---- LongRoPE searched-cache extension (additive; env-resolved) ---- + // "" (default) => off, byte-identical. When set, a searched per-dim rescale + // vector + n_hat build a LongRoPE cos/sin cache (context_parallel/LongRoPEOps.h). + // LongRoPE-alone installs the cache directly; composed with CREAM it becomes the + // gather source (must satisfy longrope_search_len == cream_scaled_max, s==factor). + std::string longrope_factors = ""; // CP_LONGROPE_FACTORS (path) + std::vector longrope_lambda; // parsed per-dim factors (len head_dim/2) + int64_t longrope_n_hat = 0; // initial-token threshold + int longrope_s = 0; // extension ratio recorded in file + int64_t longrope_search_len = 0; // S_search recorded in file + float longrope_mscale = 1.0f; // CP_LONGROPE_MSCALE: 1.0=paper (no temp), + // else YaRN m=0.1*ln(s)+1 baked in (fair-vs-YaRN) + + // ---- Ms-PoE (Zhang et al. 2024, arXiv:2403.04797) ----------------------- + // Per-KV-GROUP position rescaling: group g rotates at + // lambda_g[i] = 1 + (lambda_base[i] - 1) * w_g + // so w_g=1 reproduces the base LongRoPE/NTK vector (full reach) and w_g->0 + // approaches original, un-extended positions (short-range precision). That + // is the paper's "some heads keep near-original positions, others get + // heavily compressed ones", expressed in our lambda parameterisation. + // + // The ratio is per KV GROUP, not per head: RoPE's relative-position + // property needs Q and K rotated at the SAME rate, and under GQA one K head + // serves q_heads/kv_heads query heads. So the number of distinct ratios a + // checkpoint can express is kv_heads (an MQA checkpoint cannot express + // Ms-PoE at all). See LongRoPEOps.h. + // + // CP_MSPOE_WEIGHTS="w0,w1,..." explicit per-group weights (count==kv_heads) + // CP_MSPOE_WMIN= with no explicit weights: kv_heads weights + // evenly spaced over [WMIN, 1] (default 0.25) + // Requires CP_LONGROPE_FACTORS (the base lambda) and is FORWARD-ONLY. + std::vector mspoe_weights; // empty => weight mode off + + // FAITHFUL Ms-PoE (matches VITA-Group/Ms-PoE utils/modify_arch/llama.py): + // that implementation divides the POSITION uniformly, `t = t / ratio_h`, then + // builds freqs from the untouched inv_freq -- so the ratio scales every + // frequency dimension identically (PI-style), and needs NO base lambda: it + // compresses positions BELOW the native window to fight RoPE decay inside + // it. In our per-dim parameterisation that is lambda_g[i] = r_g, constant + // over i. Ratios come from CP_MSPOE_RATIOS (explicit) or the repo's ladder + // r_g = rmin + (rmax - rmin) * (g / n_groups) + // via CP_MSPOE_RMIN / CP_MSPOE_RMAX (repo defaults 1.2 / 1.8; note the repo + // divides by n_groups, not n_groups-1, so rmax is approached, never reached). + // Mutually exclusive with mspoe_weights. + std::vector mspoe_ratios; // empty => ratio mode off + // Layers Ms-PoE applies to. The repo passes apply_layers "2,...,31" for a + // 32-layer model, i.e. the first two layers keep unscaled positions. Empty + // => every layer. + std::vector mspoe_layers; + + // ---- checkpointing ---- + bool checkpointing = true; + int ckpt_freq = 50; + int ckpt_keep = 6; + std::string ckpt_dir = "checkpoints_bluscriptcp"; + std::string ckpt_prefix = "blumodelcp"; +}; + +// ============================================================================= +// Shared advancing RNG for custom weight init (unchanged from bluscript.cpp). +// ============================================================================= +struct InitRng { + std::mt19937 gen; + explicit InitRng(uint64_t seed) : gen(static_cast(seed)) {} + Tensor normal(const Shape& shape, float std) { + Tensor t(shape, TensorOptions().with_dtype(Dtype::Float32)); + std::normal_distribution dist(0.0f, std); + float* d = t.data(); + const size_t n = t.numel(); + for (size_t i = 0; i < n; ++i) d[i] = dist(gen); + return t; + } +}; + +static void copy_into_param(Tensor& dst, const Tensor& src) { + const bool was = dst.requires_grad(); + if (was) dst.set_requires_grad(false); + // Cast the fp32 initialiser draw to the parameter's dtype. Required once + // CP_BF16=1 makes the Linears bf16: the draw is always fp32 (InitRng fills a + // Float32 tensor on the host), and a bare copy_ across dtypes is not a + // conversion. bluscript_zero.cpp:156 guards this the same way. + dst.copy_(src.dtype() == dst.dtype() ? src : src.as_type(dst.dtype())); + if (was) dst.set_requires_grad(true); +} +// Mixed-precision boundary helpers. Both are exact no-ops when bf16 is false, so +// the fp32 graph is unchanged shape-for-shape and node-for-node -- CP_BF16=0 does +// not merely compute the same values, it builds the same autograd graph. +static inline Tensor to_compute(const Tensor& t, bool bf16) { + return bf16 ? autograd::to_dtype(t, Dtype::Bfloat16) : t; +} +static inline Tensor to_master(const Tensor& t, bool bf16) { + return bf16 ? autograd::to_dtype(t, Dtype::Float32) : t; +} +static inline Dtype linear_dtype(bool bf16) { + return bf16 ? Dtype::Bfloat16 : Dtype::Float32; +} +// Per-site gate: bf16 at this site only when CP_BF16=1 AND its mask bit is set. +enum : int { BF16_QKV = 1, BF16_WO = 2, BF16_MLP = 4, BF16_HEAD = 8 }; +static inline bool bf16_at(bool bf16, int mask, int bit) { return bf16 && (mask & bit); } + +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()) { + 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 -> CP fused (QK-norm + RoPE + causal GQA) -> Wo -> +x +// The attention (RoPE + QK-norm + GQA + causal) runs INSIDE ContextParallel via +// the fused ring (enable_rope) or Ulysses (enable_ulysses_fused) path. x is the +// LOCAL sequence shard [B, T_local, d_model]; forward_cp keeps the output local. +// ============================================================================= +class CausalGQA : public nn::Module { +public: + nn::RMSNorm norm; + nn::Linear wQ, wK, wV, wO; + Tensor q_gamma, k_gamma; + Tensor cos_sin_cache; + ModelConfig cfg; + std::shared_ptr cp_; + + // CP_GENERATE (additive): during decoding GPT::forward hands this block the + // FULL [B,T,d_model] sequence instead of the local shard, so CP must shard + // q/k/v itself and gather the output back. Default false => training and + // CP_EVAL_PPL are byte-identical to before this flag existed. + bool generation_mode_ = false; + void set_generation_mode(bool g) { generation_mode_ = g; } + + // cp_pg is the CP process group. In hybrid mode it is the RING sub-group and + // ulysses_pg carries the inner all-to-all; otherwise ulysses_pg is null. + CausalGQA(const ModelConfig& c, DeviceIndex dev, const Tensor& cache, InitRng& rng, + const DeviceMesh& mesh, std::shared_ptr cp_pg, + std::shared_ptr ulysses_pg = nullptr) + : norm(c.d_model, c.rms_eps), + wQ(c.d_model, c.q_heads * c.head_dim, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_QKV))), + wK(c.d_model, c.kv_heads * c.head_dim, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_QKV))), + wV(c.d_model, c.kv_heads * c.head_dim, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_QKV))), + wO(c.q_heads * c.head_dim, c.d_model, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_WO))), + cos_sin_cache(cache), + cfg(c) + { + const float resid_scale = 1.0f / std::sqrt(2.0f * static_cast(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); + + auto gopts = TensorOptions().with_dtype(Dtype::Float32).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); + + // ---- Context-parallel attention layer ---- + // ring => HeadTail load-balance (matches shard_sequence_pre_embed's perm); + // ulysses=> contiguous shard; + // hybrid => Ulysses-inner (ulysses_pg) nested in Ring-outer (cp_pg==ring pg), + // HeadTail on the ring axis. RoPE/QK-norm/GQA applied INSIDE CP. + const float attn_scale = 1.0f / std::sqrt(static_cast(c.head_dim)); + const bool lb = (c.ring_mode || c.hybrid_mode); + cp_ = std::make_shared( + mesh, cp_pg, attn_scale, /*is_causal=*/true, c.rotator, + /*load_balance=*/lb); + if (c.hybrid_mode) { + cp_->enable_hybrid(ulysses_pg, c.ulysses_size, cos_sin_cache, + q_gamma, k_gamma, c.qk_eps); + } else if (c.ring_mode) { + cp_->enable_rope(cos_sin_cache, q_gamma, k_gamma, c.qk_eps, + /*split_attn=*/c.attn_fusion == "split"); + } else { + cp_->enable_ulysses_fused(cos_sin_cache, q_gamma, k_gamma, c.qk_eps, + c.rope_interleaved, + /*cplocal=*/c.ulysses_cplocal); + } + } + + // x: LOCAL sequence shard [B, T_local, d_model] + Tensor forward(const Tensor& x) override { + const int64_t B = x.shape().dims[0]; + const int64_t Tl = x.shape().dims[1]; + + // Norm in fp32, projections in the compute dtype, then back to fp32: + // forward_cp's QK-norm gammas and the RoPE cos/sin cache are fp32, and + // FusedRoPESDPA.h:110 does its own as_type(Bfloat16) at the kernel edge. + const bool bf_qkv = bf16_at(cfg.bf16, cfg.bf16_mask, BF16_QKV); + const bool bf_wo = bf16_at(cfg.bf16, cfg.bf16_mask, BF16_WO); + Tensor hn = norm.forward(x); + Tensor h = to_compute(hn, bf_qkv); + Tensor q = to_master(wQ.forward(h), bf_qkv); // [B, Tl, Nq*hd] + Tensor k = to_master(wK.forward(h), bf_qkv); // [B, Tl, Nkv*hd] + Tensor v = to_master(wV.forward(h), bf_qkv); // [B, Tl, Nkv*hd] + + // -> head-major [B, heads, T_local, hd] + q = autograd::transpose(autograd::reshape(q, Shape{{B, Tl, cfg.q_heads, cfg.head_dim}}), 1, 2); + k = autograd::transpose(autograd::reshape(k, Shape{{B, Tl, cfg.kv_heads, cfg.head_dim}}), 1, 2); + v = autograd::transpose(autograd::reshape(v, Shape{{B, Tl, cfg.kv_heads, cfg.head_dim}}), 1, 2); + + // CP fused attention: RoPE + QK-norm + causal GQA inside the ring/ulysses + // path. Training / CP_EVAL_PPL: pre_sharded=true (q/k/v are already the + // local shard); unshard=false (output stays [B, Nq, T_local, hd], no + // all-gather). + // + // CP_GENERATE: q/k/v are the FULL sequence, so CP shards them itself + // (pre_sharded=false) and all-gathers the output (unshard=true). RoPE + // stays correct because RopeDeltas.h derives global positions from + // (rank, ring step, world_size, T_local) -- its documented "contiguous + // CP" degenerate case -- not from the caller's HeadTail perm. When + // T is not divisible by 2*world_size, forward_cp's lb_div_ok check + // (ContextParallel.h:461) drops to contiguous sharding for that call and + // the deltas follow, so short decode steps stay correct too. + Tensor attn = cp_->forward_cp(q, k, v, + /*unshard=*/generation_mode_, + /*pre_sharded=*/!generation_mode_); + + Tensor merged = autograd::reshape(autograd::transpose(attn, 1, 2), + Shape{{B, Tl, cfg.q_heads * cfg.head_dim}}); + // to_compute here rather than assuming forward_cp's output dtype: it + // returns fp32 today, and if that ever changes to bf16 the cast is a + // no-op instead of a silent dtype mismatch at wO. + Tensor proj = to_master(wO.forward(to_compute(merged, bf_wo)), bf_wo); + return autograd::add(x, proj); // fp32 residual + } +}; + +// ============================================================================= +// SwiGLU feed-forward (unchanged from bluscript.cpp; operates on the local shard) +// ============================================================================= +class SwiGLUMLP : public nn::Module { +public: + nn::RMSNorm norm; + nn::Linear gate_up; + nn::Linear down; + bool bf16_ = false; // SwiGLUMLP does not keep a ModelConfig copy + + SwiGLUMLP(const ModelConfig& c, DeviceIndex dev, InitRng& rng) + : norm(c.d_model, c.rms_eps), + gate_up(c.d_model, 2 * c.ffn_hidden, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_MLP))), + down(c.ffn_hidden, c.d_model, /*bias=*/false, linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_MLP))), + bf16_(bf16_at(c.bf16,c.bf16_mask,BF16_MLP)) + { + const float resid_scale = 1.0f / std::sqrt(2.0f * static_cast(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 { + // norm fp32 -> bf16 gate_up -> back to fp32 -> swiglu in FP32 -> bf16 down + // -> fp32 residual. Only the two MATMULS are bf16. + // + // swiglu MUST run in fp32. Leaving it in bf16 (which is what I did first) + // gives NaN/Inf loss at step 0: swiglu is swish(gate)*up, so it carries a + // sigmoid and a product of two activations, and bf16's 8 mantissa bits + // blow that up immediately. bluscript_zero.cpp casts back to fp32 for + // exactly this reason. + Tensor h = to_compute(norm.forward(x), bf16_); + Tensor gu = to_master(gate_up.forward(h), bf16_); + Tensor act = autograd::swiglu(gu); + Tensor out = to_master(down.forward(to_compute(act, bf16_)), bf16_); + return autograd::add(x, out); + } +}; + +// ============================================================================= +// Full model. GPT::forward shards the SEQUENCE pre-embed (over the CP axis), +// then runs all blocks on the local [B, T_local] shard. No wpe (RoPE is applied +// inside CP). Weight-tied lm_head. +// ============================================================================= +class GPT : public nn::Module { +public: + ModelConfig cfg; +private: + InitRng init_rng_; +public: + nn::Embedding wte; + Tensor cos_sin_cache; + // ADDITIVE (CREAM): full YaRN-scaled cos/sin cache of length scaled_max, built + // ONLY when cream_mode != "off". Used solely as the gather SOURCE for per-step + // relabeled caches; never passed to the kernel directly. Invalid/empty when off. + Tensor cream_full_cache_; + std::vector> attn; + std::vector> mlp; + nn::RMSNorm norm_f; + std::shared_ptr lm_head; + // CP sharding coordinate (over the CP mesh axis). In hybrid mode the CP axis + // is factored into (ring_rank_, ulysses_rank_); cp_size_/cp_rank_ hold the + // composite CP degree/coord and drive the non-hybrid single-axis sharder. + int cp_size_, cp_rank_; + int ring_rank_, ulysses_rank_; + bool generation_mode_ = false; // CP_GENERATE; see set_generation_mode() + + GPT(const ModelConfig& c, DeviceIndex dev, const DeviceMesh& mesh, + std::shared_ptr cp_pg, int cp_size, int cp_rank, + std::shared_ptr ulysses_pg = nullptr, + int ring_rank = 0, int ulysses_rank = 0, + uint64_t seed = 1234) + : cfg(c), + init_rng_(seed), + wte(static_cast(c.vocab_size), static_cast(c.d_model)), + cos_sin_cache(autograd::build_rope_cache(c.context_length, c.head_dim, + static_cast(c.rope_theta), dev)), + norm_f(c.d_model, c.rms_eps), + cp_size_(cp_size), cp_rank_(cp_rank), + ring_rank_(ring_rank), ulysses_rank_(ulysses_rank) + { + init_embedding(wte, 0.02f, init_rng_); + wte.to(dev); + norm_f.to(dev); + + // CREAM (additive): build the full YaRN-scaled cache (length scaled_max) as + // the gather source. build_rope_cache reads YARN_SCALE from env, so this row + // `pos` holds the YaRN-interpolated cos/sin for absolute position `pos`. Only + // built when CREAM is on; the base cos_sin_cache above is left untouched. + if (c.cream_mode != "off") { + if (!c.longrope_factors.empty()) { + // COMPOSITION: gather source is the searched LongRoPE cache (built at + // cream_scaled_max; coupling asserted in main). Replaces the YaRN source. + cream_full_cache_ = autograd::build_rope_cache_longrope( + c.cream_scaled_max, c.head_dim, static_cast(c.rope_theta), + c.longrope_lambda, c.longrope_n_hat, dev, c.longrope_mscale); + } else { + cream_full_cache_ = autograd::build_rope_cache( + c.cream_scaled_max, c.head_dim, + static_cast(c.rope_theta), dev); + } + } + + for (int i = 0; i < c.n_layers; ++i) { + auto a = std::make_shared(c, dev, cos_sin_cache, init_rng_, mesh, cp_pg, ulysses_pg); + auto m = std::make_shared(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(static_cast(c.d_model), + static_cast(c.vocab_size), /*bias=*/false, + linear_dtype(bf16_at(c.bf16,c.bf16_mask,BF16_HEAD))); + 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()); + + // LongRoPE-ALONE (no CREAM): the searched cache is STATIC (no per-step + // relabeling), so build it once at context_length and install on every layer + // here. Composed-with-CREAM is handled by the gather source above + the + // per-step install. Off (empty factors) leaves the base cache untouched. + const bool mspoe_on = !c.mspoe_weights.empty() || !c.mspoe_ratios.empty(); + if (c.cream_mode == "off" && (mspoe_on || !c.longrope_factors.empty())) { + const int64_t half = c.head_dim / 2; + + // The non-Ms-PoE plane every UNSELECTED layer keeps: the searched + // LongRoPE cache when factors were given, otherwise the base cache + // the constructor already installed (nothing to do for those layers). + Tensor shared; + if (!c.longrope_factors.empty()) + shared = autograd::build_rope_cache_longrope( + c.context_length, c.head_dim, static_cast(c.rope_theta), + c.longrope_lambda, c.longrope_n_hat, dev, c.longrope_mscale); + + Tensor grouped; + if (mspoe_on) { + std::vector> lam_g; + if (!c.mspoe_ratios.empty()) { + // FAITHFUL: r_g constant over dims == the repo's t/ratio. + for (float r : c.mspoe_ratios) + lam_g.emplace_back(static_cast(half), r); + } else { + // COMPOSED variant: scale the base lambda toward 1 per group. + for (float w : c.mspoe_weights) { + std::vector lam(c.longrope_lambda.size()); + for (size_t i = 0; i < lam.size(); ++i) + lam[i] = 1.0f + (c.longrope_lambda[i] - 1.0f) * w; + lam_g.push_back(std::move(lam)); + } + } + grouped = autograd::build_rope_cache_longrope_grouped( + c.context_length, c.head_dim, static_cast(c.rope_theta), + lam_g, c.longrope_n_hat, dev, c.longrope_mscale); + } + + auto layer_selected = [&](int li) { + if (c.mspoe_layers.empty()) return true; + return std::find(c.mspoe_layers.begin(), c.mspoe_layers.end(), li) + != c.mspoe_layers.end(); + }; + for (size_t li = 0; li < attn.size(); ++li) { + const bool use_grouped = mspoe_on && layer_selected(static_cast(li)); + if (use_grouped) attn[li]->cp_->set_rope_cache(grouped); + else if (shared.is_valid()) attn[li]->cp_->set_rope_cache(shared); + } + } + } + + // CP_GENERATE (additive): switch every CP block to full-sequence mode so the + // decode loop gets logits for all T positions. Call with true before the + // generation loop and false after; training / CP_EVAL_PPL never touch this. + void set_generation_mode(bool g) { + generation_mode_ = g; + for (auto& a : attn) a->set_generation_mode(g); + } + + // idx [B, T] (full, UInt16/Int) -> logits [B, T_local, vocab] on the CP shard, + // or [B, T, vocab] (full) when generation_mode_ is set. + Tensor forward(const Tensor& idx) override { + Tensor x; + if (generation_mode_) { + // Decoding: idx grows one token per step, so read its ACTUAL length + // rather than cfg.T (which the sharded path below hard-codes because + // training windows are always exactly T). No pre-embed shard: each + // CausalGQA shards q/k/v itself and gathers the output back. + x = wte.forward(idx); // [B, T, d_model] + } else { + // Shard the sequence over the CP axis (HeadTail for ring, contiguous for + // ulysses, composed HeadTail-over-ring + contiguous-over-ulysses for + // hybrid). Targets are sharded with the SAME perm in the loss (main). + Tensor empty_y; + ShardedInputs sh; + if (cfg.hybrid_mode) { + const int ring_size = cfg.cp_size / cfg.ulysses_size; + sh = shard_sequence_pre_embed_hybrid( + idx, empty_y, cfg.T, ring_size, ring_rank_, + cfg.ulysses_size, ulysses_rank_, idx.device()); + } else { + sh = shard_sequence_pre_embed( + idx, empty_y, cfg.T, cp_size_, cp_rank_, /*load_balance=*/cfg.ring_mode, + idx.device()); + } + + x = wte.forward(sh.idx_local); // [B, T_local, d_model] (NO wpe) + } + for (int i = 0; i < cfg.n_layers; ++i) { + x = attn[i]->forward(x); + x = mlp[i]->forward(x); + } + x = norm_f.forward(x); + + // Output projection is the model's largest matmul (d_model x vocab), so it + // is the biggest single beneficiary of bf16. The embedding itself stays + // fp32 -- it is a gather, not a matmul, and the tied weight is the + // parameter AdamW updates, so keeping it fp32 keeps the autograd edge in + // fp32 and lets the master-weight path stay untouched for this tensor. + // Logits come back to fp32 for sparse_cross_entropy_loss. + if (cfg.weight_tying) { + Tensor w_t = autograd::transpose(wte.weight, 0, 1); // fp32 [d_model, vocab] + const bool bf_head = bf16_at(cfg.bf16, cfg.bf16_mask, BF16_HEAD); + Tensor logits = autograd::matmul(to_compute(x, bf_head), + to_compute(w_t, bf_head)); + return to_master(logits, bf_head); // [B, T_local, vocab] + } + const bool bf_h2 = bf16_at(cfg.bf16, cfg.bf16_mask, BF16_HEAD); + return to_master(lm_head->forward(to_compute(x, bf_h2)), bf_h2); + } + + // CREAM (additive): generate manipulated RoPE position LABELS for this step and + // install a gathered cos/sin cache on EVERY CP layer. The gather is done ONCE + // here and fanned out to all layers -- do NOT move it inside a per-layer loop + // (it would cost n_layers x for nothing). No-op when cream_mode=="off"; call + // before forward() when CREAM is active. Single-GPU (CP_SIZE=1): local_idx==pos, + // so cream_cache[local_idx] rotates each token at its label L[local_idx]. + void install_cream_positions(uint64_t seed) { + if (cfg.cream_mode == "off") return; + const int T = static_cast(cfg.T); + const int factor = cfg.cream_factor; + std::vector L; + if (cfg.cream_mode == "cream") + L = OwnTensor::cp::cream::generate_cream_positions(T, factor, cfg.cream_sigma, seed); + else if (cfg.cream_mode == "pose") + L = OwnTensor::cp::cream::generate_pose_positions(T, factor, seed); + else if (cfg.cream_mode == "randpos") + L = OwnTensor::cp::cream::generate_randpos_positions(T, factor, seed); + else + throw std::runtime_error("unknown CP_CREAM_MODE: " + cfg.cream_mode); + + // [T, hd] gather index: row i = L[i] broadcast across the hd columns so + // gather(dim=0) selects whole cache rows: out[i][j] = full_cache[L[i]][j]. + const int64_t hd = cfg.head_dim; + Tensor idx_cpu(Shape{{static_cast(T), hd}}, + TensorOptions().with_dtype(Dtype::Int64)); + int64_t* ip = static_cast(idx_cpu.data()); + for (int i = 0; i < T; ++i) + for (int64_t j = 0; j < hd; ++j) + ip[static_cast(i) * hd + j] = L[static_cast(i)]; + Tensor idx_dev = idx_cpu.to(cream_full_cache_.device()); + + Tensor cream_cache = OwnTensor::gather(cream_full_cache_, /*dim=*/0, idx_dev); + for (auto& a : attn) a->cp_->set_rope_cache(cream_cache); + } + + // CREAM diagnostics: the middle-block START position sampled for a seed, so the + // training loop can log its distribution (verification step 4 -- confirm the + // truncated-Gaussian shape). head==tail, so head = the leading run [0,1,2,...]; + // the middle block starts at the first label that breaks that run. -1 if not + // cream mode. + int64_t cream_middle_start(uint64_t seed) const { + if (cfg.cream_mode != "cream") return -1; + auto L = OwnTensor::cp::cream::generate_cream_positions( + static_cast(cfg.T), cfg.cream_factor, cfg.cream_sigma, seed); + size_t head = 0; + while (head < L.size() && L[head] == static_cast(head)) ++head; + return (head < L.size()) ? L[head] : -1; + } +}; + +// ============================================================================= +// Cosine LR schedule (unchanged). +// ============================================================================= +// WSD LR: warmup -> stable (held at max_lr) -> cosine decay to min_lr. +// `warmup_start` anchors the warmup ramp: 0 for a fresh run, or the resume step +// for a 2-stage re-warmup on context extension (ramp +// [warmup_start, warmup_start+warmup), then stable, then decay by max_steps). +// `max_lr` is the (possibly reduced) peak reached at the top of the ramp. +// +// `stable` (CP_STABLE) DEFAULTS TO 0, which collapses this to the pure +// warmup+cosine this function computed before -- same value at every step, so +// existing runs and the recorded run100 / step-500 probe baselines stay +// comparable. Ported from bluscript_zero.cpp:395. +// +// Why WSD matters here beyond final loss: cosine bakes the horizon into every +// step, so extending a run (4k -> 16k context) means either re-warming from a +// decayed LR or discarding the schedule. A stable phase gives a flat region to +// branch from -- extend at constant LR, then decay once at the end. +static float get_lr(int step, float max_lr, float min_lr, int warmup, int max_steps, + int warmup_start = 0, int stable = 0) { + if (max_lr < min_lr) max_lr = min_lr; // guard: absurdly low peak -> constant min + if (stable < 0) stable = 0; + const int local = step - warmup_start; // steps since the (re)warmup anchor + if (warmup > 0 && local < warmup) + return max_lr * static_cast(local + 1) / static_cast(warmup); + if (step >= max_steps) return min_lr; + const int decay_start = warmup_start + warmup + stable; + if (step < decay_start) return max_lr; // stable phase (no-op when stable==0) + float denom = static_cast(max_steps - decay_start); + if (denom <= 0.0f) denom = 1.0f; // guard: warmup+stable past max_steps + const float decay = static_cast(step - decay_start) / denom; + const float coeff = 0.5f * (1.0f + std::cos(M_PI * decay)); + return min_lr + coeff * (max_lr - min_lr); +} + +static int64_t env_i64(const char* k, int64_t d) { + const char* v = std::getenv(k); return v ? std::atoll(v) : d; +} + +// Parse a LongRoPE factors file. Format (whitespace-separated, keys in any order): +// n_hat +// s (extension ratio the search used) +// S_search (target length the search targeted; = factor*T when composed) +// lambda ... (head_dim/2 monotone-nondecreasing values >= 1) +// '#'-prefixed tokens skip to end of line (comments). Throws on any malformation. +static void parse_longrope_factors(const std::string& path, int64_t half, + std::vector& lambda, int64_t& n_hat, + int& s, int64_t& search_len) { + std::ifstream f(path); + if (!f) throw std::runtime_error("CP_LONGROPE_FACTORS: cannot open '" + path + "'"); + lambda.clear(); n_hat = -1; s = -1; search_len = -1; + std::string tok; + while (f >> tok) { + if (tok == "n_hat") { if (!(f >> n_hat)) throw std::runtime_error("CP_LONGROPE_FACTORS: bad n_hat"); } + else if (tok == "s") { if (!(f >> s)) throw std::runtime_error("CP_LONGROPE_FACTORS: bad s"); } + else if (tok == "S_search") { if (!(f >> search_len)) throw std::runtime_error("CP_LONGROPE_FACTORS: bad S_search"); } + else if (tok == "lambda") { + lambda.resize(static_cast(half)); + for (int64_t i = 0; i < half; ++i) + if (!(f >> lambda[static_cast(i)])) + throw std::runtime_error("CP_LONGROPE_FACTORS: expected " + + std::to_string(half) + " lambda values"); + } + else if (!tok.empty() && tok[0] == '#') { std::string rest; std::getline(f, rest); } + else throw std::runtime_error("CP_LONGROPE_FACTORS: unexpected token '" + tok + "'"); + } + if (static_cast(lambda.size()) != half) + throw std::runtime_error("CP_LONGROPE_FACTORS: missing/short 'lambda' (need " + + std::to_string(half) + ")"); + if (n_hat < 0) throw std::runtime_error("CP_LONGROPE_FACTORS: missing 'n_hat'"); + for (int64_t i = 0; i < half; ++i) { + if (!(lambda[static_cast(i)] >= 1.0f)) + throw std::runtime_error("CP_LONGROPE_FACTORS: lambda must be >= 1"); + if (i > 0 && lambda[static_cast(i)] < lambda[static_cast(i - 1)]) + throw std::runtime_error("CP_LONGROPE_FACTORS: lambda must be non-decreasing"); + } +} + +// ============================================================================= +// Training +// ============================================================================= +int main(int argc, char** argv) { + MPI_Init(&argc, &argv); + 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); + + // Only rank 0 writes to stdout. Redirect std::cout on every other rank to a + // discard buffer so un-gated library chatter (CheckpointManager auto-load, CP + // ring/overlap notices, etc.) is emitted once instead of world_size times. + // std::cerr is left intact so genuine errors/warnings from any rank still show. + static NullBuf global_null_buf; + if (!is_master) std::cout.rdbuf(&global_null_buf); + + if (is_master) + std::cout << "=== Llama Context Parallel Training Script (bluscriptCP) ===" + << std::endl; + + // CP_DEV_OVERRIDE pins every rank to one physical device, so a rank-dependent + // bug can be separated from a device-index-dependent one. + int cp_dev = rank; + if (const char* dv = std::getenv("CP_DEV_OVERRIDE")) cp_dev = std::atoi(dv); + cudaSetDevice(cp_dev); + DeviceIndex device(Device::CUDA, cp_dev); + + ModelConfig cfg; + + // ---- context-parallel topology from env ---- + cfg.cp_size = static_cast(env_i64("CP_SIZE", world_size)); // default CP-only + if (const char* m = std::getenv("CP_ATTN_MODE")) cfg.attn_mode = m; + cfg.ring_mode = (cfg.attn_mode == "ring"); + cfg.hybrid_mode = (cfg.attn_mode == "hybrid"); + cfg.ulysses_size = static_cast(env_i64("CP_ULYSSES_SIZE", 1)); + if (const char* f = std::getenv("CP_ATTN_FUSION")) cfg.attn_fusion = f; + // Ulysses kernel dispatch is decided by the ARCH, not left to an env var. + // forward_ulysses_fused's default leg calls libtensor's + // gqa_fused_flash_attn_forward, which after the TI upgrade ships only + // sm_100a/sm_103a images (GQA_fused_*_sm103.cu). On anything below Blackwell + // that launch can only fail with "no kernel image is available", and because + // TI reports the sticky error at the NEXT check it surfaced as a bogus + // "convert_cuda_bf16_to_fp32_sm90 kernel launch failed" -- which is why + // ulysses appeared broken in every configuration, single-GPU included. + // The CP-local leg is a WMMA fork of the same flash core, compiles hd 64 and + // 128, and is arch-portable, so it is the correct default on sm_90. + // CP_ULYSSES_CPLOCAL still overrides in both directions (=0 forces the + // libtensor leg, for Blackwell parity testing). + { + int dev = 0, major = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev); + const bool blackwell_or_newer = (major >= 10); + cfg.ulysses_cplocal = !blackwell_or_newer; + if (const char* u = std::getenv("CP_ULYSSES_CPLOCAL")) + cfg.ulysses_cplocal = (u[0] == '1'); + if (is_master && !cfg.ring_mode) + std::cout << "[CP ulysses kernel] sm_" << major << "0 -> " + << (cfg.ulysses_cplocal + ? "CP-LOCAL fused (gqa_fused_rope_cp_forward, deltas=0), arch-portable" + : "libtensor gqa_fused_flash_attn (Blackwell-only images)") + << std::endl; + } + if (is_master && cfg.ring_mode) + std::cout << "[CP attn fusion] " << cfg.attn_fusion + << (cfg.attn_fusion == "split" + ? " (Stage-2 rope-once + Stage-3 Hopper causal-attn kernels)" + : " (single-kernel WMMA fused RoPE+QKnorm+attention)") + << "\n"; + if (is_master) { + if (cfg.hybrid_mode) + std::cout << "[CP attn mode] HYBRID (Ulysses-inner all-to-all + Ring-outer " + "P2P); HeadTail load-balance on the ring axis\n"; + else if (!cfg.ring_mode) + std::cout << "[CP attn mode] ULYSSES (all-to-all); load_balancing forced " + "off (contiguous sharding)\n"; + } + std::string rotator_label = "p2p"; + if (const char* e = std::getenv("CP_ROTATOR")) { + std::string rv = e; + if (rv == "p2p" || rv == "P2P") { cfg.rotator = RotatorType::P2P; rotator_label = "p2p"; } + else if (rv == "allgather" || rv == "AllGather") { cfg.rotator = RotatorType::AllGather; rotator_label = "allgather"; } + else { cfg.rotator = RotatorType::AlltoAll; rotator_label = "alltoall"; } + } + // Under pure Ulysses the ring rotator is unused (all-to-all attention); tag the + // snapshot with "ulysses" so it matches the sweep harness's filename. Hybrid + // DOES use the ring rotator (outer axis), so it keeps the rotator label. + if (!cfg.ring_mode && !cfg.hybrid_mode) rotator_label = "ulysses"; + + // ---- memory-occupancy probe (mirrors gpt2_cp_test.cpp) ---- + // CP_MEM_PROBE=1 : run CP_MEM_PROBE_STEPS steps (grad_accum forced to 1), + // skip validation + checkpointing, snapshot nvidia-smi + + // cudaMemGetInfo + allocator stats after the last step, + // then exit. Used by Tests/bluscriptcp/mem_scaling_sweep.sh. + // CP_MEM_PROBE_STEPS : steps in probe mode (default 2; must be >=2 so the + // optimizer-state allocation is counted). + // CP_MODEL_LABEL : free-text label embedded in the snapshot (e.g. "48M"). + const bool mem_probe = std::getenv("CP_MEM_PROBE") != nullptr && + std::atoi(std::getenv("CP_MEM_PROBE")) != 0; + + // ---- long-context PPL-vs-position eval mode (additive, forward-only) ---- + // CP_EVAL_PPL=1 : load the checkpoint (CP_CKPT_RESUME=), teacher-force + // over CP_EVAL_WINDOWS windows of length T from the given + // split, write pos,count,mean_nll,ppl to CP_EVAL_OUT, exit. + // Requires NP=1 / CP_SIZE=1 so T_local==T (positions are + // global -> no cross-rank NLL gather). Skips training, + // checkpoint saving, and training-log/config writes. + // Build the base-vs-extended comparison by running it per checkpoint+YaRN: + // base ckpt + YARN_SCALE=1 (naive extrapolation @ T) + // extended ckpt + YARN_SCALE=4 ORIG=... (YaRN @ T) + const bool eval_ppl_mode = std::getenv("CP_EVAL_PPL") != nullptr && + std::atoi(std::getenv("CP_EVAL_PPL")) != 0; + const int mem_probe_steps = + std::getenv("CP_MEM_PROBE_STEPS") ? std::atoi(std::getenv("CP_MEM_PROBE_STEPS")) : 2; + const std::string mem_label = + std::getenv("CP_MODEL_LABEL") ? std::getenv("CP_MODEL_LABEL") : "bluscriptcp"; + // Optional config overrides (smoke tests / scaling). + cfg.T = cfg.context_length = env_i64("CP_T", cfg.context_length); + cfg.d_model = env_i64("CP_N_EMBD", cfg.d_model); + cfg.n_layers = env_i64("CP_N_LAYER", cfg.n_layers); + cfg.q_heads = env_i64("CP_N_HEAD", cfg.q_heads); + cfg.kv_heads = env_i64("CP_N_KVHEAD", cfg.kv_heads); + cfg.head_dim = cfg.d_model / cfg.q_heads; + // SwiGLU inner dim: honor CP_FFN, else scale with d_model (8/3*d_model). Was + // previously frozen at the d=384 default (1024) regardless of CP_N_EMBD, which + // silently undersized the FFN for any wider model in the scaling sweeps. + cfg.ffn_hidden = env_i64("CP_FFN", (8 * cfg.d_model) / 3); + // rope_theta is BAKED INTO trained weights but is set by env, so a mismatch + // between train and eval is silent (unlike d_model/heads, which fail loudly on + // a shape mismatch). The value is persisted to a per-run sidecar and checked on + // load -- see the CP_ROPE_THETA precedence block near run_prefix below. + cfg.rope_theta = static_cast(env_i64("CP_ROPE_THETA", cfg.rope_theta)); + if (const char* e = std::getenv("CP_WEIGHT_TYING")) cfg.weight_tying = (e[0] == '1'); + cfg.B = env_i64("CP_B", cfg.B); + cfg.global_batch = env_i64("CP_GLOBAL_BATCH", cfg.global_batch); + cfg.max_steps = static_cast(env_i64("CP_MAX_STEPS", cfg.max_steps)); + cfg.warmup_steps = static_cast(env_i64("CP_WARMUP", cfg.warmup_steps)); + cfg.stable_steps = static_cast(env_i64("CP_STABLE", cfg.stable_steps)); + cfg.bf16 = env_i64("CP_BF16", cfg.bf16 ? 1 : 0) != 0; + cfg.bf16_mask = static_cast(env_i64("CP_BF16_MASK", cfg.bf16_mask)); + // Defaults ON when bf16 is requested at world_size>1, because DDP is broken + // there. Explicit CP_ZERO=0 still forces DDP, for bisecting. + cfg.zero = env_i64("CP_ZERO", (cfg.bf16 && world_size > 1) ? 1 : 0) != 0; + cfg.zero_v2 = env_i64("CP_ZERO_V2", 0) != 0; + cfg.zero_fp32_reduce = env_i64("CP_ZERO_FP32_REDUCE", 0) != 0; + if (cfg.zero && world_size == 1 && is_master) + std::cout << "[ZeRO-1] world_size=1: sharding is a no-op, running for parity" + << std::endl; + + // ---- checkpoint / logging run controls (env) ---- + if (const char* e = std::getenv("CP_CKPT")) cfg.checkpointing = (e[0] == '1'); + // CP_CKPT_DIR: where checkpoints are read/written. Previously hardcoded, which + // made CP_CKPT_WEIGHTS_ONLY unusable -- that flag's whole purpose is seeding a + // new run from a base checkpoint held somewhere else. + if (const char* e = std::getenv("CP_CKPT_DIR")) cfg.ckpt_dir = e; + if (const char* e = std::getenv("CP_CKPT_FREQ")) cfg.ckpt_freq = std::atoi(e); + // CP_CKPT_KEEP: rotation depth. Was hardcoded 6, which for the 600M is 6 x + // 8.18 GB = 49 GB and does not fit beside anything else on a 96 GB volume. + // Frequent checkpoints (FREQ=25) matter more than deep history on a rented + // box, so make the depth tunable rather than trading away the frequency. + if (const char* e = std::getenv("CP_CKPT_KEEP")) cfg.ckpt_keep = std::atoi(e); + // Force a brand-new run number even if resumable checkpoints exist. + const bool ckpt_new_run = + std::getenv("CP_CKPT_NEW_RUN") && std::getenv("CP_CKPT_NEW_RUN")[0] == '1'; + // Resume a specific run number (>=0); -1 = auto-resume latest incomplete run. + int ckpt_resume_run = -1; + if (const char* e = std::getenv("CP_CKPT_RESUME")) ckpt_resume_run = std::atoi(e); + + // 2-stage re-warmup for context extension: on a RESUMING run, re-anchor a fresh + // warmup ramp at the resume step (the base-context schedule's warmup is absolute + // from step 0, so it cannot ramp again). CP_REWARMUP= (0 = off; only + // applies when start_step>0). CP_REWARMUP_PEAK= gives a + // reduced peak for the long-context phase (default 1.0 = same peak). + const int rewarmup_steps = static_cast(env_i64("CP_REWARMUP", 0)); + float rewarmup_peak_frac = 1.0f; + if (const char* e = std::getenv("CP_REWARMUP_PEAK")) rewarmup_peak_frac = std::atof(e); + + // ---- LOUD invariant asserts (silent-failure guards) ---- + auto die = [&](const std::string& msg) { + if (is_master) std::cerr << "ERROR: " << msg << std::endl; + MPI_Abort(MPI_COMM_WORLD, 1); + }; + if (world_size % cfg.cp_size != 0) + die("world_size=" + std::to_string(world_size) + " % CP_SIZE=" + std::to_string(cfg.cp_size) + " != 0"); + // WSD needs a non-empty decay tail. Without this the schedule silently holds + // max_lr to the last step and never anneals, which looks like a converging + // run right up to the point the checkpoint is unusable. bluscript_zero.cpp:486 + // aborts on the same condition. + if (cfg.stable_steps > 0 && cfg.warmup_steps + cfg.stable_steps >= cfg.max_steps - 1) + die("CP_WARMUP(" + std::to_string(cfg.warmup_steps) + ") + CP_STABLE(" + + std::to_string(cfg.stable_steps) + ") leaves no cosine decay tail before " + "CP_MAX_STEPS(" + std::to_string(cfg.max_steps) + ")"); + if (cfg.attn_fusion != "fused" && cfg.attn_fusion != "split") + die("CP_ATTN_FUSION must be 'fused' or 'split' (got '" + cfg.attn_fusion + "')"); + if (cfg.attn_fusion == "split" && !cfg.ring_mode) + die("CP_ATTN_FUSION=split only wires the ring path (CP_ATTN_MODE=ring); Ulysses " + "still calls the separate gqa_fused_flash_attn_forward kernel (out of scope, " + "see .mdfiles/HOPPER_MIGRATION_PLAN.md section 4). Run CP_ATTN_MODE=ring."); + const int cp_size = cfg.cp_size; + const int dp_size = world_size / cp_size; + + // ---- CREAM context-extension (additive; env-resolved; single-GPU only) ---- + if (const char* e = std::getenv("CP_CREAM_MODE")) cfg.cream_mode = e; + if (const char* e = std::getenv("CP_CREAM_SIGMA")) cfg.cream_sigma = std::atof(e); + if (const char* e = std::getenv("CP_CREAM_SEED")) + cfg.cream_seed = static_cast(std::strtoull(e, nullptr, 10)); + const char* cream_log_path = std::getenv("CP_CREAM_LOG"); // optional dist-check CSV + const bool cream_on = (cfg.cream_mode != "off"); + if (cream_on) { + // factor == YARN_SCALE (locked); scaled_max == factor * T. YaRN must be ON. + const double yarn_scale = + std::getenv("YARN_SCALE") ? std::atof(std::getenv("YARN_SCALE")) : 1.0; + cfg.cream_factor = static_cast(yarn_scale + 0.5); + cfg.cream_scaled_max = static_cast(cfg.cream_factor) * cfg.T; + if (cfg.cream_mode != "cream" && cfg.cream_mode != "pose" && + cfg.cream_mode != "randpos") + die("CP_CREAM_MODE must be off|cream|pose|randpos, got '" + cfg.cream_mode + "'"); + if (cp_size != 1) + die("CREAM is single-GPU only for now: requires CP_SIZE=1 (got " + + std::to_string(cp_size) + ")"); + if (yarn_scale <= 1.0) + die("CREAM requires YARN_SCALE > 1 (cache must be YaRN-scaled, got " + + std::to_string(yarn_scale) + ")"); + if (cfg.cream_factor < 2) + die("CREAM requires integer YARN_SCALE >= 2 (got factor=" + + std::to_string(cfg.cream_factor) + ")"); + if (cfg.cream_scaled_max % cfg.T != 0) + die("CREAM: scaled_max must be a multiple of T"); + if (cfg.cream_scaled_max < cfg.T) + die("CREAM: scaled_max < T"); + } + + // ---- LongRoPE searched-cache extension (additive; env-resolved) ---- + // CP_LONGROPE_FACTORS= installs a searched per-dim rescale cache. Unset => + // byte-identical. Composed with CREAM (both set): the LongRoPE cache is CREAM's + // gather source and MUST have been searched at the same target length/ratio. + if (const char* e = std::getenv("CP_LONGROPE_FACTORS")) cfg.longrope_factors = e; + const bool longrope_on = !cfg.longrope_factors.empty(); + if (longrope_on) { + const int64_t half = cfg.head_dim / 2; + try { + parse_longrope_factors(cfg.longrope_factors, half, cfg.longrope_lambda, + cfg.longrope_n_hat, cfg.longrope_s, cfg.longrope_search_len); + } catch (const std::exception& ex) { die(ex.what()); } + // CP_LONGROPE_MSCALE=1 bakes YaRN's attention temperature m=0.1*ln(s)+1 into + // the LongRoPE cache (a LongRoPE+m hybrid) for an apples-to-apples fight with + // YaRN. Default off => paper-faithful LongRoPE (mscale=1, no temperature). + if (std::getenv("CP_LONGROPE_MSCALE") && std::atoi(std::getenv("CP_LONGROPE_MSCALE")) != 0) { + cfg.longrope_mscale = (cfg.longrope_s > 1) + ? (0.1f * std::log(static_cast(cfg.longrope_s)) + 1.0f) : 1.0f; + if (is_master) + std::cout << "[LongRoPE] mscale ON: m=" << cfg.longrope_mscale + << " (YaRN temperature baked in)" << std::endl; + } + if (cream_on) { + // Composition coupling: CREAM's tail labels reach cream_scaled_max, so the + // LongRoPE gather source must be that long AND tuned for that target. + if (cfg.longrope_search_len != cfg.cream_scaled_max) + die("LongRoPE S_search (" + std::to_string(cfg.longrope_search_len) + + ") != cream_scaled_max (" + std::to_string(cfg.cream_scaled_max) + + ") -- search the LongRoPE cache at factor*T"); + if (cfg.longrope_s != cfg.cream_factor) + die("LongRoPE s (" + std::to_string(cfg.longrope_s) + + ") != cream factor (" + std::to_string(cfg.cream_factor) + ")"); + } + if (is_master) + std::cout << "[LongRoPE] factors=" << cfg.longrope_factors + << " n_hat=" << cfg.longrope_n_hat << " s=" << cfg.longrope_s + << " S_search=" << cfg.longrope_search_len + << (cream_on ? " (composed with CREAM)" : " (alone)") << std::endl; + } + + // ───────────────────────────────────────────────────────────────────────── + // Ms-PoE (arXiv:2403.04797) gates + weights. Per-KV-group position rescale + // composed on top of the base LongRoPE/NTK lambda: + // lambda_g[i] = 1 + (lambda_base[i] - 1) * w_g + // All w_g == 1 makes every plane identical, which must reproduce the + // non-Ms-PoE numbers exactly -- that is the runtime parity check. + // + // Every gate below guards a path whose kernel still assumes ONE shared cache + // plane and would therefore rotate groups 1..n at group 0's positions while + // reporting a clean run. Silence is the failure mode being prevented. + // ───────────────────────────────────────────────────────────────────────── + { + const char* mw = std::getenv("CP_MSPOE_WEIGHTS"); + const char* wmin_s = std::getenv("CP_MSPOE_WMIN"); + const char* mr = std::getenv("CP_MSPOE_RATIOS"); + const char* rmin_s = std::getenv("CP_MSPOE_RMIN"); + const char* rmax_s = std::getenv("CP_MSPOE_RMAX"); + const bool ratio_mode = (mr != nullptr) || (rmin_s != nullptr) || (rmax_s != nullptr); + const bool weight_mode = (mw != nullptr) || (wmin_s != nullptr); + if (ratio_mode && weight_mode) + die("Ms-PoE: CP_MSPOE_RATIOS/RMIN/RMAX (faithful, uniform over dims, no " + "base lambda) and CP_MSPOE_WEIGHTS/WMIN (composed onto a base " + "lambda) are different parameterisations -- set one, not both"); + if (ratio_mode || weight_mode) { + // Ratio mode needs no base lambda: the faithful method compresses + // positions below the native window and does not extend context. + if (weight_mode && !longrope_on) + die("Ms-PoE weight mode requires CP_LONGROPE_FACTORS: the per-group " + "weights rescale a BASE lambda vector, they do not replace it. " + "For the paper's parameterisation use CP_MSPOE_RMIN/RMAX, which " + "needs no base lambda."); + if (!eval_ppl_mode) + die("Ms-PoE is forward-only (the method is training-free): run it " + "under CP_EVAL_PPL=1. The fused RoPE BACKWARD still assumes a " + "single shared cache plane and would emit wrong gradients."); + if (cream_on) + die("Ms-PoE + CREAM is not supported: CREAM installs its own " + "per-step relabelled cache, which is 2-D and would silently " + "overwrite the per-group planes"); + if (cfg.attn_fusion == "split") + die("Ms-PoE requires CP_ATTN_FUSION=fused: the split path " + "(sdpa_split_rope_once_forward) has no per-group cache stride"); + if (!(cfg.attn_mode == "ring" || + (cfg.attn_mode == "ulysses" && cfg.ulysses_cplocal))) + die("Ms-PoE requires CP_ATTN_MODE=ring, or ulysses with " + "CP_ULYSSES_CPLOCAL=1 -- only those two reach the patched " + "sdpa_fused_forward_rope. Got attn_mode=" + cfg.attn_mode); + if (OwnTensor::cp::string_cfg_from_env().enabled) + die("Ms-PoE + CP_STRING is not supported: string_attention indexes " + "the cache itself and has no per-group stride"); + if (cfg.kv_heads < 2) + die("Ms-PoE needs kv_heads >= 2 to express more than one ratio " + "(got kv_heads=" + std::to_string(cfg.kv_heads) + "). Under GQA " + "the ratio must be constant within a KV group, so an MQA " + "checkpoint degenerates to a single uniform lambda -- which is " + "the plain LongRoPE arm, already measured."); + + if (ratio_mode) { + if (mr) { + std::string spec(mr); + for (char& ch : spec) if (ch == ',') ch = ' '; + std::istringstream is(spec); + float r; + while (is >> r) cfg.mspoe_ratios.push_back(r); + if (static_cast(cfg.mspoe_ratios.size()) != cfg.kv_heads) + die("CP_MSPOE_RATIOS has " + + std::to_string(cfg.mspoe_ratios.size()) + " values but the " + "model has " + std::to_string(cfg.kv_heads) + " KV heads"); + } else { + // Repo ladder, including its g/n (not g/(n-1)) divisor, so rmax + // is approached and never reached -- kept for faithfulness. + const float rmin = rmin_s ? std::strtof(rmin_s, nullptr) : 1.2f; + const float rmax = rmax_s ? std::strtof(rmax_s, nullptr) : 1.8f; + if (!(rmax >= rmin)) die("CP_MSPOE_RMAX must be >= CP_MSPOE_RMIN"); + const int64_t n = cfg.kv_heads; + for (int64_t g = 0; g < n; ++g) + cfg.mspoe_ratios.push_back( + rmin + (rmax - rmin) * (static_cast(g) / + static_cast(n))); + } + for (float r : cfg.mspoe_ratios) + if (!(r >= 1.0f)) + die("Ms-PoE ratios must be >= 1 (got " + std::to_string(r) + + "): the method COMPRESSES positions, and build_rope_cache_" + "longrope rejects lambda < 1 anyway"); + } else if (mw) { + std::string spec(mw); + for (char& ch : spec) if (ch == ',') ch = ' '; + std::istringstream is(spec); + float w; + while (is >> w) cfg.mspoe_weights.push_back(w); + if (static_cast(cfg.mspoe_weights.size()) != cfg.kv_heads) + die("CP_MSPOE_WEIGHTS has " + + std::to_string(cfg.mspoe_weights.size()) + " values but the " + "model has " + std::to_string(cfg.kv_heads) + " KV heads " + "(one weight per KV group is required)"); + } else { + const float wmin = std::strtof(wmin_s, nullptr); + // Evenly spaced over [wmin, 1]; group 0 is the most precise + // (least-extended) view, the last group the longest-reach one. + const int64_t n = cfg.kv_heads; + for (int64_t g = 0; g < n; ++g) + cfg.mspoe_weights.push_back( + wmin + (1.0f - wmin) * (static_cast(g) / + static_cast(n - 1))); + } + for (float w : cfg.mspoe_weights) + if (!(w > 0.0f) || !(w <= 1.0f)) + die("Ms-PoE weights must lie in (0, 1] (got " + + std::to_string(w) + "). w=1 is the base lambda (full " + "reach); smaller w moves that group toward original, " + "un-extended positions. w>1 would over-extend beyond the " + "searched vector and w<=0 is meaningless."); + + // CP_MSPOE_LAYERS: comma list and/or a-b ranges, e.g. "2-11" or + // "2,3,7". The repo passes apply_layers "2,...,31" so layers 0 and 1 + // keep unscaled positions. Unset => all layers. + if (const char* ml = std::getenv("CP_MSPOE_LAYERS")) { + std::string spec(ml); + for (char& ch : spec) if (ch == ',') ch = ' '; + std::istringstream is(spec); + std::string tok; + while (is >> tok) { + const size_t dash = tok.find('-'); + if (dash == std::string::npos) { + cfg.mspoe_layers.push_back(std::stoi(tok)); + } else { + const int a = std::stoi(tok.substr(0, dash)); + const int b = std::stoi(tok.substr(dash + 1)); + if (b < a) die("CP_MSPOE_LAYERS: bad range '" + tok + "'"); + for (int i = a; i <= b; ++i) cfg.mspoe_layers.push_back(i); + } + } + for (int li : cfg.mspoe_layers) + if (li < 0 || li >= static_cast(cfg.n_layers)) + die("CP_MSPOE_LAYERS: layer " + std::to_string(li) + + " out of range for n_layers=" + std::to_string(cfg.n_layers)); + if (cfg.mspoe_layers.empty()) + die("CP_MSPOE_LAYERS was set but parsed to no layers"); + } + + if (is_master) { + std::cout << "[Ms-PoE] mode=" + << (ratio_mode ? "RATIO (faithful: lambda_g[i]=r_g, uniform " + "over dims, no base lambda)" + : "WEIGHT (composed: lambda_g=1+(lambda_base-1)*w)") + << "\n[Ms-PoE] groups=" << cfg.kv_heads + << " (kv_heads) q_heads_per_group=" + << (cfg.q_heads / cfg.kv_heads) << " "; + if (ratio_mode) { + std::cout << "ratios="; + for (size_t g = 0; g < cfg.mspoe_ratios.size(); ++g) + std::cout << (g ? "," : "") << cfg.mspoe_ratios[g]; + } else { + std::cout << "weights="; + for (size_t g = 0; g < cfg.mspoe_weights.size(); ++g) + std::cout << (g ? "," : "") << cfg.mspoe_weights[g]; + std::cout << " base lambda from " << cfg.longrope_factors + << " (lam_last=" << cfg.longrope_lambda.back() << ")"; + } + std::cout << "\n[Ms-PoE] layers="; + if (cfg.mspoe_layers.empty()) std::cout << "ALL(" << cfg.n_layers << ")"; + else for (size_t i = 0; i < cfg.mspoe_layers.size(); ++i) + std::cout << (i ? "," : "") << cfg.mspoe_layers[i]; + std::cout << "\n[Ms-PoE] NOTE: ratios are assigned by KV-GROUP INDEX. " + "The reference implementation permutes them by a MEASURED " + "position-awareness statistic (outlier.argsort on the last " + "attention row); that ordering is NOT implemented here, so " + "this is the repo's --head_type normal ablation, not the " + "full method." << std::endl; + } + } + } + + // ───────────────────────────────────────────────────────────────────────── + // STRING (arXiv:2410.18745) gates + banner. See context_parallel/StringRoPE.h. + // STRING shifts a QUERY's cos/sin CACHE ROW, so its relative-distance + // semantics are exact only while the angle is LINEAR in position: + // theta_i(p) = p * base_freq_i / lambda_i (lambda_i a per-dim constant) + // True for PI / NTK / YaRN unconditionally. LongRoPE with n_hat > 0 leaves + // positions below n_hat unscaled, making theta PIECEWISE-linear -- the shift + // is then silently wrong near the threshold. Every searched vector we hold has + // n_hat 0, so this gate costs nothing and catches the one case that would + // produce plausible-but-wrong numbers. + // ───────────────────────────────────────────────────────────────────────── + { + const OwnTensor::cp::StringConfig& scfg = OwnTensor::cp::string_cfg_from_env(); + if (scfg.enabled) { + if (world_size != 1) + die("CP_STRING=1 requires CP_SIZE==1 (got world_size=" + + std::to_string(world_size) + "). The block decomposition assumes a " + "cache row IS the absolute position, which is false under CP " + "sharding. Refusing rather than returning baseline numbers " + "labelled STRING."); + // The block decomposition needs Q and K passed as SEPARATE tensors of + // DIFFERENT lengths, with independent q/k position deltas -- an interface + // only sdpa_fused_forward_rope (the ring body of ContextParallel::forward_cp) + // exposes. The ulysses path routes to gqa_fused_flash_attn_forward, which + // takes a PACKED qkv tensor at one shared T and a single scalar pos_offset, + // so it cannot express "Qi vs K[0:near_lo] with Q's rows shifted". Without + // this gate CP_STRING=1 under the DEFAULT ulysses mode parses, prints this + // banner, and then silently computes plain baseline attention -- which is + // exactly what the W=S identity test cannot distinguish from a real pass. + // CONSEQUENCE FOR THE EXPERIMENT: the STRING-off baseline of every + // comparison must ALSO be run with CP_ATTN_MODE=ring, or the two arms + // differ by attention path as well as by STRING. + if (cfg.attn_mode != "ring") + die("CP_STRING=1 requires CP_ATTN_MODE=ring (got '" + cfg.attn_mode + + "'). Only the ring body calls sdpa_fused_forward_rope, the one entry " + "point taking separate Q/K tensors with independent position deltas. " + "Under ulysses/hybrid the STRING branch is never reached and the run " + "would report baseline numbers labelled STRING. Run BOTH arms with " + "CP_ATTN_MODE=ring."); + if (longrope_on && cfg.longrope_n_hat != 0) + die("CP_STRING=1 with a LongRoPE cache whose n_hat=" + + std::to_string(cfg.longrope_n_hat) + " (!= 0). STRING shifts a cache " + "ROW and needs the angle linear in position; n_hat>0 makes it " + "piecewise-linear, so the shift is silently wrong near the " + "threshold. Use an n_hat=0 vector, or disable CP_STRING."); + // Validate S/B/W and surface the resolved operating point. string_plan + // throws on a bad combination (including the no-op refusal), so a + // misconfiguration dies HERE at startup rather than mid-eval. + const OwnTensor::cp::StringPlan sp = + OwnTensor::cp::string_plan(static_cast(cfg.T), scfg); + if (is_master) + std::cout << "[STRING] S=" << sp.S << " B=" << sp.B << " W=" << sp.W + << " -> G=" << sp.G << " k=" << sp.k + << " shift=" << sp.shift + << " far=" << (sp.far_exists ? "yes" : "NO (no-op, allowed)") + << " threshold fuzzy over [" << sp.S << "," << (sp.S + sp.B) + << ")" << std::endl; + } + } + + // ---- CP_DUMP_ROPE: dump the YaRN rope cache for the parity gate (plan 0.3) ---- + // Writes autograd::build_rope_cache(seq_len, head_dim, rope_theta) as raw float32 + // [seq_len, head_dim] (cols [0:half]=cos*mscale, [half:head_dim]=sin*mscale, the + // YARNOps layout) so scripts/convert_ckpt/rope_parity.py --cp-cache can diff it + // against HF Qwen3 YaRN. YARN_SCALE / YARN_ORIG_MAXPOS are read by build_rope_cache + // from env. Needs no model or checkpoint; forward-only, exits. Placed BEFORE the + // T!=context_length guard so CP_T need not be set for a dump. + if (const char* rope_path = std::getenv("CP_DUMP_ROPE")) { + const int64_t seqlen = env_i64("CP_DUMP_ROPE_SEQLEN", 8192); + Tensor cache = autograd::build_rope_cache(seqlen, cfg.head_dim, + static_cast(cfg.rope_theta), device); + if (is_master) { + Tensor c = cache.as_type(Dtype::Float32).to_cpu(); + std::ofstream rf(rope_path, std::ios::binary); + rf.write(reinterpret_cast(c.data()), + static_cast(c.numel()) * sizeof(float)); + std::cout << "[dump-rope] wrote YaRN cache [" << seqlen << "," << cfg.head_dim + << "] raw float32 (YARN_SCALE via env) -> " << rope_path << std::endl; + } + MPI_Finalize(); + return 0; + } + + // The base RoPE cache is built at context_length. With CREAM the kernel-facing + // cache is a per-step gather of the full scaled cache, so T is legitimately + // decoupled from context_length -- but ONLY under CREAM. The original guard is + // preserved verbatim for the default (CREAM-off) path. + if (!cream_on) { + if (cfg.T != cfg.context_length) + die("T != context_length (RoPE cache is built at context_length)"); + } + // Hybrid: CP_SIZE = ring_size * ulysses_size. ring_size is the OUTER (P2P + // rotation) degree, ulysses_size the INNER (all-to-all) degree. + int ring_size = cp_size; // non-hybrid: ring axis == whole CP + if (cfg.hybrid_mode) { + if (cfg.ulysses_size < 1) + die("CP_ULYSSES_SIZE must be >= 1"); + if (cp_size % cfg.ulysses_size != 0) + die("CP_SIZE=" + std::to_string(cp_size) + " % CP_ULYSSES_SIZE=" + + std::to_string(cfg.ulysses_size) + " != 0"); + ring_size = cp_size / cfg.ulysses_size; + // Ring HeadTail on the ring axis: T divisible by 2*ring_size. + if (cfg.T % (2 * ring_size) != 0) + die("T=" + std::to_string(cfg.T) + " not divisible by 2*ring_size=" + + std::to_string(2 * ring_size) + " (hybrid HeadTail)"); + // Contiguous inner Ulysses slice of each ring block. + if ((cfg.T / ring_size) % cfg.ulysses_size != 0) + die("(T/ring_size)=" + std::to_string(cfg.T / ring_size) + + " not divisible by CP_ULYSSES_SIZE=" + std::to_string(cfg.ulysses_size) + + " (hybrid inner slice)"); + // Fused-RoPE kernel tile: the HeadTail sub-chunk reaching the kernel is + // T/(2*ring_size); it MUST be a multiple of 32 or gqa_fused_rope_cp aborts. + if ((cfg.T / (2 * ring_size)) % 32 != 0) + die("T/(2*ring_size)=" + std::to_string(cfg.T / (2 * ring_size)) + + " not a multiple of 32 (fused-RoPE tile); increase CP_T so " + "T % (64*ring_size) == 0"); + // Ulysses head split needs both q and kv heads divisible by U. The U>kv + // (kv-replication) path is NOT supported in hybrid v1. + if (cfg.q_heads % cfg.ulysses_size != 0) + die("q_heads=" + std::to_string(cfg.q_heads) + + " not divisible by CP_ULYSSES_SIZE=" + std::to_string(cfg.ulysses_size)); + if (cfg.kv_heads % cfg.ulysses_size != 0) + die("kv_heads=" + std::to_string(cfg.kv_heads) + + " not divisible by CP_ULYSSES_SIZE=" + std::to_string(cfg.ulysses_size) + + " (hybrid does not yet support U > kv_heads replication)"); + // Advisory fast-domain guard (perf, not correctness): warn if a Ulysses + // group would straddle a fast-interconnect island. F defaults to + // NO_GPUS_PER_NODE / cudaGetDeviceCount; set it to a NUMA/NVLink-island + // size to place Ulysses intra-island and Ring across it. NOT fatal. + int fast_domain = 0; + if (const char* e = std::getenv("NO_GPUS_PER_NODE")) fast_domain = std::atoi(e); + if (fast_domain <= 0) { int gc = 1; cudaGetDeviceCount(&gc); fast_domain = gc; } + if (fast_domain > 0 && (fast_domain % cfg.ulysses_size) != 0 && is_master) + std::cerr << "[CP WARN] fast-domain size (" << fast_domain + << ") % CP_ULYSSES_SIZE (" << cfg.ulysses_size + << ") != 0: a Ulysses group straddles a fast-interconnect " + "boundary; the intra-node all-to-all may land on a slower " + "link. This is a performance hint, not a correctness error." + << std::endl; + } else if (cfg.ring_mode) { + if (cfg.T % (2 * cp_size) != 0) + die("T=" + std::to_string(cfg.T) + " not divisible by 2*CP_SIZE=" + std::to_string(2*cp_size) + " (HeadTail)"); + } else { + // Pure Ulysses: contiguous sequence shard + all-to-all HEAD split across + // the CP group. Both q and kv heads must divide cp_size, else the split is + // ill-defined and (without this guard) throws late from UlyssesAttention.h. + if (cfg.T % cp_size != 0) + die("T=" + std::to_string(cfg.T) + " not divisible by CP_SIZE=" + std::to_string(cp_size) + " (contiguous)"); + if (cp_size > 1 && cfg.q_heads % cp_size != 0) + die("CP_ATTN_MODE=ulysses: q_heads=" + std::to_string(cfg.q_heads) + + " not divisible by CP_SIZE=" + std::to_string(cp_size) + + " (heads are split across CP ranks; use CP_ATTN_MODE=ring for this head/CP combo)"); + if (cp_size > 1 && cfg.kv_heads % cp_size != 0) + die("CP_ATTN_MODE=ulysses: kv_heads=" + std::to_string(cfg.kv_heads) + + " not divisible by CP_SIZE=" + std::to_string(cp_size) + + " (heads are split across CP ranks; use CP_ATTN_MODE=ring for this head/CP combo)"); + } + if (cfg.head_dim != 32 && cfg.head_dim != 64 && cfg.head_dim != 128) + // hd=32 is ADDITIVE and FUSED-ONLY: the CP WMMA kernels instantiate + // D=32 (D%16==0 holds), but the split/sm90 path cannot -- its 128 B + // swizzle atom is 64 bf16, so hd=32 is half an atom. CP_ATTN_FUSION=split + // therefore still requires hd==128 (enforced by gqa_causal_cp_fwd_shape_ok). + die("head_dim must be 32, 64 or 128 (fused kernel; split requires 128)"); + if (cfg.q_heads % cfg.kv_heads != 0) + die("q_heads must be a multiple of kv_heads"); + + // Probe mode: one micro-batch of B*T*dp_size tokens (grad_accum=1) captures the + // per-microstep peak faithfully (the accum loop is sequential), runs fast, and + // skips checkpointing/validation. max_steps>=2 so optimizer state is allocated. + if (mem_probe) { + cfg.global_batch = cfg.B * cfg.T * static_cast(dp_size); + cfg.max_steps = mem_probe_steps; + cfg.warmup_steps = 1; + cfg.checkpointing = false; + } + + // Tokens per optimizer step: CP shares the batch (no data factor); DP multiplies. + const int64_t tokens_per_micro = cfg.B * cfg.T * static_cast(dp_size); + if (cfg.global_batch % tokens_per_micro != 0) + die("global_batch not divisible by B*T*dp_size (" + std::to_string(tokens_per_micro) + ")"); + const int GRAD_ACCUM = static_cast(cfg.global_batch / tokens_per_micro); + + // ---- process groups: the mesh builds the WORLD PG (ncclCommSplit parent) + // + per-axis sub-groups. Reuse the mesh's world PG for DataParallel — a second + // init_process_group would throw (BluTrain's PG registry is a singleton). + // non-hybrid: 2-D {dp, cp} -> cp_pg = axis 1. + // hybrid (default, dp outer): 3-D {dp, ring, ulysses}. Ulysses is the fastest + // axis (U consecutive ranks); RING stride = U; DP stride = ring*U. Good when + // you want DP to cross the slow boundary (e.g. inter-node) and ring+ulysses + // to stay on the fast domain. + // hybrid + CP_RING_OUTER=1: 3-D {ring, dp, ulysses}. RING stride = dp*U so a + // ring group spans across the coarse boundary while Ulysses AND DDP stay on + // the fast domain. Use this to put the RING (P2P, overlaps) across NUMA/nodes + // and keep the bulk all-to-all + all-reduce intra-domain. Grouping is by rank + // index; rank r binds to GPU (r % gpus_per_node), so on one node the logical + // groups land on physical GPUs in order. ---- + const bool ring_outer = + cfg.hybrid_mode && std::getenv("CP_RING_OUTER") && + std::getenv("CP_RING_OUTER")[0] == '1'; + std::vector ranks_vec(world_size); + for (int i = 0; i < world_size; ++i) ranks_vec[i] = i; + // Axis indices into the mesh (depend on the chosen ordering). + int dp_axis = 0, ring_axis = 1, uly_axis = 2; + std::vector mesh_shape; + if (!cfg.hybrid_mode) { + mesh_shape = {dp_size, cp_size}; // {dp, cp} + } else if (ring_outer) { + mesh_shape = {ring_size, dp_size, cfg.ulysses_size}; // {ring, dp, ulysses} + ring_axis = 0; dp_axis = 1; uly_axis = 2; + } else { + mesh_shape = {dp_size, ring_size, cfg.ulysses_size}; // {dp, ring, ulysses} + dp_axis = 0; ring_axis = 1; uly_axis = 2; + } + DeviceMesh mesh(mesh_shape, ranks_vec); + std::shared_ptr pg = mesh.world_pg(); // global (all ranks), for DDP + std::shared_ptr cp_pg; // CP ring/ulysses sub-group (has ring surface) + std::shared_ptr ulysses_pg; // inner ulysses sub-group (base methods only) + int cp_rank = 0, dp_rank = mesh.get_dim_rank(dp_axis); + int ring_rank = 0, ulysses_rank = 0; + if (cfg.hybrid_mode) { + cp_pg = mesh.get_process_group(ring_axis); // RING sub-group (ContextParallel pg_) + ulysses_pg = mesh.get_process_group(uly_axis); // Ulysses sub-group (inner all-to-all) + ring_rank = mesh.get_dim_rank(ring_axis); + ulysses_rank = mesh.get_dim_rank(uly_axis); + cp_rank = ring_rank * cfg.ulysses_size + ulysses_rank; // composite CP coord + if (is_master) + std::cout << "[CP hybrid mesh] order=" + << (ring_outer ? "{ring,dp,ulysses} (RING outer -> crosses coarse boundary)" + : "{dp,ring,ulysses} (DP outer)") + << " ring=" << ring_size << " dp=" << dp_size + << " ulysses=" << cfg.ulysses_size << std::endl; + } else { + cp_pg = mesh.get_process_group(1); // CP ring/ulysses sub-group + cp_rank = mesh.get_dim_rank(1); + ring_rank = cp_rank; // whole CP axis is the ring axis + if (std::getenv("CP_RANK_NORM")) { std::fprintf(stderr, "[rank %d] cp_pg: worldsize=%d rank=%d | dp_size=%d dp_rank=%d cp_size=%d cp_rank=%d\n", rank, cp_pg ? cp_pg->get_worldsize() : -1, cp_pg ? cp_pg->get_rank() : -1, dp_size, dp_rank, cp_size, cp_rank); std::fflush(stderr); } + } + + if (is_master) { + std::cout << "================ Configuration (bluscriptCP) ================\n" + << " d_model/layers : " << cfg.d_model << " / " << cfg.n_layers << "\n" + << " q/kv heads : " << cfg.q_heads << " / " << cfg.kv_heads + << " (G=" << (cfg.q_heads / cfg.kv_heads) << ", hd=" << cfg.head_dim << ")\n" + << " ffn_hidden : " << cfg.ffn_hidden << " (SwiGLU)\n" + << " context_length : " << cfg.context_length << "\n" + << " vocab_size : " << cfg.vocab_size << " tying=" << cfg.weight_tying << "\n" + << " B x T : " << cfg.B << " x " << cfg.T << "\n" + << " world_size : " << world_size + << " (dp=" << dp_size << " x cp=" << cp_size + << (cfg.hybrid_mode ? " [ring=" + std::to_string(ring_size) + + " x ulysses=" + std::to_string(cfg.ulysses_size) + "]" + : "") + << ")\n" + << " CP_ATTN_MODE : " << cfg.attn_mode + << " (load_balance=" << ((cfg.ring_mode || cfg.hybrid_mode) ? "HeadTail" : "contiguous") << ")\n" + << " global_batch : " << cfg.global_batch << " (grad_accum=" << GRAD_ACCUM << ")\n" + << " precision : " << (cfg.bf16 ? "bf16 Linears + fp32 norms/residual/logits" + " (fp32 master weights in AdamW)" + : "fp32 everywhere") + << (cfg.bf16 ? " [mask=" + std::to_string(cfg.bf16_mask) + + std::string(cfg.bf16_mask & 1 ? " QKV" : "") + + std::string(cfg.bf16_mask & 2 ? " WO" : "") + + std::string(cfg.bf16_mask & 4 ? " MLP" : "") + + std::string(cfg.bf16_mask & 8 ? " HEAD": "") + "]" + : "") + << "\n" + << " max_lr/min_lr : " << cfg.max_lr << " / " << cfg.min_lr << "\n" + << " warmup/steps : " << cfg.warmup_steps << " / " << cfg.max_steps + << (cfg.stable_steps > 0 + ? " (WSD: stable=" + std::to_string(cfg.stable_steps) + ", decay tail=" + + std::to_string(cfg.max_steps - cfg.warmup_steps - cfg.stable_steps) + ")" + : " (cosine, no stable phase)") + << "\n" + << " checkpointing : " << (cfg.checkpointing ? "true" : "false") + << " (ckpt_freq=" << cfg.ckpt_freq << ", dir=" << cfg.ckpt_dir << ")\n" + << "=============================================================" << std::endl; + } + + if (is_master) std::cout << "Initializing model on CUDA device " << rank << "..." << std::endl; + GPT model(cfg, device, mesh, cp_pg, cp_size, cp_rank, + ulysses_pg, ring_rank, ulysses_rank, /*seed=*/1234); + + // CP_SHARE_FWD_ROTATOR=1 (additive, default OFF): share ONE forward ring holder + // (rotator + send_buf[2] + persistent exch_work) across all layers instead of + // one set per layer. Forward is sequential (layer i finishes before i+1), so a + // single set suffices and saves ~kv_staging * n_layers resident memory. Unset => + // per-layer path, byte-identical to before. Value-aware, mirrors CP_NO_OVERLAP. + { + const char* _share = std::getenv("CP_SHARE_FWD_ROTATOR"); + if (_share && !(_share[0] == '0' && _share[1] == '\0')) { + auto holder = std::make_shared(); + for (auto& a : model.attn) a->cp_->set_shared_fwd_ring(holder); + if (is_master) + std::cout << "[CP] CP_SHARE_FWD_ROTATOR=1: one forward ring holder shared across " + << model.attn.size() << " layers." << std::endl; + } + } + + auto params = model.parameters(); + + // ---- CP_INIT_FROM= -------------------------------- + // BRANCH a new run from a trained base, taking ONLY the model weights. + // This is what a long-context (YaRN / LongRoPE) extension is, and it is + // deliberately NOT the same operation as a resume: + // + // resume every piece of state is restored together, so ordering + // does not matter -- load_state overwrites the optimizer's + // fp32 masters after the fact. + // initialize-from there IS no saved optimizer state. The masters must be + // DERIVED from the weights, so the weights have to land + // BEFORE any optimizer is constructed. + // + // Hence this hook sits here, after model.parameters() and before BOTH + // optimizer construction sites (the ZeRO instances below, and the plain + // nn::AdamW further down). Every one of them snapshots fp32 masters at + // construction; loading afterwards leaves the masters holding the random + // init, and the first step() publishes them straight back over the weights + // you just loaded. Measured, before this hook existed: + // step 1000 loss 2.481 <- weights loaded + // step 1001 loss 11.084 <- masters clobbered them + // + // Because only replicated, T-independent PARAMETERS move, this is portable + // across every configuration change that matters: any world size, ZeRO or + // no ZeRO, dp<->cp resharding, and any context length (cos_sin_cache is a + // plain member, not a registered parameter, so T and YARN_SCALE do not + // change the parameter count). + bool init_from_active = false; // consumed by the LR schedule below + int init_from_base_step = 0; + if (const char* init_path = std::getenv("CP_INIT_FROM")) { + if (!std::filesystem::exists(init_path)) + die(std::string("CP_INIT_FROM: no such checkpoint file: ") + init_path); + int base_step = 0; + float base_loss = 0.0f; + load_checkpoint_weights_only(init_path, model, base_step, base_loss); + init_from_active = true; + init_from_base_step = base_step; + if (is_master) + std::cout << "[InitFrom] weights loaded from " << init_path + << " (base step " << base_step << ", loss " << base_loss << ")" + << " -- optimizer state and RNG intentionally NOT restored;" + << " pair with CP_REWARMUP" << std::endl; + } + int64_t num_params = 0; + for (auto& p : params) num_params += p.numel(); + if (is_master) { + std::cout << "Parameters: " << num_params << "\n"; + std::cout << "Parameters per GPU: " << num_params << "\n"; + std::cout << "max_steps: " << cfg.max_steps << "\n"; + std::cout << "warmup_steps: " << cfg.warmup_steps << std::endl; + } + + // ---- DDP wrapper KEPT over the GLOBAL pg: init-broadcast + param-grad + // all-reduce over ALL dp*cp ranks (correct group for globally-replicated + // params). The CP ring adds only the attention-activation grad sum over + // cp_pg (disjoint from DDP's leaf-param reduction). ---- + std::unique_ptr ddp_model; + if (world_size > 1 && !cfg.zero && !std::getenv("CP_NO_DDP")) { + DDP_Options ddp_opts; + ddp_opts = ddp_opts + .with_process_group(pg) + .with_world_size(world_size) + .with_bucket_data(std::getenv("CP_DDP_NO_BUCKET") == nullptr, 25 * 1024 * 1024) + .with_grad_view(std::getenv("CP_DDP_NO_GRADVIEW") == nullptr); + // CP_DDP_INIT_SYNC=0 skips DataParallel's startup parameter broadcast. + // Safe here BY CONSTRUCTION: GPT is built with a hardcoded rank-independent + // seed (1234, bluscriptCP.cpp:1442), so every rank already holds identical + // parameters and the broadcast is a redundant safety net. + // + // It is also a BUG WORKAROUND. With bf16 params (CP_BF16=1) and world_size>1 + // the broadcast corrupts weights: validation loss is nan BEFORE the first + // training step, while the identical model at NP=1 gives 11.1653. Gradient + // buckets are forced fp32 (DataParallel.cpp:300) and Dtype::Bfloat16 does map + // to ncclBfloat16 (ProcessGroupNCCL.h:103), so the fault is in the mixed-dtype + // parameter path of sync_model_parameter/broadcast_coalesced, not the dtype + // mapping. Defaults to 1 so fp32 runs are untouched. + // DEFAULTS OFF UNDER bf16, on otherwise: bf16 + init_sync is MEASURED broken + // (val loss nan at NP=4 vs 11.1509 with the broadcast skipped, everything else + // identical), and the broadcast is redundant anyway given the fixed seed. An + // explicit CP_DDP_INIT_SYNC=1 still forces it on, for anyone bisecting this. + const bool ddp_init_sync = env_i64("CP_DDP_INIT_SYNC", cfg.bf16 ? 0 : 1) != 0; + if (is_master && cfg.bf16) + std::cout << "[DDP] init_sync=" << (ddp_init_sync ? "1 (FORCED -- expect nan" + " with bf16)" + : "0 (skipped: redundant with" + " the fixed init seed, and" + " broken for bf16 params)") + << std::endl; + ddp_model = std::make_unique(&model, ddp_opts, ddp_init_sync); + } + + // ---- ZeRO-1 optimizer-state sharding (owt::zero), mirroring bluscript_zero ---- + // ShardedOptimizer requires every parameter it manages to share ONE dtype (the + // fused Adam kernel and the reduce-scatter wire format are single-dtype). This + // model is deliberately mixed under CP_BF16=1 -- bf16 Linears, fp32 embedding / + // RMSNorm scales / QK-norm gammas -- so group by dtype and give each group its + // own instance. In practice exactly two groups. Stepping them back-to-back + // (never concurrently) is safe: each local_adam_step() stream-syncs after its + // fused-Adam launch, so the multi-tensor kernel's shared launch-metadata buffer + // is not clobbered by the next instance. + std::vector> zopts; + std::vector zgroup_label; + if (cfg.zero) { + std::map> params_by_dtype; + for (auto& p : params) params_by_dtype[p.dtype()].push_back(p); + + owt::zero::ZeroConfig zcfg; + zcfg.variant = owt::zero::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 + zcfg.max_grad_norm = cfg.grad_clip; + // Every rank draws identical init weights from the same seed (GPT ctor, + // seed=1234), so the broadcast would only copy data already equal. Same + // reasoning bluscript_zero.cpp uses to disable it. + zcfg.init_broadcast = false; + zcfg.verbose = is_master; + if (cfg.zero_fp32_reduce) { + // 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. Reduce in fp32 when that is suspected. + zcfg.reduce_dtype = owt::zero::ReduceDtype::FP32; + } + for (auto& kv : params_by_dtype) { + if (cfg.zero_v2) + zopts.push_back(std::make_unique(kv.second, pg, zcfg)); + else + zopts.push_back(std::make_unique(kv.second, pg, zcfg)); + int64_t n = 0; for (auto& p : kv.second) n += p.numel(); + zgroup_label.push_back(OwnTensor::get_dtype_name(kv.first) + "/" + std::to_string(n)); + } + if (is_master) { + std::cout << "[ZeRO-1] variant=" << (cfg.zero_v2 ? "v2 reduce-scatter/all-gather" + : "v1 reduce/broadcast") + << " groups=" << zopts.size(); + for (auto& g : zgroup_label) std::cout << " [" << g << "]"; + std::cout << "\n[ZeRO-1] DataParallel is DISABLED (both would double-reduce)" + << std::endl; + } + } + + // ---- optimizer (single AdamW, decoupled wd applied manually to dim>=2) ---- + std::vector decay_params; + for (const auto& p : params) if (p.shape().dims.size() >= 2) decay_params.push_back(p); + nn::AdamW optimizer(params, cfg.max_lr, 0.9f, 0.95f, 1e-8f, /*weight_decay=*/0.0f); + + // Mixed precision depends on Optim.cpp:28 having built an fp32 master weight for + // every non-fp32 parameter, because AdamW::step (Optim.cpp:779) updates through + // it. bluscript_zero gets its masters from ZeRO's ShardedOptimizer instead, so + // this is the ONE part of the bf16 scheme its validation does not cover -- assert + // it here rather than discover it as a run that trains but does not learn (a + // bf16 param updated in place loses every step smaller than ~2^-9 relative). + if (cfg.bf16) { + int64_t n_lowp = 0, n_master = 0; + for (auto& p : params) { + if (!p.requires_grad() || p.dtype() == Dtype::Float32) continue; + ++n_lowp; + if (optimizer.get_master_weight(p) != nullptr) ++n_master; + } + if (n_lowp == 0) + die("CP_BF16=1 but no parameter is non-fp32 -- the Linears were built fp32"); + if (n_master != n_lowp) + die("CP_BF16=1: only " + std::to_string(n_master) + " of " + + std::to_string(n_lowp) + " bf16 params have an fp32 master weight. " + "Adam would update the bf16 tensors in place and the run would not learn."); + if (is_master) + std::cout << "[precision] bf16 params=" << n_lowp + << ", all with fp32 master weights" << std::endl; + } + + // ---- data: shard the BATCH over the DP axis; all CP ranks in a group share it ---- + const std::string data_root = std::getenv("CP_DATA_ROOT") + ? std::getenv("CP_DATA_ROOT") : "Data_Loader/Data"; + DataLoaderLite train_loader(static_cast(cfg.B), static_cast(cfg.T), + std::getenv("CP_SAME_DATA") ? 0 : dp_rank, std::getenv("CP_SAME_DATA") ? 1 : dp_size, "train", data_root, is_master, 100000000, rank); + DataLoaderLite val_loader(static_cast(cfg.B), static_cast(cfg.T), + std::getenv("CP_SAME_DATA") ? 0 : dp_rank, std::getenv("CP_SAME_DATA") ? 1 : dp_size, "val", data_root, is_master, 100000000, rank); + + Tensor grad_scale = Tensor::full(Shape{{1}}, TensorOptions().with_device(device), + 1.0f / static_cast(GRAD_ACCUM)); + Tensor loss_accum_gpu = Tensor::zeros(Shape{{1}}, TensorOptions().with_device(device)); + float last_val_loss = -1.0f; + + // ── Run-numbered logging + checkpoints (mirrors gpt2_cp_test.cpp) ───────── + // run_number governs BOTH the CSV log index (CP_BluScreipt_Training_logs/CP_Training_log + // .csv) AND the checkpoint filename prefix (blumodelcp_run_step_.ckpt). + // Resolved on rank 0 (needs filesystem scans), then broadcast so every rank + // builds the same prefix. When resuming, run_number is kept (log is appended). + std::string log_filename, config_filename; + std::ofstream log_file; + int run_number = 0; + bool ckpt_resuming = false; // rank-0 decision; broadcast below + + if (is_master) { + std::filesystem::create_directories("CP_BluScreipt_Training_logs"); + int next_free_log = 1; + while (std::filesystem::exists("CP_BluScreipt_Training_logs/CP_Training_log" + + std::to_string(next_free_log) + ".csv")) + next_free_log++; + + // Scan the checkpoint dir for _run_step_.ckpt, tracking each + // run's highest step; a run is "complete" once its top step >= max_steps. + int resume_run = -1, resume_top_step = -1; + if (cfg.checkpointing && !ckpt_new_run) { + std::map run_top_step; // run K -> highest step S + // Scan BOTH layouts. Under ZeRO the CheckpointManager is constructed + // with shard_dir, so the .ckpt files live in /rank_/ -- + // scanning only the parent finds nothing and every resume silently + // becomes a fresh run. Checking the parent as well keeps checkpoints + // written before that change (and every non-ZeRO run) resumable. + auto scan_dir = [&](const std::string& dir) { + if (!std::filesystem::exists(dir)) return; + for (const auto& de : std::filesystem::directory_iterator(dir)) { + if (!de.is_regular_file()) continue; + const std::string fn = de.path().filename().string(); + const std::string p1 = cfg.ckpt_prefix + "_run", p2 = "_step_", p3 = ".ckpt"; + if (fn.rfind(p1, 0) != 0) continue; + size_t spos = fn.find(p2); + if (spos == std::string::npos) continue; + if (fn.size() < p3.size() || + fn.compare(fn.size() - p3.size(), p3.size(), p3) != 0) + continue; + try { + int k = std::stoi(fn.substr(p1.size(), spos - p1.size())); + int s = std::stoi(fn.substr(spos + p2.size(), + fn.size() - p3.size() - (spos + p2.size()))); + auto it = run_top_step.find(k); + if (it == run_top_step.end() || s > it->second) run_top_step[k] = s; + } catch (const std::exception&) { continue; } + } + }; + scan_dir(cfg.ckpt_dir); + if (cfg.zero) scan_dir(cfg.ckpt_dir + "/rank_" + std::to_string(rank)); + if (ckpt_resume_run >= 0) { + auto it = run_top_step.find(ckpt_resume_run); + if (it != run_top_step.end()) { resume_run = ckpt_resume_run; resume_top_step = it->second; } + else std::cerr << "[Resume] CP_CKPT_RESUME=" << ckpt_resume_run + << " has no checkpoint; starting a new run." << std::endl; + } else if (!run_top_step.empty()) { + auto last = std::prev(run_top_step.end()); // highest run key + if (last->second < cfg.max_steps) { resume_run = last->first; resume_top_step = last->second; } + else std::cout << "[Resume] latest run " << last->first + << " already complete at step " << last->second + << "; starting a new run." << std::endl; + } + } + if (resume_run >= 0) { run_number = resume_run; ckpt_resuming = true; } + else { run_number = next_free_log; } + + log_filename = "CP_BluScreipt_Training_logs/CP_Training_log" + std::to_string(run_number) + ".csv"; + const bool log_exists = std::filesystem::exists(log_filename); + std::cout << "Saving logs to: " << log_filename + << (ckpt_resuming ? " (resume run " : " (run ") << run_number << ")\n"; + + config_filename = "CP_BluScreipt_Training_logs/CP_Training_log" + std::to_string(run_number) + "_config.txt"; + std::ofstream config_file(eval_ppl_mode ? "/dev/null" : config_filename.c_str(), + (ckpt_resuming && log_exists) ? std::ios::app : std::ios::out); + if (ckpt_resuming && log_exists) + config_file << "\n# resumed run " << run_number << " from step " << resume_top_step << "\n"; + config_file << "Configuration (bluscriptCP):\n"; + config_file << " d_model: " << cfg.d_model << "\n"; + config_file << " n_layers: " << cfg.n_layers << "\n"; + config_file << " q_heads: " << cfg.q_heads << "\n"; + config_file << " kv_heads: " << cfg.kv_heads << " (G=" << (cfg.q_heads / cfg.kv_heads) + << ", head_dim=" << cfg.head_dim << ")\n"; + config_file << " ffn_hidden: " << cfg.ffn_hidden << "\n"; + config_file << " context_length: " << cfg.context_length << "\n"; + config_file << " vocab_size: " << cfg.vocab_size << " (tying=" << cfg.weight_tying << ")\n"; + config_file << " B: " << cfg.B << "\n"; + config_file << " T: " << cfg.T << "\n"; + config_file << " world_size: " << world_size << " (dp=" << dp_size << " x cp=" << cp_size << ")\n"; + if (cfg.hybrid_mode) + config_file << " hybrid: ring_size=" << ring_size + << " ulysses_size=" << cfg.ulysses_size << "\n"; + config_file << " CP_ATTN_MODE: " << cfg.attn_mode + << " (load_balance=" << ((cfg.ring_mode || cfg.hybrid_mode) ? "HeadTail" : "contiguous") << ")\n"; + config_file << " global_batch: " << cfg.global_batch << "\n"; + config_file << " grad_accum_steps: " << GRAD_ACCUM << "\n"; + config_file << " Parameters: " << num_params << "\n"; + config_file << " Parameters per GPU: " << num_params << "\n"; + config_file << " Max Learning Rate: " << cfg.max_lr << "\n"; + config_file << " Min Learning Rate: " << cfg.min_lr << "\n"; + config_file << " max_steps: " << cfg.max_steps << "\n"; + config_file << " warmup_steps: " << cfg.warmup_steps << "\n"; + config_file << " checkpointing: " << (cfg.checkpointing ? 1 : 0) << "\n"; + if (cfg.checkpointing) config_file << " ckpt_freq: " << cfg.ckpt_freq << "\n"; + config_file << " run: " << run_number << "\n"; + size_t free_mem = 0, total_mem = 0; cudaMemGetInfo(&free_mem, &total_mem); + double used_mb = static_cast(total_mem - free_mem) / (1024.0 * 1024.0); + double total_mb = static_cast(total_mem) / (1024.0 * 1024.0); + config_file << " GPU Memory Used (rank 0): " << std::fixed << std::setprecision(1) + << used_mb << " MB / " << total_mb << " MB\n"; + config_file.close(); + + // Resume -> append to the same run's CSV; fresh run -> truncate + header. + // Eval mode must not touch the training CSV -> write to /dev/null. + const bool append_log = ckpt_resuming && log_exists; + log_file.open(eval_ppl_mode ? "/dev/null" : log_filename.c_str(), + append_log ? std::ios::app : std::ios::out); + if (!log_file.is_open()) { + std::cerr << "ERROR: Could not open log file " << log_filename << "\n"; + std::exit(1); + } + if (append_log) + log_file << "# resumed run " << run_number << " from step " << resume_top_step << "\n"; + else + log_file << "step,train_loss,val_loss,lr,grad_norm,dt_ms,tok_per_sec," + "timer_data,timer_fwd,timer_loss,timer_bwd,timer_clip,timer_optim,mem_gpu_mb\n"; + log_file << std::fixed << std::setprecision(6); + } + + // Broadcast the run number so every CheckpointManager builds the same prefix. + int run_number_bc = is_master ? run_number : 0; + MPI_Bcast(&run_number_bc, 1, MPI_INT, 0, MPI_COMM_WORLD); + run_number = run_number_bc; + + // ---- checkpointing (master save, all-rank load; params replicated) ---- + const std::string run_prefix = cfg.ckpt_prefix + "_run" + std::to_string(run_number); + + // ───────────────────────────────────────────────────────────────────────── + // rope_theta train/eval agreement. rope_theta is baked into trained weights, + // but it is an env knob, so evaluating a theta=10000 run WITHOUT exporting + // CP_ROPE_THETA=10000 would silently rebuild the cache at the compiled default + // and produce plausible-but-wrong numbers for an entire sweep row. Same + // "two values must agree" class as the CREAM S_search == cream_scaled_max + // coupling. Persisted to a sidecar (no checkpoint format change) and resolved + // by an explicit precedence table: + // + // sidecar env behaviour + // present unset use the sidecar (authoritative) + // present agrees proceed + // present disagrees DIE, naming both values + // absent unset WARN, keep the compiled default, proceed <-- legacy + // absent set use env, warn provenance unverified, write sidecar + // + // The legacy row is load-bearing: every checkpoint predating this sidecar + // (run1, run100, run200, ...) would otherwise be bricked. + // ───────────────────────────────────────────────────────────────────────── + { + const std::string meta_path = cfg.ckpt_dir + "/" + run_prefix + ".meta"; + const bool env_set = std::getenv("CP_ROPE_THETA") != nullptr; + long long meta_theta = -1; + if (std::ifstream mf(meta_path); mf) { + std::string key; + while (mf >> key) { + if (key == "rope_theta") { mf >> meta_theta; break; } + } + } + if (meta_theta > 0) { + if (!env_set) { + cfg.rope_theta = static_cast(meta_theta); + if (is_master) + std::cout << "[rope_theta] " << meta_theta << " from " << meta_path + << " (authoritative)\n"; + } else if (meta_theta != static_cast(cfg.rope_theta)) { + die("rope_theta MISMATCH: " + meta_path + " records " + + std::to_string(meta_theta) + " but CP_ROPE_THETA=" + + std::to_string(cfg.rope_theta) + ". The checkpoint was trained at " + + std::to_string(meta_theta) + "; evaluating it at another base is " + "silently wrong. Unset CP_ROPE_THETA, or set it to " + + std::to_string(meta_theta) + "."); + } + } else if (is_master) { + // No sidecar: a legacy run. Warn, do not die. + std::cout << "[rope_theta] " << cfg.rope_theta + << (env_set ? " from CP_ROPE_THETA (no sidecar -- provenance UNVERIFIED" + " against the checkpoint)" + : " compiled default (no sidecar -- legacy run, assuming" + " it was trained at this base)") + << "\n"; + } + // Record it for this run so the legacy case is self-extinguishing -- but ONLY + // from a path that actually TRAINS. A forward-only run (eval / search / probe) + // of a legacy checkpoint has merely ASSUMED this theta; writing it would + // launder that assumption into an authoritative record that later runs trust, + // and would then `die` against a correct env value. Training is the only + // context where cfg.rope_theta is what the weights were actually built with. + const bool forward_only = + eval_ppl_mode + || (std::getenv("CP_LONGROPE_SEARCH") && std::atoi(std::getenv("CP_LONGROPE_SEARCH")) != 0) + || (std::getenv("CP_MEM_PROBE") && std::atoi(std::getenv("CP_MEM_PROBE")) != 0) + || std::getenv("CP_DUMP_ROPE") != nullptr + || std::getenv("CP_DUMP_ACTS") != nullptr; + if (is_master && cfg.checkpointing && !forward_only && meta_theta <= 0) { + std::error_code ec; + std::filesystem::create_directories(cfg.ckpt_dir, ec); + if (std::ofstream mf(meta_path); mf) + mf << "rope_theta " << cfg.rope_theta << "\n" + << "d_model " << cfg.d_model << "\n" + << "q_heads " << cfg.q_heads << "\n" + << "kv_heads " << cfg.kv_heads << "\n" + << "head_dim " << cfg.head_dim << "\n"; + } + } + + int start_step = 0; + // shard_dir under ZeRO-1. Every rank holds a DIFFERENT slice of the optimizer + // state, so a master-only save persists rank 0's slice and nothing else -- + // and because loads run on every rank, each rank then reads rank 0's bytes + // as though they were its own. For v1/v2 (flat contiguous shards) that is + // SILENT: rank r loads [0,S) into the buffer that owns [r*S,(r+1)*S). + // ShardedMuon happens to catch it, because load_state verifies the owned + // record indices and throws. Either way resume at W>1 was broken. + // + // With shard_dir the manager appends rank_, so every rank saves and loads + // its own shard. The cost is W copies of the (replicated) model weights; + // correctness of the optimizer state is worth more than the bytes. Non-ZeRO + // stays master-only: there the optimizer state is replicated, so one copy is + // the whole truth. + // CP_CKPT_WEIGHTS_ONLY=1: branch a NEW run from a trained base, taking only + // the model weights. This is what a long-context (YaRN / LongRoPE) extension + // needs, and it is the ONLY way to seed from a base checkpoint on a box with + // a different GPU count: ZeRO state is sharded to the world size that wrote + // it, so a normal resume throws on the layout check. Pair it with + // CP_REWARMUP, which is the schedule half of the same idea. + // + // Reads FLAT (not rank_/): a weights-only load is world-size agnostic by + // construction, and at a larger np the rank_ directories would not exist. + // Point CP_CKPT_DIR at a directory holding the flat base .ckpt. + const bool ckpt_weights_only = (std::getenv("CP_CKPT_WEIGHTS_ONLY") != nullptr); + CheckpointManager ckpt(cfg.ckpt_dir, run_prefix, cfg.ckpt_keep, rank, + /*shard_dir=*/(cfg.zero && !ckpt_weights_only), + /*use_async=*/false); + // Weights-only is meaningless with checkpointing disabled: the resume block + // below is gated on cfg.checkpointing, so CP_CKPT=0 would silently skip the + // load and train from scratch -- the exact silent failure this flag exists to + // prevent. Fail loudly instead. + // CP_CKPT_WEIGHTS_ONLY is retired: it loaded the weights HERE, long after the + // optimizers had already snapshotted their fp32 masters from the random init, + // so the first step() overwrote the loaded weights. Superseded by + // CP_INIT_FROM, which loads before any optimizer exists. + if (ckpt_weights_only) { + die("CP_CKPT_WEIGHTS_ONLY is retired -- it silently discarded the loaded " + "weights on the first optimizer step. Use CP_INIT_FROM= " + "instead (loads before the optimizer is built). Pair it with CP_REWARMUP."); + } + if (is_master && ckpt_weights_only) + std::cout << "[Resume] CP_CKPT_WEIGHTS_ONLY: model weights only; optimizer " + "state and RNG are NOT restored (flat checkpoint dir)\n"; + // Where the g>=1 ZeRO sidecars live. Must follow the .ckpt into the rank + // subdirectory, or they would collide across ranks in the parent. + const std::string zero_side_dir = + cfg.zero ? (cfg.ckpt_dir + "/rank_" + std::to_string(rank)) : cfg.ckpt_dir; + if (is_master) { + if (cfg.checkpointing) + std::cout << "Checkpointing: ON -> " << cfg.ckpt_dir << "/" << run_prefix + << "_step_.ckpt (every " << cfg.ckpt_freq << " steps, keep " + << cfg.ckpt_keep << ")\n"; + else + std::cout << "Checkpointing: OFF\n"; + } + // CP_INIT_FROM is EXCLUSIVE with resume. Without this the branch loaded its base + // weights at line ~1612 and then this block ran anyway, found whatever checkpoint + // happened to sit in ckpt_dir, and loaded it ON TOP -- silently evaluating a + // DIFFERENT model. Observed: CP_INIT_FROM=run7_step_750 followed by "[Resume] + // latest run 8 already complete ... Auto-loading blumodelcp_run8_step_1000.ckpt". + // It surfaced only because run8 is bf16 and the probe runs fp32, so the second + // load threw copy_: dtype mismatch; with matching dtypes it would have produced a + // plausible reach number for the wrong checkpoint. A branch defines its own + // starting weights and nothing may override them afterwards. + if (cfg.checkpointing && !init_from_active) { + float resume_loss = 0.0f; + // CheckpointManager's API takes ONE nn::Optimizer&, so under ZeRO-1 group 0 + // rides the normal path alongside the model's full (unsharded) params, and + // every further dtype group round-trips through a sidecar keyed off the same + // step. Without the sidecar a bf16 run resumes with group 1's Adam m/v at + // zero -- it trains on, so the loss curve only looks slightly wrong. + // Loads run on EVERY rank (concurrent reads are safe) so each rank restores + // its own shard; saves are master-only to avoid a write race. + nn::Optimizer& ckpt_opt = cfg.zero + ? static_cast(*zopts[0]) + : static_cast(optimizer); + const bool loaded = ckpt_weights_only + ? ckpt.load_latest_weights_only(model, start_step, resume_loss) + : ckpt.load_latest(model, ckpt_opt, start_step, resume_loss); + // A weights-only load that finds nothing is ALWAYS a mistake: the caller + // asked to branch from a specific base, and silently training a randomly + // initialised model instead looks like a working run for hours. The normal + // resume path may legitimately find no checkpoint (a fresh run), so this + // check is deliberately confined to the weights-only case. + if (ckpt_weights_only && !loaded) { + die("CP_CKPT_WEIGHTS_ONLY=1 but no checkpoint matching prefix '" + + run_prefix + "' was found in '" + cfg.ckpt_dir + "'. Point " + "CP_CKPT_DIR at the directory holding the flat base .ckpt and make " + "sure CP_CKPT_RESUME names its run number. Refusing to train from a " + "random initialisation."); + } + if (loaded) { + // Sidecars carry ZeRO state, which weights-only deliberately drops. + if (cfg.zero && !ckpt_weights_only) { + for (size_t g = 1; g < zopts.size(); ++g) { + const std::string side = zero_side_dir + "/" + run_prefix + "_step_" + + std::to_string(start_step) + ".zero_g" + std::to_string(g); + std::ifstream zis(side, std::ios::binary); + if (zis) zopts[g]->load_state(zis); + else if (is_master) + std::cerr << "[WARN] missing ZeRO sidecar state: " << side << std::endl; + } + } + if (is_master) + std::cout << "[Resume] run " << run_number << " from step " << start_step + << " (loss " << resume_loss << ")" << std::endl; + const size_t batches_to_skip = + static_cast(start_step) * static_cast(GRAD_ACCUM); + // Forward-only eval never reads train_loader: CP_EVAL_PPL scores either + // CP_EVAL_TOKENS_BIN or a val window, so replaying the training stream + // to its resume position is pure cost. skip_batches walks one iteration + // per batch and re-maps a shard on every crossing (tens of thousands of + // iterations at step 1300+), which is a large slice of the fixed + // per-eval overhead in a sweep that launches one process per eval. + // Cannot change eval output -- it touches only the loader's read + // position, never weights, RNG, or the eval's token source. + if (batches_to_skip > 0 && !eval_ppl_mode) { + // Silence the per-shard-crossing debug spam during the replay, then + // emit a single line summarizing where the loader was fast-forwarded to. + static NullBuf null_buf; + std::streambuf* prev = std::cout.rdbuf(&null_buf); + train_loader.skip_batches(batches_to_skip); + std::cout.rdbuf(prev); + if (is_master) + std::cout << "[resume] fast-forwarded data loader " << batches_to_skip + << " batches (" << (static_cast(start_step) * + static_cast(cfg.global_batch)) + << " tokens)" << std::endl; + } + } + } + + // Effective LR schedule: on a resuming run with CP_REWARMUP>0, ramp a fresh + // warmup to a (possibly reduced) peak, then cosine-decay to min_lr by + // max_steps. Otherwise the original absolute-from-0 schedule. + // + // The ramp anchor MUST be a fixed absolute step (the step the extension + // warmup first began at), NOT the current resume step -- otherwise every + // resume of the extension re-anchors a fresh ramp and warmup restarts. + // CP_REWARMUP_ANCHOR= pins it; default (-1) falls back to start_step + // (correct only on the very first extension resume). Read it once, then + // reuse the SAME value on every subsequent resume of the extension run. + const int rewarmup_anchor_env = static_cast(env_i64("CP_REWARMUP_ANCHOR", -1)); + // CP_REWARMUP applies to a BRANCH as well as a resume. The original condition + // required start_step > 0, which is true for longrope_progressive.sh (it COPIES + // the base .ckpt to a new run number and RESUMES it, inheriting prev_step) but + // never for CP_INIT_FROM, which legitimately starts a new run at step 0. Without + // this, CP_REWARMUP_PEAK=0.3 silently did nothing and an extension trained at the + // full max_lr -- 1e-3 on a converged base, which is how you wash out the model you + // were trying to adapt. Same root cause as the CP_CKPT_DIR and master-weight + // ordering bugs: machinery written for the resume path did not fire for the + // initialize-from path, because the two were never distinguished. + const bool rewarmup_active = (rewarmup_steps > 0 && (start_step > 0 || init_from_active)); + const int lr_warmup_start = rewarmup_active + ? (rewarmup_anchor_env >= 0 ? rewarmup_anchor_env : start_step) + : 0; + (void)init_from_base_step; // reported at load; the branch's own schedule starts at 0 + const int lr_warmup = rewarmup_active ? rewarmup_steps : cfg.warmup_steps; + const float lr_peak = rewarmup_active ? (rewarmup_peak_frac * cfg.max_lr) : cfg.max_lr; + if (is_master && rewarmup_active) { + const int steps_into_ramp = start_step - lr_warmup_start; + std::cout << "[re-warmup] anchored at step " << lr_warmup_start << " for " << rewarmup_steps + << " steps -> peak " << lr_peak << " (" << rewarmup_peak_frac + << " x max_lr), then cosine-decay to " << cfg.min_lr + << " | resuming at step " << start_step << " (" + << (steps_into_ramp >= rewarmup_steps ? "past ramp, in cosine-decay" + : (steps_into_ramp < 0 ? "BEFORE anchor -- check CP_REWARMUP_ANCHOR" + : "mid-ramp")) + << ")" << std::endl; + } + + // Per-phase event timers (data/fwd/loss/bwd/clip/optim) + whole-step timer, + // to produce gpt2_cp_test-style per-step [TIMING] breakdown lines. + CudaTimer timer_step, timer_data, timer_fwd, timer_loss, timer_bwd, timer_clip, timer_optim; + + // Shard the targets with the SAME perm GPT::forward applies to the inputs + // (single-axis for ring/ulysses, composed HeadTail-over-ring + contiguous- + // over-ulysses for hybrid). + auto shard_targets = [&](const Tensor& in, const Tensor& tgt) -> ShardedInputs { + if (cfg.hybrid_mode) + return shard_sequence_pre_embed_hybrid( + in, tgt, cfg.T, ring_size, ring_rank, cfg.ulysses_size, ulysses_rank, device); + return shard_sequence_pre_embed( + in, tgt, cfg.T, cp_size, cp_rank, cfg.ring_mode, device); + }; + + // ───────────────────────────────────────────────────────────────────────── + // CP_DUMP_NAMED: identity manifest for the ckpt2hf converter gate. + // Writes a value fingerprint of EACH weight read from its LIVE NAMED MEMBER + // (model.attn[i]->wK.weight, model.mlp[i]->gate_up.weight, ...), NOT by + // re-walking model.parameters()/the serialization order. The converter + // (scripts/convert_ckpt/ckpt2hf.py --verify-manifest) independently + // fingerprints each raw positional .ckpt slot under the same name and + // cross-checks, so the C++ ACTUAL parameter order is verified against the + // Python ASSUMED order -- this is the only check that catches a same-shape + // permutation (q/k_gamma, wK/wV, attn/mlp norm) that shape-validation can't. + // Requires the base checkpoint to be loaded (CP_CKPT_RESUME). Forward-only; + // exits after writing the manifest. Names are CP-native (layers.i.wK, ...) so + // they map 1:1 to positional slots on both sides. + // ───────────────────────────────────────────────────────────────────────── + if (const char* dump_path = std::getenv("CP_DUMP_NAMED")) { + auto write_fp = [](std::ostream& os, const Tensor& t) { + const Tensor c = t.as_type(Dtype::Float32).to_cpu(); + const float* p = c.data(); + const size_t n = c.numel(); + double sum = 0.0, sumsq = 0.0; + for (size_t i = 0; i < n; ++i) { const double v = p[i]; sum += v; sumsq += v * v; } + // Fields (must match the header): n,sum,sumsq,e_first,e_mid,e_last. + os << n << "," << std::setprecision(17) + << sum << "," << sumsq << "," + << (n ? static_cast(p[0]) : 0.0) << "," + << (n ? static_cast(p[n / 2]) : 0.0) << "," + << (n ? static_cast(p[n - 1]) : 0.0); + }; + if (is_master) { + std::ofstream mf(dump_path); + mf << "name,n,sum,sumsq,e_first,e_mid,e_last\n"; + for (int i = 0; i < cfg.n_layers; ++i) { + const std::string pre = "layers." + std::to_string(i) + "."; + mf << pre << "q_gamma,"; write_fp(mf, model.attn[i]->q_gamma); mf << "\n"; + mf << pre << "k_gamma,"; write_fp(mf, model.attn[i]->k_gamma); mf << "\n"; + mf << pre << "attn_norm,"; write_fp(mf, model.attn[i]->norm.weight); mf << "\n"; + mf << pre << "wQ,"; write_fp(mf, model.attn[i]->wQ.weight); mf << "\n"; + mf << pre << "wK,"; write_fp(mf, model.attn[i]->wK.weight); mf << "\n"; + mf << pre << "wV,"; write_fp(mf, model.attn[i]->wV.weight); mf << "\n"; + mf << pre << "wO,"; write_fp(mf, model.attn[i]->wO.weight); mf << "\n"; + mf << pre << "mlp_norm,"; write_fp(mf, model.mlp[i]->norm.weight); mf << "\n"; + mf << pre << "gate_up,"; write_fp(mf, model.mlp[i]->gate_up.weight); mf << "\n"; + mf << pre << "down,"; write_fp(mf, model.mlp[i]->down.weight); mf << "\n"; + } + mf << "wte,"; write_fp(mf, model.wte.weight); mf << "\n"; + mf << "norm_f,"; write_fp(mf, model.norm_f.weight); mf << "\n"; + if (!cfg.weight_tying && model.lm_head) { mf << "lm_head,"; write_fp(mf, model.lm_head->weight); mf << "\n"; } + std::cout << "[dump-named] wrote identity manifest (live named members) -> " + << dump_path << std::endl; + } + MPI_Finalize(); + return 0; + } + + // ───────────────────────────────────────────────────────────────────────── + // CP_DUMP_ACTS: per-layer hidden-state dump for the HF forward-parity debug. + // Runs a FIXED short input (first CP_DUMP_ACTS_T tokens of CP_EVAL_TOKENS_BIN, + // default 32 == fused-kernel tile) through the SAME path the eval uses + // (CP_ATTN_MODE=ulysses, CP_SIZE=1), tapping x after embedding, after each + // transformer block (attn+mlp), after final norm, and the logits. Each tap is + // written raw float32 to /cp_.bin. The HF side dumps the + // same taps (output_hidden_states) so a layer-by-layer diff localizes the first + // divergent op. Forward-only; exits. + // ───────────────────────────────────────────────────────────────────────── + if (const char* acts_dir = std::getenv("CP_DUMP_ACTS")) { + const char* tb = std::getenv("CP_EVAL_TOKENS_BIN"); + if (!tb) throw std::runtime_error("CP_DUMP_ACTS requires CP_EVAL_TOKENS_BIN (fixed input)"); + const int Ta = static_cast(env_i64("CP_DUMP_ACTS_T", 32)); // multiple of 32 (fused tile) + std::ifstream tf(tb, std::ios::binary); + if (!tf) throw std::runtime_error(std::string("CP_DUMP_ACTS: cannot open ") + tb); + std::vector toks(static_cast(Ta)); + tf.read(reinterpret_cast(toks.data()), static_cast(Ta) * 4); + std::vector xh(static_cast(Ta)); + for (int t = 0; t < Ta; ++t) { + const int32_t v = toks[static_cast(t)]; + if (v < 0 || v >= cfg.vocab_size) + throw std::runtime_error("CP_DUMP_ACTS: input token out of [0,vocab)"); + xh[static_cast(t)] = static_cast(v); + } + Tensor idx(Shape{{1, static_cast(Ta)}}, + TensorOptions().with_dtype(Dtype::UInt16).with_device(device)); + cudaMemcpy(idx.data(), xh.data(), static_cast(Ta) * sizeof(uint16_t), + cudaMemcpyHostToDevice); + + auto dump = [&](const Tensor& t, const std::string& name) { + if (!is_master) return; + const Tensor c = t.as_type(Dtype::Float32).to_cpu(); + std::ofstream f(std::string(acts_dir) + "/cp_" + name + ".bin", std::ios::binary); + f.write(reinterpret_cast(c.data()), + static_cast(c.numel()) * sizeof(float)); + }; + + autograd::NoGradGuard ng; + Tensor x = model.wte.forward(idx); // [1, Ta, d_model] + dump(x, "embed"); + { // layer-0 attention sub-taps: split "before kernel" (input-norm + Q/K proj) + // from "inside kernel" (QK-norm+RoPE+scores). Pre-rope, pre-QK-norm. + Tensor h0 = model.attn[0]->norm.forward(x); + dump(h0, "l0_norm"); + dump(model.attn[0]->wQ.forward(h0), "l0_qproj"); + dump(model.attn[0]->wK.forward(h0), "l0_kproj"); + } + for (int i = 0; i < cfg.n_layers; ++i) { + x = model.attn[i]->forward(x); // RMSNorm -> fused(QKnorm+RoPE+GQA) -> Wo -> +x + dump(x, "attn" + std::to_string(i)); // post-attention residual (== HF input to post_attn_ln) + x = model.mlp[i]->forward(x); // RMSNorm -> SwiGLU -> +x + dump(x, "block" + std::to_string(i)); // full block output (== HF hidden_states[i+1]) + } + x = model.norm_f.forward(x); + dump(x, "final"); + Tensor logits; + if (cfg.weight_tying) { + Tensor w_t = autograd::transpose(model.wte.weight, 0, 1); + logits = autograd::matmul(x, w_t); + } else { + logits = model.lm_head->forward(x); + } + dump(logits, "logits"); + if (is_master) + std::cout << "[dump-acts] wrote per-layer activations (T=" << Ta << ") -> " + << acts_dir << std::endl; + MPI_Finalize(); + return 0; + } + + // ───────────────────────────────────────────────────────────────────────── + // Long-context PPL-vs-position eval (forward-only; exits before training). + // Requires cp_size==1 so model.forward returns full [B,T,vocab] and each + // position t maps directly to the global index (no cross-rank NLL gather). + // Teacher-forces CP_EVAL_WINDOWS contiguous windows of length T from the + // chosen split and accumulates per-position NLL; PPL(t) = exp(mean_nll(t)). + // ───────────────────────────────────────────────────────────────────────── + // CP_LONGROPE_SEARCH: resident evaluator for the LongRoPE GA driver. Loads the + // model+ckpt ONCE (above), then loops reading candidate factor-file paths from + // stdin; for each it rebuilds ONLY the cheap CPU LongRoPE cache, installs it on + // all layers, runs forward-only PPL over CP_EVAL_WINDOWS, and prints "PPL " + // to stdout. Avoids reloading the model per candidate (~2560 candidates/search). + // The Python GA feeds paths on stdin and reads PPL back. Requires CP_SIZE=1. + if (std::getenv("CP_LONGROPE_SEARCH") && std::atoi(std::getenv("CP_LONGROPE_SEARCH")) != 0) { + // CP-enabled resident evaluator. Requires a SINGLE DP replica so every MPI + // rank is a CP rank and MPI_COMM_WORLD is the reduction group (dp_size==1 => + // every rank reads the SAME val batch, then forward()/shard_targets shard it). + if (world_size != cp_size || dp_size != 1) + die("CP_LONGROPE_SEARCH requires a single DP replica: world_size == CP_SIZE and " + "dp_size == 1 (got world_size=" + std::to_string(world_size) + + ", CP_SIZE=" + std::to_string(cp_size) + ", dp_size=" + std::to_string(dp_size) + ")"); + const int64_t half = cfg.head_dim / 2; + const int n_windows = static_cast(env_i64("CP_EVAL_WINDOWS", 5)); + const int skip_win = static_cast(env_i64("CP_EVAL_SKIP_WINDOWS", 0)); + // Fitness CE is chunked over the local sequence (see the window loop below). + // 16384 x 50304 = 8.2e8 elements per call -- an order of magnitude under the + // 2^32 bound where a single whole-tensor call was measured to go wrong. + const int Tlocal_ce = static_cast(cfg.T) / cp_size; // logits shard length + const int ce_chunk = std::max(1, static_cast( + env_i64("CP_SEARCH_CE_CHUNK", 16384))); + // If CP_LONGROPE_MSCALE is set, bake YaRN's temperature into every candidate's + // cache (fair-vs-YaRN search). m is per-candidate from the file's s (below). + const bool lr_mscale_on = std::getenv("CP_LONGROPE_MSCALE") + && std::atoi(std::getenv("CP_LONGROPE_MSCALE")) != 0; + if (is_master) + std::cerr << "[longrope-search] resident CP evaluator ready: cp=" << cp_size + << " T=" << cfg.T << " windows=" << n_windows << " skip=" << skip_win + << " (feed candidate factor-file paths on stdin; 'PPL ' per line back)\n"; + // Per-candidate lockstep loop: rank 0 owns stdin + parsing, then BROADCASTS + // the parsed genome (ok flag + lambda + n_hat + s) so all ranks act together. + // A parse error broadcasts ok=0 => ALL ranks `continue` (no forward, no + // collective) — never a partial set entering NCCL. EOF broadcasts ok=-1. + while (true) { + int ok = 0; + std::string path; + std::vector lam(static_cast(half), 1.0f); + int64_t nhat = 0; int s_ = 0; + if (is_master) { + if (std::getline(std::cin, path) && !path.empty()) { + int64_t slen = 0; + try { parse_longrope_factors(path, half, lam, nhat, s_, slen); ok = 1; } + catch (const std::exception& ex) { + std::cout << "ERR " << ex.what() << std::endl << std::flush; ok = 0; + } + } else { + ok = -1; // EOF / blank -> terminate + } + } + if (cp_size > 1) MPI_Bcast(&ok, 1, MPI_INT, 0, MPI_COMM_WORLD); + if (ok == -1) break; + if (ok == 0) continue; + if (cp_size > 1) { + MPI_Bcast(lam.data(), static_cast(half), MPI_FLOAT, 0, MPI_COMM_WORLD); + long long nhat_ll = static_cast(nhat); + MPI_Bcast(&nhat_ll, 1, MPI_LONG_LONG, 0, MPI_COMM_WORLD); + nhat = static_cast(nhat_ll); + MPI_Bcast(&s_, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + // Rebuild only the cache (cheap CPU loop, GLOBAL length) and swap it in. + const float cand_mscale = (lr_mscale_on && s_ > 1) + ? (0.1f * std::log(static_cast(s_)) + 1.0f) : 1.0f; + Tensor lr = autograd::build_rope_cache_longrope( + cfg.T, cfg.head_dim, static_cast(cfg.rope_theta), lam, nhat, device, cand_mscale); + for (auto& a : model.attn) a->cp_->set_rope_cache(lr); + + // Fitness = mean NLL on the GPU (no host logits copy, no manual logsumexp). + // Each rank scores its local shard [B, T/cp_size, V]; equal token counts => + // averaging per-rank means then /world_size gives the global mean NLL. + double sum_loss = 0.0; int nwin = 0; + { + autograd::NoGradGuard no_grad; + val_loader.reset(); + if (skip_win > 0) val_loader.skip_batches(static_cast(skip_win)); + for (int w = 0; w < n_windows; ++w) { + Batch b = val_loader.next_batch(); + ShardedInputs shy = shard_targets(b.input, b.target); + Tensor logits = model.forward(b.input); // [B, T_local, V] + // CE over SEQUENCE CHUNKS, not the whole [B, T_local, V] at once. + // MEASURED: a single call is correct while B*T_local*V stays under + // 2^32 elements but silently wrong past it -- at CP=1/T=131072 + // (6.59e9 elems) it reports PPL 7.55 where the CSV path, which reads + // the SAME logits in 4096-position chunks, reports 59.87; at + // T<=65536 (3.30e9 elems) the two agree to 6 significant figures. + // Chunking keeps every CE call far below that bound, so the search + // fitness no longer depends on picking a CP degree that happens to + // make the shard small enough. + // Exactness: chunks are equal-length except possibly the last, and + // each contributes B*clen tokens, so the token-weighted mean of the + // per-chunk means IS the whole-window mean. When T_local <= chunk we + // take the single-call path unchanged, keeping the already-verified + // short-T results bit-identical. + double win_loss = 0.0; int64_t win_tok = 0; + for (int t0 = 0; t0 < Tlocal_ce; t0 += ce_chunk) { + const int clen = std::min(ce_chunk, Tlocal_ce - t0); + const bool whole = (clen == Tlocal_ce); + Tensor lc = whole ? logits : logits.narrow(1, t0, clen).contiguous(); + Tensor yc = whole ? shy.y_local : shy.y_local.narrow(1, t0, clen).contiguous(); + Tensor loss = autograd::sparse_cross_entropy_loss(lc, yc); + win_loss += static_cast(loss.to_cpu().data()[0]) + * static_cast(clen); + win_tok += clen; + } + sum_loss += (win_tok > 0) ? (win_loss / static_cast(win_tok)) : 0.0; + ++nwin; + } + } + float local_mean = nwin > 0 ? static_cast(sum_loss / nwin) : 0.0f; + float global_mean = local_mean; + if (cp_size > 1) + MPI_Allreduce(&local_mean, &global_mean, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD); + global_mean /= static_cast(world_size); + if (is_master) + std::cout << "PPL " << std::setprecision(8) << std::exp(global_mean) + << std::endl << std::flush; + } + MPI_Finalize(); + return 0; + } + + // ───────────────────────────────────────────────────────────────────────── + // CP_GENERATE: autoregressive decoding for match-scored long-context + // benchmarks (RULER / NIAH / VAL), which score GENERATED text and therefore + // cannot use the teacher-forced CP_EVAL_PPL path. + // + // Ports the proven loop from gpt2_cp_test.cpp:1237-1346 (last-position + // gather -> topk -> pick -> rank-0 broadcast -> append). Three things differ + // because bluscriptCP is Llama-style with RoPE INSIDE ContextParallel: + // + // 1. PADDING TO 256, not to world_size. CP_ATTN_FUSION=split dispatches the + // Hopper kernel whose guard is hd==128 && T_q%256==0 && T_k%128==0 + // (GQA_causal_cp_fwd_sm90.cu:579) and THROWS otherwise. Trailing pad is + // safe under causal masking -- it cannot influence position seq_len-1, + // which is the only logit we read. 256 is a multiple of 2*world_size for + // every world_size we run, so it also satisfies forward_cp's lb_div_ok. + // 2. RoPE CACHE BOUND. The cos/sin cache is built for cfg.T, so the padded + // length must not exceed it -- prompt + max_new must fit inside CP_T. + // Checked up front rather than indexing past the cache. + // 3. GREEDY BY DEFAULT. Benchmarks are scored by exact match, so sampling + // would make results irreproducible. CP_GEN_GREEDY=0 opts into top-k. + // + // Input : CP_GEN_TOKENS_BIN = int32 stream of CP_GEN_N prompts, each exactly + // CP_GEN_PROMPT_LEN tokens, laid end to end. + // Output: CP_GEN_OUT, one line per prompt ",,,..." (token ids; + // the scorer matches in token space, so no detokenizer is needed). + // ───────────────────────────────────────────────────────────────────────── + if (std::getenv("CP_GENERATE") && std::atoi(std::getenv("CP_GENERATE")) != 0) { + const char* tb = std::getenv("CP_GEN_TOKENS_BIN"); + if (!tb) die("CP_GENERATE requires CP_GEN_TOKENS_BIN"); + if (world_size != cp_size || dp_size != 1) + die("CP_GENERATE requires world_size == CP_SIZE and dp_size == 1 (got world_size=" + + std::to_string(world_size) + ", CP_SIZE=" + std::to_string(cp_size) + ")"); + + const int P_LEN = static_cast(env_i64("CP_GEN_PROMPT_LEN", 0)); + const int N_PROMPT= static_cast(env_i64("CP_GEN_N", 1)); + const int MAX_NEW = static_cast(env_i64("CP_GEN_MAX_NEW", 32)); + const int TOPK = static_cast(env_i64("CP_GEN_TOPK", 50)); + // Substantive tokens required before a newline may end generation; see the + // early-stop note in the decode loop below. + const int MIN_TOK = static_cast(env_i64("CP_GEN_MIN_TOKENS", 3)); + const bool greedy = env_i64("CP_GEN_GREEDY", 1) != 0; + const std::string out_path = std::getenv("CP_GEN_OUT") + ? std::getenv("CP_GEN_OUT") : "generations.csv"; + if (P_LEN <= 0) die("CP_GENERATE requires CP_GEN_PROMPT_LEN > 0"); + + const int PAD = 256; // see note 1 above + auto pad_up = [&](int n) { return ((n + PAD - 1) / PAD) * PAD; }; + const int max_padded = pad_up(P_LEN + MAX_NEW); + if (max_padded > static_cast(cfg.T)) + die("CP_GENERATE: prompt(" + std::to_string(P_LEN) + ") + max_new(" + + std::to_string(MAX_NEW) + ") pads to " + std::to_string(max_padded) + + " which exceeds CP_T=" + std::to_string(cfg.T) + + " (the RoPE cache length). Raise CP_T or shorten the prompt."); + + std::ifstream tf(tb, std::ios::binary); + if (!tf) die(std::string("CP_GENERATE: cannot open ") + tb); + + std::ofstream of; + if (is_master) { + of.open(out_path); + if (!of) die("CP_GENERATE: cannot write " + out_path); + } + if (is_master) + std::cout << "[gen] " << N_PROMPT << " prompts x " << P_LEN + << " tok, max_new=" << MAX_NEW + << (greedy ? " greedy" : " topk") << " -> " << out_path << std::endl; + + autograd::NoGradGuard ng; + model.set_generation_mode(true); + + for (int pi = 0; pi < N_PROMPT; ++pi) { + std::vector raw(static_cast(P_LEN)); + tf.read(reinterpret_cast(raw.data()), + static_cast(P_LEN) * 4); + if (!tf) die("CP_GENERATE: short read on prompt " + std::to_string(pi)); + + std::vector seq; + seq.reserve(static_cast(P_LEN + MAX_NEW)); + for (int t = 0; t < P_LEN; ++t) { + const int32_t v = raw[static_cast(t)]; + if (v < 0 || v >= cfg.vocab_size) + die("CP_GENERATE: prompt token out of [0,vocab)"); + seq.push_back(static_cast(v)); + } + + std::vector produced; + // Early stop: answers are 3-5 tokens but the loop would otherwise run + // all MAX_NEW forwards, and each forward is a full T=16384 pass. This + // counter MUST live inside the per-prompt loop: hoisted out, it would + // stay satisfied after the first prompt and every later prompt would + // break at its first newline -- invisible except as a low score. + int n_sub = 0; + for (int step = 0; step < MAX_NEW; ++step) { + const int cur = static_cast(seq.size()); + const int padded = pad_up(cur); + + std::vector buf(static_cast(padded), 0); + std::copy(seq.begin(), seq.end(), buf.begin()); + Tensor idx(Shape{{1, static_cast(padded)}}, + TensorOptions().with_dtype(Dtype::UInt16).with_device(device)); + cudaMemcpy(idx.data(), buf.data(), + static_cast(padded) * sizeof(uint16_t), + cudaMemcpyHostToDevice); + + Tensor logits = model.forward(idx); // [1, padded, vocab] + const int64_t Vg = logits.shape().dims[2]; + + // Logit row that predicts the NEXT token is at cur-1. + Tensor gidx = Tensor::full( + Shape{{1, 1, Vg}}, + TensorOptions().with_dtype(Dtype::Int64).with_device(device), + static_cast(cur - 1)); + Tensor last = OwnTensor::gather(logits, 1, gidx); // [1,1,V] + + // Pick on rank 0 only, then broadcast. Every rank MUST append the + // same token or their shards diverge and the next step is garbage. + // Greedy is deterministic given identical logits, and unshard=true + // makes the ranks' logits identical -- but deciding on one rank and + // broadcasting is what actually guarantees it, so a future + // non-deterministic reduction cannot silently desynchronise them. + int64_t next = 0; + if (rank == 0) { + if (greedy) { + auto tk = last.topk(1, -1); + next = tk.second.to_cpu().data()[0]; + } else { + Tensor probs = autograd::softmax(last, -1); + auto tk = probs.topk(TOPK, -1); + Tensor p2 = autograd::reshape(tk.first, Shape{{1, TOPK}}); + Tensor i2 = autograd::reshape(tk.second, Shape{{1, TOPK}}); + Tensor pick = Tensor::multinomial(p2, 1); + next = OwnTensor::gather(i2, 1, pick).to_cpu().data()[0]; + } + } + if (world_size > 1) + MPI_Bcast(&next, 1, MPI_INT64_T, 0, MPI_COMM_WORLD); + + produced.push_back(next); + seq.push_back(static_cast(next)); + + // Stop at a newline, but only once the model has emitted MIN_TOK + // real tokens. The answer prefix ends "... is" with no newline, so + // an untuned base model may emit "\n" FIRST and the answer after; + // breaking on that would truncate the row to nothing and turn a HIT + // into a MISS. MIN_TOK is per task (14 for niah_multivalue, which + // needs three 7-digit values) -- see MINTOK in run_longeval.sh. + // + // The break sits AFTER the MPI_Bcast deliberately: every rank must + // observe the same token and reach the same decision, or the ranks + // desynchronise. (Moot at NP=1, but this block is shared.) + if (next != 198 && next != 628) ++n_sub; // 198="\n", 628="\n\n" + else if (n_sub >= MIN_TOK) break; + } + + if (is_master) { + of << pi; + for (int64_t t : produced) of << "," << t; + of << "\n"; + of.flush(); + if ((pi + 1) % 10 == 0 || pi + 1 == N_PROMPT) + std::cout << "[gen] " << (pi + 1) << "/" << N_PROMPT << std::endl; + } + } + + model.set_generation_mode(false); + if (is_master) { of.close(); std::cout << "[gen] wrote " << out_path << std::endl; } + MPI_Finalize(); + return 0; + } + + if (eval_ppl_mode) { + // CP-enabled per-position PPL. Single DP replica only (world==cp, dp==1) so + // MPI_COMM_WORLD is the reduction group; hybrid perm unsupported in v1. + if (world_size != cp_size || dp_size != 1) + die("CP_EVAL_PPL requires a single DP replica: world_size == CP_SIZE and dp_size == 1 " + "(got world_size=" + std::to_string(world_size) + ", CP_SIZE=" + + std::to_string(cp_size) + ", dp_size=" + std::to_string(dp_size) + ")"); + // Hybrid (ring x ulysses) IS supported: the per-position global map is read + // from shy.pos_local (the shard's own composed perm), not hand-derived below. + + int n_windows = static_cast(env_i64("CP_EVAL_WINDOWS", 10)); + const int skip_win = static_cast(env_i64("CP_EVAL_SKIP_WINDOWS", 0)); + const std::string out_csv = std::getenv("CP_EVAL_OUT") + ? std::getenv("CP_EVAL_OUT") : "ppl_vs_pos.csv"; + const int Tpos = static_cast(cfg.T); // GLOBAL length (CSV rows) + const int Tlocal = static_cast(cfg.T) / cp_size; // local shard length (logits) + const int Bloc = static_cast(cfg.B); + const int64_t V = cfg.vocab_size; + + // ─── Eval-faithfulness: optional shared token stream ──────────────────── + // When CP_EVAL_TOKENS_BIN is set we bypass val_loader and window a raw + // int32 little-endian token file directly (input=toks[w*T .. w*T+T-1], + // target=toks[w*T+1 .. w*T+T]) with B forced to 1, so the exact same + // [pos] windows are scored here AND by the HF scorer (scripts/eval/ + // ppl_vs_pos.py, which also uses B=1). This makes the CP-vs-HF PPL + // comparison read byte-identical token IDs. When UNSET, behavior below is + // byte-identical to the previous val_loader path (additive change). + const char* tok_bin = std::getenv("CP_EVAL_TOKENS_BIN"); + std::vector eval_toks; + const bool bypass = (tok_bin != nullptr); + if (bypass) { + std::ifstream tf(tok_bin, std::ios::binary | std::ios::ate); + if (!tf) throw std::runtime_error(std::string("CP_EVAL_TOKENS_BIN not readable: ") + tok_bin); + const std::streamsize nbytes = tf.tellg(); + tf.seekg(0); + if (nbytes < 0 || (nbytes % 4) != 0) + throw std::runtime_error("CP_EVAL_TOKENS_BIN size not a multiple of 4 (expect int32 tokens)"); + eval_toks.resize(static_cast(nbytes) / 4); + tf.read(reinterpret_cast(eval_toks.data()), nbytes); + // Windows AVAILABLE after skipping the calibration windows (held-out). + const long long total_w = (static_cast(eval_toks.size()) - 1) / Tpos; + if (total_w - skip_win <= 0) + throw std::runtime_error("CP_EVAL_TOKENS_BIN too short for skip+one window at T=CP_T"); + if (n_windows > total_w - skip_win) { + if (is_master) + std::cout << "[eval] clamping windows " << n_windows << " -> " << (total_w - skip_win) + << " (stream has " << eval_toks.size() << " tokens at T=" << Tpos + << ", skip=" << skip_win << ")\n"; + n_windows = static_cast(total_w - skip_win); + } + if (is_master) + std::cout << "[eval] shared token stream: " << eval_toks.size() << " int32 tokens from " + << tok_bin << " (B forced to 1 for HF-scorer parity)" << std::endl; + } + + if (is_master) + std::cout << "[eval] PPL-vs-position: " << n_windows << " windows (skip=" << skip_win + << ") x B=" << (bypass ? 1 : Bloc) << " Tglobal=" << Tpos << " Tlocal=" << Tlocal + << " cp=" << cp_size << " (YARN_SCALE via env) -> " << out_csv << std::endl; + + // Per-GLOBAL-position accumulators (each rank fills only its shard's global + // slots; MPI_Reduce combines to the master which writes the full [0,T) CSV). + std::vector sum_nll(static_cast(Tpos), 0.0); + std::vector cnt(static_cast(Tpos), 0); + + // CP_EVAL_NLL_HOST=1 forces the legacy host-side logsumexp. Default is the + // on-GPU path: the host path copied the full logits D2H (~13 GB/eval at + // T=16384) and was CPU-bound (GPU 0%, 124 W of 700 W), making eval ~36x the + // cost of training. Keeping both is what makes the numeric A/B gate possible. + const bool nll_host = std::getenv("CP_EVAL_NLL_HOST") != nullptr && + std::atoi(std::getenv("CP_EVAL_NLL_HOST")) != 0; + double* d_sum = nullptr; + long long* d_cnt = nullptr; + if (!nll_host) { + cudaMalloc(&d_sum, static_cast(Tpos) * sizeof(double)); + cudaMalloc(&d_cnt, static_cast(Tpos) * sizeof(long long)); + cudaMemset(d_sum, 0, static_cast(Tpos) * sizeof(double)); + cudaMemset(d_cnt, 0, static_cast(Tpos) * sizeof(long long)); + } + if (is_master) + std::cout << "[eval] per-position NLL: " + << (nll_host ? "HOST (legacy logsumexp)" : "GPU") << std::endl; + // Host-copy chunk over LOCAL positions -> peak host RAM is CHUNK*V floats + // per rank, NOT Tlocal*V (which is 13.2 GB/rank at s=64). MANDATORY (round 3). + const int CHUNK = 4096; + std::vector host_buf; + + autograd::NoGradGuard no_grad; + if (!bypass) { val_loader.reset(); if (skip_win > 0) val_loader.skip_batches(static_cast(skip_win)); } + // CREAM (additive): install ONE fixed labeling for the whole eval. No-op when + // off, and CREAM is single-GPU (dies earlier under CP), so inert here at cp>1. + if (cream_on) model.install_cream_positions(cfg.cream_seed); + for (int w = 0; w < n_windows; ++w) { + // Build the FULL [Bcur, Tglobal] input; forward() shards it internally and + // returns local [Bcur, Tlocal, V] logits (y_local is the matching shard). + Tensor input, y_local, pos_local; // pos_local[1,Tlocal] = global pos of each local idx + int Bcur; + if (bypass) { + Bcur = 1; + std::vector xh(static_cast(Tpos)), yh(static_cast(Tpos)); + const size_t base = static_cast(skip_win + w) * static_cast(Tpos); + for (int t = 0; t < Tpos; ++t) { + const int32_t xi = eval_toks[base + t]; + const int32_t yi = eval_toks[base + t + 1]; + if (xi < 0 || xi >= V) + throw std::runtime_error("CP_EVAL_TOKENS_BIN: input token out of [0,vocab)"); + if (yi < 0 || yi >= 65536) + throw std::runtime_error("CP_EVAL_TOKENS_BIN: target token out of uint16 range"); + xh[static_cast(t)] = static_cast(xi); + yh[static_cast(t)] = static_cast(yi); + } + input = Tensor(Shape{{1, static_cast(Tpos)}}, + TensorOptions().with_dtype(Dtype::UInt16).with_device(device)); + Tensor tgt(Shape{{1, static_cast(Tpos)}}, + TensorOptions().with_dtype(Dtype::UInt16).with_device(device)); + cudaMemcpy(input.data(), xh.data(), static_cast(Tpos) * sizeof(uint16_t), + cudaMemcpyHostToDevice); + cudaMemcpy(tgt.data(), yh.data(), static_cast(Tpos) * sizeof(uint16_t), + cudaMemcpyHostToDevice); + ShardedInputs shy = shard_targets(input, tgt); // shards to [1, Tlocal] + y_local = shy.y_local; pos_local = shy.pos_local; + } else { + Batch b = val_loader.next_batch(); + Bcur = Bloc; + input = b.input; + ShardedInputs shy = shard_targets(b.input, b.target); // shards to [Bloc, Tlocal] + y_local = shy.y_local; pos_local = shy.pos_local; + } + Tensor logits = model.forward(input); // [Bcur, Tlocal, V] + Tensor logits_dev = (logits.dtype() == Dtype::Float32) ? logits + : logits.as_type(Dtype::Float32); + const float* dptr = reinterpret_cast(logits_dev.data()); // DEVICE ptr + + // Targets: y_local inherits the loader dtype (UInt16) -- read in actual dtype. + Tensor y_cpu = y_local.to_cpu(); // [Bcur, Tlocal] (small) + const Dtype ydt = y_cpu.dtype(); + const uint16_t* yp16 = (ydt == Dtype::UInt16) ? y_cpu.data() : nullptr; + const int32_t* yp32 = (ydt == Dtype::Int32) ? y_cpu.data() : nullptr; + const int64_t* yp64 = (ydt == Dtype::Int64) ? y_cpu.data() : nullptr; + if (!yp16 && !yp32 && !yp64) + throw std::runtime_error("CP_EVAL_PPL: unexpected target dtype for y_local"); + + // Global position of each LOCAL index, straight from the shard's own perm + // (pos_local). Works for contiguous, HeadTail, AND hybrid with no map math. + Tensor pos_cpu = pos_local.to_cpu(); // [1, Tlocal] int64 + const int64_t* posp = pos_cpu.data(); + + // Per-LOCAL-position NLL. Each local tl maps to a unique GLOBAL position g + // (bijective per rank), which is why neither the OMP writes below nor the + // GPU path's block writes can collide -- and why the GPU path needs no + // atomics. The permutation holds PER ROW, so both paths loop over bi. + if (!nll_host) { + // One launch per batch row: within a row the scatter is a permutation, + // across rows accumulation is sequential -> deterministic for any Bcur. + Tensor y32 = y_local.as_type(Dtype::Int32).contiguous(); + Tensor pos32 = pos_local.contiguous(); + const int32_t* yg = y32.data(); + const int64_t* pg = pos32.data(); + for (int bi = 0; bi < Bcur; ++bi) { + OwnTensor::cp::cuda::eval_nll_accumulate_row( + dptr + static_cast(bi) * Tlocal * V, + yg + static_cast(bi) * Tlocal, + pg, d_sum, d_cnt, + static_cast(Tlocal), V, Tpos); + } + } else + for (int bi = 0; bi < Bcur; ++bi) { + for (int tl0 = 0; tl0 < Tlocal; tl0 += CHUNK) { + const int clen = std::min(CHUNK, Tlocal - tl0); + host_buf.resize(static_cast(clen) * static_cast(V)); + cudaMemcpy(host_buf.data(), + dptr + (static_cast(bi) * Tlocal + tl0) * static_cast(V), + static_cast(clen) * static_cast(V) * sizeof(float), + cudaMemcpyDeviceToHost); + #pragma omp parallel for schedule(dynamic) + for (int c = 0; c < clen; ++c) { + const int tl = tl0 + c; + const float* row = host_buf.data() + static_cast(c) * static_cast(V); + float mx = row[0]; + for (int64_t v = 1; v < V; ++v) if (row[v] > mx) mx = row[v]; + double se = 0.0; + for (int64_t v = 0; v < V; ++v) se += std::exp(static_cast(row[v] - mx)); + const double lse = static_cast(mx) + std::log(se); + const size_t idx = static_cast(bi) * Tlocal + tl; + const int64_t tgt = yp16 ? static_cast(yp16[idx]) + : yp64 ? yp64[idx] + : static_cast(yp32[idx]); + if (tgt < 0 || tgt >= V) continue; // skip padding/ignore ids + const int g = static_cast(posp[tl]); // shy.pos_local[tl] + sum_nll[static_cast(g)] += lse - static_cast(row[tgt]); + cnt[static_cast(g)] += 1; + } + } + } + if (is_master) + std::cout << "[eval] window " << (w + 1) << "/" << n_windows << " done" << std::endl; + } + + // Device -> host once per eval (not per window): the whole point is that the + // logits never leave the GPU. Frees the accumulators too. + if (!nll_host) { + cudaDeviceSynchronize(); + cudaMemcpy(sum_nll.data(), d_sum, + static_cast(Tpos) * sizeof(double), cudaMemcpyDeviceToHost); + cudaMemcpy(cnt.data(), d_cnt, + static_cast(Tpos) * sizeof(long long), cudaMemcpyDeviceToHost); + cudaFree(d_sum); + cudaFree(d_cnt); + } + + // Combine per-position sums across CP ranks onto the master. + if (cp_size > 1) { + std::vector g_nll(static_cast(Tpos), 0.0); + std::vector g_cnt(static_cast(Tpos), 0); + MPI_Reduce(sum_nll.data(), g_nll.data(), Tpos, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD); + MPI_Reduce(cnt.data(), g_cnt.data(), Tpos, MPI_LONG_LONG, MPI_SUM, 0, MPI_COMM_WORLD); + if (is_master) { sum_nll.swap(g_nll); cnt.swap(g_cnt); } + } + + if (is_master) { + std::ofstream f(out_csv); + f << "pos,count,mean_nll,ppl\n"; + f << std::fixed << std::setprecision(6); + for (int t = 0; t < Tpos; ++t) { + const double mn = cnt[static_cast(t)] > 0 + ? sum_nll[static_cast(t)] / static_cast(cnt[static_cast(t)]) + : 0.0; + f << t << "," << cnt[static_cast(t)] << "," << mn << "," << std::exp(mn) << "\n"; + } + std::cout << "[eval] wrote " << out_csv << " (" << Tpos << " positions, " + << n_windows << " windows)" << std::endl; + } + MPI_Finalize(); + return 0; + } + + if (is_master) std::cout << "\nStarting training...\n" << std::endl; + + for (int step = start_step; step < cfg.max_steps; ++step) { + timer_step.start_timer(); + + // ---------- periodic validation (local shard loss + global avg) ---------- + if (!mem_probe && !std::getenv("CP_NO_VAL") && (step % cfg.val_freq == 0 || step == cfg.max_steps - 1)) { + autograd::NoGradGuard no_grad; + val_loader.reset(); + float val_accum = 0.0f; + for (int v = 0; v < cfg.val_loss_steps; ++v) { + Batch b = val_loader.next_batch(); + ShardedInputs shy = shard_targets(b.input, b.target); + Tensor logits = model.forward(b.input); // [B, T_local, vocab] + Tensor loss = autograd::sparse_cross_entropy_loss(logits, shy.y_local); + val_accum += loss.to_cpu().data()[0] / static_cast(cfg.val_loss_steps); + } + float global_val = val_accum; + MPI_Allreduce(&val_accum, &global_val, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD); + global_val /= world_size; // avg over all dp*cp ranks + last_val_loss = global_val; + if (is_master) + std::cout << "validation loss: " << std::fixed + << std::setprecision(4) << global_val << std::endl; + } + + // ---------- training step ---------- + double time_data = 0, time_forward = 0, time_loss = 0; + double time_backward = 0, time_clip = 0, time_optim = 0; + if (cfg.zero) { for (auto& o : zopts) o->zero_grad(); } + else optimizer.zero_grad(); + loss_accum_gpu *= 0.0f; + + // CREAM (additive): resample position labels once per optimizer step (shared + // across micro-batches) and install the gathered cache on all layers. No-op + // when off. Optional CP_CREAM_LOG= records the middle-block start per + // step for the truncated-Gaussian distribution sanity check. + if (cream_on) { + const uint64_t step_seed = cfg.cream_seed + static_cast(step); + model.install_cream_positions(step_seed); + if (is_master && cream_log_path) { + std::ofstream cf(cream_log_path, std::ios::app); + cf << step << "," << model.cream_middle_start(step_seed) << "\n"; + } + } + + // no_sync()/sync() gate ONLY DDP's param-grad hooks; the CP ring's cp_pg + // comm runs every microstep (inside model.forward/backward) regardless. + if (ddp_model) ddp_model->no_sync(); + for (int micro = 0; micro < GRAD_ACCUM; ++micro) { + timer_data.start_timer(); + Batch batch = train_loader.next_batch(); + // Shard targets with the SAME perm GPT::forward uses on the inputs. + ShardedInputs shy = shard_targets(batch.input, batch.target); + time_data += timer_data.get_elapsed_seconds(); + + timer_fwd.start_timer(); + Tensor logits = ddp_model ? ddp_model->forward(batch.input) + : model.forward(batch.input); // [B, T_local, vocab] + time_forward += timer_fwd.get_elapsed_seconds(); + + timer_loss.start_timer(); + Tensor loss = autograd::sparse_cross_entropy_loss(logits, shy.y_local); + loss_accum_gpu += loss.detach(); + // CP_DIAG=1: per-micro-batch loss. The reported train loss is the MEAN + // over GRAD_ACCUM micros, so one pathological micro is diluted 32x while + // its gradient dominates the accumulated norm almost undiluted -- which is + // exactly how a modest loss bump sits next to a 15x norm spike (observed + // at steps 108/137 of the bf16 run). Printing the values separates "one + // bad micro-batch" from "the whole step is worse". + // Costs a device sync per micro, so it is off by default. Does NOT perturb + // the result: the loss is already fully determined at this point. + // bluscript_zero also records micro_where (shard index + token offset); + // not ported because CP's Batch carries neither field. The per-micro + // vector + printer is likewise not ported here -- this file has no + // diag_micro/micro_loss plumbing -- so only the probe below is kept. + if (std::getenv("CP_FWD_PROBE")) { + // Bitwise forward comparison across ranks: raw loss bits plus a + // logits checksum. A 6-decimal match is not proof of identity. + const float lv = loss.to_cpu().data()[0]; + unsigned int lb; std::memcpy(&lb, &lv, sizeof(lb)); + const Tensor lg = logits.detach().as_type(Dtype::Float32).to_cpu(); + const float* lp = lg.data(); + double s0 = 0.0, s1 = 0.0; + for (size_t i = 0; i < lg.numel(); ++i) { s0 += lp[i]; s1 += (double)lp[i] * lp[i]; } + std::fprintf(stderr, "[rank %d] step %d micro %d loss=0x%08x (%.9g) logits n=%zu sum=%.17g sumsq=%.17g\n", + rank, step, micro, lb, lv, (size_t)lg.numel(), s0, s1); + std::fflush(stderr); + } + time_loss += timer_loss.get_elapsed_seconds(); + + // Gradient reduction fires on the LAST micro-step only. DDP gates its + // hooks with sync(); ZeRO-1 arms its overlapped reduce the same way, so + // the earlier micro-steps accumulate locally and exactly one reduction + // happens per optimizer step. + if (micro == GRAD_ACCUM - 1) { + if (ddp_model && !std::getenv("CP_DDP_NO_REDUCE")) ddp_model->sync(); + if (cfg.zero) { for (auto& o : zopts) o->arm_grad_sync(); } + } + + timer_bwd.start_timer(); + loss.backward(&grad_scale); + time_backward += timer_bwd.get_elapsed_seconds(); + } + + // Logged loss: local mean over micros, then avg over all dp*cp ranks. + float loss_accum; + { + float local = loss_accum_gpu.to_cpu().data()[0] / static_cast(GRAD_ACCUM); + if (std::getenv("CP_RANK_NORM")) { std::fprintf(stderr, "[rank %d] step %d LOCAL loss = %.6f\n", rank, step, local); std::fflush(stderr); } + MPI_Allreduce(&local, &loss_accum, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD); + loss_accum /= world_size; + } + if (std::isnan(loss_accum) || std::isinf(loss_accum)) { + if (is_master) { std::cerr << "ERROR: NaN/Inf loss at step " << step << std::endl; log_file.close(); } + MPI_Abort(MPI_COMM_WORLD, 2); + } + + // Clipping and decoupled weight decay are the OPTIMIZER's job under ZeRO-1 + // (zcfg.max_grad_norm + zcfg.weight_decay + exclude_1d_from_weight_decay), + // and it applies them to its fp32 shard. Doing them here as well would + // clip twice and decay twice -- and the manual `p *= wd_factor` would write + // directly into bf16 parameters, discarding the fp32 master state ZeRO + // maintains. So both are skipped when zero is on. + timer_clip.start_timer(); + // Under ZeRO the norm comes from the optimizer (it clips on its fp32 shard), + // not from a second clip here. Reporting 0.0 instead would throw away the + // early-warning signal that caught the last divergence -- grad norm went + // 1.29 -> 14.17 several steps BEFORE the loss blew up. Max over groups. + float grad_norm = 0.0f; + // CP_GRAD_PROBE: per-parameter grad norm on EVERY rank, to localise a + // rank-dependent backward divergence. Uses clip_grad_norm_ with an + // unreachable max so it measures without scaling. + if (std::getenv("CP_GRAD_PROBE")) { + for (size_t pi = 0; pi < params.size(); ++pi) { + std::vector one{params[pi]}; + const float n = nn::clip_grad_norm_(one, 1e30f); + if (n > 1e3f || pi < 4) + std::fprintf(stderr, "[rank %d] step %d param[%zu] gnorm=%.6g\n", + rank, step, pi, n); + } + std::fflush(stderr); + } + if (!cfg.zero) grad_norm = nn::clip_grad_norm_(params, cfg.grad_clip); + if (std::getenv("CP_RANK_NORM")) { int _cd=-1; cudaGetDevice(&_cd); std::fprintf(stderr, "[rank %d] step %d cudaGetDevice=%d norm=%.4f loss=%.6f\n", rank, step, _cd, grad_norm, loss_accum); std::fflush(stderr); } + time_clip = timer_clip.get_elapsed_seconds(); + + const float lr = get_lr(step, lr_peak, cfg.min_lr, lr_warmup, cfg.max_steps, + lr_warmup_start, cfg.stable_steps); + timer_optim.start_timer(); + if (cfg.zero) { + // Groups here are a dtype split of ONE parameter set (see the ZeRO setup + // above), so they all ride the same WSD curve -- unlike the Muon variant + // in bluscript_zero, where each group carries its own (peak, floor) pair. + // + // The reported norm is the L2 over groups, sqrt(sum g^2) -- NOT the max, + // which would under-report a spike confined to one group. + float gsq = 0.0f; + for (auto& o : zopts) { + o->set_lr(lr); + o->step(); + const float g = o->last_grad_norm(); + gsq += g * g; + } + grad_norm = std::sqrt(gsq); + if (std::getenv("CP_RANK_NORM")) { std::fprintf(stderr, "[rank %d] step %d ZERO norm=%.4f\n", rank, step, grad_norm); std::fflush(stderr); } + } else { + optimizer.set_lr(lr); + const float wd_factor = 1.0f - lr * cfg.weight_decay; + for (auto& p : decay_params) p *= wd_factor; + optimizer.step(); + } + time_optim = timer_optim.get_elapsed_seconds(); + + const double dt = timer_step.get_elapsed_seconds(); + const int64_t tokens = cfg.B * cfg.T * GRAD_ACCUM * static_cast(dp_size); + const double tok_per_sec = static_cast(tokens) / dt; + + // Throughput + time left + GPU memory (gpt2_cp_test-style per-step report). + long long total_sec = static_cast((cfg.max_steps - step) * dt); + int h = static_cast(total_sec / 3600); + int m = static_cast((total_sec % 3600) / 60); + size_t free_mem = 0, total_mem = 0; + cudaMemGetInfo(&free_mem, &total_mem); + const double used_mb = static_cast(total_mem - free_mem) / (1024.0 * 1024.0); + + if (is_master) { + std::cout << "step " << std::setw(5) << step + << " | loss: " << std::fixed << std::setprecision(6) << loss_accum + << " | lr " << std::scientific << std::setprecision(4) << lr + << " | norm: " << std::fixed << std::setprecision(4) << grad_norm + << " | dt: " << std::fixed << std::setprecision(2) << (dt * 1000.0) << "ms" + << " | tok/sec: " << std::fixed << std::setprecision(1) << tok_per_sec + << " | mem: " << std::fixed << std::setprecision(0) << used_mb << "MB" + << " | Time Left: " << std::setfill('0') << std::setw(2) << h << " hrs : " + << std::setw(2) << m << " mins" << std::setfill(' ') << "\n"; + + std::cout << " [TIMING] data: " << std::fixed << std::setprecision(1) + << (time_data * 1000.0) << "ms" + << " | fwd: " << (time_forward * 1000.0) << "ms" + << " | loss: " << (time_loss * 1000.0) << "ms" + << " | bwd: " << (time_backward * 1000.0) << "ms" + << " | clip: " << (time_clip * 1000.0) << "ms" + << " | optim: " << (time_optim * 1000.0) << "ms\n"; + + log_file << step << "," << loss_accum << "," << last_val_loss << "," << lr << "," + << grad_norm << "," << (dt * 1000.0) << "," << tok_per_sec << "," + << (time_data * 1000.0) << "," << (time_forward * 1000.0) << "," + << (time_loss * 1000.0) << "," << (time_backward * 1000.0) << "," + << (time_clip * 1000.0) << "," << (time_optim * 1000.0) << "," + << used_mb << "\n"; + log_file.flush(); + } + + // ── Memory-probe snapshot: after the final probe step, capture nvidia-smi + // + cudaMemGetInfo + allocator stats (live, before teardown), then stop. + // Same snapshot format as gpt2_cp_test.cpp so mem_scaling_table.py parses + // both. tag = CPP_