271 lines
8.9 KiB
C
271 lines
8.9 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_TERNARIZE_H
|
|
#define FOUNDRY_TERNARIZE_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Post-training ternarization using Run 47 gradient-based scoring.
|
|
*
|
|
* Converts pre-trained FP16/BF16 models to ternary {-1, 0, +1} weights
|
|
* with per-group float32 scales. Layer-by-layer streaming so a 72B model
|
|
* never needs to be fully in memory. Each layer is an independent job.
|
|
*
|
|
* Four-phase pipeline:
|
|
* 1. split - Extract per-layer FP16 weights, initial absmean ternary
|
|
* 2. sweep - Stream forward+backward for activation/gradient snapshots
|
|
* 3. flip - Run 47 scoring on one layer (embarrassingly parallel)
|
|
* 4. merge - Assemble final safetensors with packed ternary + scales
|
|
*
|
|
* Prefix: ft_ (foundry ternarize)
|
|
*/
|
|
|
|
/* ============================================================= */
|
|
/* Architecture */
|
|
/* ============================================================= */
|
|
|
|
typedef enum {
|
|
FT_ARCH_LLAMA,
|
|
FT_ARCH_QWEN2,
|
|
FT_ARCH_QWEN3,
|
|
FT_ARCH_PHI3,
|
|
FT_ARCH_GEMMA2,
|
|
FT_ARCH_RWKVX,
|
|
FT_ARCH_UNKNOWN,
|
|
} ft_arch;
|
|
|
|
/* ============================================================= */
|
|
/* Configuration */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
/* Model architecture */
|
|
ft_arch arch;
|
|
int n_layers;
|
|
int n_embd;
|
|
int n_heads;
|
|
int n_kv_heads;
|
|
int head_dim;
|
|
int ffn_dim;
|
|
int vocab_size;
|
|
float rope_theta;
|
|
float rms_eps;
|
|
int has_qkv_bias;
|
|
int has_qk_norm;
|
|
int tie_word_embeddings;
|
|
|
|
/* Run 47 parameters */
|
|
int group_size; /* Per-group scale granularity (default 128) */
|
|
float k_frac_start; /* Initial flip rate (default 0.00008) */
|
|
float k_frac_end; /* Final flip rate (default 0.000005) */
|
|
float snr_weight; /* SNR exponent in score blend (default 0.3) */
|
|
float hess_power; /* Hessian dampening exponent (default 0.15) */
|
|
float momentum_bonus; /* Momentum bonus/penalty (default 0.2) */
|
|
float boltz_frac_start; /* Initial Boltzmann temperature (default 0.17) */
|
|
float boltz_frac_end; /* Final Boltzmann temperature (default 0.0) */
|
|
float temporal_weight; /* Temporal consistency weight (default 1.8) */
|
|
int phase1_iters; /* Sign-only phase iterations (default 60) */
|
|
int phase2_iters; /* Mask phase iterations (default 140) */
|
|
int batch_size; /* Calibration batch size (default 128) */
|
|
|
|
/* Paths */
|
|
char workdir[512];
|
|
char model_path[512];
|
|
char calib_path[512];
|
|
} ft_config;
|
|
|
|
/* ============================================================= */
|
|
/* Weight Registry (per-layer tensor metadata) */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
char name[128]; /* Tensor name (e.g., "self_attn.q_proj.weight") */
|
|
int rows;
|
|
int cols;
|
|
int64_t offset; /* Offset into the flat sign/mask/momentum arrays */
|
|
int64_t numel; /* rows * cols */
|
|
} ft_weight_entry;
|
|
|
|
/* ============================================================= */
|
|
/* Layer State */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
int layer_idx;
|
|
int64_t numel_total; /* Total weight elements in this layer */
|
|
|
|
/* Ternary state: sign + mask = effective weight {-1, 0, +1} */
|
|
uint8_t *sign; /* 1 = +1, 0 = -1 (bit array, packed 8 per byte) */
|
|
uint8_t *mask; /* 1 = active, 0 = zero (bit array, packed 8 per byte) */
|
|
uint8_t *momentum; /* Last flip direction (bit array, packed 8 per byte) */
|
|
|
|
/* Per-group scales */
|
|
float *scales;
|
|
int n_groups;
|
|
int group_size;
|
|
|
|
/* Scoring state */
|
|
float *prev_scores; /* Previous iteration scores for temporal */
|
|
|
|
/* Weight registry */
|
|
ft_weight_entry *weights;
|
|
int n_weights;
|
|
|
|
/* Norm tensors (kept FP32, not ternarized) */
|
|
float *norm_attn; /* input_layernorm.weight [n_embd] */
|
|
float *norm_ffn; /* post_attention_layernorm.weight [n_embd] */
|
|
|
|
/* Status */
|
|
int iteration;
|
|
int phase; /* 0 = sign, 1 = mask */
|
|
int64_t total_flips;
|
|
float best_loss;
|
|
} ft_layer_state;
|
|
|
|
/* ============================================================= */
|
|
/* Status */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
int n_layers;
|
|
int layers_split;
|
|
int layers_swept;
|
|
int layers_flipped;
|
|
int current_sweep_iter;
|
|
float overall_loss;
|
|
char model_name[128];
|
|
} ft_status;
|
|
|
|
/* ============================================================= */
|
|
/* Gradient Statistics (per-layer, saved by sweep) */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
float *grad_sum; /* Sum of gradients over batch [numel_total] */
|
|
float *grad_sq_sum; /* Sum of squared gradients [numel_total] */
|
|
float *hessian_diag; /* Hessian diagonal approximation [numel_total] */
|
|
int64_t numel;
|
|
int n_samples; /* Number of samples accumulated */
|
|
} ft_grad_stats;
|
|
|
|
/* ============================================================= */
|
|
/* Public API */
|
|
/* ============================================================= */
|
|
|
|
/* Get default config with Run 47 parameters */
|
|
ft_config ft_config_defaults(void);
|
|
|
|
/* Parse model config.json and populate architecture fields */
|
|
bool ft_config_from_json(ft_config *cfg, const char *config_json_path);
|
|
|
|
/* ---- Split ---- */
|
|
|
|
/* Extract model into per-layer workdir. Initial ternary via absmean.
|
|
* Creates workdir/layer_NNN/ folders with ternary.bin, scales.bin, etc. */
|
|
bool ft_split(const ft_config *cfg);
|
|
|
|
/* ---- Sweep ---- */
|
|
|
|
/* Stream forward+backward through all layers. Saves activation snapshots
|
|
* and gradient statistics per layer. One layer in memory at a time. */
|
|
bool ft_sweep(const ft_config *cfg);
|
|
|
|
/* ---- Flip ---- */
|
|
|
|
/* Run 47 scoring on a single layer. Fully self-contained.
|
|
* iterations: number of flip iterations to run (0 = use config defaults).
|
|
* Returns number of flips performed. */
|
|
int64_t ft_flip(const ft_config *cfg, int layer_idx, int iterations);
|
|
|
|
/* ---- Merge ---- */
|
|
|
|
/* Assemble final safetensors with packed ternary weights + scales.
|
|
* output_path: path for the output .safetensors file. */
|
|
bool ft_merge(const ft_config *cfg, const char *output_path);
|
|
|
|
/* ---- Layer State I/O ---- */
|
|
|
|
/* Load layer state from workdir */
|
|
ft_layer_state* ft_layer_load(const ft_config *cfg, int layer_idx);
|
|
|
|
/* Save layer state to workdir */
|
|
bool ft_layer_save(const ft_layer_state *state, const ft_config *cfg);
|
|
|
|
/* Free layer state */
|
|
void ft_layer_free(ft_layer_state *state);
|
|
|
|
/* ---- Gradient Stats I/O ---- */
|
|
|
|
/* Load gradient stats from workdir */
|
|
ft_grad_stats* ft_grad_stats_load(const ft_config *cfg, int layer_idx);
|
|
|
|
/* Save gradient stats to workdir */
|
|
bool ft_grad_stats_save(const ft_grad_stats *stats, const ft_config *cfg, int layer_idx);
|
|
|
|
/* Free gradient stats */
|
|
void ft_grad_stats_free(ft_grad_stats *stats);
|
|
|
|
/* ---- Status ---- */
|
|
|
|
/* Read/write pipeline status */
|
|
ft_status ft_status_read(const ft_config *cfg);
|
|
bool ft_status_write(const ft_status *status, const ft_config *cfg);
|
|
|
|
/* ---- Utilities ---- */
|
|
|
|
/* Get architecture name */
|
|
const char* ft_arch_name(ft_arch arch);
|
|
|
|
/* Get effective ternary value for element i */
|
|
static inline int8_t ft_get_ternary(const ft_layer_state *s, int64_t i) {
|
|
int64_t byte_idx = i / 8;
|
|
int bit = (int)(i % 8);
|
|
int m = (s->mask[byte_idx] >> bit) & 1;
|
|
if (!m) return 0;
|
|
int sgn = (s->sign[byte_idx] >> bit) & 1;
|
|
return sgn ? 1 : -1;
|
|
}
|
|
|
|
/* Set ternary value for element i */
|
|
static inline void ft_set_ternary(ft_layer_state *s, int8_t val, int64_t i) {
|
|
int64_t byte_idx = i / 8;
|
|
int bit = (int)(i % 8);
|
|
if (val == 0) {
|
|
s->mask[byte_idx] &= ~(1 << bit);
|
|
} else {
|
|
s->mask[byte_idx] |= (1 << bit);
|
|
if (val > 0)
|
|
s->sign[byte_idx] |= (1 << bit);
|
|
else
|
|
s->sign[byte_idx] &= ~(1 << bit);
|
|
}
|
|
}
|
|
|
|
/* Get momentum bit for element i */
|
|
static inline int ft_get_momentum(const ft_layer_state *s, int64_t i) {
|
|
return (s->momentum[i / 8] >> (i % 8)) & 1;
|
|
}
|
|
|
|
/* Set momentum bit for element i */
|
|
static inline void ft_set_momentum(ft_layer_state *s, int val, int64_t i) {
|
|
int64_t byte_idx = i / 8;
|
|
int bit = (int)(i % 8);
|
|
if (val)
|
|
s->momentum[byte_idx] |= (1 << bit);
|
|
else
|
|
s->momentum[byte_idx] &= ~(1 << bit);
|
|
}
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_TERNARIZE_H */
|