// =============================================================================
// 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;
}