// =============================================================================
// 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 <cstdint>
#include <cstdlib>
#include <cstring> // memcpy: CP_FWD_PROBE raw-loss bit dump
#include <iostream>
#include <iomanip>
#include <chrono>
#include <cmath>
#include <vector>
#include <string>
#include <algorithm>
#include <fstream>
#include <sstream> // Ms-PoE weight list parsing
#include <filesystem>
#include <map>
#include <random>
#include <mpi.h>
#include <cuda_runtime.h>
// 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<float> 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=<w> 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<float> 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<float> 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<int> 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<uint32_t>(seed)) {}
Tensor normal(const Shape& shape, float std) {
Tensor t(shape, TensorOptions().with_dtype(Dtype::Float32));
std::normal_distribution<float> dist(0.0f, std);
float* d = t.data<float>();
const size_t n = t.numel();
for (size_t i = 0; i < n; ++i) d[i] = dist(gen);
return t;
}
};
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<ContextParallel> 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<CPProcessGroupNCCL> cp_pg,
std::shared_ptr<ProcessGroupNCCL> 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<float>(c.n_layers));
init_linear(wQ, 0.02f, rng);
init_linear(wK, 0.02f, rng);
init_linear(wV, 0.02f, rng);
init_linear(wO, 0.02f * resid_scale, rng);
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<float>(c.head_dim));
const bool lb = (c.ring_mode || c.hybrid_mode);
cp_ = std::make_shared<ContextParallel>(
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<float>(c.n_layers));
init_linear(gate_up, 0.02f, rng);
init_linear(down, 0.02f * resid_scale, rng);
norm.to(dev); gate_up.to(dev); down.to(dev);
register_module(norm);
register_module(gate_up);
register_module(down);
}
Tensor forward(const Tensor& x) override {
// 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<std::shared_ptr<CausalGQA>> attn;
std::vector<std::shared_ptr<SwiGLUMLP>> mlp;
nn::RMSNorm norm_f;
std::shared_ptr<nn::Linear> 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<CPProcessGroupNCCL> cp_pg, int cp_size, int cp_rank,
std::shared_ptr<ProcessGroupNCCL> ulysses_pg = nullptr,
int ring_rank = 0, int ulysses_rank = 0,
uint64_t seed = 1234)
: cfg(c),
init_rng_(seed),
wte(static_cast<int>(c.vocab_size), static_cast<int>(c.d_model)),
cos_sin_cache(autograd::build_rope_cache(c.context_length, c.head_dim,
static_cast<float>(c.rope_theta), dev)),
norm_f(c.d_model, c.rms_eps),
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<float>(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<float>(c.rope_theta), dev);
}
}
for (int i = 0; i < c.n_layers; ++i) {
auto a = std::make_shared<CausalGQA>(c, dev, cos_sin_cache, init_rng_, mesh, cp_pg, ulysses_pg);
auto m = std::make_shared<SwiGLUMLP>(c, dev, init_rng_);
attn.push_back(a);
mlp.push_back(m);
register_module(a.get());
register_module(m.get());
}
if (!c.weight_tying) {
lm_head = std::make_shared<nn::Linear>(static_cast<int>(c.d_model),
static_cast<int>(c.vocab_size), /*bias=*/false,
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<float>(c.rope_theta),
c.longrope_lambda, c.longrope_n_hat, dev, c.longrope_mscale);
Tensor grouped;
if (mspoe_on) {
std::vector<std::vector<float>> 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<size_t>(half), r);
} else {
// COMPOSED variant: scale the base lambda toward 1 per group.
for (float w : c.mspoe_weights) {
std::vector<float> 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<float>(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<int>(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<int>(cfg.T);
const int factor = cfg.cream_factor;
std::vector<int64_t> 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<int64_t>(T), hd}},
TensorOptions().with_dtype(Dtype::Int64));
int64_t* ip = static_cast<int64_t*>(idx_cpu.data());
for (int i = 0; i < T; ++i)
for (int64_t j = 0; j < hd; ++j)
ip[static_cast<int64_t>(i) * hd + j] = L[static_cast<size_t>(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<int>(cfg.T), cfg.cream_factor, cfg.cream_sigma, seed);
size_t head = 0;
while (head < L.size() && L[head] == static_cast<int64_t>(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<float>(local + 1) / static_cast<float>(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<float>(max_steps - decay_start);
if (denom <= 0.0f) denom = 1.0f; // guard: warmup+stable past max_steps
const float decay = static_cast<float>(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 <int>
// s <int> (extension ratio the search used)
// S_search <int> (target length the search targeted; = factor*T when composed)
// lambda <f0> <f1> ... <f_{half-1}> (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<float>& 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<size_t>(half));
for (int64_t i = 0; i < half; ++i)
if (!(f >> lambda[static_cast<size_t>(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<int64_t>(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<size_t>(i)] >= 1.0f))
throw std::runtime_error("CP_LONGROPE_FACTORS: lambda must be >= 1");
if (i > 0 && lambda[static_cast<size_t>(i)] < lambda[static_cast<size_t>(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<int>(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<int>(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=<run>), 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<int>(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<int>(env_i64("CP_MAX_STEPS", cfg.max_steps));
cfg.warmup_steps = static_cast<int>(env_i64("CP_WARMUP", cfg.warmup_steps));
cfg.stable_steps = static_cast<int>(env_i64("CP_STABLE", cfg.stable_steps));
cfg.bf16 = env_i64("CP_BF16", cfg.bf16 ? 1 : 0) != 0;
cfg.bf16_mask = static_cast<int>(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=<steps> (0 = off; only
// applies when start_step>0). CP_REWARMUP_PEAK=<fraction of max_lr> gives a
// reduced peak for the long-context phase (default 1.0 = same peak).
const int rewarmup_steps = static_cast<int>(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<uint64_t>(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<int>(yarn_scale + 0.5);
cfg.cream_scaled_max = static_cast<int64_t>(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=<file> 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<float>(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<int64_t>(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<float>(g) /
static_cast<float>(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<int64_t>(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<float>(g) /
static_cast<float>(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<int>(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<int>(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<float>(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<const char*>(c.data<float>()),
static_cast<std::streamsize>(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<int64_t>(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<int64_t>(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<int>(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<int> 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<int> 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<ProcessGroupNCCL> pg = mesh.world_pg(); // global (all ranks), for DDP
std::shared_ptr<CPProcessGroupNCCL> cp_pg; // CP ring/ulysses sub-group (has ring surface)
std::shared_ptr<ProcessGroupNCCL> 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<ContextParallel::SharedFwdRing>();
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=<path/to/base.ckpt> --------------------------------
// 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<DataParallel> 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<DataParallel>(&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<std::unique_ptr<owt::zero::ZeroOptimizerBase>> zopts;
std::vector<std::string> zgroup_label;
if (cfg.zero) {
std::map<Dtype, std::vector<Tensor>> 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<owt::zero::SymmetricShardedOptimizer>(kv.second, pg, zcfg));
else
zopts.push_back(std::make_unique<owt::zero::ShardedOptimizer>(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<Tensor> 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<int>(cfg.B), static_cast<int>(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<int>(cfg.B), static_cast<int>(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<float>(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
// <N>.csv) AND the checkpoint filename prefix (blumodelcp_run<N>_step_<S>.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 <prefix>_run<K>_step_<S>.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<int, int> 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 <ckpt_dir>/rank_<N>/ --
// 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<double>(total_mem - free_mem) / (1024.0 * 1024.0);
double total_mb = static_cast<double>(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<int>(meta_theta);
if (is_master)
std::cout << "[rope_theta] " << meta_theta << " from " << meta_path
<< " (authoritative)\n";
} else if (meta_theta != static_cast<long long>(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_<N>, 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_<N>/): a weights-only load is world-size agnostic by
// construction, and at a larger np the rank_<N> 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=<path/to/base.ckpt> "
"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_<N>.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<nn::Optimizer&>(*zopts[0])
: static_cast<nn::Optimizer&>(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<size_t>(start_step) * static_cast<size_t>(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<size_t>(start_step) *
static_cast<size_t>(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=<step> 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<int>(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<float>();
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<double>(p[0]) : 0.0) << ","
<< (n ? static_cast<double>(p[n / 2]) : 0.0) << ","
<< (n ? static_cast<double>(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_DUMP_ACTS>/cp_<name>.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<int>(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<int32_t> toks(static_cast<size_t>(Ta));
tf.read(reinterpret_cast<char*>(toks.data()), static_cast<std::streamsize>(Ta) * 4);
std::vector<uint16_t> xh(static_cast<size_t>(Ta));
for (int t = 0; t < Ta; ++t) {
const int32_t v = toks[static_cast<size_t>(t)];
if (v < 0 || v >= cfg.vocab_size)
throw std::runtime_error("CP_DUMP_ACTS: input token out of [0,vocab)");
xh[static_cast<size_t>(t)] = static_cast<uint16_t>(v);
}
Tensor idx(Shape{{1, static_cast<int64_t>(Ta)}},
TensorOptions().with_dtype(Dtype::UInt16).with_device(device));
cudaMemcpy(idx.data(), xh.data(), static_cast<size_t>(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<const char*>(c.data<float>()),
static_cast<std::streamsize>(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 <val>"
// 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<int>(env_i64("CP_EVAL_WINDOWS", 5));
const int skip_win = static_cast<int>(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<int>(cfg.T) / cp_size; // logits shard length
const int ce_chunk = std::max(1, static_cast<int>(
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 <v>' 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<float> lam(static_cast<size_t>(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<int>(half), MPI_FLOAT, 0, MPI_COMM_WORLD);
long long nhat_ll = static_cast<long long>(nhat);
MPI_Bcast(&nhat_ll, 1, MPI_LONG_LONG, 0, MPI_COMM_WORLD);
nhat = static_cast<int64_t>(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<float>(s_)) + 1.0f) : 1.0f;
Tensor lr = autograd::build_rope_cache_longrope(
cfg.T, cfg.head_dim, static_cast<float>(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<size_t>(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<double>(loss.to_cpu().data<float>()[0])
* static_cast<double>(clen);
win_tok += clen;
}
sum_loss += (win_tok > 0) ? (win_loss / static_cast<double>(win_tok)) : 0.0;
++nwin;
}
}
float local_mean = nwin > 0 ? static_cast<float>(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<float>(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 "<i>,<tok>,<tok>,..." (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<int>(env_i64("CP_GEN_PROMPT_LEN", 0));
const int N_PROMPT= static_cast<int>(env_i64("CP_GEN_N", 1));
const int MAX_NEW = static_cast<int>(env_i64("CP_GEN_MAX_NEW", 32));
const int TOPK = static_cast<int>(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<int>(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<int>(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<int32_t> raw(static_cast<size_t>(P_LEN));
tf.read(reinterpret_cast<char*>(raw.data()),
static_cast<std::streamsize>(P_LEN) * 4);
if (!tf) die("CP_GENERATE: short read on prompt " + std::to_string(pi));
std::vector<uint16_t> seq;
seq.reserve(static_cast<size_t>(P_LEN + MAX_NEW));
for (int t = 0; t < P_LEN; ++t) {
const int32_t v = raw[static_cast<size_t>(t)];
if (v < 0 || v >= cfg.vocab_size)
die("CP_GENERATE: prompt token out of [0,vocab)");
seq.push_back(static_cast<uint16_t>(v));
}
std::vector<int64_t> 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<int>(seq.size());
const int padded = pad_up(cur);
std::vector<uint16_t> buf(static_cast<size_t>(padded), 0);
std::copy(seq.begin(), seq.end(), buf.begin());
Tensor idx(Shape{{1, static_cast<int64_t>(padded)}},
TensorOptions().with_dtype(Dtype::UInt16).with_device(device));
cudaMemcpy(idx.data(), buf.data(),
static_cast<size_t>(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<float>(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<int64_t>()[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<int64_t>()[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<uint16_t>(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<int>(env_i64("CP_EVAL_WINDOWS", 10));
const int skip_win = static_cast<int>(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<int>(cfg.T); // GLOBAL length (CSV rows)
const int Tlocal = static_cast<int>(cfg.T) / cp_size; // local shard length (logits)
const int Bloc = static_cast<int>(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<int32_t> 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<size_t>(nbytes) / 4);
tf.read(reinterpret_cast<char*>(eval_toks.data()), nbytes);
// Windows AVAILABLE after skipping the calibration windows (held-out).
const long long total_w = (static_cast<long long>(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<int>(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<double> sum_nll(static_cast<size_t>(Tpos), 0.0);
std::vector<long long> cnt(static_cast<size_t>(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<size_t>(Tpos) * sizeof(double));
cudaMalloc(&d_cnt, static_cast<size_t>(Tpos) * sizeof(long long));
cudaMemset(d_sum, 0, static_cast<size_t>(Tpos) * sizeof(double));
cudaMemset(d_cnt, 0, static_cast<size_t>(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<float> host_buf;
autograd::NoGradGuard no_grad;
if (!bypass) { val_loader.reset(); if (skip_win > 0) val_loader.skip_batches(static_cast<size_t>(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<uint16_t> xh(static_cast<size_t>(Tpos)), yh(static_cast<size_t>(Tpos));
const size_t base = static_cast<size_t>(skip_win + w) * static_cast<size_t>(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<size_t>(t)] = static_cast<uint16_t>(xi);
yh[static_cast<size_t>(t)] = static_cast<uint16_t>(yi);
}
input = Tensor(Shape{{1, static_cast<int64_t>(Tpos)}},
TensorOptions().with_dtype(Dtype::UInt16).with_device(device));
Tensor tgt(Shape{{1, static_cast<int64_t>(Tpos)}},
TensorOptions().with_dtype(Dtype::UInt16).with_device(device));
cudaMemcpy(input.data(), xh.data(), static_cast<size_t>(Tpos) * sizeof(uint16_t),
cudaMemcpyHostToDevice);
cudaMemcpy(tgt.data(), yh.data(), static_cast<size_t>(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<const float*>(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<uint16_t>() : nullptr;
const int32_t* yp32 = (ydt == Dtype::Int32) ? y_cpu.data<int32_t>() : nullptr;
const int64_t* yp64 = (ydt == Dtype::Int64) ? y_cpu.data<int64_t>() : 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<int64_t>();
// 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<int32_t>();
const int64_t* pg = pos32.data<int64_t>();
for (int bi = 0; bi < Bcur; ++bi) {
OwnTensor::cp::cuda::eval_nll_accumulate_row(
dptr + static_cast<size_t>(bi) * Tlocal * V,
yg + static_cast<size_t>(bi) * Tlocal,
pg, d_sum, d_cnt,
static_cast<int>(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<size_t>(clen) * static_cast<size_t>(V));
cudaMemcpy(host_buf.data(),
dptr + (static_cast<size_t>(bi) * Tlocal + tl0) * static_cast<size_t>(V),
static_cast<size_t>(clen) * static_cast<size_t>(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<size_t>(c) * static_cast<size_t>(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<double>(row[v] - mx));
const double lse = static_cast<double>(mx) + std::log(se);
const size_t idx = static_cast<size_t>(bi) * Tlocal + tl;
const int64_t tgt = yp16 ? static_cast<int64_t>(yp16[idx])
: yp64 ? yp64[idx]
: static_cast<int64_t>(yp32[idx]);
if (tgt < 0 || tgt >= V) continue; // skip padding/ignore ids
const int g = static_cast<int>(posp[tl]); // shy.pos_local[tl]
sum_nll[static_cast<size_t>(g)] += lse - static_cast<double>(row[tgt]);
cnt[static_cast<size_t>(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<size_t>(Tpos) * sizeof(double), cudaMemcpyDeviceToHost);
cudaMemcpy(cnt.data(), d_cnt,
static_cast<size_t>(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<double> g_nll(static_cast<size_t>(Tpos), 0.0);
std::vector<long long> g_cnt(static_cast<size_t>(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<size_t>(t)] > 0
? sum_nll[static_cast<size_t>(t)] / static_cast<double>(cnt[static_cast<size_t>(t)])
: 0.0;
f << t << "," << cnt[static_cast<size_t>(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<float>()[0] / static_cast<float>(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=<csv> 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<uint64_t>(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<float>()[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<float>();
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<float>()[0] / static_cast<float>(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<Tensor> 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<int64_t>(dp_size);
const double tok_per_sec = static_cast<double>(tokens) / dt;
// Throughput + time left + GPU memory (gpt2_cp_test-style per-step report).
long long total_sec = static_cast<long long>((cfg.max_steps - step) * dt);
int h = static_cast<int>(total_sec / 3600);
int m = static_cast<int>((total_sec % 3600) / 60);
size_t free_mem = 0, total_mem = 0;
cudaMemGetInfo(&free_mem, &total_mem);
const double used_mb = static_cast<double>(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_<label>_<rotator>_T<T>_ws<ws>.txt. ────────────────────
if (mem_probe && step == cfg.max_steps - 1) {
cudaDeviceSynchronize();
size_t pf = 0, pt = 0;
cudaMemGetInfo(&pf, &pt);
const double probe_used_mb = static_cast<double>(pt - pf) / (1024.0 * 1024.0);
auto mstats = CachingCUDAAllocator::instance().get_stats(rank);
const double MB = 1024.0 * 1024.0;
const double st_reserved_mb = mstats.reserved_current / MB;
const double st_active_mb = mstats.active_current / MB;
const double st_requested_mb = mstats.allocated_current / MB;
const double st_frag_pct = mstats.fragmentation_ratio();
const double st_active_peak_mb = mstats.active_peak / MB;
const double st_reserved_peak_mb = mstats.reserved_peak / MB;
const double st_cached_free_mb =
(mstats.reserved_current > mstats.active_current)
? (mstats.reserved_current - mstats.active_current) / MB : 0.0;
const std::string tag = "CPP_" + mem_label + "_" + rotator_label + "_T" +
std::to_string(cfg.T) + "_ws" + std::to_string(world_size);
std::cout << "[MEM_PROBE rank=" << rank << "] tag=" << tag
<< " used_mb=" << std::fixed << std::setprecision(1)
<< probe_used_mb << "\n";
if (is_master) {
const char* sd = std::getenv("MEM_SNAPSHOT_DIR");
const std::string snap_dir = sd ? sd : "mem_scaling_runs";
std::filesystem::create_directories(snap_dir);
const std::string snap_path = snap_dir + "/" + tag + ".txt";
std::ofstream sf(snap_path);
sf << "# MEM PROBE SNAPSHOT tag=" << tag << "\n";
sf << "# impl=Cpp label=" << mem_label << " rotator=" << rotator_label
<< " n_embd=" << cfg.d_model << " n_layer=" << cfg.n_layers
<< " n_head=" << cfg.q_heads
<< " weight_tying=" << (cfg.weight_tying ? 1 : 0) << "\n";
sf << "# B=" << cfg.B << " T=" << cfg.T
<< " cp_world_size=" << world_size << " params=" << num_params << "\n";
sf << "# cudaMemGetInfo used_mb(rank0)=" << std::fixed
<< std::setprecision(1) << probe_used_mb << "\n";
sf << "# ALLOC reserved_mb=" << std::fixed << std::setprecision(1)
<< st_reserved_mb << " active_mb=" << st_active_mb
<< " requested_mb=" << st_requested_mb
<< " cached_free_mb=" << st_cached_free_mb
<< " internal_frag_pct=" << std::setprecision(2) << st_frag_pct << "\n";
sf << "# ALLOC_PEAK active_peak_mb=" << std::fixed << std::setprecision(1)
<< st_active_peak_mb << " reserved_peak_mb=" << st_reserved_peak_mb << "\n";
// Machine-readable live per-GPU used MiB (parsed by the table gen).
// nvidia-smi ignores CUDA_VISIBLE_DEVICES; keep only this run's GPUs.
{
std::string smi_used;
std::vector<std::string> vis;
if (const char* cvd = std::getenv("CUDA_VISIBLE_DEVICES")) {
std::string s(cvd), cur;
for (char c : s) {
if (c == ',') { if (!cur.empty()) vis.push_back(cur); cur.clear(); }
else if (c != ' ') cur += c;
}
if (!cur.empty()) vis.push_back(cur);
}
auto allowed = [&](const std::string& idx) {
if (vis.empty()) return true;
for (auto& v : vis) if (v == idx) return true;
return false;
};
FILE* qp = popen("nvidia-smi --query-gpu=index,memory.used "
"--format=csv,noheader,nounits 2>/dev/null", "r");
if (qp) {
char buf[256];
while (fgets(buf, sizeof(buf), qp)) {
std::string ln(buf);
while (!ln.empty() && (ln.back() == '\n' || ln.back() == '\r')) ln.pop_back();
std::string cleaned;
for (char c : ln) if (c != ' ') cleaned += c;
if (cleaned.empty()) continue;
const std::string idx = cleaned.substr(0, cleaned.find(','));
if (!allowed(idx)) continue;
if (!smi_used.empty()) smi_used += ";";
smi_used += cleaned;
}
pclose(qp);
}
sf << "# SMI_USED_MB_PER_GPU=" << smi_used << "\n";
}
sf << "# ---- nvidia-smi ----\n";
sf.close();
const std::string cmd = "nvidia-smi >> '" + snap_path + "' 2>&1";
int rc = std::system(cmd.c_str());
(void)rc;
std::cout << "[MEM_PROBE] wrote snapshot " << snap_path << "\n";
}
MPI_Barrier(MPI_COMM_WORLD);
break;
}
// Under ZeRO every rank must write its own shard (see shard_dir above);
// each writes into its own rank_<N> directory, so there is no race.
if (cfg.checkpointing && (cfg.zero || is_master) &&
((step + 1) % cfg.ckpt_freq == 0 || step == cfg.max_steps - 1)) {
try {
ckpt.save(step + 1, model,
cfg.zero ? static_cast<nn::Optimizer&>(*zopts[0])
: static_cast<nn::Optimizer&>(optimizer),
loss_accum);
if (cfg.zero) {
for (size_t g = 1; g < zopts.size(); ++g) {
const std::string side = zero_side_dir + "/" + run_prefix + "_step_" +
std::to_string(step + 1) + ".zero_g" + std::to_string(g);
std::ofstream zos(side, std::ios::binary);
zopts[g]->save_state(zos);
}
}
}
catch (const std::exception& e) {
std::cerr << "[WARN] checkpoint save at step " << (step + 1) << " failed: " << e.what() << std::endl;
}
}
if (step == 0 || step == 1) CachingCUDAAllocator::instance().empty_cache();
}
if (cfg.checkpointing) {
try { ckpt.wait_for_completion(); }
catch (const std::exception& e) {
if (is_master) std::cerr << "[WARN] checkpoint flush: " << e.what() << std::endl;
}
}
if (is_master) { log_file.close(); std::cout << "\n=== Training complete ===" << std::endl; }
MPI_Finalize();
return 0;
}
// =============================================================================
// 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
// =============================================================================
// Tensor library
// Data + distributed
// 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).
// Context parallelism
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 {