1313 lines
48 KiB
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;
|
|
}
|