bonsai.cpp/examples/training/ternary.cpp

#include "ternary-graph.h"
#include "ternary-data.h"
#include "ternary-timing.h"
#include "build-info.h"
#include "ggml-cpu.h"
#include "ggml-cuda.h"
#include "ggml-impl.h"
#include "llama-context.h"

#include <openssl/evp.h>

#include <chrono>
#include <cmath>
#include <filesystem>
#include <iomanip>
#include <iostream>
#include <set>
#include <sstream>

using steady_clock = std::chrono::steady_clock;

static void check(bool value, const std::string & message) {
    if (!value) { throw std::runtime_error(message); }
}

static double seconds(steady_clock::time_point start) {
    return std::chrono::duration<double>(steady_clock::now() - start).count();
}

struct options {
    std::string model, data, output, resume, targets = "q,v,ffn_up,ffn_down", layers;
    std::string prompt = "prompt", response = "response", text, split, mode = "ternary", generate;
    int rank = 128, epochs = 1, batch = 1, micro = 1, context = 64, steps = 0, max_train = 0, max_eval = 0, validate_block = -1, validate_depth = 0, data_version = 2;
    unsigned seed = 42;
    double lr = 0.001, eval_split = 0.1;
    float scale = 0.0001f;
    bool all_tokens = false, chat = true, pack = true, eval_only = false, int8_activations = false, verbose = false, base_only = false, validate_forward = false, profile = false, gpu_update = true, verify_reload = false, validate_batch = false;
};

static options parse(int argc, char ** argv) {
    options o;
    for (int i = 1; i < argc; ++i) {
        const std::string key = argv[i];
        auto value = [&]() { check(++i < argc, "missing value for " + key); return std::string(argv[i]); };
        if (key == "--model") { o.model = value(); }
        else if (key == "--train-data") { o.data = value(); }
        else if (key == "--output") { o.output = value(); }
        else if (key == "--resume" || key == "--adapter") { o.resume = value(); }
        else if (key == "--target") { o.targets = value(); }
        else if (key == "--layers") { o.layers = value(); }
        else if (key == "--prompt-column") { o.prompt = value(); }
        else if (key == "--response-column") { o.response = value(); }
        else if (key == "--text-column") { o.text = value(); }
        else if (key == "--split-column") { o.split = value(); }
        else if (key == "--mode") { o.mode = value(); }
        else if (key == "--generate") { o.generate = value(); }
        else if (key == "--rank") { o.rank = std::stoi(value()); }
        else if (key == "--epochs") { o.epochs = std::stoi(value()); }
        else if (key == "--batch-size") { o.batch = std::stoi(value()); }
        else if (key == "--micro-batch") { o.micro = std::stoi(value()); }
        else if (key == "--context") { o.context = std::stoi(value()); }
        else if (key == "--max-steps") { o.steps = std::stoi(value()); }
        else if (key == "--max-train-samples") { o.max_train = std::stoi(value()); }
        else if (key == "--max-eval-samples") { o.max_eval = std::stoi(value()); }
        else if (key == "--validate-block") { o.validate_block = std::stoi(value()); }
        else if (key == "--validate-depth") { o.validate_depth = std::stoi(value()); }
        else if (key == "--validate-forward") { o.validate_forward = true; }
        else if (key == "--validate-batch") { o.validate_batch = true; }
        else if (key == "--seed") { o.seed = std::stoul(value()); }
        else if (key == "--lr") { o.lr = std::stod(value()); }
        else if (key == "--eval-split") { o.eval_split = std::stod(value()); }
        else if (key == "--adapter-scale") { o.scale = std::stof(value()); }
        else if (key == "--train-all-tokens") { o.all_tokens = true; }
        else if (key == "--no-chat-template") { o.chat = false; }
        else if (key == "--no-pack") { o.pack = false; }
        else if (key == "--eval-only") { o.eval_only = true; }
        else if (key == "--verify-reload") { o.verify_reload = true; }
        else if (key == "--int8-activations") { o.int8_activations = true; }
        else if (key == "--verbose") { o.verbose = true; }
        else if (key == "--profile") { o.profile = true; }
        else if (key == "--cpu-update") { o.gpu_update = false; }
        else if (key == "--base-only") { o.base_only = true; }
        else if (key == "--help") {
            std::puts("bonsai-train-ternary-lora --model MODEL --train-data CSV --output ADAPTER.gguf\n"
                    "  --rank 128 --target q,v,ffn_up,ffn_down --layers 0,3,63 (default all blocks)\n"
                    "  --prompt-column prompt --response-column response | --text-column text\n"
                    "  --epochs 1 --batch-size 1 --micro-batch 1 --context 64 --seed 42 --eval-split .1\n"
                    "  --lr .001 (ternary: fraction of trits selected by gradient magnitude per update)\n"
                    "  --mode bf16 (continuous BF16 SGD baseline, lr is a gradient step size)\n"
                    "  --adapter-scale .0001 --int8-activations --train-all-tokens --no-pack\n"
                    "  --split-column split --max-steps N --max-train-samples N --max-eval-samples N\n"
                    "  --resume ADAPTER.gguf --eval-only --verify-reload --base-only --verbose\n"
                    "  --validate-block N (real block forward/backward diagnostic)\n"
                    "  --validate-depth N (prefix causal-LM backward diagnostic, no training run)\n"
                    "  --validate-forward (compare all blocks/logits with normal inference)\n"
                    "  --validate-batch (compare joined micro-batch loss with independent sequences)\n"
                    "  --profile (CUDA event timing per node; introduces synchronization overhead)\n"
                    "  --cpu-update (CPU reference optimizer; default accumulates/updates on GPU)\n"
                    "  --adapter ADAPTER.gguf --generate PROMPT (fresh-process greedy inference)");
            std::exit(0);
        } else { throw std::runtime_error("unknown option: " + key); }
    }
    check(!o.model.empty(), "--model is required");
    check(o.rank > 0 && o.context > 1 && o.batch > 0 && o.micro > 0 && o.batch % o.micro == 0, "invalid rank/context/batch or micro-batch does not divide batch-size");
    check(o.epochs >= 0 && std::isfinite(o.lr) && o.lr > 0 && (o.mode == "bf16" || o.lr <= 1) && o.eval_split >= 0 && o.eval_split < 1 && std::isfinite(o.scale) && o.scale > 0, "invalid epochs/lr/eval-split/scale");
    check(o.mode == "ternary" || o.mode == "bf16", "--mode must be ternary or bf16");
    if (o.base_only || o.epochs == 0) { o.eval_only = true; }
    check(o.generate.empty() || !o.resume.empty(), "--generate requires --adapter");
    check(!o.verify_reload || (o.eval_only && !o.resume.empty()), "--verify-reload requires --eval-only and --adapter");
    check(o.validate_block >= 0 || !o.generate.empty() || !o.data.empty(), "--train-data is required");
    check(o.steps >= 0 && o.max_train >= 0 && o.max_eval >= 0 && o.validate_depth >= 0 && o.validate_depth <= 64, "invalid diagnostic depth or sample/step limit");
    check(o.eval_only || o.validate_forward || o.validate_batch || o.validate_depth || o.validate_block >= 0 || !o.generate.empty() || !o.output.empty(), "--output is required for training");
    if (!o.output.empty()) {
        for (const auto & destination : {o.output, o.output + ".tmp"}) {
            std::error_code error;
            const bool same_file = std::filesystem::equivalent(o.model, destination, error);
            check(!same_file && std::filesystem::weakly_canonical(o.model) != std::filesystem::weakly_canonical(destination),
                    "adapter output or temporary path would overwrite the frozen base model");
        }
    }
    return o;
}

static std::string hash_file(const std::string & path) {
    std::ifstream f(path, std::ios::binary);
    check((bool) f, "cannot hash model: " + path);
    std::unique_ptr<EVP_MD_CTX, decltype(&EVP_MD_CTX_free)> state(EVP_MD_CTX_new(), EVP_MD_CTX_free);
    check(state && EVP_DigestInit_ex(state.get(), EVP_sha256(), nullptr) == 1, "SHA256 initialization failed");
    std::vector<unsigned char> buffer(8 * 1024 * 1024);
    while (f) {
        f.read((char *) buffer.data(), buffer.size());
        check(EVP_DigestUpdate(state.get(), buffer.data(), f.gcount()) == 1, "SHA256 update failed");
    }
    check(f.eof(), "error reading model hash");
    unsigned char digest[32]; unsigned length = 0;
    check(EVP_DigestFinal_ex(state.get(), digest, &length) == 1 && length == 32, "SHA256 finalization failed");
    std::ostringstream result;
    for (unsigned char byte : digest) { result << std::hex << std::setw(2) << std::setfill('0') << (unsigned) byte; }
    return result.str();
}

static std::string hash_values(const std::vector<float> & values) {
    unsigned char digest[32]; unsigned length = 0;
    check(EVP_Digest(values.data(), values.size() * sizeof(float), digest, &length, EVP_sha256(), nullptr) == 1 && length == 32, "logit SHA256 failed");
    std::ostringstream result;
    for (unsigned char byte : digest) { result << std::hex << std::setw(2) << std::setfill('0') << (unsigned) byte; }
    return result.str();
}

static std::string hash_model(const llama_model * model) {
    std::unique_ptr<EVP_MD_CTX, decltype(&EVP_MD_CTX_free)> state(EVP_MD_CTX_new(), EVP_MD_CTX_free);
    check(state && EVP_DigestInit_ex(state.get(), EVP_sha256(), nullptr) == 1, "tensor hash initialization failed");
    std::vector<unsigned char> buffer(8 * 1024 * 1024);
    for (const auto & entry : llama_internal_get_tensor_map(model)) {
        const auto * t = entry.second;
        EVP_DigestUpdate(state.get(), entry.first.data(), entry.first.size());
        for (size_t offset = 0; offset < ggml_nbytes(t); offset += buffer.size()) {
            const size_t n = std::min(buffer.size(), ggml_nbytes(t) - offset);
            ggml_backend_tensor_get(t, buffer.data(), offset, n); EVP_DigestUpdate(state.get(), buffer.data(), n);
        }
    }
    unsigned char digest[32]; unsigned length = 0;
    check(EVP_DigestFinal_ex(state.get(), digest, &length) == 1 && length == 32, "tensor hash finalization failed");
    std::ostringstream result;
    for (unsigned char byte : digest) { result << std::hex << std::setw(2) << std::setfill('0') << (unsigned) byte; }
    return result.str();
}

struct factor {
    int64_t in, out;
    std::vector<uint8_t> bytes;
    std::vector<float> gradient;
    ggml_tensor * tensor = nullptr;
    ggml_tensor * gpu_gradient = nullptr;
    bool gradient_ready = false;
    std::array<uint64_t, 10> last_stats{};
};

struct adapter_weight {
    factor a, b;
    float scale;
    int layer;
    double base_squared = 0, delta_squared = 0;
    size_t output_entries = 0;
};

struct timing {
    double graph = 0, forward = 0, backward = 0, transfer = 0, download = 0, update = 0;
    size_t peak_compute = 0, peak_used = 0, peak_adapter = 0, peak_gradients = 0, peak_boundaries = 0;
    std::map<std::string, double> profile;
};

struct trainer {
    options o;
    llama_model * model;
    ggml_backend_t gpu, cpu;
    ggml_backend_sched_t sched;
    size_t graph_capacity = 65536;
    std::map<std::string, adapter_weight> weights;
    std::string model_hash;
    uint64_t step = 0;
    uint64_t epoch = 0, cursor = 0;
    std::string dataset_hash;
    std::string last_logits_hash, eval_logits_hash, saved_logits_hash;
    timing times;
    size_t initial_used = 0, packed_gpu = 0;
    void * timer = nullptr;
    void * node_timer = nullptr;
    ggml_context_ptr gradient_context;
    ggml_backend_buffer_ptr gradient_buffer;

    trainer(options o, llama_model * model, ggml_backend_t gpu, ggml_backend_t cpu, std::string hash) :
        o(o), model(model), gpu(gpu), cpu(cpu), model_hash(std::move(hash)) {
        ggml_backend_t backends[] = {gpu, cpu};
        sched = ggml_backend_sched_new(backends, nullptr, 2, graph_capacity, false, true);
        timer = bonsai_cuda_timer_new(gpu); node_timer = bonsai_cuda_timer_new(gpu);
        if (o.profile) { ggml_backend_sched_set_eval_callback(sched, profile_node, this); }
        size_t free, total;
        ggml_backend_dev_memory(ggml_backend_get_device(gpu), &free, &total);
        initial_used = total - free;
        for (const auto & entry : llama_internal_get_tensor_map(model)) {
            const auto * tensor = entry.second;
            check(!(tensor->flags & GGML_TENSOR_FLAG_PARAM), "backbone unexpectedly marked trainable");
            if (tensor->type == GGML_TYPE_PQ2_0 && !ggml_backend_buffer_is_host(tensor->buffer)) { packed_gpu += ggml_nbytes(tensor); }
        }
        if (!o.resume.empty()) { load(); } else if (!o.base_only) { initialize(); }
        check(this->o.mode != "bf16" || !this->o.int8_activations, "--int8-activations requires ternary factors");
    }

    ~trainer() { ggml_backend_sched_free(sched); bonsai_cuda_timer_free(timer); bonsai_cuda_timer_free(node_timer); }

    static bool profile_node(ggml_tensor * t, bool ask, void * opaque) {
        auto & train = *(trainer *) opaque;
        if (ask) { bonsai_cuda_timer_start(train.node_timer); return true; }
        const double duration = bonsai_cuda_timer_stop(train.node_timer);
        std::string category = "other";
        const bool factor = t->src[0] && std::string(t->src[0]->name).find(".lora_") != std::string::npos;
        const bool base = t->src[0] && (t->src[0]->type == GGML_TYPE_PQ2_0 || (!factor && t->src[0]->type == GGML_TYPE_BF16));
        if (t->op == GGML_OP_MUL_MAT && base) { category = "base_forward"; }
        else if (t->op == GGML_OP_OUT_PROD && base) { category = "base_backward"; }
        else if (t->op == GGML_OP_MUL_MAT && std::string(t->src[0]->name).find(".lora_") != std::string::npos) { category = "adapter_forward"; }
        else if (t->op == GGML_OP_OUT_PROD && (factor || std::string(t->name).find(".lora_") != std::string::npos)) { category = "adapter_backward"; }
        train.times.profile[category] += duration;
        return true;
    }

    static std::set<std::string> list(const std::string & value) {
        std::set<std::string> result;
        std::stringstream stream(value); std::string part;
        while (std::getline(stream, part, ',')) { if (!part.empty()) { result.insert(part); } }
        return result;
    }

    void prepare_gradients() {
        if (!o.gpu_update || gradient_buffer) { return; }
        gradient_context.reset(ggml_init({(weights.size() * 2 + 8) * ggml_tensor_overhead(), nullptr, true}));
        for (auto & item : weights) {
            for (auto * f : {&item.second.a, &item.second.b}) {
                f->gpu_gradient = ggml_new_tensor_2d(gradient_context.get(), GGML_TYPE_F32, f->in, f->out);
            }
        }
        gradient_buffer.reset(ggml_backend_alloc_ctx_tensors(gradient_context.get(), gpu));
        check(gradient_buffer != nullptr, "GPU gradient accumulation allocation failed; use --cpu-update if memory is limited");
        ggml_backend_buffer_clear(gradient_buffer.get(), 0);
    }

    void clear_gradients() {
        if (gradient_buffer) { ggml_backend_buffer_clear(gradient_buffer.get(), 0); }
        for (auto & item : weights) { for (auto * f : {&item.second.a, &item.second.b}) { f->gradient.clear(); f->gradient_ready = false; } }
    }

    void initialize() {
        auto targets = list(o.targets);
        if (targets.count("all-linear")) { targets = {"q", "k", "v", "o", "qkv", "attn_gate", "ssm_out", "ffn_up", "ffn_gate", "ffn_down"}; }
        if (targets.erase("attention")) { targets.insert({"q", "k", "v", "o", "qkv", "attn_gate", "ssm_out"}); }
        if (targets.erase("mlp")) { targets.insert({"ffn_up", "ffn_gate", "ffn_down"}); }
        const std::map<std::string, std::string> names = {{"q", "attn_q"}, {"k", "attn_k"}, {"v", "attn_v"}, {"o", "attn_output"}, {"attention_output", "attn_output"},
            {"qkv", "attn_qkv"}, {"attn_gate", "attn_gate"}, {"ssm_out", "ssm_out"}, {"ffn_gate", "ffn_gate"}, {"ffn_up", "ffn_up"}, {"ffn_down", "ffn_down"}};
        const auto layers = list(o.layers);
        std::mt19937 rng(o.seed);
        for (int il = 0; il < (int) model->hparams.n_layer(); ++il) {
            if (!layers.empty() && !layers.count(std::to_string(il))) { continue; }
            for (const auto & target : targets) {
                auto name = names.find(target); check(name != names.end(), "unknown adapter target: " + target);
                const std::string tensor_name = "blk." + std::to_string(il) + "." + name->second + ".weight";
                auto * base = model->get_tensor(tensor_name.c_str());
                if (!base) { continue; }
                check(base->type == GGML_TYPE_PQ2_0, "adapter target is not packed PQ2_0: " + tensor_name);
                adapter_weight w{{base->ne[0], o.rank, {}, {}, nullptr}, {o.rank, base->ne[1], {}, {}, nullptr}, o.scale, il};
                for (auto * f : {&w.a, &w.b}) {
                    f->bytes.resize(f->in * f->out * (o.mode == "bf16" ? 2 : 1));
                    for (int64_t i = 0; i < f->in * f->out; ++i) {
                        const int random = rng() % 8;
                        const int trit = random == 0 ? -1 : random == 1 ? 1 : 0;
                        if (o.mode == "bf16") {
                            const auto v = ggml_fp32_to_bf16((float) trit);
                            std::memcpy(f->bytes.data() + 2 * i, &v, 2);
                        } else { f->bytes[i] = (uint8_t) trit; }
                    }
                }
                weights.emplace(tensor_name, std::move(w));
            }
        }
        check(!weights.empty(), "no adapter targets matched the selected model/layers");
        std::printf("initialized projections=%zu rank=%d state=%s basis=bonsai-projection-input-v1 fixed_scale=%.8g nonzero_init=.25\n", weights.size(), o.rank, o.mode.c_str(), o.scale);
    }

    void load() {
        ggml_context * meta = nullptr;
        gguf_context_ptr file{gguf_init_from_file(o.resume.c_str(), {true, &meta})};
        ggml_context_ptr ctx{meta}; check(file != nullptr, "cannot load adapter metadata");
        auto string = [&](const char * key) { int id = gguf_find_key(file.get(), key); check(id >= 0, std::string("missing adapter key: ") + key); return std::string(gguf_get_val_str(file.get(), id)); };
        check(string("adapter.ternary.base_sha256") == model_hash, "adapter base model SHA256 mismatch");
        check(string("adapter.ternary.basis") == "bonsai-projection-input-v1", "unsupported adapter basis");
        const int version = gguf_find_key(file.get(), "adapter.ternary.version");
        check(version >= 0 && gguf_get_val_u32(file.get(), version) == 1, "unsupported adapter version");
        const int base_type = gguf_find_key(file.get(), "adapter.ternary.base_type");
        check(base_type >= 0 && gguf_get_val_u32(file.get(), base_type) == GGML_TYPE_PQ2_0, "adapter requires PQ2_0 base weights");
        o.mode = string("adapter.ternary.mode");
        check(o.mode == "ternary" || o.mode == "bf16", "invalid saved adapter mode");
        int id = gguf_find_key(file.get(), "adapter.ternary.step");
        step = id >= 0 ? gguf_get_val_u64(file.get(), id) : 0;
        id = gguf_find_key(file.get(), "adapter.ternary.int8_activations");
        o.int8_activations = id >= 0 && gguf_get_val_bool(file.get(), id);
        auto u32 = [&](const char * key, int & value) { const int k = gguf_find_key(file.get(), key); if (k >= 0) { value = gguf_get_val_u32(file.get(), k); } };
        auto boolean = [&](const char * key, bool & value) { const int k = gguf_find_key(file.get(), key); if (k >= 0) { value = gguf_get_val_bool(file.get(), k); } };
        auto str = [&](const char * key, std::string & value) { const int k = gguf_find_key(file.get(), key); if (k >= 0) { value = gguf_get_val_str(file.get(), k); } };
        u32("adapter.ternary.rank", o.rank); u32("adapter.ternary.context", o.context);
        o.data_version = 1; u32("adapter.ternary.data_version", o.data_version);
        check(o.data_version == 1 || o.data_version == 2, "unsupported saved preprocessing version");
        u32("adapter.ternary.batch", o.batch); u32("adapter.ternary.micro", o.micro);
        id = gguf_find_key(file.get(), "adapter.ternary.seed"); if (id >= 0) { o.seed = gguf_get_val_u32(file.get(), id); }
        id = gguf_find_key(file.get(), "adapter.ternary.lr_f64"); if (id >= 0) { o.lr = gguf_get_val_f64(file.get(), id); }
        id = gguf_find_key(file.get(), "adapter.ternary.eval_split"); if (id >= 0) { o.eval_split = gguf_get_val_f64(file.get(), id); }
        str("adapter.ternary.targets", o.targets); str("adapter.ternary.prompt_column", o.prompt); str("adapter.ternary.response_column", o.response);
        str("adapter.ternary.text_column", o.text); str("adapter.ternary.split_column", o.split); str("adapter.ternary.dataset_sha256", dataset_hash);
        str("adapter.ternary.eval_logits_sha256", saved_logits_hash);
        boolean("adapter.ternary.all_tokens", o.all_tokens); boolean("adapter.ternary.chat", o.chat); boolean("adapter.ternary.pack", o.pack);
        boolean("adapter.ternary.gpu_update", o.gpu_update);
        u32("adapter.ternary.max_train", o.max_train);
        check(o.rank > 0 && o.context > 1 && o.batch > 0 && o.micro > 0 && o.batch % o.micro == 0 &&
                std::isfinite(o.lr) && o.lr > 0 && (o.mode == "bf16" || o.lr <= 1) &&
                o.eval_split >= 0 && o.eval_split < 1, "invalid saved training settings");
        id = gguf_find_key(file.get(), "adapter.ternary.epoch"); if (id >= 0) { epoch = gguf_get_val_u64(file.get(), id); }
        id = gguf_find_key(file.get(), "adapter.ternary.cursor"); if (id >= 0) { cursor = gguf_get_val_u64(file.get(), id); }
        std::ifstream input(o.resume, std::ios::binary);
        for (int64_t i = 0; i < gguf_get_n_tensors(file.get()); ++i) {
            const std::string name = gguf_get_tensor_name(file.get(), i);
            check(name.size() > 7 && (name.substr(name.size() - 7) == ".lora_a" || name.substr(name.size() - 7) == ".lora_b"), "unexpected adapter tensor");
            const std::string base_name = name.substr(0, name.size() - 7);
            auto * base = model->get_tensor(base_name.c_str()); check(base != nullptr && base->type == GGML_TYPE_PQ2_0, "adapter target missing from packed base");
            auto * tensor = ggml_get_tensor(ctx.get(), name.c_str());
            check(tensor->ne[2] == 1 && tensor->ne[3] == 1, "adapter factor must be two-dimensional");
            check(tensor->type == (o.mode == "ternary" ? GGML_TYPE_I8 : GGML_TYPE_BF16), "incorrect adapter factor storage type");
            auto & w = weights[base_name];
            auto & f = name.back() == 'a' ? w.a : w.b;
            f.in = tensor->ne[0]; f.out = tensor->ne[1]; f.bytes.resize(ggml_nbytes(tensor));
            input.seekg(gguf_get_data_offset(file.get()) + gguf_get_tensor_offset(file.get(), i));
            input.read((char *) f.bytes.data(), f.bytes.size()); check(input.good(), "truncated adapter factor");
            if (o.mode == "ternary") { for (uint8_t byte : f.bytes) { check((int8_t) byte >= -1 && (int8_t) byte <= 1, "adapter has an invalid trit"); } }
            w.layer = std::stoi(base_name.substr(4));
            const std::string scale_key = "adapter.ternary.scale." + base_name;
            id = gguf_find_key(file.get(), scale_key.c_str()); check(id >= 0, "missing saved adapter scale");
            w.scale = gguf_get_val_f32(file.get(), id); check(std::isfinite(w.scale) && w.scale > 0, "invalid saved scale");
        }
        check(!weights.empty(), "adapter has no projections");
        for (const auto & item : weights) {
            const auto & w = item.second;
            auto * base = model->get_tensor(item.first.c_str());
            check(w.a.in == base->ne[0] && w.b.out == base->ne[1] && w.a.out == w.b.in && !w.a.bytes.empty() && !w.b.bytes.empty(), "saved adapter shape mismatch");
        }
        std::fprintf(o.generate.empty() ? stdout : stderr, "resumed adapter=%s projections=%zu optimizer_step=%llu mode=%s\n", o.resume.c_str(), weights.size(), (unsigned long long) step, o.mode.c_str());
    }

    void save() {
        gguf_context_ptr file{gguf_init_empty()};
        gguf_set_val_str(file.get(), "general.type", "adapter");
        gguf_set_val_str(file.get(), "general.architecture", "qwen35");
        gguf_set_val_str(file.get(), "adapter.type", "lora");
        gguf_set_val_f32(file.get(), "adapter.lora.alpha", 0.0f);
        gguf_set_val_u32(file.get(), "adapter.ternary.version", 1);
        gguf_set_val_u32(file.get(), "adapter.ternary.base_type", GGML_TYPE_PQ2_0);
        gguf_set_val_str(file.get(), "adapter.ternary.base_sha256", model_hash.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.base_path", o.model.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.basis", "bonsai-projection-input-v1");
        gguf_set_val_str(file.get(), "adapter.ternary.mode", o.mode.c_str());
        gguf_set_val_bool(file.get(), "adapter.ternary.int8_activations", o.int8_activations);
        gguf_set_val_u64(file.get(), "adapter.ternary.step", step);
        gguf_set_val_u32(file.get(), "adapter.ternary.seed", o.seed);
        gguf_set_val_f32(file.get(), "adapter.ternary.lr", o.lr);
        gguf_set_val_f64(file.get(), "adapter.ternary.lr_f64", o.lr);
        gguf_set_val_u64(file.get(), "adapter.ternary.epoch", epoch);
        gguf_set_val_u64(file.get(), "adapter.ternary.cursor", cursor);
        gguf_set_val_str(file.get(), "adapter.ternary.dataset_sha256", dataset_hash.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.eval_logits_sha256", eval_logits_hash.c_str());
        gguf_set_val_u32(file.get(), "adapter.ternary.rank", o.rank);
        gguf_set_val_u32(file.get(), "adapter.ternary.data_version", o.data_version);
        gguf_set_val_u32(file.get(), "adapter.ternary.batch", o.batch);
        gguf_set_val_u32(file.get(), "adapter.ternary.micro", o.micro);
        gguf_set_val_f64(file.get(), "adapter.ternary.eval_split", o.eval_split);
        gguf_set_val_str(file.get(), "adapter.ternary.prompt_column", o.prompt.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.response_column", o.response.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.text_column", o.text.c_str());
        gguf_set_val_str(file.get(), "adapter.ternary.split_column", o.split.c_str());
        gguf_set_val_bool(file.get(), "adapter.ternary.all_tokens", o.all_tokens);
        gguf_set_val_bool(file.get(), "adapter.ternary.chat", o.chat);
        gguf_set_val_bool(file.get(), "adapter.ternary.pack", o.pack);
        gguf_set_val_bool(file.get(), "adapter.ternary.gpu_update", o.gpu_update);
        gguf_set_val_u32(file.get(), "adapter.ternary.max_train", o.max_train);
        gguf_set_val_str(file.get(), "adapter.ternary.optimizer", o.mode == "ternary" ? "adjacent-top-gradient-v1" : "bf16-sgd-v1");
        gguf_set_val_str(file.get(), "adapter.ternary.commit", llama_commit());
        gguf_set_val_str(file.get(), "adapter.ternary.targets", o.targets.c_str());
        gguf_set_val_u32(file.get(), "adapter.ternary.context", o.context);
        ggml_context_ptr ctx{ggml_init({(weights.size() * 2 + 8) * ggml_tensor_overhead(), nullptr, true})};
        for (auto & item : weights) {
            auto & w = item.second;
            gguf_set_val_f32(file.get(), ("adapter.ternary.scale." + item.first).c_str(), w.scale);
            const std::array<std::pair<factor *, const char *>, 2> factors = {{{&w.a, ".lora_a"}, {&w.b, ".lora_b"}}};
            for (const auto & pair : factors) {
                auto * tensor = ggml_new_tensor_2d(ctx.get(), o.mode == "ternary" ? GGML_TYPE_I8 : GGML_TYPE_BF16, pair.first->in, pair.first->out);
                ggml_set_name(tensor, (item.first + pair.second).c_str());
                tensor->data = pair.first->bytes.data();
                gguf_add_tensor(file.get(), tensor);
            }
        }
        const std::string temporary = o.output + ".tmp";
        check(gguf_write_to_file(file.get(), temporary.c_str(), false), "failed writing adapter");
        check(std::rename(temporary.c_str(), o.output.c_str()) == 0, "failed installing adapter file");
        std::printf("saved adapter=%s step=%llu native_factor_bytes=%zu\n", o.output.c_str(), (unsigned long long) step, parameter_bytes());
    }

    size_t parameter_bytes() const {
        size_t bytes = 0; for (const auto & item : weights) { bytes += item.second.a.bytes.size() + item.second.b.bytes.size(); } return bytes;
    }

    llm_graph_params params(llm_graph_result & result, int nt, llama_adapter_loras & loras) {
        llm_graph_params p{};
        p.arch = model->arch; p.hparams = model->hparams;
        p.cparams.n_ctx = o.context; p.cparams.n_ctx_seq = o.context;
        p.cparams.n_ctx_orig_yarn = model->hparams.n_ctx_orig_yarn;
        p.cparams.rope_freq_base = model->hparams.rope_freq_base_train;
        p.cparams.rope_freq_scale = model->hparams.rope_freq_scale_train;
        p.cparams.yarn_ext_factor = 0; p.cparams.yarn_attn_factor = 1;
        p.cparams.yarn_beta_fast = 32; p.cparams.yarn_beta_slow = 1;
        p.cparams.offload_kqv = true; p.cparams.causal_attn = true;
        p.ubatch.n_tokens = nt; p.ubatch.n_seq_tokens = nt; p.ubatch.n_seqs = 1; p.ubatch.n_seqs_unq = 1;
        p.n_outputs = nt; p.sched = sched; p.backend_cpu = cpu;
        p.loras = &loras; p.res = &result;
        p.hadamard_rotations = &model->hadamard_rotations;
        p.hadamard_inverses = &model->hadamard_inverses;
        return p;
    }

    std::vector<float> graph_run(int il, const std::vector<float> & input, const std::vector<int> & segments,
            const std::vector<float> * incoming, const bonsai_sample * sample, double * loss_value = nullptr) {
        const int nt = segments.size(), width = model->hparams.n_embd;
        const size_t required_capacity = std::max<size_t>(65536, (size_t) nt * 256 + 16384);
        if (required_capacity > graph_capacity) {
            ggml_backend_sched_free(sched);
            graph_capacity = required_capacity;
            ggml_backend_t backends[] = {gpu, cpu};
            sched = ggml_backend_sched_new(backends, nullptr, 2, graph_capacity, false, true);
            if (o.profile) { ggml_backend_sched_set_eval_callback(sched, profile_node, this); }
        }
        const bool backward = incoming || (sample && loss_value == nullptr && il == (int) model->hparams.n_layer());
        if (backward && !weights.empty()) { prepare_gradients(); }
        ggml_backend_sched_reset(sched);
        const auto graph_start = steady_clock::now();
        llm_graph_result result(graph_capacity);
        llama_adapter_lora adapter(model); adapter.alpha = 0; adapter.transformed_basis = true; adapter.int8_activations = o.int8_activations;
        ggml_context_ptr actx{ggml_init({(weights.size() * 2 + 8) * ggml_tensor_overhead(), nullptr, true})};
        for (auto & item : weights) {
            auto & w = item.second;
            if (w.layer != il || o.base_only) { continue; }
            for (auto f : {&w.a, &w.b}) {
                f->tensor = ggml_new_tensor_2d(actx.get(), o.mode == "ternary" ? GGML_TYPE_I8 : GGML_TYPE_BF16, f->in, f->out);
                if (backward) { ggml_set_param(f->tensor); }
            }
            ggml_set_name(w.a.tensor, (item.first + ".lora_a").c_str());
            ggml_set_name(w.b.tensor, (item.first + ".lora_b").c_str());
            adapter.ab_map[item.first] = {w.a.tensor, w.b.tensor};
            adapter.ab_map[item.first].ternary_scale = w.scale;
        }
        ggml_backend_buffer_ptr abuf{ggml_backend_alloc_ctx_tensors(actx.get(), gpu)};
        if (abuf) { times.peak_adapter = std::max(times.peak_adapter, ggml_backend_buffer_get_size(abuf.get())); }
        if (!adapter.ab_map.empty()) { check(abuf != nullptr, "adapter GPU allocation failed"); ggml_backend_buffer_set_usage(abuf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); }
        llama_adapter_loras loras{{&adapter, 1.0f}};
        auto p = params(result, nt, loras);
        bonsai_training_graph build(*model, p, segments);
        auto * ctx = result.get_ctx();
        auto * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, nt);
        ggml_set_input(x); ggml_set_name(x, "block_input");
        if (backward) { ggml_set_param(x); }
        ggml_tensor * output;
        ggml_tensor * ids = nullptr;
        if (il == -1) {
            ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, nt); ggml_set_input(ids);
            output = build.build_embd_rows(model->tok_embd, ids);
        } else if (il == (int) model->hparams.n_layer()) {
            output = build.build_norm(x, model->output_norm, nullptr, LLM_NORM_RMS, -1);
            output = build.build_lora_mm(model->output, output, model->output_s);
        } else { output = build.block(x, il); }
        ggml_set_output(output);
        auto * graph = ggml_graph_dup(ctx, build.gf, true);
        ggml_build_forward_expand(graph, output);
        ggml_tensor * loss = nullptr, * labels = nullptr, * loss_ids = nullptr, * dy = nullptr, * dx = nullptr;
        std::vector<int32_t> selected;
        if (incoming) {
            dy = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, nt); ggml_set_input(dy);
            loss = ggml_sum(ctx, ggml_mul(ctx, output, dy));
        } else if (sample && il == (int) model->hparams.n_layer()) {
            for (int t = 0; t < nt; ++t) {
                if (sample->loss[t + 1] && sample->segments[t] == sample->segments[t + 1]) { selected.push_back(t); }
            }
            check(!selected.empty(), "sequence has no causal loss tokens");
            loss_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, selected.size()); ggml_set_input(loss_ids);
            auto * logits = ggml_get_rows(ctx, output, loss_ids);
            labels = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, model->vocab.n_tokens(), selected.size()); ggml_set_input(labels);
            loss = ggml_cross_entropy_loss(ctx, logits, labels);
        }
        if (loss) {
            ggml_set_loss(loss); ggml_set_output(loss); ggml_build_forward_expand(graph, loss);
        }
        if (backward) {
            ggml_build_backward_expand(ctx, graph, nullptr);
            dx = ggml_graph_get_grad(graph, x); ggml_set_output(dx);
            for (auto & item : adapter.ab_map) {
                ggml_set_output(ggml_graph_get_grad(graph, item.second.a));
                ggml_set_output(ggml_graph_get_grad(graph, item.second.b));
            }
        }
        struct diagnostic { adapter_weight * w; ggml_tensor * base, * delta; size_t entries; };
        std::vector<diagnostic> diagnostics;
        if (!backward) {
            const int nodes = ggml_graph_n_nodes(graph);
            for (int i = 0; i < nodes; ++i) {
                auto * node = ggml_graph_node(graph, i);
                if (node->op != GGML_OP_ADD || !node->src[1] || node->src[1]->op != GGML_OP_SCALE) { continue; }
                auto * delta = node->src[1], * mm = delta->src[0];
                if (mm->op != GGML_OP_MUL_MAT) { continue; }
                for (auto & item : weights) {
                    auto & w = item.second;
                    if (w.layer != il || w.b.tensor != mm->src[0] || w.output_entries) { continue; }
                    auto * base_sum = ggml_sum(ctx, ggml_sqr(ctx, node->src[0]));
                    auto * delta_sum = ggml_sum(ctx, ggml_sqr(ctx, delta));
                    ggml_set_output(base_sum); ggml_set_output(delta_sum);
                    ggml_build_forward_expand(graph, base_sum); ggml_build_forward_expand(graph, delta_sum);
                    diagnostics.push_back({&w, base_sum, delta_sum, (size_t) ggml_nelements(delta)});
                }
            }
        }
        for (auto * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) {
            if (t->type != GGML_TYPE_F32 && t->type != GGML_TYPE_F16 && t->type != GGML_TYPE_BF16) { continue; }
            if (t->ne[2] != 1 || t->ne[3] != 1 || ggml_nelements(t) < 10000000) { continue; }
            for (const auto & item : llama_internal_get_tensor_map(model)) {
                if (item.second->type == GGML_TYPE_PQ2_0 && ggml_are_same_shape(t, item.second)) { throw std::runtime_error("packed-backbone diagnostic: floating tensor has base weight shape"); }
            }
        }
        check(ggml_backend_sched_alloc_graph(sched, graph), "training graph allocation failed");
        times.graph += seconds(graph_start);
        const auto transfer_start = steady_clock::now();
        if (il != -1) { ggml_backend_tensor_set(x, input.data(), 0, input.size() * 4); }
        if (ids) { ggml_backend_tensor_set(ids, sample->tokens.data(), 0, nt * sizeof(llama_token)); }
        if (dy) { ggml_backend_tensor_set(dy, incoming->data(), 0, incoming->size() * 4); }
        for (auto & item : weights) {
            auto & w = item.second;
            if (w.layer != il || o.base_only) { continue; }
            for (auto f : {&w.a, &w.b}) { ggml_backend_tensor_set(f->tensor, f->bytes.data(), 0, f->bytes.size()); }
        }
        for (auto & constant : build.constants) { ggml_backend_tensor_set(constant.first, constant.second.data(), 0, constant.second.size() * 4); }
        if (build.positions) {
            std::vector<int32_t> positions(ggml_nelements(build.positions));
            for (size_t i = 0; i < positions.size(); ++i) {
                const int t = i % nt; int start = t;
                while (start > 0 && segments[start - 1] == segments[t]) { --start; }
                positions[i] = t - start;
            }
            ggml_backend_tensor_set(build.positions, positions.data(), 0, positions.size() * 4);
        }
        if (labels) {
            ggml_backend_tensor_memset(labels, 0, 0, ggml_nbytes(labels));
            const float one = 1;
            ggml_backend_tensor_set(loss_ids, selected.data(), 0, selected.size() * 4);
            for (size_t j = 0; j < selected.size(); ++j) {
                ggml_backend_tensor_set(labels, &one, (j * labels->ne[0] + sample->tokens[selected[j] + 1]) * 4, 4);
            }
        }
        if (backward) { ggml_graph_reset(graph); }
        ggml_backend_sched_synchronize(sched);
        times.transfer += seconds(transfer_start);
        const auto compute_start = steady_clock::now();
        bonsai_cuda_timer_start(timer);
        check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, "training graph compute failed");
        const double gpu_seconds = bonsai_cuda_timer_stop(timer);
        ggml_backend_sched_synchronize(sched);
        if (backward) { times.backward += gpu_seconds; } else { times.forward += gpu_seconds; }
        size_t free, total; ggml_backend_dev_memory(ggml_backend_get_device(gpu), &free, &total);
        times.peak_used = std::max(times.peak_used, total - free);
        times.peak_compute = std::max(times.peak_compute, ggml_backend_sched_get_buffer_size(sched, gpu));
        const auto download_start = steady_clock::now();
        for (auto & d : diagnostics) {
            float b, a; ggml_backend_tensor_get(d.base, &b, 0, 4); ggml_backend_tensor_get(d.delta, &a, 0, 4);
            d.w->base_squared = b; d.w->delta_squared = a; d.w->output_entries = d.entries;
        }
        std::vector<float> values(loss_value ? 0 : ggml_nelements(backward ? dx : output));
        if (!values.empty()) { ggml_backend_tensor_get(backward ? dx : output, values.data(), 0, values.size() * 4); }
        if (loss) {
            float value; ggml_backend_tensor_get(loss, &value, 0, sizeof(value));
            check(std::isfinite(value), "nonfinite language-model loss");
            if (loss_value) { *loss_value = value; }
            if (sample && !loss_value) { last_loss = value; }
        }
        if (loss_value) {
            std::vector<float> row(output->ne[0]);
            ggml_backend_tensor_get(output, row.data(), (output->ne[1] - 1) * output->nb[1], row.size() * sizeof(float));
            last_logits_hash = hash_values(row);
        }
        if (backward) {
            size_t gradient_bytes = ggml_nbytes(dx);
            for (auto & item : weights) {
                auto & w = item.second;
                if (w.layer != il || o.base_only) { continue; }
                for (auto f : {&w.a, &w.b}) {
                    auto * grad = ggml_graph_get_grad(graph, f->tensor);
                    check(grad != nullptr, "missing adapter gradient");
                    gradient_bytes += ggml_nbytes(grad);
                    if (o.gpu_update) {
                        bonsai_cuda_accumulate(gpu, f->gpu_gradient, grad, loss_weight); f->gradient_ready = true;
                        continue;
                    }
                    std::vector<float> current(ggml_nelements(grad));
                    ggml_backend_tensor_get(grad, current.data(), 0, current.size() * 4);
                    if (f->gradient.empty()) { f->gradient.assign(current.size(), 0); }
                    for (size_t i = 0; i < current.size(); ++i) {
                        check(std::isfinite(current[i]), "nonfinite adapter gradient");
                        f->gradient[i] += current[i] * (float) loss_weight;
                    }
                }
            }
            times.peak_gradients = std::max(times.peak_gradients, gradient_bytes);
            if (o.gpu_update) { ggml_backend_synchronize(gpu); }
        }
        times.download += seconds(download_start);
        if (o.verbose) { std::printf("graph block=%d tokens=%d backward=%d nodes=%d compute_ms=%.3f GPU_graph_MiB=%.2f\n", il, nt, backward, ggml_graph_n_nodes(graph), seconds(compute_start) * 1000, times.peak_compute / 1048576.0); }
        ggml_backend_sched_reset(sched);
        return values;
    }

    double last_loss = 0, loss_weight = 1;

    std::pair<double, size_t> sequence(const bonsai_sample & sample, bool train) {
        for (auto & item : weights) { item.second.output_entries = 0; }
        const int nt = sample.tokens.size() - 1;
        std::vector<int> segments(sample.segments.begin(), sample.segments.begin() + nt);
        size_t nloss = 0;
        for (int t = 0; t < nt; ++t) { nloss += sample.loss[t + 1] && sample.segments[t] == sample.segments[t + 1]; }
        std::vector<std::vector<float>> boundaries;
        auto hidden = graph_run(-1, {}, segments, nullptr, &sample);
        const int depth = o.validate_depth ? o.validate_depth : model->hparams.n_layer();
        for (int il = 0; il < depth; ++il) {
            if (train) { boundaries.push_back(hidden); }
            hidden = graph_run(il, hidden, segments, nullptr, nullptr);
        }
        times.peak_boundaries = std::max(times.peak_boundaries, boundaries.size() * hidden.size() * sizeof(float));
        if (!train) {
            double loss = 0; graph_run(model->hparams.n_layer(), hidden, segments, nullptr, &sample, &loss);
            return {loss, nloss};
        }
        auto gradient = graph_run(model->hparams.n_layer(), hidden, segments, nullptr, &sample);
        const double loss = last_loss;
        for (int il = depth - 1; il >= 0; --il) {
            gradient = graph_run(il, boundaries[il], segments, &gradient, nullptr);
            if (o.verbose) {
                double rms = 0; for (float v : gradient) { rms += v * (double) v; }
                std::printf("gradient block=%d input_rms=%.8g\n", il, std::sqrt(rms / gradient.size()));
            }
        }
        return {loss, nloss};
    }

    void validate_inference(const bonsai_sample & sample) {
        check(std::count(sample.segments.begin(), sample.segments.end(), 0) == (int) sample.tokens.size(), "forward reference requires --no-pack");
        const int nt = sample.tokens.size() - 1, width = model->hparams.n_embd;
        auto cp = llama_context_default_params();
        cp.n_ctx = std::max(256, nt); cp.n_batch = nt; cp.n_ubatch = nt;
        cp.type_k = cp.type_v = GGML_TYPE_F32;
        cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
        llama_context * reference = llama_init_from_model(model, cp); check(reference != nullptr, "reference context failed");
        for (int il = 0; il < 64; ++il) { reference->set_embeddings_layer_inp(il, true); }
        llama_adapter_lora * adapter = nullptr;
        if (!o.base_only) {
            check(!o.resume.empty(), "adapter forward reference requires --adapter or --base-only");
            adapter = llama_adapter_lora_init(model, o.resume.c_str()); check(adapter != nullptr, "reference adapter load failed");
            llama_adapter_lora * list[] = {adapter}; float scales[] = {1}; reference->set_adapters_lora(list, 1, scales);
        }
        auto batch = llama_batch_init(nt, 0, 1); batch.n_tokens = nt;
        for (int t = 0; t < nt; ++t) {
            batch.token[t] = sample.tokens[t]; batch.pos[t] = t; batch.n_seq_id[t] = 1; batch.seq_id[t][0] = 0; batch.logits[t] = 1;
        }
        check(llama_decode(reference, batch) == 0, "reference decode failed"); llama_synchronize(reference);
        auto relative = [](const std::vector<float> & a, const float * b) {
            double error = 0, norm = 0;
            for (size_t i = 0; i < a.size(); ++i) { error += (a[i] - (double) b[i]) * (a[i] - (double) b[i]); norm += b[i] * (double) b[i]; }
            return std::sqrt(error / std::max(norm, 1e-30));
        };
        std::vector<int> segments(nt, 0);
        auto hidden = graph_run(-1, {}, segments, nullptr, &sample);
        double maximum = 0;
        for (int il = 0; il < 64; ++il) {
            const auto error = relative(hidden, reference->get_embeddings_layer_inp(il));
            maximum = std::max(maximum, error);
            if (o.verbose || il == 0 || il == 3 || il == 63) { std::printf("INFERENCE_REFERENCE layer_input=%d relative_L2=%.9g entries=%d\n", il, error, nt * width); }
            hidden = graph_run(il, hidden, segments, nullptr, nullptr);
        }
        const auto logits = graph_run(64, hidden, segments, nullptr, nullptr);
        const double error = relative(logits, llama_get_logits(reference));
        std::printf("INFERENCE_REFERENCE logits_relative_L2=%.9g max_layer_relative_L2=%.9g tokens=%d\n", error, maximum, nt);
        llama_batch_free(batch); llama_free(reference); if (adapter) { llama_adapter_lora_free(adapter); }
        check(error < 0.025 && maximum < 0.025, "training/inference forward mismatch exceeds tolerance");
    }

    void validate_micro_batch(const std::vector<bonsai_sample> & samples) {
        check(samples.size() >= 2, "batch validation needs at least two sequences");
        const auto first = sequence(samples[0], false), second = sequence(samples[1], false);
        const double independent = (first.first * first.second + second.first * second.second) / (first.second + second.second);
        auto batch = bonsai_micro_batch(samples, {0, 1}, 0, 2);
        const auto joined = sequence(batch, false);
        const auto original_hash = last_logits_hash;
        for (size_t i = 0; i < samples[0].tokens.size(); ++i) { batch.tokens[i] = samples[0].tokens.back(); }
        sequence(batch, false);
        check(last_logits_hash == original_hash, "changing an earlier document changes the later document's logits");
        std::printf("MICRO_BATCH_REFERENCE independent_loss=%.9f joined_loss=%.9f absolute_error=%.9g supervised_tokens=%zu\n", independent, joined.first, std::fabs(joined.first - independent), joined.second);
        std::puts("MICRO_BATCH_DOCUMENT_ISOLATION_BIT_IDENTICAL");
        check(joined.second == first.second + second.second && std::fabs(joined.first - independent) < .02, "micro-batch changes causal loss beyond the low-bit projection tolerance");
    }

    void update() {
        const auto start = steady_clock::now();
        size_t changed = 0, selected = 0, saturated = 0, positive = 0, negative = 0, count = 0, trits[3] = {};
        for (auto & item : weights) {
            for (auto f : {&item.second.a, &item.second.b}) {
                if (o.gpu_update && f->gradient_ready) {
                    auto * stats = f->last_stats.data();
                    bonsai_cuda_update(gpu, f->bytes.data(), o.mode == "ternary" ? GGML_TYPE_I8 : GGML_TYPE_BF16, f->gpu_gradient, o.lr, stats);
                    check(stats[9] == 0, "nonfinite adapter gradient");
                    count += stats[0]; changed += stats[1]; selected += stats[2]; saturated += stats[3];
                    negative += stats[4]; positive += stats[5];
                    for (int i = 0; i < 3; ++i) { trits[i] += stats[6 + i]; }
                    f->gradient_ready = false; continue;
                }
                if (f->gradient.empty()) { continue; }
                const std::array<size_t, 9> before = {count, changed, selected, saturated, negative, positive, trits[0], trits[1], trits[2]};
                count += f->gradient.size();
                if (o.mode == "ternary") {
                    std::vector<float> magnitude(f->gradient.size());
                    for (size_t i = 0; i < magnitude.size(); ++i) { magnitude[i] = std::fabs(f->gradient[i]); }
                    const size_t budget = std::max<size_t>(1, std::ceil(o.lr * magnitude.size()));
                    std::nth_element(magnitude.begin(), magnitude.end() - budget, magnitude.end());
                    const float threshold = magnitude[magnitude.size() - budget];
                    size_t above = 0, ties = 0;
                    for (float g : f->gradient) { above += std::fabs(g) > threshold; }
                    for (size_t i = 0; i < f->gradient.size(); ++i) {
                        const float g = f->gradient[i]; positive += g > 0; negative += g < 0;
                        int8_t current = (int8_t) f->bytes[i];
                        const bool chosen = std::fabs(g) > threshold || (std::fabs(g) == threshold && ties++ < budget - above);
                        if (g != 0 && chosen) {
                            ++selected;
                            const int next = current + (g > 0 ? -1 : 1);
                            saturated += next < -1 || next > 1;
                            const int8_t updated = std::max(-1, std::min(1, next));
                            changed += updated != current; current = updated;
                            f->bytes[i] = (uint8_t) current;
                        }
                        ++trits[current + 1];
                    }
                } else {
                    for (size_t i = 0; i < f->gradient.size(); ++i) {
                        negative += f->gradient[i] < 0; positive += f->gradient[i] > 0;
                        ggml_bf16_t old; std::memcpy(&old, f->bytes.data() + i * 2, 2);
                        auto next = ggml_fp32_to_bf16(ggml_bf16_to_fp32(old) - (float) o.lr * f->gradient[i]);
                        changed += std::memcmp(&old, &next, 2) != 0;
                        std::memcpy(f->bytes.data() + i * 2, &next, 2);
                    }
                }
                const std::array<size_t, 9> after = {count, changed, selected, saturated, negative, positive, trits[0], trits[1], trits[2]};
                for (size_t i = 0; i < before.size(); ++i) { f->last_stats[i] = after[i] - before[i]; }
                f->gradient.clear();
            }
        }
        ++step; times.update += seconds(start);
        size_t free, total; ggml_backend_dev_memory(ggml_backend_get_device(gpu), &free, &total); times.peak_used = std::max(times.peak_used, total - free);
        std::printf("update step=%llu parameters=%zu changed=%zu fraction=%.8g selected=%zu saturated=%zu grad_sign_neg=%.5f grad_sign_pos=%.5f trits_neg=%.5f trits_zero=%.5f trits_pos=%.5f update_ms=%.3f\n",
            (unsigned long long) step, count, changed, count ? (double) changed / count : 0, selected, saturated,
            count ? (double) negative / count : 0, count ? (double) positive / count : 0,
            count ? (double) trits[0] / count : 0, count ? (double) trits[1] / count : 0, count ? (double) trits[2] / count : 0, seconds(start) * 1000);
        if (!changed) { std::puts("optimizer_no_changes: zero gradients, BF16 rounding, or saturated discrete transitions"); }
    }

    double evaluate(const std::vector<bonsai_sample> & eval, const char * label) {
        if (eval.empty()) { std::printf("%s unavailable: no held-out sequences\n", label); return NAN; }
        double total = 0; size_t count = 0; std::string first_logits_hash;
        const bool heldout = std::strncmp(label, "heldout", 7) == 0;
        for (size_t i = 0; i < eval.size() && (!o.max_eval || i < (size_t) o.max_eval); ++i) {
            auto result = sequence(eval[i], false); total += result.first * result.second; count += result.second;
            if (i == 0) { first_logits_hash = last_logits_hash; }
        }
        const double loss = total / count;
        std::printf("%s loss=%.9f perplexity=%.9f supervised_tokens=%zu\n", label, loss, std::exp(loss), count);
        std::printf("%s first_sequence_last_logits_sha256=%s\n", label, first_logits_hash.c_str());
        if (heldout) { eval_logits_hash = first_logits_hash; }
        if (heldout && o.eval_only && !saved_logits_hash.empty()) {
            if (first_logits_hash == saved_logits_hash) { std::puts("SAVED_RELOADED_LOGITS_BIT_IDENTICAL"); }
            else { check(!o.verify_reload, "saved/reloaded logits differ; verify identical GPU/runtime and dataset settings"); std::puts("saved/reloaded logits differ at bit level"); }
        }
        if (heldout && o.verify_reload) { check(!saved_logits_hash.empty(), "adapter has no saved validation logits hash"); }
        return loss;
    }

    void memory() {
        size_t free, total; ggml_backend_dev_memory(ggml_backend_get_device(gpu), &free, &total);
        size_t host_gradients = 0;
        for (const auto & item : weights) { host_gradients += (item.second.a.gradient.capacity() + item.second.b.gradient.capacity()) * sizeof(float); }
        std::printf("memory accounting MiB: packed_model_GPU=%.2f adapter_parameters_host=%.2f adapter_parameters_GPU_peak=%.2f adapter_auxiliary_state=0 activation_boundaries_host_peak=%.2f gradients_host=%.2f gradient_accumulation_GPU=%.2f retained_gradients_GPU_peak=%.2f activations_and_temporary_gradients_GPU_peak=%.2f CUDA_cuBLAS_workspace=%.2f KV=0 total_GPU_global=%.2f initial_GPU_global=%.2f peak_GPU_global=%.2f\n",
                packed_gpu / 1048576.0, parameter_bytes() / 1048576.0, times.peak_adapter / 1048576.0, times.peak_boundaries / 1048576.0,
                host_gradients / 1048576.0, gradient_buffer ? ggml_backend_buffer_get_size(gradient_buffer.get()) / 1048576.0 : 0, times.peak_gradients / 1048576.0, (times.peak_compute - std::min(times.peak_compute, times.peak_gradients)) / 1048576.0,
                bonsai_cuda_workspace_bytes(gpu) / 1048576.0, (total - free) / 1048576.0, initial_used / 1048576.0, times.peak_used / 1048576.0);
        const size_t reserved = ggml_backend_sched_get_buffer_size(sched, gpu);
        const size_t known = initial_used + reserved + (gradient_buffer ? ggml_backend_buffer_get_size(gradient_buffer.get()) : 0) + bonsai_cuda_workspace_bytes(gpu);
        std::printf("memory allocator accounting MiB: graph_reserved_GPU=%.2f CUDA_allocator_cache_and_other_over_startup=%.2f startup_other_than_packed_model=%.2f (component peaks occur at different times)\n", reserved / 1048576.0, (total - free > known ? total - free - known : 0) / 1048576.0, (initial_used - packed_gpu) / 1048576.0);
        for (const auto & item : times.profile) { std::printf("profile CUDA_event_seconds category=%s seconds=%.6f\n", item.first.c_str(), item.second); }
    }

    void diagnostics() {
        double base = 0, delta = 0; size_t entries = 0;
        for (const auto & item : weights) {
            const auto & w = item.second;
            base += w.base_squared; delta += w.delta_squared; entries += w.output_entries;
            if (o.verbose && w.output_entries) {
                std::printf("adapter_output projection=%s base_RMS=%.8g adapter_RMS=%.8g ratio=%.8g\n", item.first.c_str(), std::sqrt(w.base_squared / w.output_entries), std::sqrt(w.delta_squared / w.output_entries), std::sqrt(w.delta_squared / std::max(w.base_squared, 1e-30)));
                for (const auto * f : {&w.a, &w.b}) {
                    const double n = f->in * f->out; size_t distribution[3] = {};
                    if (o.mode == "ternary") { for (uint8_t v : f->bytes) { ++distribution[(int8_t) v + 1]; } }
                    std::printf("adapter_state projection=%s factor=%s trits_neg=%.6g trits_zero=%.6g trits_pos=%.6g changed=%.6g gradient_neg=%.6g gradient_pos=%.6g saturation_rate=%.6g\n", item.first.c_str(), f == &w.a ? "A" : "B", distribution[0] / n, distribution[1] / n, distribution[2] / n, f->last_stats[1] / n, f->last_stats[4] / n, f->last_stats[5] / n, f->last_stats[2] ? (double) f->last_stats[3] / f->last_stats[2] : 0);
                }
            }
        }
        if (entries) { std::printf("adapter_output optimizer_step=%llu base_RMS=%.8g adapter_RMS=%.8g ratio=%.8g entries=%zu\n", (unsigned long long) step, std::sqrt(base / entries), std::sqrt(delta / entries), std::sqrt(delta / std::max(base, 1e-30)), entries); }
    }
};

static void generate(llama_model * model, const options & o) {
    llama_adapter_lora * adapter = llama_adapter_lora_init(model, o.resume.c_str()); check(adapter != nullptr, "inference adapter load failed");
    auto cp = llama_context_default_params(); cp.n_ctx = std::max(256, o.context); cp.n_batch = cp.n_ctx; cp.n_ubatch = cp.n_ctx; cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
    llama_context * ctx = llama_init_from_model(model, cp); check(ctx != nullptr, "inference context failed");
    llama_adapter_lora * adapters[] = {adapter}; float scales[] = {1}; llama_set_adapters_lora(ctx, adapters, 1, scales);
    auto templates = common_chat_templates_init(model, ""); common_chat_templates_inputs inputs;
    common_chat_msg user; user.role = "user"; user.content = o.generate; inputs.messages = {user}; inputs.enable_thinking = false;
    auto tokens = common_tokenize(llama_model_get_vocab(model), common_chat_templates_apply(templates.get(), inputs).prompt, true, true);
    check(tokens.size() + 64 < cp.n_ctx, "generation prompt exceeds context");
    check(llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size())) == 0, "generation prefill failed");
    for (int i = 0; i < 64; ++i) {
        const float * logits = llama_get_logits_ith(ctx, -1);
        const auto token = (llama_token) (std::max_element(logits, logits + llama_vocab_n_tokens(llama_model_get_vocab(model))) - logits);
        if (llama_vocab_is_eog(llama_model_get_vocab(model), token)) { break; }
        std::printf("%s", common_token_to_piece(ctx, token).c_str()); std::fflush(stdout);
        llama_token next = token; check(llama_decode(ctx, llama_batch_get_one(&next, 1)) == 0, "generation decode failed");
    }
    std::puts(""); llama_free(ctx); llama_adapter_lora_free(adapter);
}

int main(int argc, char ** argv) {
    try {
        const auto run_start = steady_clock::now();
        const options o = parse(argc, argv);
        auto log_callback = [](ggml_log_level level, const char * message, void * data) {
            if (level != GGML_LOG_LEVEL_DEBUG || *(const bool *) data) { std::fputs(message, stderr); }
        };
        bool verbose_logs = o.verbose;
        llama_log_set(log_callback, &verbose_logs); ggml_log_set(log_callback, &verbose_logs);
        setenv("GGML_PQ2_NO_FP_COPY", "1", 1);
        llama_backend_init();
        ggml_backend_t gpu = ggml_backend_init_by_name("CUDA0", nullptr), cpu = ggml_backend_cpu_init();
        check(gpu && cpu, "CUDA0 and CPU backends are required");
        ggml_backend_dev_t devices[] = {ggml_backend_get_device(gpu), nullptr};
        auto mp = llama_model_default_params(); mp.devices = devices; mp.n_gpu_layers = 99;
        llama_model * model = llama_model_load_from_file(o.model.c_str(), mp); check(model != nullptr, "model load failed");
        check(model->arch == LLM_ARCH_QWEN35 && model->hparams.n_layer() == 64 && model->hparams.n_embd == 5120 && model->output->type == GGML_TYPE_PQ2_0, "trainer requires the Bonsai 2 27B packed Qwen3.5 checkpoint");
        FILE * log = o.generate.empty() ? stdout : stderr;
        std::fprintf(log, "environment GPU=%s compiler=%s commit=%s CUDA_device=CUDA0 model=%s\n", ggml_backend_dev_description(devices[0]), llama_compiler(), llama_commit(), o.model.c_str());
        std::fprintf(log, "environment %s\n", bonsai_cuda_environment(gpu).c_str());
        std::fflush(stdout);
        const std::string hash = hash_file(o.model); std::fprintf(log, "base_sha256=%s\n", hash.c_str());
        const std::string tensor_hash = hash_model(model); std::fprintf(log, "frozen_loaded_tensor_sha256=%s\n", tensor_hash.c_str()); std::fflush(log);
        if (!o.generate.empty()) {
            { trainer verify(o, model, gpu, cpu, hash); }
            generate(model, o);
        } else {
            trainer train(o, model, gpu, cpu, hash);
            const options & o = train.o;
            std::printf("settings rank=%d target=%s context=%d batch_size=%d micro_batch=%d epochs=%d lr=%.17g seed=%u eval_split=%.6g mode=%s pack=%d response_only=%d chat=%d data_version=%d int8_activations=%d GPU_update=%d\n",
                    o.rank, o.targets.c_str(), o.context, o.batch, o.micro, o.epochs, o.lr, o.seed, o.eval_split, o.mode.c_str(), o.pack, o.text.empty() && !o.all_tokens, o.chat, o.data_version, o.int8_activations, o.gpu_update);
            train.memory();
            if (o.validate_block >= 0) {
                check(o.validate_block < 64, "block validation index out of range");
                std::vector<float> input(model->hparams.n_embd * o.context), incoming(input.size(), 0.001f);
                for (size_t i = 0; i < input.size(); ++i) { input[i] = 0.1f * std::sin(i * 0.017); }
                std::vector<int> segments(o.context, 0);
                for (int warmup = 0; warmup < 3; ++warmup) {
                    train.graph_run(o.validate_block, input, segments, nullptr, nullptr);
                    train.graph_run(o.validate_block, input, segments, &incoming, nullptr);
                    train.clear_gradients();
                }
                train.times = {};
                train.graph_run(o.validate_block, input, segments, nullptr, nullptr);
                auto dx = train.graph_run(o.validate_block, input, segments, &incoming, nullptr);
                size_t nz = 0; for (float v : dx) { check(std::isfinite(v), "nonfinite real block input gradient"); nz += v != 0; }
                std::printf("REAL_BLOCK_BACKWARD block=%d input_grad_nonzero=%zu/%zu\n", o.validate_block, nz, dx.size());
                train.diagnostics(); train.update(); train.memory();
                check(hash_file(o.model) == hash, "frozen base GGUF changed");
                std::puts("FROZEN_BASE_SHA256_UNCHANGED");
            } else {
                const auto tokenize_start = steady_clock::now();
                const std::string dataset_hash = hash_file(o.data);
                check(train.dataset_hash.empty() || train.dataset_hash == dataset_hash, "resume dataset SHA256 mismatch");
                train.dataset_hash = dataset_hash;
                auto dataset = bonsai_dataset(model, o.data, o.prompt, o.response, o.text, o.split, o.all_tokens, o.chat, o.context, o.seed, o.eval_split, o.data_version);
                if (o.max_train && dataset.first.size() > (size_t) o.max_train) { dataset.first.resize(o.max_train); }
                const auto packing_start = steady_clock::now();
                if (o.pack) { dataset.first = bonsai_pack(dataset.first, o.context); dataset.second = bonsai_pack(dataset.second, o.context); }
                std::printf("packing_seconds=%.6f\n", seconds(packing_start));
                std::printf("tokenization_and_packing_seconds=%.6f train_sequences=%zu eval_sequences=%zu\n", seconds(tokenize_start), dataset.first.size(), dataset.second.size());
                if (o.validate_forward) { train.validate_inference(dataset.first.front()); }
                else if (o.validate_batch) { train.validate_micro_batch(dataset.first); }
                else if (o.validate_depth) {
                    auto result = train.sequence(dataset.first.front(), true);
                    std::printf("PREFIX_CAUSAL_LM_BACKWARD depth=%d loss=%.9f supervised_tokens=%zu\n", o.validate_depth, result.first, result.second);
                    train.diagnostics(); train.update(); train.memory();
                } else { train.evaluate(dataset.second, "heldout_initial"); train.evaluate(dataset.first, "train_evaluation"); }
                if (!o.eval_only && !o.validate_forward && !o.validate_batch && !o.validate_depth) {
                    check(train.epoch < (size_t) o.epochs, "resume has completed the requested total epochs; increase --epochs or use --eval-only");
                    auto warm = train.sequence(dataset.first[0], true);
                    train.clear_gradients();
                    std::printf("warmup train_loss=%.9f\n", warm.first); train.diagnostics(); train.times = {};
                    size_t tokens = 0, samples = 0, updates = 0; double training_seconds = 0;
                    check(train.cursor <= dataset.first.size(), "saved sampler cursor exceeds dataset size");
                    while (train.epoch < (size_t) o.epochs && (!o.steps || updates < (size_t) o.steps)) {
                        std::vector<size_t> order(dataset.first.size()); std::iota(order.begin(), order.end(), 0);
                        std::mt19937 rng(o.seed + (unsigned) train.epoch); std::shuffle(order.begin(), order.end(), rng);
                        for (size_t begin = train.cursor; begin < dataset.first.size() && (!o.steps || updates < (size_t) o.steps); begin += o.batch) {
                            const size_t end = std::min(dataset.first.size(), begin + o.batch);
                            size_t supervised = 0;
                            for (size_t j = begin; j < end; ++j) {
                                const auto & s = dataset.first[order[j]];
                                for (size_t t = 1; t < s.tokens.size(); ++t) { supervised += s.loss[t] && s.segments[t] == s.segments[t - 1]; }
                            }
                            double loss = 0; const auto step_start = steady_clock::now();
                            for (size_t j = begin; j < end; j += o.micro) {
                                const auto s = bonsai_micro_batch(dataset.first, order, j, std::min(end, j + o.micro)); size_t nloss = 0;
                                for (size_t t = 1; t < s.tokens.size(); ++t) { nloss += s.loss[t] && s.segments[t] == s.segments[t - 1]; }
                                train.loss_weight = (double) nloss / supervised;
                                auto result = train.sequence(s, true); loss += result.first * result.second;
                                tokens += s.tokens.size() - 1; samples += std::min(end, j + o.micro) - j;
                            }
                            train.update(); ++updates;
                            const double step_seconds = seconds(step_start); training_seconds += step_seconds; train.cursor = end;
                            std::printf("train epoch=%llu step=%llu loss=%.9f supervised_tokens=%zu step_seconds=%.6f\n", (unsigned long long) train.epoch + 1, (unsigned long long) train.step, loss / supervised, supervised, step_seconds);
                            std::printf("GPU_utilization_last_NVML_interval_percent=%d\n", bonsai_cuda_utilization(gpu));
                            train.diagnostics(); train.memory(); std::fflush(stdout);
                        }
                        if (train.cursor == dataset.first.size()) { ++train.epoch; train.cursor = 0; }
                        const auto measured = train.times;
                        train.evaluate(dataset.second, "heldout_epoch"); train.evaluate(dataset.first, "train_epoch"); train.save(); train.times = measured;
                        std::printf("checkpoint_progress completed_epochs=%llu optimizer_step=%llu warmed_training_seconds=%.6f elapsed_wall_seconds=%.6f\n",
                                (unsigned long long) train.epoch, (unsigned long long) train.step, training_seconds, seconds(run_start));
                    }
                    std::printf("performance measured_after_warmup tokens_per_second=%.6f samples_per_second=%.6f training_seconds=%.6f graph_CPU_seconds=%.6f forward_CUDA_event_seconds=%.6f recompute_plus_backward_CUDA_event_seconds=%.6f upload_sync_seconds=%.6f download_and_gradient_accumulation_seconds=%.6f update_seconds=%.6f\n",
                            tokens / training_seconds, samples / training_seconds, training_seconds, train.times.graph, train.times.forward, train.times.backward, train.times.transfer, train.times.download, train.times.update);
                }
                train.memory();
                check(hash_file(o.model) == hash, "frozen base GGUF changed");
                std::puts("FROZEN_BASE_SHA256_UNCHANGED");
            }
        }
        check(hash_model(model) == tensor_hash, "frozen model tensor bytes changed in memory");
        std::fprintf(log, "FROZEN_LOADED_TENSORS_SHA256_UNCHANGED\n");
        llama_model_free(model); ggml_backend_free(cpu); ggml_backend_free(gpu); llama_backend_free();
        return 0;
    } catch (const std::exception & e) {
        std::fprintf(stderr, "ERROR: %s\n", e.what()); return 1;
    }
}