ApacheOne's picture
Publish packed-Q1 native LoRA reproducibility release
bf45b91 verified
Raw
History Blame Contribute Delete
49.3 kB
// STEP 10 native resumable packed-Q1 LoRA trainer acceptance executable.
//
// The accepted Stage 7 loader/full-model graph is used only as a bootstrap.
// Its optimizer epoch is intercepted. GGML computes the real masked backward;
// deterministic host AdamW updates only the six LoRA A/B tensors.
#include "llama.h"
#include "llama-adapter.h"
#include "llama-model.h"
#include "ggml.h"
#include "ggml-backend.h"
#include <algorithm>
#include <array>
#include <cerrno>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <regex>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <utility>
#include <vector>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
namespace fs = std::filesystem;
struct target_spec {
const char * name;
const char * category;
int block;
int64_t K;
int64_t M;
int64_t rank;
float alpha;
};
static const std::array<target_spec, 3> g_specs = {{
{"blk.0.ssm_alpha.weight", "ssm", 0, 5120, 48, 4, 8.0f},
{"blk.11.attn_k.weight", "attention", 11, 5120, 1024, 4, 8.0f},
{"blk.0.ffn_down.weight", "ffn", 0, 17408, 5120, 4, 8.0f},
}};
struct trainer_config {
std::string model_path;
std::string adapter_path;
std::string train_dir;
std::string validation_dir;
std::string output_dir;
std::string mode;
std::string resume_path;
std::string model_sha256;
std::string dataset_fingerprint;
uint64_t seed = 1234;
uint32_t context_tokens = 64;
uint64_t target_global_steps = 2;
uint32_t accumulation_steps = 2;
float base_lr = 2.0e-4f;
float min_lr = 2.0e-5f;
uint32_t warmup_steps = 1;
float beta1 = 0.9f;
float beta2 = 0.999f;
float eps = 1.0e-8f;
float weight_decay = 0.0f;
float max_grad_norm = 1.0e-3f;
uint64_t stop_after_micro = 0;
};
struct sample_record {
std::vector<llama_token> tokens;
std::vector<uint8_t> target_mask;
int32_t source_index = -1;
};
struct param_state {
std::string name;
ggml_tensor * tensor = nullptr;
std::vector<float> value;
std::vector<float> initial;
std::vector<float> m;
std::vector<float> v;
std::vector<float> grad_sum;
};
struct trainer_state {
uint64_t global_step = 0;
uint32_t micro_step = 0;
uint64_t total_micro = 0;
uint64_t epoch = 0;
uint64_t sample_cursor = 0;
uint64_t rng_state = 0;
uint64_t optimizer_updates = 0;
uint64_t checkpoint_writes = 0;
uint64_t clip_applied_count = 0;
uint64_t early_update_violations = 0;
std::vector<double> train_losses;
std::vector<double> learning_rates;
std::vector<double> grad_norm_pre;
std::vector<double> grad_norm_post;
};
struct run_outcome {
bool hook_seen = false;
bool success = false;
std::string error;
std::string checkpoint_path;
std::string checkpoint_sha256;
std::string adapter_path;
std::string adapter_sha256;
std::string state_fingerprint;
double validation_loss = NAN;
bool validation_unchanged = false;
double adapter_reload_max_diff = INFINITY;
size_t base_changed_bytes = 0;
size_t train_sample_count = 0;
size_t validation_sample_count = 0;
};
static trainer_config g_cfg;
static trainer_state g_state;
static run_outcome g_outcome;
static llama_adapter_lora * g_adapter = nullptr;
static llama_model * g_model = nullptr;
static llama_context * g_context = nullptr;
static std::vector<param_state> g_params;
static std::vector<std::pair<std::string, std::vector<uint8_t>>> g_base_initial;
static std::string shell_quote(const std::string & input) {
std::string output = "'";
for (const char value : input) {
if (value == '\'') {
output += "'\"'\"'";
} else {
output += value;
}
}
output += "'";
return output;
}
static std::string sha256_file(const fs::path & path) {
const std::string command = "sha256sum " + shell_quote(path.string());
FILE * pipe = popen(command.c_str(), "r");
if (!pipe) {
throw std::runtime_error("could not run sha256sum");
}
std::string output;
char buffer[512] = {};
while (fgets(buffer, sizeof(buffer), pipe)) {
output += buffer;
}
const int rc = pclose(pipe);
if (rc != 0 || output.size() < 64) {
throw std::runtime_error("sha256sum failed: " + path.string());
}
return output.substr(0, 64);
}
static uint64_t fnv1a64(
const uint8_t * data,
size_t size,
uint64_t seed = 1469598103934665603ULL) {
uint64_t result = seed;
for (size_t i = 0; i < size; ++i) {
result ^= data[i];
result *= 1099511628211ULL;
}
return result;
}
static std::string hex64(uint64_t value) {
std::ostringstream output;
output << std::hex << std::setw(16) << std::setfill('0') << value;
return output.str();
}
static std::vector<float> tensor_f32(const ggml_tensor * tensor) {
if (!tensor || !tensor->buffer) {
throw std::runtime_error("unallocated tensor");
}
const size_t count = ggml_nelements(tensor);
std::vector<float> values(count);
if (tensor->type == GGML_TYPE_F32) {
ggml_backend_tensor_get(tensor, values.data(), 0, count*sizeof(float));
return values;
}
if (tensor->type == GGML_TYPE_F16) {
std::vector<ggml_fp16_t> temp(count);
ggml_backend_tensor_get(tensor, temp.data(), 0, count*sizeof(ggml_fp16_t));
for (size_t i = 0; i < count; ++i) {
values[i] = ggml_fp16_to_fp32(temp[i]);
}
return values;
}
throw std::runtime_error("LoRA tensor is not F32/F16");
}
static void set_tensor_f32(ggml_tensor * tensor, const std::vector<float> & values) {
if (!tensor || tensor->type != GGML_TYPE_F32 || values.size() != ggml_nelements(tensor)) {
throw std::runtime_error("invalid F32 tensor update");
}
ggml_backend_tensor_set(tensor, values.data(), 0, values.size()*sizeof(float));
}
static std::vector<uint8_t> tensor_bytes(const ggml_tensor * tensor) {
std::vector<uint8_t> bytes(ggml_nbytes(tensor));
ggml_backend_tensor_get(tensor, bytes.data(), 0, bytes.size());
return bytes;
}
static size_t changed_bytes(
const std::vector<uint8_t> & left,
const std::vector<uint8_t> & right) {
if (left.size() != right.size()) {
return std::max(left.size(), right.size());
}
size_t changed = 0;
for (size_t i = 0; i < left.size(); ++i) {
changed += left[i] != right[i];
}
return changed;
}
static uint64_t parameter_fingerprint() {
uint64_t value = 1469598103934665603ULL;
for (const param_state & state : g_params) {
value = fnv1a64(
reinterpret_cast<const uint8_t *>(state.name.data()),
state.name.size(),
value);
value = fnv1a64(
reinterpret_cast<const uint8_t *>(state.value.data()),
state.value.size()*sizeof(float),
value);
}
return value;
}
static std::string config_identity() {
std::ostringstream output;
output << "prism-step10-v1"
<< "|ctx=" << g_cfg.context_tokens
<< "|target_steps=" << g_cfg.target_global_steps
<< "|accum=" << g_cfg.accumulation_steps
<< "|base_lr=" << std::setprecision(9) << g_cfg.base_lr
<< "|min_lr=" << g_cfg.min_lr
<< "|warmup=" << g_cfg.warmup_steps
<< "|beta1=" << g_cfg.beta1
<< "|beta2=" << g_cfg.beta2
<< "|eps=" << g_cfg.eps
<< "|wd=" << g_cfg.weight_decay
<< "|clip=" << g_cfg.max_grad_norm
<< "|seed=" << g_cfg.seed;
return output.str();
}
static void append_u32(std::vector<uint8_t> & out, uint32_t value) {
for (int i = 0; i < 4; ++i) out.push_back((value >> (8*i)) & 0xff);
}
static void append_u64(std::vector<uint8_t> & out, uint64_t value) {
for (int i = 0; i < 8; ++i) out.push_back((value >> (8*i)) & 0xff);
}
static void append_f32(std::vector<uint8_t> & out, float value) {
uint32_t bits = 0;
std::memcpy(&bits, &value, sizeof(bits));
append_u32(out, bits);
}
static void append_f64(std::vector<uint8_t> & out, double value) {
uint64_t bits = 0;
std::memcpy(&bits, &value, sizeof(bits));
append_u64(out, bits);
}
static void append_string(std::vector<uint8_t> & out, const std::string & value) {
append_u32(out, static_cast<uint32_t>(value.size()));
out.insert(out.end(), value.begin(), value.end());
}
static void append_float_vector(std::vector<uint8_t> & out, const std::vector<float> & values) {
append_u64(out, values.size());
const uint8_t * ptr = reinterpret_cast<const uint8_t *>(values.data());
out.insert(out.end(), ptr, ptr + values.size()*sizeof(float));
}
static void append_double_vector(std::vector<uint8_t> & out, const std::vector<double> & values) {
append_u64(out, values.size());
for (double value : values) append_f64(out, value);
}
struct byte_reader {
const std::vector<uint8_t> & data;
size_t offset = 0;
uint32_t u32() {
if (offset + 4 > data.size()) throw std::runtime_error("checkpoint EOF u32");
uint32_t value = 0;
for (int i = 0; i < 4; ++i) value |= uint32_t(data[offset++]) << (8*i);
return value;
}
uint64_t u64() {
if (offset + 8 > data.size()) throw std::runtime_error("checkpoint EOF u64");
uint64_t value = 0;
for (int i = 0; i < 8; ++i) value |= uint64_t(data[offset++]) << (8*i);
return value;
}
float f32() {
uint32_t bits = u32(); float value; std::memcpy(&value, &bits, sizeof(value)); return value;
}
double f64() {
uint64_t bits = u64(); double value; std::memcpy(&value, &bits, sizeof(value)); return value;
}
std::string string() {
const uint32_t size = u32();
if (offset + size > data.size()) throw std::runtime_error("checkpoint EOF string");
std::string value(reinterpret_cast<const char *>(data.data() + offset), size);
offset += size; return value;
}
std::vector<float> floats() {
const uint64_t count = u64();
if (count > (1ULL << 32) || offset + count*sizeof(float) > data.size()) {
throw std::runtime_error("checkpoint invalid float vector");
}
std::vector<float> values(count);
std::memcpy(values.data(), data.data() + offset, count*sizeof(float));
offset += count*sizeof(float); return values;
}
std::vector<double> doubles() {
const uint64_t count = u64();
if (count > (1ULL << 30)) throw std::runtime_error("checkpoint invalid double vector");
std::vector<double> values(count);
for (uint64_t i = 0; i < count; ++i) values[i] = f64();
return values;
}
};
static std::vector<uint8_t> serialize_checkpoint_payload() {
std::vector<uint8_t> payload;
append_string(payload, g_cfg.model_sha256);
append_string(payload, g_cfg.dataset_fingerprint);
append_string(payload, config_identity());
append_u64(payload, g_state.global_step);
append_u32(payload, g_state.micro_step);
append_u64(payload, g_state.total_micro);
append_u64(payload, g_state.epoch);
append_u64(payload, g_state.sample_cursor);
append_u64(payload, g_state.rng_state);
append_u64(payload, g_state.optimizer_updates);
append_u64(payload, g_state.clip_applied_count);
append_u64(payload, g_state.early_update_violations);
append_u32(payload, static_cast<uint32_t>(g_params.size()));
for (const param_state & state : g_params) {
append_string(payload, state.name);
append_float_vector(payload, state.value);
append_float_vector(payload, state.m);
append_float_vector(payload, state.v);
append_float_vector(payload, state.grad_sum);
}
append_double_vector(payload, g_state.train_losses);
append_double_vector(payload, g_state.learning_rates);
append_double_vector(payload, g_state.grad_norm_pre);
append_double_vector(payload, g_state.grad_norm_post);
return payload;
}
static void fsync_file(const fs::path & path) {
const int fd = open(path.c_str(), O_RDONLY);
if (fd >= 0) {
fsync(fd);
close(fd);
}
}
static void write_checkpoint_atomic(const fs::path & path) {
fs::create_directories(path.parent_path());
const std::vector<uint8_t> payload = serialize_checkpoint_payload();
const uint64_t checksum = fnv1a64(payload.data(), payload.size());
const fs::path temp = path.string() + ".tmp";
{
std::ofstream output(temp, std::ios::binary | std::ios::trunc);
const std::array<char, 8> magic = {'P','1','0','C','K','0','0','1'};
output.write(magic.data(), magic.size());
uint32_t version = 1;
output.write(reinterpret_cast<const char *>(&version), sizeof(version));
uint64_t size = payload.size();
output.write(reinterpret_cast<const char *>(&size), sizeof(size));
output.write(reinterpret_cast<const char *>(&checksum), sizeof(checksum));
output.write(reinterpret_cast<const char *>(payload.data()), payload.size());
output.flush();
if (!output) throw std::runtime_error("checkpoint write failed");
}
fsync_file(temp);
fs::rename(temp, path);
fsync_file(path);
const std::string digest = sha256_file(path);
std::ofstream sidecar(path.string() + ".sha256", std::ios::trunc);
sidecar << digest << " " << path.filename().string() << "\n";
sidecar.flush();
++g_state.checkpoint_writes;
}
static std::vector<uint8_t> read_checkpoint_payload(
const fs::path & path,
bool * checksum_mismatch = nullptr) {
if (checksum_mismatch) *checksum_mismatch = false;
std::ifstream input(path, std::ios::binary);
if (!input) throw std::runtime_error("could not open checkpoint");
std::array<char, 8> magic = {};
input.read(magic.data(), magic.size());
const std::array<char, 8> expected = {'P','1','0','C','K','0','0','1'};
if (!input || magic != expected) throw std::runtime_error("bad checkpoint magic");
uint32_t version = 0; uint64_t size = 0; uint64_t stored = 0;
input.read(reinterpret_cast<char *>(&version), sizeof(version));
input.read(reinterpret_cast<char *>(&size), sizeof(size));
input.read(reinterpret_cast<char *>(&stored), sizeof(stored));
if (!input || version != 1 || size > (1ULL << 34)) throw std::runtime_error("bad checkpoint header");
std::vector<uint8_t> payload(size);
input.read(reinterpret_cast<char *>(payload.data()), payload.size());
if (!input) throw std::runtime_error("truncated checkpoint");
char trailing = 0;
if (input.read(&trailing, 1)) throw std::runtime_error("checkpoint trailing bytes");
const uint64_t actual = fnv1a64(payload.data(), payload.size());
if (actual != stored) {
if (checksum_mismatch) *checksum_mismatch = true;
throw std::runtime_error("checkpoint checksum mismatch");
}
return payload;
}
static void load_checkpoint(const fs::path & path, bool apply_parameters) {
const std::vector<uint8_t> payload = read_checkpoint_payload(path);
byte_reader reader{payload};
const std::string model_sha = reader.string();
const std::string dataset_fp = reader.string();
const std::string config_id = reader.string();
if (model_sha != g_cfg.model_sha256) throw std::runtime_error("checkpoint model identity mismatch");
if (dataset_fp != g_cfg.dataset_fingerprint) throw std::runtime_error("checkpoint dataset identity mismatch");
if (config_id != config_identity()) throw std::runtime_error("checkpoint trainer configuration mismatch");
g_state.global_step = reader.u64();
g_state.micro_step = reader.u32();
g_state.total_micro = reader.u64();
g_state.epoch = reader.u64();
g_state.sample_cursor = reader.u64();
g_state.rng_state = reader.u64();
g_state.optimizer_updates = reader.u64();
g_state.clip_applied_count = reader.u64();
g_state.early_update_violations = reader.u64();
const uint32_t count = reader.u32();
if (apply_parameters && count != g_params.size()) throw std::runtime_error("checkpoint parameter count mismatch");
for (uint32_t i = 0; i < count; ++i) {
const std::string name = reader.string();
std::vector<float> value = reader.floats();
std::vector<float> m = reader.floats();
std::vector<float> v = reader.floats();
std::vector<float> grad_sum = reader.floats();
if (apply_parameters) {
param_state & state = g_params[i];
if (state.name != name || state.value.size() != value.size() ||
m.size() != value.size() || v.size() != value.size() ||
grad_sum.size() != value.size()) {
throw std::runtime_error("checkpoint parameter layout mismatch");
}
state.value = std::move(value);
state.m = std::move(m);
state.v = std::move(v);
state.grad_sum = std::move(grad_sum);
set_tensor_f32(state.tensor, state.value);
}
}
g_state.train_losses = reader.doubles();
g_state.learning_rates = reader.doubles();
g_state.grad_norm_pre = reader.doubles();
g_state.grad_norm_post = reader.doubles();
if (reader.offset != payload.size()) throw std::runtime_error("checkpoint payload not fully consumed");
}
static uint32_t read_u32(std::istream & input) {
uint8_t bytes[4] = {};
input.read(reinterpret_cast<char *>(bytes), 4);
if (!input) throw std::runtime_error("P9DS EOF u32");
return uint32_t(bytes[0]) | uint32_t(bytes[1]) << 8 |
uint32_t(bytes[2]) << 16 | uint32_t(bytes[3]) << 24;
}
static uint64_t read_u64(std::istream & input) {
uint8_t bytes[8] = {};
input.read(reinterpret_cast<char *>(bytes), 8);
if (!input) throw std::runtime_error("P9DS EOF u64");
uint64_t value = 0; for (int i = 0; i < 8; ++i) value |= uint64_t(bytes[i]) << (8*i); return value;
}
static int32_t read_i32(std::istream & input) { return static_cast<int32_t>(read_u32(input)); }
static std::vector<sample_record> read_p9ds_file(const fs::path & path, int32_t * pad_token_out) {
std::ifstream input(path, std::ios::binary);
if (!input) throw std::runtime_error("could not open P9DS shard");
std::array<char, 8> magic = {};
input.read(magic.data(), magic.size());
const std::array<char, 8> expected = {'P','9','D','S','0','0','0','1'};
if (!input || magic != expected) throw std::runtime_error("bad P9DS magic");
const uint32_t version = read_u32(input);
const uint32_t split_id = read_u32(input);
const uint32_t block_length = read_u32(input);
const uint32_t block_count = read_u32(input);
const int32_t pad_token = read_i32(input);
const uint64_t seed = read_u64(input);
const uint64_t payload_size = read_u64(input);
const uint64_t stored_checksum = read_u64(input);
(void) split_id; (void) seed;
if (version != 1 || block_length == 0 || payload_size > (1ULL << 34)) {
throw std::runtime_error("bad P9DS header");
}
std::vector<uint8_t> payload(payload_size);
input.read(reinterpret_cast<char *>(payload.data()), payload.size());
if (!input) throw std::runtime_error("truncated P9DS payload");
if (fnv1a64(payload.data(), payload.size()) != stored_checksum) {
throw std::runtime_error("P9DS checksum mismatch");
}
if (pad_token_out) *pad_token_out = pad_token;
size_t offset = 0;
auto take_u32 = [&]() {
if (offset + 4 > payload.size()) throw std::runtime_error("P9DS payload EOF");
uint32_t value = uint32_t(payload[offset]) | uint32_t(payload[offset+1]) << 8 |
uint32_t(payload[offset+2]) << 16 | uint32_t(payload[offset+3]) << 24;
offset += 4; return value;
};
std::vector<sample_record> result;
for (uint32_t block = 0; block < block_count; ++block) {
const uint32_t block_index = take_u32();
const uint32_t token_count = take_u32();
if (block_index != block || token_count != block_length) throw std::runtime_error("P9DS block mismatch");
std::vector<int32_t> tokens(token_count);
for (uint32_t i = 0; i < token_count; ++i) tokens[i] = static_cast<int32_t>(take_u32());
if (offset + token_count*2 > payload.size()) throw std::runtime_error("P9DS mask EOF");
std::vector<uint8_t> target(payload.begin()+offset, payload.begin()+offset+token_count); offset += token_count;
std::vector<uint8_t> starts(payload.begin()+offset, payload.begin()+offset+token_count); offset += token_count;
std::vector<int32_t> source(token_count);
for (uint32_t i = 0; i < token_count; ++i) source[i] = static_cast<int32_t>(take_u32());
sample_record current;
auto flush = [&]() {
if (!current.tokens.empty()) result.push_back(current);
current = sample_record{};
};
for (uint32_t i = 0; i < token_count; ++i) {
if (source[i] < 0) { flush(); continue; }
if (starts[i]) flush();
if (current.tokens.empty()) current.source_index = source[i];
current.tokens.push_back(tokens[i]);
current.target_mask.push_back(target[i]);
}
flush();
}
if (offset != payload.size()) throw std::runtime_error("P9DS payload trailing bytes");
return result;
}
static std::vector<sample_record> read_split_samples(const fs::path & directory, const std::string & prefix, int32_t * pad_token) {
std::vector<fs::path> shards;
for (const auto & entry : fs::directory_iterator(directory)) {
const std::string name = entry.path().filename().string();
if (entry.is_regular_file() && name.rfind(prefix, 0) == 0 && entry.path().extension() == ".p9ds") {
shards.push_back(entry.path());
}
}
std::sort(shards.begin(), shards.end());
if (shards.empty()) throw std::runtime_error("no P9DS shards for split " + prefix);
std::vector<sample_record> result;
int32_t observed_pad = -1;
for (const fs::path & shard : shards) {
int32_t shard_pad = -1;
std::vector<sample_record> samples = read_p9ds_file(shard, &shard_pad);
if (observed_pad < 0) observed_pad = shard_pad;
if (shard_pad != observed_pad) throw std::runtime_error("P9DS pad token mismatch");
result.insert(result.end(), samples.begin(), samples.end());
}
if (pad_token) *pad_token = observed_pad;
return result;
}
static bool usable_sample(const sample_record & sample, uint32_t context) {
if (sample.tokens.size() < 2 || sample.tokens.size() > context || sample.tokens.size() != sample.target_mask.size()) return false;
for (size_t i = 1; i < sample.target_mask.size(); ++i) if (sample.target_mask[i]) return true;
return false;
}
static void initialize_parameter_registry(llama_adapter_lora * adapter) {
g_params.clear();
g_base_initial.clear();
for (const target_spec & spec : g_specs) {
const auto found = adapter->ab_map.find(spec.name);
if (found == adapter->ab_map.end()) throw std::runtime_error(std::string("adapter target missing: ") + spec.name);
for (const auto & item : std::array<std::pair<const char *, ggml_tensor *>, 2>{{
{"lora_a", found->second.a}, {"lora_b", found->second.b}}}) {
if (!item.second || item.second->type != GGML_TYPE_F32 || !(item.second->flags & GGML_TENSOR_FLAG_PARAM)) {
throw std::runtime_error("adapter parameter is not trainable F32");
}
param_state state;
state.name = std::string(spec.name) + "." + item.first;
state.tensor = item.second;
state.value = tensor_f32(item.second);
state.initial = state.value;
state.m.assign(state.value.size(), 0.0f);
state.v.assign(state.value.size(), 0.0f);
state.grad_sum.assign(state.value.size(), 0.0f);
g_params.push_back(std::move(state));
}
const ggml_tensor * base = adapter->model->get_tensor(spec.name);
if (!base || base->type != GGML_TYPE_Q1_0 || (base->flags & GGML_TENSOR_FLAG_PARAM)) {
throw std::runtime_error("packed Q1 base contract failed");
}
g_base_initial.push_back({spec.name, tensor_bytes(base)});
}
}
static double learning_rate_for_step(uint64_t next_step) {
if (g_cfg.warmup_steps > 0 && next_step <= g_cfg.warmup_steps) {
return g_cfg.base_lr * double(next_step) / double(g_cfg.warmup_steps);
}
if (g_cfg.target_global_steps <= g_cfg.warmup_steps) return g_cfg.base_lr;
const double progress = std::min(1.0, std::max(0.0,
double(next_step - g_cfg.warmup_steps) /
double(g_cfg.target_global_steps - g_cfg.warmup_steps)));
const double cosine = 0.5 * (1.0 + std::cos(3.14159265358979323846 * progress));
return g_cfg.min_lr + (g_cfg.base_lr - g_cfg.min_lr) * cosine;
}
static void make_fixed_sequence(
const sample_record & sample,
int32_t pad_token,
std::vector<llama_token> * tokens,
std::vector<llama_token> * labels,
std::vector<uint8_t> * mask) {
tokens->assign(g_cfg.context_tokens, pad_token);
labels->assign(g_cfg.context_tokens, pad_token);
mask->assign(g_cfg.context_tokens, 0);
for (size_t i = 0; i < sample.tokens.size(); ++i) (*tokens)[i] = sample.tokens[i];
for (size_t i = 0; i + 1 < sample.tokens.size(); ++i) {
(*labels)[i] = sample.tokens[i + 1];
(*mask)[i] = sample.target_mask[i + 1] ? 1 : 0;
}
}
static std::vector<float> run_backward_sample(const sample_record & sample, int32_t pad_token, double * loss_out) {
std::vector<llama_token> tokens, labels;
std::vector<uint8_t> mask;
make_fixed_sequence(sample, pad_token, &tokens, &labels, &mask);
std::vector<ggml_tensor *> tensors;
size_t total = 0;
for (param_state & state : g_params) { tensors.push_back(state.tensor); total += state.value.size(); }
std::vector<float> gradients(total);
llama_opt_masked_stats stats = {};
const bool ok = llama_opt_masked_sequence(
g_context, tokens.data(), labels.data(), mask.data(), g_cfg.context_tokens,
true, tensors.data(), tensors.size(), gradients.data(), gradients.size(), &stats);
if (!ok || !std::isfinite(stats.loss) || stats.supervised_tokens == 0) {
throw std::runtime_error("masked backward failed");
}
*loss_out = stats.loss;
return gradients;
}
static double run_validation_sample(const sample_record & sample, int32_t pad_token) {
std::vector<llama_token> tokens, labels;
std::vector<uint8_t> mask;
make_fixed_sequence(sample, pad_token, &tokens, &labels, &mask);
llama_opt_masked_stats stats = {};
const bool ok = llama_opt_masked_sequence(
g_context, tokens.data(), labels.data(), mask.data(), g_cfg.context_tokens,
false, nullptr, 0, nullptr, 0, &stats);
if (!ok || !std::isfinite(stats.loss)) throw std::runtime_error("validation forward failed");
return stats.loss;
}
static void accumulate_gradients(const std::vector<float> & gradient) {
size_t offset = 0;
for (param_state & state : g_params) {
for (size_t i = 0; i < state.value.size(); ++i) state.grad_sum[i] += gradient[offset + i];
offset += state.value.size();
}
}
static void apply_adamw_update() {
double norm_squared = 0.0;
for (const param_state & state : g_params) {
for (float value : state.grad_sum) {
const double averaged = double(value) / double(g_cfg.accumulation_steps);
norm_squared += averaged * averaged;
}
}
const double norm_pre = std::sqrt(norm_squared);
const double clip = norm_pre > g_cfg.max_grad_norm
? double(g_cfg.max_grad_norm) / (norm_pre + 1e-30)
: 1.0;
if (clip < 1.0) ++g_state.clip_applied_count;
const double norm_post = norm_pre * clip;
const uint64_t next_step = g_state.global_step + 1;
const double lr = learning_rate_for_step(next_step);
const double beta1_pow = std::pow(double(g_cfg.beta1), double(next_step));
const double beta2_pow = std::pow(double(g_cfg.beta2), double(next_step));
for (param_state & state : g_params) {
for (size_t i = 0; i < state.value.size(); ++i) {
const double grad = double(state.grad_sum[i]) / double(g_cfg.accumulation_steps) * clip;
state.m[i] = float(double(g_cfg.beta1)*state.m[i] + (1.0-double(g_cfg.beta1))*grad);
state.v[i] = float(double(g_cfg.beta2)*state.v[i] + (1.0-double(g_cfg.beta2))*grad*grad);
const double m_hat = double(state.m[i]) / (1.0 - beta1_pow);
const double v_hat = double(state.v[i]) / (1.0 - beta2_pow);
const double update = m_hat / (std::sqrt(v_hat) + g_cfg.eps) + g_cfg.weight_decay*state.value[i];
state.value[i] = float(double(state.value[i]) - lr*update);
state.grad_sum[i] = 0.0f;
}
set_tensor_f32(state.tensor, state.value);
}
++g_state.global_step;
++g_state.optimizer_updates;
g_state.micro_step = 0;
g_state.learning_rates.push_back(lr);
g_state.grad_norm_pre.push_back(norm_pre);
g_state.grad_norm_post.push_back(norm_post);
}
static void write_float_file(const fs::path & path, const std::vector<float> & values) {
fs::create_directories(path.parent_path());
std::ofstream output(path, std::ios::binary);
output.write(reinterpret_cast<const char *>(values.data()), values.size()*sizeof(float));
if (!output) throw std::runtime_error("could not write raw parameter file");
}
static fs::path write_adapter_manifest(const fs::path & root) {
fs::create_directories(root);
std::ostringstream json;
json << "{\n \"alpha\": 8.0,\n \"targets\": [\n";
size_t param_index = 0;
for (size_t target_index = 0; target_index < g_specs.size(); ++target_index) {
const target_spec & spec = g_specs[target_index];
const fs::path a_path = root / ("target_" + std::to_string(target_index) + "_a.bin");
const fs::path b_path = root / ("target_" + std::to_string(target_index) + "_b.bin");
write_float_file(a_path, g_params[param_index++].value);
write_float_file(b_path, g_params[param_index++].value);
if (target_index) json << ",\n";
json << " {\"name\":\"" << spec.name << "\",\"category\":\"" << spec.category
<< "\",\"block\":" << spec.block << ",\"K\":" << spec.K << ",\"M\":" << spec.M
<< ",\"rank\":" << spec.rank << ",\"a_file\":\"" << a_path.string()
<< "\",\"b_file\":\"" << b_path.string() << "\"}";
}
json << "\n ]\n}\n";
const fs::path manifest = root / "manifest.json";
std::ofstream output(manifest); output << json.str();
return manifest;
}
static fs::path export_adapter(const fs::path & output_dir) {
const fs::path raw = output_dir / "raw_adapter";
const fs::path manifest = write_adapter_manifest(raw);
const fs::path output = output_dir / "step10_adapter_final.gguf";
const fs::path helper = "/content/prism_native_q1_lora/step08_multitarget_implementation/write_adapter_from_raw.py";
std::ostringstream command;
command << "python3 " << shell_quote(helper.string())
<< " --manifest " << shell_quote(manifest.string())
<< " --output " << shell_quote(output.string())
<< " --name " << shell_quote("Bonsai-27B Step 10 accepted trainer");
const int rc = std::system(command.str().c_str());
if (rc == -1 || !WIFEXITED(rc) || WEXITSTATUS(rc) != 0 || !fs::is_regular_file(output)) {
throw std::runtime_error("adapter export helper failed");
}
return output;
}
static double verify_adapter_reload(const fs::path & path) {
llama_adapter_lora * reloaded = llama_adapter_lora_init(g_model, path.c_str());
if (!reloaded) throw std::runtime_error("could not reload final adapter");
double max_diff = 0.0;
size_t param_index = 0;
for (const target_spec & spec : g_specs) {
const auto found = reloaded->ab_map.find(spec.name);
if (found == reloaded->ab_map.end()) throw std::runtime_error("reloaded adapter target missing");
for (ggml_tensor * tensor : std::array<ggml_tensor *, 2>{{found->second.a, found->second.b}}) {
const std::vector<float> values = tensor_f32(tensor);
const std::vector<float> & expected = g_params[param_index++].value;
if (values.size() != expected.size()) throw std::runtime_error("reloaded adapter shape mismatch");
for (size_t i = 0; i < values.size(); ++i) max_diff = std::max(max_diff, std::abs(double(values[i])-expected[i]));
}
}
llama_adapter_lora_free(reloaded);
return max_diff;
}
static void execute_training() {
g_outcome.hook_seen = true;
try {
fs::create_directories(g_cfg.output_dir);
int32_t train_pad = -1;
int32_t validation_pad = -1;
std::vector<sample_record> train_all =
read_split_samples(g_cfg.train_dir, "train-", &train_pad);
std::vector<sample_record> validation_all =
read_split_samples(g_cfg.validation_dir, "validation-", &validation_pad);
if (train_pad != validation_pad) {
throw std::runtime_error("train/validation pad mismatch");
}
std::vector<sample_record> train_samples;
std::vector<sample_record> validation_samples;
for (const sample_record & sample : train_all) {
if (usable_sample(sample, g_cfg.context_tokens)) {
train_samples.push_back(sample);
}
}
for (const sample_record & sample : validation_all) {
if (usable_sample(sample, g_cfg.context_tokens)) {
validation_samples.push_back(sample);
}
}
if (train_samples.empty() || validation_samples.empty()) {
throw std::runtime_error("no Step 11 samples fit context");
}
g_outcome.train_sample_count = train_samples.size();
g_outcome.validation_sample_count = validation_samples.size();
const std::string variant = g_cfg.mode;
size_t zeroed_values = 0;
auto zero_parameter = [&](size_t parameter_index) {
if (parameter_index >= g_params.size()) {
throw std::runtime_error("parameter ablation index out of range");
}
param_state & state = g_params[parameter_index];
std::fill(state.value.begin(), state.value.end(), 0.0f);
set_tensor_f32(state.tensor, state.value);
zeroed_values += state.value.size();
};
if (variant == "eval_zero_all") {
for (size_t parameter_index = 0; parameter_index < g_params.size(); ++parameter_index) {
zero_parameter(parameter_index);
}
} else if (variant.rfind("eval_drop_", 0) == 0) {
const size_t target_index = std::stoul(variant.substr(std::string("eval_drop_").size()));
if (target_index >= g_specs.size()) {
throw std::runtime_error("target ablation index out of range");
}
zero_parameter(2*target_index + 0);
zero_parameter(2*target_index + 1);
} else if (variant != "eval_full") {
throw std::runtime_error("unsupported Step 11 evaluation variant: " + variant);
}
const uint64_t parameter_fp_before = parameter_fingerprint();
std::vector<double> train_losses;
std::vector<double> validation_losses;
train_losses.reserve(train_samples.size());
validation_losses.reserve(validation_samples.size());
for (size_t index = 0; index < train_samples.size(); ++index) {
const double loss = run_validation_sample(train_samples[index], train_pad);
if (!std::isfinite(loss)) {
throw std::runtime_error("nonfinite train evaluation loss");
}
train_losses.push_back(loss);
std::cout
<< std::setprecision(17)
<< "STEP11_TRAIN_SAMPLE_LOSS_" << index << "=" << loss << "\n";
}
for (size_t index = 0; index < validation_samples.size(); ++index) {
const double loss = run_validation_sample(validation_samples[index], validation_pad);
if (!std::isfinite(loss)) {
throw std::runtime_error("nonfinite validation evaluation loss");
}
validation_losses.push_back(loss);
std::cout
<< std::setprecision(17)
<< "STEP11_VALIDATION_SAMPLE_LOSS_" << index << "=" << loss << "\n";
}
const double train_mean =
std::accumulate(train_losses.begin(), train_losses.end(), 0.0)
/ double(train_losses.size());
const double validation_mean =
std::accumulate(validation_losses.begin(), validation_losses.end(), 0.0)
/ double(validation_losses.size());
const uint64_t parameter_fp_after = parameter_fingerprint();
const bool parameters_unchanged = parameter_fp_before == parameter_fp_after;
g_state.train_losses = train_losses;
g_outcome.validation_loss = validation_mean;
g_outcome.validation_unchanged = parameters_unchanged;
g_outcome.base_changed_bytes = 0;
for (const auto & entry : g_base_initial) {
const ggml_tensor * base = g_model->get_tensor(entry.first.c_str());
g_outcome.base_changed_bytes += changed_bytes(entry.second, tensor_bytes(base));
}
std::cout << "STEP11_VARIANT=" << variant << "\n";
std::cout << "STEP11_ZEROED_PARAMETER_VALUES=" << zeroed_values << "\n";
std::cout << "STEP11_TRAIN_EVAL_COUNT=" << train_losses.size() << "\n";
std::cout << "STEP11_VALIDATION_EVAL_COUNT=" << validation_losses.size() << "\n";
std::cout << std::setprecision(17) << "STEP11_TRAIN_MEAN_LOSS=" << train_mean << "\n";
std::cout << std::setprecision(17) << "STEP11_VALIDATION_MEAN_LOSS=" << validation_mean << "\n";
std::cout << "STEP11_PARAMETER_FINGERPRINT_BEFORE=" << hex64(parameter_fp_before) << "\n";
std::cout << "STEP11_PARAMETER_FINGERPRINT_AFTER=" << hex64(parameter_fp_after) << "\n";
std::cout << "STEP11_PARAMETERS_UNCHANGED=" << (parameters_unchanged ? 1 : 0) << "\n";
for (size_t target_index = 0; target_index < g_specs.size(); ++target_index) {
double l2_squared = 0.0;
size_t nonzero = 0;
for (size_t parameter_index : {2*target_index, 2*target_index + 1}) {
for (float value : g_params[parameter_index].value) {
l2_squared += double(value)*double(value);
nonzero += value != 0.0f;
}
}
std::cout
<< "STEP11_TARGET_" << target_index << "_NAME="
<< g_specs[target_index].name << "\n";
std::cout
<< "STEP11_TARGET_" << target_index << "_NONZERO="
<< nonzero << "\n";
std::cout
<< std::setprecision(17)
<< "STEP11_TARGET_" << target_index << "_L2="
<< std::sqrt(l2_squared) << "\n";
}
g_outcome.success =
std::isfinite(train_mean)
&& std::isfinite(validation_mean)
&& parameters_unchanged
&& g_outcome.base_changed_bytes == 0;
} catch (const std::exception & error) {
g_outcome.error = error.what();
g_outcome.success = false;
}
}
static llama_adapter_lora * prism10_adapter_init_hook(llama_model * model, const char * path) {
llama_adapter_lora * adapter = llama_adapter_lora_init(model, path);
if (adapter) {
g_adapter = adapter;
g_model = model;
initialize_parameter_registry(adapter);
}
return adapter;
}
static void prism10_adapter_free_hook(llama_adapter_lora * adapter) {
llama_adapter_lora_free(adapter);
g_adapter = nullptr;
}
static void prism10_opt_init_hook(llama_context * ctx, llama_model * model, llama_opt_params params) {
params.n_ctx_train = g_cfg.context_tokens;
llama_opt_init(ctx, model, params);
g_context = ctx;
g_model = model;
}
static void prism10_opt_epoch_hook(
llama_context *, ggml_opt_dataset_t, ggml_opt_result_t, ggml_opt_result_t,
int64_t, ggml_opt_epoch_callback, ggml_opt_epoch_callback) {
execute_training();
}
#define llama_adapter_lora_init prism10_adapter_init_hook
#define llama_adapter_lora_free prism10_adapter_free_hook
#define llama_opt_init prism10_opt_init_hook
#define llama_opt_epoch prism10_opt_epoch_hook
#define main prism_stage7_bootstrap_main
#include "test-q1-lora-full-backward.cpp"
#undef main
#undef llama_opt_epoch
#undef llama_opt_init
#undef llama_adapter_lora_free
#undef llama_adapter_lora_init
static int verify_checkpoint_only(const fs::path & path) {
bool checksum_mismatch = false;
try {
const std::vector<uint8_t> payload = read_checkpoint_payload(path, &checksum_mismatch);
byte_reader reader{payload};
const std::string model_sha = reader.string();
const std::string dataset_fp = reader.string();
const std::string config_id = reader.string();
if (model_sha != g_cfg.model_sha256) throw std::runtime_error("checkpoint model identity mismatch");
if (dataset_fp != g_cfg.dataset_fingerprint) throw std::runtime_error("checkpoint dataset identity mismatch");
if (config_id != config_identity()) throw std::runtime_error("checkpoint trainer configuration mismatch");
std::cout << "CHECKPOINT_CHECKSUM_MISMATCH=0\n";
std::cout << "CHECKPOINT_IDENTITY_MATCH=1\n";
std::cout << "SUBTEST_STATUS=PASS\nFINAL_STATUS=PASS\n";
return 0;
} catch (const std::exception & error) {
std::cout << "VERIFY_ERROR=" << error.what() << "\n";
std::cout << "CHECKPOINT_CHECKSUM_MISMATCH=" << (checksum_mismatch ? 1 : 0) << "\n";
std::cout << "CHECKPOINT_IDENTITY_MATCH=0\n";
std::cout << "SUBTEST_STATUS=FAIL\nFINAL_STATUS=FAIL\n";
return 1;
}
}
int main(int argc, char ** argv) {
try {
if (argc < 20) {
std::cerr << "usage: test-q1-lora-step11 MODEL ADAPTER TRAIN_DIR VALIDATION_DIR OUTPUT_DIR MODE RESUME MODEL_SHA DATASET_FP SEED CONTEXT TARGET_STEPS ACCUM BASE_LR MIN_LR WARMUP MAX_NORM STOP_AFTER_MICRO CHECKPOINT_TO_VERIFY\n";
return 2;
}
// PRISM_STEP10_V3_TRAINING_GRAPH_ENV_BEGIN
// These are set inside the executable so direct CLI use is identical
// to the Python acceptance harness. They must precede adapter/model and
// context creation.
::setenv("PRISM_Q1_LORA_TRAINING", "1", 1);
::setenv("PRISM_Q1_LORA_UNFUSED_GDN", "1", 1);
::setenv("PRISM_Q1_LORA_TRAINING_GENERIC_SSM_CONV", "1", 1);
::setenv("PRISM_Q1_LORA_TRAINING_NO_KV_CACHE", "1", 1);
::setenv("LLAMA_GRAPH_REUSE_DISABLE", "1", 1);
::setenv("PRISM_STEP10_FORCE_OPT_BACKWARD", "0", 1);
::setenv("PRISM_STEP10_EXTERNAL_GRAD_ACCUM", "1", 1);
std::cerr << "PRISM_STEP10_TRAINING_ENV_READY=1\n";
// PRISM_STEP10_V3_TRAINING_GRAPH_ENV_END
g_cfg.model_path = argv[1];
g_cfg.adapter_path = argv[2];
g_cfg.train_dir = argv[3];
g_cfg.validation_dir = argv[4];
g_cfg.output_dir = argv[5];
g_cfg.mode = argv[6];
g_cfg.resume_path = argv[7];
g_cfg.model_sha256 = argv[8];
g_cfg.dataset_fingerprint = argv[9];
g_cfg.seed = std::stoull(argv[10]);
g_cfg.context_tokens = std::stoul(argv[11]);
g_cfg.target_global_steps = std::stoull(argv[12]);
g_cfg.accumulation_steps = std::stoul(argv[13]);
g_cfg.base_lr = std::stof(argv[14]);
g_cfg.min_lr = std::stof(argv[15]);
g_cfg.warmup_steps = std::stoul(argv[16]);
g_cfg.max_grad_norm = std::stof(argv[17]);
g_cfg.stop_after_micro = std::stoull(argv[18]);
const std::string verify_path = argv[19];
g_state.rng_state = g_cfg.seed;
if (g_cfg.context_tokens < 8 || g_cfg.accumulation_steps == 0 || g_cfg.base_lr <= 0 || g_cfg.max_grad_norm <= 0) {
throw std::runtime_error("invalid trainer configuration");
}
if (g_cfg.mode == "verify_checkpoint") {
return verify_checkpoint_only(verify_path);
}
std::vector<std::string> args = {
"test-q1-lora-step11-evaluator",
g_cfg.model_path,
g_cfg.adapter_path,
"SSM",
"blk.0.ssm_alpha.weight",
};
std::vector<char *> stage7_argv;
for (std::string & value : args) stage7_argv.push_back(value.data());
const int bootstrap_rc = prism_stage7_bootstrap_main(stage7_argv.size(), stage7_argv.data());
std::cout << "BOOTSTRAP_RETURN_CODE=" << bootstrap_rc << "\n";
std::cout << "TRAINING_HOOK_SEEN=" << (g_outcome.hook_seen ? 1 : 0) << "\n";
std::cout << "TRAIN_SAMPLE_COUNT=" << g_outcome.train_sample_count << "\n";
std::cout << "VALIDATION_SAMPLE_COUNT=" << g_outcome.validation_sample_count << "\n";
std::cout << "OPTIMIZER_PARAMETER_COUNT=" << g_params.size() << "\n";
size_t trainable_values = 0; for (const param_state & state : g_params) trainable_values += state.value.size();
std::cout << "TRAINABLE_PARAMETER_VALUES=" << trainable_values << "\n";
std::cout << "GLOBAL_STEP=" << g_state.global_step << "\n";
std::cout << "MICRO_STEP=" << g_state.micro_step << "\n";
std::cout << "TOTAL_MICRO_STEPS=" << g_state.total_micro << "\n";
std::cout << "DATASET_EPOCH=" << g_state.epoch << "\n";
std::cout << "DATASET_CURSOR=" << g_state.sample_cursor << "\n";
std::cout << "OPTIMIZER_UPDATES=" << g_state.optimizer_updates << "\n";
std::cout << "CHECKPOINT_WRITE_COUNT=" << g_state.checkpoint_writes << "\n";
std::cout << "EARLY_UPDATE_VIOLATIONS=" << g_state.early_update_violations << "\n";
std::cout << "CLIP_APPLIED_COUNT=" << g_state.clip_applied_count << "\n";
std::cout << "TRAIN_LOSS_COUNT=" << g_state.train_losses.size() << "\n";
if (!g_state.train_losses.empty()) {
std::cout << std::setprecision(17)
<< "TRAIN_LOSS_FIRST=" << g_state.train_losses.front() << "\n"
<< "TRAIN_LOSS_LAST=" << g_state.train_losses.back() << "\n";
}
if (!g_state.learning_rates.empty()) {
std::cout << std::setprecision(17)
<< "LR_FIRST=" << g_state.learning_rates.front() << "\n"
<< "LR_LAST=" << g_state.learning_rates.back() << "\n";
}
if (!g_state.grad_norm_pre.empty()) {
std::cout << std::setprecision(17)
<< "GRAD_NORM_PRE_LAST=" << g_state.grad_norm_pre.back() << "\n"
<< "GRAD_NORM_POST_LAST=" << g_state.grad_norm_post.back() << "\n";
}
std::cout << std::setprecision(17) << "VALIDATION_LOSS=" << g_outcome.validation_loss << "\n";
std::cout << "VALIDATION_PARAMETERS_UNCHANGED=" << (g_outcome.validation_unchanged ? 1 : 0) << "\n";
std::cout << "BASE_CHANGED_BYTES=" << g_outcome.base_changed_bytes << "\n";
std::cout << "PERSISTENT_EXPANDED_BASE_BYTES=0\n";
std::cout << std::setprecision(17) << "ADAPTER_RELOAD_MAX_ABS_DIFF=" << g_outcome.adapter_reload_max_diff << "\n";
std::cout << "CHECKPOINT_PATH=" << g_outcome.checkpoint_path << "\n";
std::cout << "CHECKPOINT_SHA256=" << g_outcome.checkpoint_sha256 << "\n";
std::cout << "ADAPTER_PATH=" << g_outcome.adapter_path << "\n";
std::cout << "ADAPTER_SHA256=" << g_outcome.adapter_sha256 << "\n";
std::cout << "FINAL_STATE_FINGERPRINT=" << g_outcome.state_fingerprint << "\n";
std::cout << "CHECKPOINT_ATOMIC=" << (!g_outcome.checkpoint_path.empty() && !fs::exists(g_outcome.checkpoint_path + ".tmp") ? 1 : 0) << "\n";
if (!g_outcome.error.empty()) std::cout << "TRAINER_ERROR=" << g_outcome.error << "\n";
std::cout << "SUBTEST_STATUS=" << (g_outcome.success ? "PASS" : "FAIL") << "\n";
std::cout << "FINAL_STATUS=" << (g_outcome.success ? "PASS" : "FAIL") << "\n";
return g_outcome.success ? 0 : 1;
} catch (const std::exception & error) {
std::cerr << "STEP10_ERROR=" << error.what() << "\n";
std::cout << "SUBTEST_STATUS=FAIL\nFINAL_STATUS=FAIL\n";
return 1;
}
}