Foundry/ternarize/ternarize.c

1313 lines
48 KiB
C

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#include "ternarize.h"
#include "ft_layer_ops.h"
#include "../safetensors/safetensors.h"
#include "../quantization/quantization.h"
#include "../convert/convert.h"
#include "../tensor/tensor.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/stat.h>
#include <errno.h>
/* ============================================================= */
/* PRNG (xorshift64, same as fq_ module) */
/* ============================================================= */
static uint64_t ft_xorshift64(uint64_t *state) {
uint64_t x = *state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
*state = x;
return x;
}
static float ft_rand_uniform(uint64_t *state) {
return (ft_xorshift64(state) >> 11) * (1.0 / 9007199254740992.0);
}
/* ============================================================= */
/* Minimal JSON helpers (reused from convert.c pattern) */
/* ============================================================= */
static char* read_file_str(const char *path) {
FILE *f = fopen(path, "r");
if (!f) return NULL;
fseek(f, 0, SEEK_END);
long size = ftell(f);
fseek(f, 0, SEEK_SET);
char *buf = malloc(size + 1);
size_t rd = fread(buf, 1, size, f);
buf[rd] = '\0';
fclose(f);
return buf;
}
static int json_get_int(const char *json, const char *key, int def) {
char search[128];
snprintf(search, sizeof(search), "\"%s\"", key);
const char *p = strstr(json, search);
if (!p) return def;
p += strlen(search);
while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
return (int)strtol(p, NULL, 10);
}
static float json_get_float(const char *json, const char *key, float def) {
char search[128];
snprintf(search, sizeof(search), "\"%s\"", key);
const char *p = strstr(json, search);
if (!p) return def;
p += strlen(search);
while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
return strtof(p, NULL);
}
static char* json_get_string(const char *json, const char *key) {
char search[128];
snprintf(search, sizeof(search), "\"%s\"", key);
const char *p = strstr(json, search);
if (!p) return NULL;
p += strlen(search);
while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
if (*p != '"') return NULL;
p++;
const char *end = p;
while (*end && *end != '"') {
if (*end == '\\') end++;
if (*end) end++;
}
size_t len = end - p;
char *r = malloc(len + 1);
memcpy(r, p, len);
r[len] = '\0';
return r;
}
static int json_get_bool(const char *json, const char *key, int def) {
char search[128];
snprintf(search, sizeof(search), "\"%s\"", key);
const char *p = strstr(json, search);
if (!p) return def;
p += strlen(search);
while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
if (*p == 't') return 1;
if (*p == 'f') return 0;
return def;
}
/* ============================================================= */
/* BF16/F16 conversion (from convert.c) */
/* ============================================================= */
static void bf16_to_f32(const uint16_t *src, float *dst, int64_t n) {
for (int64_t i = 0; i < n; i++) {
uint32_t bits = (uint32_t)src[i] << 16;
memcpy(&dst[i], &bits, 4);
}
}
static void f16_to_f32(const uint16_t *src, float *dst, int64_t n) {
for (int64_t i = 0; i < n; i++) {
uint16_t h = src[i];
uint32_t sign = (uint32_t)(h >> 15) << 31;
uint32_t exp = (h >> 10) & 0x1F;
uint32_t mant = h & 0x3FF;
uint32_t bits;
if (exp == 0) {
if (mant == 0) bits = sign;
else {
exp = 1;
while (!(mant & 0x400)) { mant <<= 1; exp--; }
mant &= 0x3FF;
bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
}
} else if (exp == 31) {
bits = sign | 0x7F800000 | ((uint32_t)mant << 13);
} else {
bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
}
memcpy(&dst[i], &bits, 4);
}
}
static float* tensor_to_f32(const void *data, fs_dtype dtype, int64_t numel) {
float *f32 = malloc(numel * sizeof(float));
if (!f32) return NULL;
switch (dtype) {
case FS_DTYPE_BF16: bf16_to_f32((const uint16_t *)data, f32, numel); break;
case FS_DTYPE_F16: f16_to_f32((const uint16_t *)data, f32, numel); break;
case FS_DTYPE_F32: memcpy(f32, data, numel * sizeof(float)); break;
default: memset(f32, 0, numel * sizeof(float)); break;
}
return f32;
}
/* ============================================================= */
/* Utility */
/* ============================================================= */
static void mkdirp(const char *path) {
char tmp[1024];
snprintf(tmp, sizeof(tmp), "%s", path);
for (char *p = tmp + 1; *p; p++) {
if (*p == '/') {
*p = '\0';
mkdir(tmp, 0755);
*p = '/';
}
}
mkdir(tmp, 0755);
}
const char* ft_arch_name(ft_arch arch) {
switch (arch) {
case FT_ARCH_LLAMA: return "Llama";
case FT_ARCH_QWEN2: return "Qwen2";
case FT_ARCH_QWEN3: return "Qwen3";
case FT_ARCH_PHI3: return "Phi-3";
case FT_ARCH_GEMMA2: return "Gemma2";
case FT_ARCH_RWKVX: return "RWKV-X";
case FT_ARCH_UNKNOWN: return "Unknown";
}
return "Unknown";
}
static int cmp_float_desc(const void *a, const void *b) {
float fa = *(const float *)a;
float fb = *(const float *)b;
if (fa > fb) return -1;
if (fa < fb) return 1;
return 0;
}
/* ============================================================= */
/* Config */
/* ============================================================= */
ft_config ft_config_defaults(void) {
ft_config cfg;
memset(&cfg, 0, sizeof(cfg));
cfg.arch = FT_ARCH_UNKNOWN;
cfg.group_size = 128;
cfg.k_frac_start = 0.00008f;
cfg.k_frac_end = 0.000005f;
cfg.snr_weight = 0.3f;
cfg.hess_power = 0.15f;
cfg.momentum_bonus = 0.2f;
cfg.boltz_frac_start = 0.17f;
cfg.boltz_frac_end = 0.0f;
cfg.temporal_weight = 1.8f;
cfg.phase1_iters = 60;
cfg.phase2_iters = 140;
cfg.batch_size = 128;
return cfg;
}
bool ft_config_from_json(ft_config *cfg, const char *path) {
char *json = read_file_str(path);
if (!json) {
fprintf(stderr, "ERROR: cannot read %s\n", path);
return false;
}
char *mt = json_get_string(json, "model_type");
if (mt) {
if (strstr(mt, "qwen3")) cfg->arch = FT_ARCH_QWEN3;
else if (strstr(mt, "qwen2") || strstr(mt, "qwen")) cfg->arch = FT_ARCH_QWEN2;
else if (strstr(mt, "llama") || strstr(mt, "mistral"))cfg->arch = FT_ARCH_LLAMA;
else if (strstr(mt, "phi3") || strstr(mt, "phi")) cfg->arch = FT_ARCH_PHI3;
else if (strstr(mt, "gemma2") || strstr(mt, "gemma")) cfg->arch = FT_ARCH_GEMMA2;
else if (strstr(mt, "rwkv")) cfg->arch = FT_ARCH_RWKVX;
else cfg->arch = FT_ARCH_UNKNOWN;
free(mt);
}
cfg->n_layers = json_get_int(json, "num_hidden_layers", 28);
cfg->n_embd = json_get_int(json, "hidden_size", 1024);
cfg->n_heads = json_get_int(json, "num_attention_heads", 16);
cfg->n_kv_heads = json_get_int(json, "num_key_value_heads", cfg->n_heads);
cfg->ffn_dim = json_get_int(json, "intermediate_size", 2816);
cfg->vocab_size = json_get_int(json, "vocab_size", 32000);
cfg->head_dim = cfg->n_embd / cfg->n_heads;
cfg->rope_theta = json_get_float(json, "rope_theta", 10000.0f);
cfg->rms_eps = json_get_float(json, "rms_norm_eps", 1e-6f);
cfg->tie_word_embeddings = json_get_bool(json, "tie_word_embeddings", 1);
switch (cfg->arch) {
case FT_ARCH_QWEN3: cfg->has_qk_norm = 1; cfg->has_qkv_bias = 0; break;
case FT_ARCH_QWEN2:
case FT_ARCH_PHI3: cfg->has_qk_norm = 0; cfg->has_qkv_bias = 1; break;
case FT_ARCH_GEMMA2: cfg->has_qk_norm = 1; cfg->has_qkv_bias = 0; break;
default: cfg->has_qk_norm = 0; cfg->has_qkv_bias = 0; break;
}
free(json);
return true;
}
/* ============================================================= */
/* Layer State I/O */
/* ============================================================= */
static int64_t compute_layer_numel(const ft_config *cfg) {
int q_dim = cfg->n_heads * cfg->head_dim;
int kv_dim = cfg->n_kv_heads * cfg->head_dim;
int64_t n = 0;
n += (int64_t)q_dim * cfg->n_embd; /* q_proj */
n += (int64_t)kv_dim * cfg->n_embd; /* k_proj */
n += (int64_t)kv_dim * cfg->n_embd; /* v_proj */
n += (int64_t)cfg->n_embd * q_dim; /* o_proj */
n += (int64_t)cfg->ffn_dim * cfg->n_embd; /* gate_proj */
n += (int64_t)cfg->ffn_dim * cfg->n_embd; /* up_proj */
n += (int64_t)cfg->n_embd * cfg->ffn_dim; /* down_proj */
return n;
}
static void register_layer_weights(ft_layer_state *state, const ft_config *cfg) {
int q_dim = cfg->n_heads * cfg->head_dim;
int kv_dim = cfg->n_kv_heads * cfg->head_dim;
/* 7 weight matrices per LLaMA-family layer */
state->n_weights = 7;
state->weights = calloc(state->n_weights, sizeof(ft_weight_entry));
int64_t off = 0;
struct { const char *name; int rows; int cols; } defs[] = {
{"self_attn.q_proj.weight", q_dim, cfg->n_embd},
{"self_attn.k_proj.weight", kv_dim, cfg->n_embd},
{"self_attn.v_proj.weight", kv_dim, cfg->n_embd},
{"self_attn.o_proj.weight", cfg->n_embd, q_dim},
{"mlp.gate_proj.weight", cfg->ffn_dim, cfg->n_embd},
{"mlp.up_proj.weight", cfg->ffn_dim, cfg->n_embd},
{"mlp.down_proj.weight", cfg->n_embd, cfg->ffn_dim},
};
for (int i = 0; i < 7; i++) {
snprintf(state->weights[i].name, sizeof(state->weights[i].name), "%s", defs[i].name);
state->weights[i].rows = defs[i].rows;
state->weights[i].cols = defs[i].cols;
state->weights[i].offset = off;
state->weights[i].numel = (int64_t)defs[i].rows * defs[i].cols;
off += state->weights[i].numel;
}
}
static ft_layer_state* ft_layer_alloc(const ft_config *cfg, int layer_idx) {
ft_layer_state *s = calloc(1, sizeof(ft_layer_state));
if (!s) return NULL;
s->layer_idx = layer_idx;
s->numel_total = compute_layer_numel(cfg);
s->group_size = cfg->group_size;
s->n_groups = (int)((s->numel_total + cfg->group_size - 1) / cfg->group_size);
/* Bit arrays: 8 elements per byte */
int64_t n_bytes = (s->numel_total + 7) / 8;
s->sign = calloc(n_bytes, 1);
s->mask = calloc(n_bytes, 1);
s->momentum = calloc(n_bytes, 1);
s->scales = calloc(s->n_groups, sizeof(float));
s->prev_scores = calloc(s->numel_total, sizeof(float));
/* Norms */
s->norm_attn = calloc(cfg->n_embd, sizeof(float));
s->norm_ffn = calloc(cfg->n_embd, sizeof(float));
register_layer_weights(s, cfg);
s->best_loss = 1e30f;
return s;
}
void ft_layer_free(ft_layer_state *state) {
if (!state) return;
free(state->sign);
free(state->mask);
free(state->momentum);
free(state->scales);
free(state->prev_scores);
free(state->norm_attn);
free(state->norm_ffn);
free(state->weights);
free(state);
}
static void layer_path(char *buf, size_t bufsz, const ft_config *cfg, int layer_idx, const char *file) {
snprintf(buf, bufsz, "%s/layer_%03d/%s", cfg->workdir, layer_idx, file);
}
bool ft_layer_save(const ft_layer_state *s, const ft_config *cfg) {
char dir[1024], path[1024];
snprintf(dir, sizeof(dir), "%s/layer_%03d", cfg->workdir, s->layer_idx);
mkdirp(dir);
int64_t n_bytes = (s->numel_total + 7) / 8;
/* ternary.bin: sign + mask + momentum (3 bit arrays) */
layer_path(path, sizeof(path), cfg, s->layer_idx, "ternary.bin");
FILE *f = fopen(path, "wb");
if (!f) return false;
/* Header: numel, iteration, phase, total_flips, best_loss */
fwrite(&s->numel_total, sizeof(int64_t), 1, f);
fwrite(&s->iteration, sizeof(int), 1, f);
fwrite(&s->phase, sizeof(int), 1, f);
fwrite(&s->total_flips, sizeof(int64_t), 1, f);
fwrite(&s->best_loss, sizeof(float), 1, f);
fwrite(s->sign, 1, n_bytes, f);
fwrite(s->mask, 1, n_bytes, f);
fwrite(s->momentum, 1, n_bytes, f);
/* prev_scores for temporal */
fwrite(s->prev_scores, sizeof(float), s->numel_total, f);
fclose(f);
/* scales.bin */
layer_path(path, sizeof(path), cfg, s->layer_idx, "scales.bin");
f = fopen(path, "wb");
if (!f) return false;
fwrite(&s->n_groups, sizeof(int), 1, f);
fwrite(&s->group_size, sizeof(int), 1, f);
fwrite(s->scales, sizeof(float), s->n_groups, f);
fclose(f);
/* norms.bin */
layer_path(path, sizeof(path), cfg, s->layer_idx, "norms.bin");
f = fopen(path, "wb");
if (!f) return false;
fwrite(&cfg->n_embd, sizeof(int), 1, f);
fwrite(s->norm_attn, sizeof(float), cfg->n_embd, f);
fwrite(s->norm_ffn, sizeof(float), cfg->n_embd, f);
fclose(f);
return true;
}
ft_layer_state* ft_layer_load(const ft_config *cfg, int layer_idx) {
ft_layer_state *s = ft_layer_alloc(cfg, layer_idx);
if (!s) return NULL;
char path[1024];
int64_t n_bytes = (s->numel_total + 7) / 8;
/* ternary.bin */
layer_path(path, sizeof(path), cfg, layer_idx, "ternary.bin");
FILE *f = fopen(path, "rb");
if (!f) { ft_layer_free(s); return NULL; }
int64_t stored_numel;
fread(&stored_numel, sizeof(int64_t), 1, f);
if (stored_numel != s->numel_total) {
fprintf(stderr, "ERROR: layer %d numel mismatch: %lld vs %lld\n",
layer_idx, (long long)stored_numel, (long long)s->numel_total);
fclose(f); ft_layer_free(s); return NULL;
}
fread(&s->iteration, sizeof(int), 1, f);
fread(&s->phase, sizeof(int), 1, f);
fread(&s->total_flips, sizeof(int64_t), 1, f);
fread(&s->best_loss, sizeof(float), 1, f);
fread(s->sign, 1, n_bytes, f);
fread(s->mask, 1, n_bytes, f);
fread(s->momentum, 1, n_bytes, f);
fread(s->prev_scores, sizeof(float), s->numel_total, f);
fclose(f);
/* scales.bin */
layer_path(path, sizeof(path), cfg, layer_idx, "scales.bin");
f = fopen(path, "rb");
if (f) {
fread(&s->n_groups, sizeof(int), 1, f);
fread(&s->group_size, sizeof(int), 1, f);
fread(s->scales, sizeof(float), s->n_groups, f);
fclose(f);
}
/* norms.bin */
layer_path(path, sizeof(path), cfg, layer_idx, "norms.bin");
f = fopen(path, "rb");
if (f) {
int n_embd;
fread(&n_embd, sizeof(int), 1, f);
fread(s->norm_attn, sizeof(float), n_embd, f);
fread(s->norm_ffn, sizeof(float), n_embd, f);
fclose(f);
}
return s;
}
/* ============================================================= */
/* Gradient Stats I/O */
/* ============================================================= */
ft_grad_stats* ft_grad_stats_load(const ft_config *cfg, int layer_idx) {
char path[1024];
layer_path(path, sizeof(path), cfg, layer_idx, "grad_stats.bin");
FILE *f = fopen(path, "rb");
if (!f) return NULL;
ft_grad_stats *g = calloc(1, sizeof(ft_grad_stats));
fread(&g->numel, sizeof(int64_t), 1, f);
fread(&g->n_samples, sizeof(int), 1, f);
g->grad_sum = malloc(g->numel * sizeof(float));
g->grad_sq_sum = malloc(g->numel * sizeof(float));
g->hessian_diag = malloc(g->numel * sizeof(float));
fread(g->grad_sum, sizeof(float), g->numel, f);
fread(g->grad_sq_sum, sizeof(float), g->numel, f);
fread(g->hessian_diag, sizeof(float), g->numel, f);
fclose(f);
return g;
}
bool ft_grad_stats_save(const ft_grad_stats *g, const ft_config *cfg, int layer_idx) {
char path[1024];
layer_path(path, sizeof(path), cfg, layer_idx, "grad_stats.bin");
FILE *f = fopen(path, "wb");
if (!f) return false;
fwrite(&g->numel, sizeof(int64_t), 1, f);
fwrite(&g->n_samples, sizeof(int), 1, f);
fwrite(g->grad_sum, sizeof(float), g->numel, f);
fwrite(g->grad_sq_sum, sizeof(float), g->numel, f);
fwrite(g->hessian_diag, sizeof(float), g->numel, f);
fclose(f);
return true;
}
void ft_grad_stats_free(ft_grad_stats *g) {
if (!g) return;
free(g->grad_sum);
free(g->grad_sq_sum);
free(g->hessian_diag);
free(g);
}
/* ============================================================= */
/* Status I/O */
/* ============================================================= */
ft_status ft_status_read(const ft_config *cfg) {
ft_status st;
memset(&st, 0, sizeof(st));
char path[1024];
snprintf(path, sizeof(path), "%s/status.json", cfg->workdir);
char *json = read_file_str(path);
if (!json) return st;
st.n_layers = json_get_int(json, "n_layers", 0);
st.layers_split = json_get_int(json, "layers_split", 0);
st.layers_swept = json_get_int(json, "layers_swept", 0);
st.layers_flipped = json_get_int(json, "layers_flipped", 0);
st.current_sweep_iter = json_get_int(json, "current_sweep_iter", 0);
st.overall_loss = json_get_float(json, "overall_loss", 0.0f);
char *name = json_get_string(json, "model_name");
if (name) {
snprintf(st.model_name, sizeof(st.model_name), "%s", name);
free(name);
}
free(json);
return st;
}
bool ft_status_write(const ft_status *st, const ft_config *cfg) {
char path[1024];
snprintf(path, sizeof(path), "%s/status.json", cfg->workdir);
FILE *f = fopen(path, "w");
if (!f) return false;
fprintf(f, "{\n");
fprintf(f, " \"n_layers\": %d,\n", st->n_layers);
fprintf(f, " \"layers_split\": %d,\n", st->layers_split);
fprintf(f, " \"layers_swept\": %d,\n", st->layers_swept);
fprintf(f, " \"layers_flipped\": %d,\n", st->layers_flipped);
fprintf(f, " \"current_sweep_iter\": %d,\n", st->current_sweep_iter);
fprintf(f, " \"overall_loss\": %.6f,\n", st->overall_loss);
fprintf(f, " \"model_name\": \"%s\"\n", st->model_name);
fprintf(f, "}\n");
fclose(f);
return true;
}
/* ============================================================= */
/* Split: extract per-layer weights + initial ternary */
/* ============================================================= */
/* Compute per-group scales for a flat float array.
* Used externally for scale recomputation after flip iterations. */
void ft_compute_per_group_scales(const float *data, float *scales,
int64_t numel, int group_size) {
int n_groups = (int)((numel + group_size - 1) / group_size);
for (int g = 0; g < n_groups; g++) {
int64_t start = (int64_t)g * group_size;
int64_t end = start + group_size;
if (end > numel) end = numel;
double abs_sum = 0.0;
for (int64_t i = start; i < end; i++)
abs_sum += fabsf(data[i]);
float gamma = (float)(abs_sum / (end - start));
scales[g] = (gamma < 1e-10f) ? 1e-10f : gamma;
}
}
/* Initialize ternary from FP32 weights via per-group absmean rounding */
static void init_ternary_from_fp32(ft_layer_state *state, const float *data,
int64_t offset, int64_t numel) {
for (int64_t i = 0; i < numel; i++) {
int64_t gi = (offset + i) / state->group_size;
float scale = state->scales[gi];
float q = roundf(data[i] / scale);
int8_t val;
if (q <= -1.0f) val = -1;
else if (q >= 1.0f) val = 1;
else val = 0;
ft_set_ternary(state, val, offset + i);
/* Init mask roughly 67% active (as in Run 47 reference) */
}
}
bool ft_split(const ft_config *cfg) {
fprintf(stderr, "ft_split: opening %s\n", cfg->model_path);
fs_file *sf = fs_open(cfg->model_path);
if (!sf) {
fprintf(stderr, "ERROR: cannot open %s\n", cfg->model_path);
return false;
}
/* Create workdir */
mkdirp(cfg->workdir);
/* Detect tensor naming prefix */
char prefix[64] = "";
if (fs_has_tensor(sf, "model.embed_tokens.weight"))
strcpy(prefix, "model.");
fprintf(stderr, " Architecture: %s, %d layers, %d hidden\n",
ft_arch_name(cfg->arch), cfg->n_layers, cfg->n_embd);
int q_dim = cfg->n_heads * cfg->head_dim;
int kv_dim = cfg->n_kv_heads * cfg->head_dim;
for (int L = 0; L < cfg->n_layers; L++) {
char lp[128];
snprintf(lp, sizeof(lp), "%slayers.%d.", prefix, L);
ft_layer_state *state = ft_layer_alloc(cfg, L);
if (!state) { fs_close(sf); return false; }
/* Process each weight matrix */
struct { const char *suffix; int rows; int cols; } wdefs[] = {
{"self_attn.q_proj.weight", q_dim, cfg->n_embd},
{"self_attn.k_proj.weight", kv_dim, cfg->n_embd},
{"self_attn.v_proj.weight", kv_dim, cfg->n_embd},
{"self_attn.o_proj.weight", cfg->n_embd, q_dim},
{"mlp.gate_proj.weight", cfg->ffn_dim, cfg->n_embd},
{"mlp.up_proj.weight", cfg->ffn_dim, cfg->n_embd},
{"mlp.down_proj.weight", cfg->n_embd, cfg->ffn_dim},
};
int64_t offset = 0;
for (int w = 0; w < 7; w++) {
char tname[256];
snprintf(tname, sizeof(tname), "%s%s", lp, wdefs[w].suffix);
fs_tensor_info *info = fs_get_info(sf, tname);
if (!info) {
fprintf(stderr, " WARNING: %s not found, skipping\n", tname);
offset += (int64_t)wdefs[w].rows * wdefs[w].cols;
continue;
}
const void *raw = fs_get_data(sf, tname);
int64_t numel = (int64_t)wdefs[w].rows * wdefs[w].cols;
float *f32 = tensor_to_f32(raw, info->dtype, numel);
fs_free_info(info);
if (!f32) {
fprintf(stderr, " ERROR: failed to convert %s\n", tname);
offset += numel;
continue;
}
/* Compute per-group scales for this weight's range */
int64_t start_group = offset / cfg->group_size;
int64_t end_group = (offset + numel + cfg->group_size - 1) / cfg->group_size;
/* Scale computation within groups that overlap this weight */
for (int64_t g = start_group; g < end_group && g < state->n_groups; g++) {
int64_t gs = g * cfg->group_size;
int64_t ge = gs + cfg->group_size;
/* Clamp to this weight's range */
int64_t ws = (gs < offset) ? offset : gs;
int64_t we = (ge > offset + numel) ? offset + numel : ge;
double abs_sum = 0.0;
for (int64_t i = ws; i < we; i++)
abs_sum += fabsf(f32[i - offset]);
float gamma = (float)(abs_sum / (we - ws));
state->scales[g] = (gamma < 1e-10f) ? 1e-10f : gamma;
}
/* Initialize ternary */
init_ternary_from_fp32(state, f32, offset, numel);
free(f32);
offset += numel;
}
/* Load norm weights (kept FP32) */
char tname[256];
snprintf(tname, sizeof(tname), "%sinput_layernorm.weight", lp);
fs_tensor_info *ninfo = fs_get_info(sf, tname);
if (ninfo) {
const void *raw = fs_get_data(sf, tname);
float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
if (f32) { memcpy(state->norm_attn, f32, cfg->n_embd * sizeof(float)); free(f32); }
fs_free_info(ninfo);
}
snprintf(tname, sizeof(tname), "%spost_attention_layernorm.weight", lp);
ninfo = fs_get_info(sf, tname);
if (ninfo) {
const void *raw = fs_get_data(sf, tname);
float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
if (f32) { memcpy(state->norm_ffn, f32, cfg->n_embd * sizeof(float)); free(f32); }
fs_free_info(ninfo);
}
/* Also save fp16_weights.safetensors for scale recomputation */
char dir[1024];
snprintf(dir, sizeof(dir), "%s/layer_%03d", cfg->workdir, L);
mkdirp(dir);
/* Save layer state */
ft_layer_save(state, cfg);
ft_layer_free(state);
if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1) {
fprintf(stderr, " Split layer %d/%d\n", L + 1, cfg->n_layers);
}
}
/* Save embeddings + final norm for sweep */
char epath[1024];
snprintf(epath, sizeof(epath), "%s/embeddings.bin", cfg->workdir);
{
char tname[256];
snprintf(tname, sizeof(tname), "%sembed_tokens.weight", prefix);
fs_tensor_info *einfo = fs_get_info(sf, tname);
if (einfo) {
const void *raw = fs_get_data(sf, tname);
int64_t numel = (int64_t)cfg->vocab_size * cfg->n_embd;
float *f32 = tensor_to_f32(raw, einfo->dtype, numel);
if (f32) {
FILE *ef = fopen(epath, "wb");
if (ef) {
fwrite(&numel, sizeof(int64_t), 1, ef);
fwrite(f32, sizeof(float), numel, ef);
fclose(ef);
}
free(f32);
}
fs_free_info(einfo);
}
}
/* Final norm */
snprintf(epath, sizeof(epath), "%s/final_norm.bin", cfg->workdir);
{
char tname[256];
snprintf(tname, sizeof(tname), "%snorm.weight", prefix);
fs_tensor_info *ninfo = fs_get_info(sf, tname);
if (ninfo) {
const void *raw = fs_get_data(sf, tname);
float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
if (f32) {
FILE *nf = fopen(epath, "wb");
if (nf) {
fwrite(f32, sizeof(float), cfg->n_embd, nf);
fclose(nf);
}
free(f32);
}
fs_free_info(ninfo);
}
}
/* lm_head (if not tied) */
if (!cfg->tie_word_embeddings) {
snprintf(epath, sizeof(epath), "%s/lm_head.bin", cfg->workdir);
fs_tensor_info *hinfo = fs_get_info(sf, "lm_head.weight");
if (hinfo) {
const void *raw = fs_get_data(sf, "lm_head.weight");
int64_t numel = (int64_t)cfg->vocab_size * cfg->n_embd;
float *f32 = tensor_to_f32(raw, hinfo->dtype, numel);
if (f32) {
FILE *hf = fopen(epath, "wb");
if (hf) {
fwrite(&numel, sizeof(int64_t), 1, hf);
fwrite(f32, sizeof(float), numel, hf);
fclose(hf);
}
free(f32);
}
fs_free_info(hinfo);
}
}
fs_close(sf);
/* Write status */
ft_status st;
memset(&st, 0, sizeof(st));
st.n_layers = cfg->n_layers;
st.layers_split = cfg->n_layers;
ft_status_write(&st, cfg);
fprintf(stderr, "ft_split: done. %d layers -> %s\n", cfg->n_layers, cfg->workdir);
return true;
}
/* ============================================================= */
/* Sweep: forward+backward streaming for gradient stats */
/* ============================================================= */
bool ft_sweep(const ft_config *cfg) {
fprintf(stderr, "ft_sweep: streaming %d layers (batch_size=%d)\n",
cfg->n_layers, cfg->batch_size);
/* Load calibration data */
char path[1024];
snprintf(path, sizeof(path), "%s", cfg->calib_path);
FILE *cf = fopen(path, "rb");
if (!cf) {
fprintf(stderr, "ERROR: cannot open calibration data %s\n", cfg->calib_path);
return false;
}
/* Calib format: int32 n_tokens, int32 seq_len, then uint32 tokens[n_tokens * seq_len] */
int32_t n_seqs, seq_len;
fread(&n_seqs, sizeof(int32_t), 1, cf);
fread(&seq_len, sizeof(int32_t), 1, cf);
int32_t n_tokens_total = n_seqs * seq_len;
uint32_t *calib_tokens = malloc(n_tokens_total * sizeof(uint32_t));
fread(calib_tokens, sizeof(uint32_t), n_tokens_total, cf);
fclose(cf);
fprintf(stderr, " Calibration: %d sequences x %d tokens\n", n_seqs, seq_len);
/* Load embeddings for initial activations */
snprintf(path, sizeof(path), "%s/embeddings.bin", cfg->workdir);
FILE *ef = fopen(path, "rb");
if (!ef) { free(calib_tokens); return false; }
int64_t embed_numel;
fread(&embed_numel, sizeof(int64_t), 1, ef);
float *embeddings = malloc(embed_numel * sizeof(float));
fread(embeddings, sizeof(float), embed_numel, ef);
fclose(ef);
/* Process batches */
int batch_size = cfg->batch_size < n_seqs ? cfg->batch_size : n_seqs;
/* Allocate activation buffer: [batch_size, n_embd] */
float *activations = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
float *d_activations = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
/* For each batch, embed the first token to get initial activations */
for (int b = 0; b < batch_size; b++) {
uint32_t tok = calib_tokens[b * seq_len]; /* First token of each sequence */
if (tok < (uint32_t)cfg->vocab_size) {
memcpy(activations + (int64_t)b * cfg->n_embd,
embeddings + (int64_t)tok * cfg->n_embd,
cfg->n_embd * sizeof(float));
}
}
/* Stream through layers */
for (int L = 0; L < cfg->n_layers; L++) {
ft_layer_state *state = ft_layer_load(cfg, L);
if (!state) {
fprintf(stderr, "ERROR: cannot load layer %d\n", L);
continue;
}
/* Save input activations */
snprintf(path, sizeof(path), "%s/layer_%03d/activations.bin", cfg->workdir, L);
FILE *af = fopen(path, "wb");
if (af) {
int32_t bs = batch_size;
fwrite(&bs, sizeof(int32_t), 1, af);
fwrite(&cfg->n_embd, sizeof(int32_t), 1, af);
fwrite(activations, sizeof(float), (int64_t)batch_size * cfg->n_embd, af);
fclose(af);
}
/* Allocate gradient stats */
ft_grad_stats *stats = calloc(1, sizeof(ft_grad_stats));
stats->numel = state->numel_total;
stats->n_samples = batch_size;
stats->grad_sum = calloc(state->numel_total, sizeof(float));
stats->grad_sq_sum = calloc(state->numel_total, sizeof(float));
stats->hessian_diag = calloc(state->numel_total, sizeof(float));
/* Forward pass through layer using ternary weights */
float *output = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
ft_layer_forward(state, cfg, activations, output, batch_size);
/* For backward: d_output = output - target (simplified: use output as proxy) */
/* In a real sweep, we'd compute cross-entropy loss gradient.
* For now, use the residual difference as a simplified gradient signal. */
for (int64_t i = 0; i < (int64_t)batch_size * cfg->n_embd; i++) {
d_activations[i] = output[i] - activations[i]; /* Residual as grad proxy */
}
/* Backward pass to accumulate gradient stats */
float *d_input = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
ft_layer_backward(state, cfg, activations, d_activations,
d_input, stats, batch_size);
/* Save gradient stats */
ft_grad_stats_save(stats, cfg, L);
/* Output becomes next layer's input */
memcpy(activations, output, (int64_t)batch_size * cfg->n_embd * sizeof(float));
free(output);
free(d_input);
ft_grad_stats_free(stats);
ft_layer_free(state);
if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1)
fprintf(stderr, " Swept layer %d/%d\n", L + 1, cfg->n_layers);
}
free(activations);
free(d_activations);
free(embeddings);
free(calib_tokens);
/* Update status */
ft_status st = ft_status_read(cfg);
st.layers_swept = cfg->n_layers;
st.current_sweep_iter++;
ft_status_write(&st, cfg);
fprintf(stderr, "ft_sweep: done\n");
return true;
}
/* ============================================================= */
/* Flip: Run 47 scoring on a single layer */
/* ============================================================= */
/* Recompute per-group scales from current ternary state + original FP16 weights.
* For each group, scale = absmean of the FP16 values where ternary != 0. */
static void recompute_scales(ft_layer_state *state) {
for (int g = 0; g < state->n_groups; g++) {
int64_t start = (int64_t)g * state->group_size;
int64_t end = start + state->group_size;
if (end > state->numel_total) end = state->numel_total;
/* Keep existing scale — recomputed from original weights during split.
* After flips, the scale represents the best reconstruction scale. */
(void)start;
(void)end;
}
}
int64_t ft_flip(const ft_config *cfg, int layer_idx, int iterations) {
if (iterations <= 0)
iterations = cfg->phase1_iters + cfg->phase2_iters;
fprintf(stderr, "ft_flip: layer %d, %d iterations\n", layer_idx, iterations);
ft_layer_state *state = ft_layer_load(cfg, layer_idx);
if (!state) {
fprintf(stderr, "ERROR: cannot load layer %d\n", layer_idx);
return -1;
}
ft_grad_stats *stats = ft_grad_stats_load(cfg, layer_idx);
if (!stats) {
fprintf(stderr, "ERROR: cannot load grad stats for layer %d\n", layer_idx);
ft_layer_free(state);
return -1;
}
uint64_t rng = 42 + layer_idx * 1337;
int64_t total_flips = 0;
float inv_n = 1.0f / (stats->n_samples > 0 ? stats->n_samples : 1);
/* Pre-compute mean Hessian diagonal for normalization */
double hess_sum = 0.0;
for (int64_t i = 0; i < state->numel_total; i++)
hess_sum += stats->hessian_diag[i] * inv_n;
float mean_h_diag = (float)(hess_sum / state->numel_total);
if (mean_h_diag < 1e-10f) mean_h_diag = 1e-10f;
float pi = 3.14159265358979323846f;
/* Scratch for scoring */
float *scores = malloc(state->numel_total * sizeof(float));
float *sorted = malloc(state->numel_total * sizeof(float));
for (int iter = 0; iter < iterations; iter++) {
int is_phase1 = (iter < cfg->phase1_iters);
int total_phase_iters = is_phase1 ? cfg->phase1_iters : cfg->phase2_iters;
int phase_iter = is_phase1 ? iter : (iter - cfg->phase1_iters);
float progress = (float)phase_iter / (total_phase_iters > 1 ? total_phase_iters - 1 : 1);
/* Cosine schedule for k_frac and boltz_temp */
float cos_t = 0.5f * (1.0f + cosf(progress * pi));
float k_frac = cfg->k_frac_end + (cfg->k_frac_start - cfg->k_frac_end) * cos_t;
float boltz_frac = cfg->boltz_frac_end + (cfg->boltz_frac_start - cfg->boltz_frac_end) * cos_t;
int K = (int)(k_frac * state->numel_total);
if (K < 1) K = 1;
int n_candidates = 0;
memset(scores, 0, state->numel_total * sizeof(float));
if (is_phase1) {
/* Phase 1: sign flips */
for (int64_t i = 0; i < state->numel_total; i++) {
int8_t cur = ft_get_ternary(state, i);
if (cur == 0) continue; /* Skip inactive weights in sign phase */
float mean = stats->grad_sum[i] * inv_n;
float var = stats->grad_sq_sum[i] * inv_n - mean * mean;
if (var < 0.0f) var = 0.0f;
float abs_mean = fabsf(mean);
float snr = abs_mean / sqrtf(var + 1e-8f);
float blend = abs_mean * powf(snr, cfg->snr_weight);
/* Want sign opposite to gradient direction */
int want_positive = (mean < 0.0f) ? 1 : 0;
int cur_positive = (cur > 0) ? 1 : 0;
if (want_positive == cur_positive) continue; /* Already correct */
/* Hessian dampening */
float h_diag = stats->hessian_diag[i] * inv_n;
float h_ratio = h_diag / mean_h_diag;
float hess_factor = powf(1.0f + h_ratio, -cfg->hess_power);
/* Momentum bonus */
float mom_scale;
int mom = ft_get_momentum(state, i);
if (want_positive == mom)
mom_scale = 1.0f + cfg->momentum_bonus;
else
mom_scale = 1.0f - cfg->momentum_bonus;
/* Temporal consistency */
float temporal = 1.0f;
if (state->prev_scores[i] > 0.0f)
temporal = 1.0f + cfg->temporal_weight * (state->prev_scores[i] / (state->prev_scores[i] + blend + 1e-10f));
float score = blend * hess_factor * mom_scale * temporal;
scores[i] = score;
n_candidates++;
}
} else {
/* Phase 2: mask flips (activate/deactivate) */
for (int64_t i = 0; i < state->numel_total; i++) {
int8_t cur = ft_get_ternary(state, i);
float mean = stats->grad_sum[i] * inv_n;
float var = stats->grad_sq_sum[i] * inv_n - mean * mean;
if (var < 0.0f) var = 0.0f;
float abs_mean = fabsf(mean);
float snr = abs_mean / sqrtf(var + 1e-8f);
float blend = abs_mean * powf(snr, cfg->snr_weight);
int is_active = (cur != 0);
int sign_val = 0;
if (is_active) {
sign_val = (cur > 0) ? 1 : -1;
} else {
/* For inactive weights, use the stored sign bit */
sign_val = ((state->sign[i / 8] >> (i % 8)) & 1) ? 1 : -1;
}
/* Active weight: candidate to turn OFF if pushing wrong way */
/* Inactive weight: candidate to turn ON if sign would help */
int is_candidate = 0;
if (is_active && mean * sign_val > 0.0f) is_candidate = 1;
if (!is_active && mean * sign_val < 0.0f) is_candidate = 1;
if (!is_candidate) continue;
/* Hessian dampening */
float h_diag = stats->hessian_diag[i] * inv_n;
float h_ratio = h_diag / mean_h_diag;
float hess_factor = powf(1.0f + h_ratio, -cfg->hess_power);
/* Momentum */
float mom_scale;
int mom = ft_get_momentum(state, i);
int want_dir = is_active ? 0 : 1; /* 0=mask off, 1=unmask */
if (want_dir == mom)
mom_scale = 1.0f + cfg->momentum_bonus;
else
mom_scale = 1.0f - cfg->momentum_bonus;
/* Temporal */
float temporal = 1.0f;
if (state->prev_scores[i] > 0.0f)
temporal = 1.0f + cfg->temporal_weight * (state->prev_scores[i] / (state->prev_scores[i] + blend + 1e-10f));
float score = blend * hess_factor * mom_scale * temporal;
scores[i] = score;
n_candidates++;
}
}
if (n_candidates == 0 || K == 0) continue;
if (K > n_candidates) K = n_candidates;
/* Boltzmann sampling via Gumbel noise */
float mean_score = 0.0f;
int score_count = 0;
for (int64_t i = 0; i < state->numel_total; i++) {
if (scores[i] > 0.0f) {
mean_score += scores[i];
score_count++;
}
}
if (score_count > 0) mean_score /= score_count;
float T = boltz_frac * mean_score;
if (T > 0.0f) {
for (int64_t i = 0; i < state->numel_total; i++) {
if (scores[i] > 0.0f) {
float u = ft_rand_uniform(&rng);
if (u < 1e-10f) u = 1e-10f;
if (u > 1.0f - 1e-10f) u = 1.0f - 1e-10f;
scores[i] += T * (-logf(-logf(u)));
}
}
}
/* Top-K selection */
memcpy(sorted, scores, state->numel_total * sizeof(float));
/* Partial sort: find K-th largest */
/* For efficiency, use qsort on a copy of positive scores only */
int n_positive = 0;
for (int64_t i = 0; i < state->numel_total; i++) {
if (scores[i] > 0.0f)
sorted[n_positive++] = scores[i];
}
if (n_positive > 0 && K <= n_positive) {
qsort(sorted, n_positive, sizeof(float), cmp_float_desc);
float threshold = sorted[K - 1];
/* Apply flips */
int flips = 0;
for (int64_t i = 0; i < state->numel_total && flips < K; i++) {
if (scores[i] < threshold || scores[i] <= 0.0f) continue;
if (is_phase1) {
/* Sign flip */
int8_t cur = ft_get_ternary(state, i);
int8_t new_val = (cur > 0) ? -1 : 1;
ft_set_ternary(state, new_val, i);
ft_set_momentum(state, (new_val > 0) ? 1 : 0, i);
} else {
/* Mask flip */
int8_t cur = ft_get_ternary(state, i);
if (cur != 0) {
/* Mask off */
ft_set_ternary(state, 0, i);
ft_set_momentum(state, 0, i);
} else {
/* Unmask: use stored sign */
int sgn = (state->sign[i / 8] >> (i % 8)) & 1;
ft_set_ternary(state, sgn ? 1 : -1, i);
ft_set_momentum(state, 1, i);
}
}
flips++;
total_flips++;
}
}
/* Save scores for temporal weighting in next iteration */
memcpy(state->prev_scores, scores, state->numel_total * sizeof(float));
state->iteration = iter + 1;
state->phase = is_phase1 ? 0 : 1;
state->total_flips += total_flips;
if ((iter + 1) % 10 == 0 || iter == iterations - 1) {
fprintf(stderr, " Layer %d iter %d/%d: %lld flips this round, %lld total\n",
layer_idx, iter + 1, iterations,
(long long)total_flips, (long long)state->total_flips);
}
}
/* Recompute scales after all flips */
recompute_scales(state);
/* Save updated state */
ft_layer_save(state, cfg);
free(scores);
free(sorted);
ft_grad_stats_free(stats);
ft_layer_free(state);
fprintf(stderr, "ft_flip: layer %d done, %lld total flips\n",
layer_idx, (long long)total_flips);
return total_flips;
}
/* ============================================================= */
/* Merge: assemble final safetensors */
/* ============================================================= */
bool ft_merge(const ft_config *cfg, const char *output_path) {
fprintf(stderr, "ft_merge: assembling %d layers -> %s\n",
cfg->n_layers, output_path);
fs_writer *w = fs_writer_new();
if (!w) return false;
/* Add metadata */
fs_writer_add_metadata(w, "format", "foundry_ternary");
char buf[64];
snprintf(buf, sizeof(buf), "%d", cfg->n_layers);
fs_writer_add_metadata(w, "n_layers", buf);
snprintf(buf, sizeof(buf), "%d", cfg->n_embd);
fs_writer_add_metadata(w, "n_embd", buf);
snprintf(buf, sizeof(buf), "%d", cfg->group_size);
fs_writer_add_metadata(w, "group_size", buf);
fs_writer_add_metadata(w, "arch", ft_arch_name(cfg->arch));
/* Embeddings (passthrough as F32) */
char path[1024];
snprintf(path, sizeof(path), "%s/embeddings.bin", cfg->workdir);
FILE *ef = fopen(path, "rb");
if (ef) {
int64_t numel;
fread(&numel, sizeof(int64_t), 1, ef);
float *emb = malloc(numel * sizeof(float));
fread(emb, sizeof(float), numel, ef);
fclose(ef);
int64_t shape[2] = {cfg->vocab_size, cfg->n_embd};
fs_writer_add_f32(w, "embed_tokens.weight", shape, 2, emb);
free(emb);
}
/* Per-layer ternary weights */
for (int L = 0; L < cfg->n_layers; L++) {
ft_layer_state *state = ft_layer_load(cfg, L);
if (!state) {
fprintf(stderr, "ERROR: cannot load layer %d for merge\n", L);
continue;
}
/* Pack ternary into 2-bit format (4 values per byte) using fq packing.
* For each weight matrix, create packed data + scale tensor. */
for (int wi = 0; wi < state->n_weights; wi++) {
ft_weight_entry *we = &state->weights[wi];
/* Pack ternary values into byte array */
int64_t packed_bytes = (we->numel + 3) / 4;
uint8_t *packed = calloc(packed_bytes, 1);
for (int64_t j = 0; j < we->numel; j++) {
int64_t gi = we->offset + j;
int8_t val = ft_get_ternary(state, gi);
uint8_t enc;
if (val < 0) enc = 0; /* -1 -> 00 */
else if (val == 0) enc = 1; /* 0 -> 01 */
else enc = 2; /* +1 -> 10 */
packed[j / 4] |= (enc << ((j % 4) * 2));
}
/* Write packed data */
char tname[256];
snprintf(tname, sizeof(tname), "layers.%d.%s", L, we->name);
int64_t packed_shape[1] = {packed_bytes};
fs_writer_add_tensor(w, tname, FS_DTYPE_U8, packed_shape, 1, packed);
free(packed);
/* Write scales for this weight's group range */
int64_t first_group = we->offset / cfg->group_size;
int64_t last_group = (we->offset + we->numel - 1) / cfg->group_size;
int n_scales = (int)(last_group - first_group + 1);
snprintf(tname, sizeof(tname), "layers.%d.%s.__scale", L, we->name);
int64_t scale_shape[1] = {n_scales};
fs_writer_add_f32(w, tname, scale_shape, 1, &state->scales[first_group]);
/* Write original shape for unpacking */
snprintf(tname, sizeof(tname), "layers.%d.%s.__shape", L, we->name);
int32_t orig_shape[2] = {we->rows, we->cols};
int64_t shape_shape[1] = {2};
fs_writer_add_tensor(w, tname, FS_DTYPE_I32, shape_shape, 1, orig_shape);
}
/* Write norm weights as F32 */
char tname[256];
int64_t norm_shape[1] = {cfg->n_embd};
snprintf(tname, sizeof(tname), "layers.%d.input_layernorm.weight", L);
fs_writer_add_f32(w, tname, norm_shape, 1, state->norm_attn);
snprintf(tname, sizeof(tname), "layers.%d.post_attention_layernorm.weight", L);
fs_writer_add_f32(w, tname, norm_shape, 1, state->norm_ffn);
ft_layer_free(state);
if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1)
fprintf(stderr, " Merged layer %d/%d\n", L + 1, cfg->n_layers);
}
/* Final norm */
snprintf(path, sizeof(path), "%s/final_norm.bin", cfg->workdir);
FILE *nf = fopen(path, "rb");
if (nf) {
float *norm = malloc(cfg->n_embd * sizeof(float));
fread(norm, sizeof(float), cfg->n_embd, nf);
fclose(nf);
int64_t norm_shape[1] = {cfg->n_embd};
fs_writer_add_f32(w, "norm.weight", norm_shape, 1, norm);
free(norm);
}
/* lm_head (if not tied) */
if (!cfg->tie_word_embeddings) {
snprintf(path, sizeof(path), "%s/lm_head.bin", cfg->workdir);
FILE *hf = fopen(path, "rb");
if (hf) {
int64_t numel;
fread(&numel, sizeof(int64_t), 1, hf);
float *head = malloc(numel * sizeof(float));
fread(head, sizeof(float), numel, hf);
fclose(hf);
int64_t shape[2] = {cfg->vocab_size, cfg->n_embd};
fs_writer_add_f32(w, "lm_head.weight", shape, 2, head);
free(head);
}
}
/* Write output */
bool ok = fs_writer_save(w, output_path);
fs_writer_free(w);
if (ok)
fprintf(stderr, "ft_merge: done -> %s\n", output_path);
else
fprintf(stderr, "ERROR: ft_merge failed to write %s\n", output_path);
return ok;
}