1342 lines
45 KiB
C
1342 lines
45 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
||
/* Copyright (C) 2026 David (Corey) Herrera */
|
||
/*
|
||
* rfa_train.c — LoRA training for RFA hybrid model (Qwen blocks only)
|
||
*
|
||
* Gradient checkpointing approach: during forward, save x_input and cache_len
|
||
* at each Qwen block. During backward, recompute all intermediates and compute
|
||
* proper gradients for all 7 LoRA layers per Qwen block.
|
||
*
|
||
* LoRA indices per Qwen block: 0=Q, 1=K, 2=V, 3=O, 4=gate, 5=up, 6=down
|
||
*
|
||
* Reuses helper patterns from rx_train.c: lora_fwd_add, lora_grad_acc,
|
||
* lora_backward_input, ensure_grad, BLAS wrappers.
|
||
*/
|
||
|
||
#include "rwkvx/rfa_train.h"
|
||
#include "rwkvx/rfa.h"
|
||
#include "tensor/tensor.h"
|
||
#include "lora/lora.h"
|
||
#include "optim/optim.h"
|
||
|
||
#include <math.h>
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
#include <string.h>
|
||
#include <time.h>
|
||
#include <sys/stat.h>
|
||
#include <pthread.h>
|
||
|
||
#ifdef USE_BLAS
|
||
#if defined(__APPLE__)
|
||
#include <Accelerate/Accelerate.h>
|
||
#else
|
||
#include <cblas.h>
|
||
#endif
|
||
#endif
|
||
|
||
/* ================================================================
|
||
* Internal BLAS helpers (same conventions as rx_train.c)
|
||
* ================================================================ */
|
||
|
||
/* out[cols] = W^T[rows,cols] @ x[rows] — i.e. out = x^T @ W */
|
||
static void t_matvec(const float *x, const float *W, float *out,
|
||
int rows, int cols) {
|
||
#ifdef USE_BLAS
|
||
cblas_sgemv(CblasRowMajor, CblasTrans, rows, cols,
|
||
1.0f, W, cols, x, 1, 0.0f, out, 1);
|
||
#else
|
||
memset(out, 0, cols * sizeof(float));
|
||
for (int i = 0; i < rows; i++) {
|
||
float xi = x[i];
|
||
const float *wi = &W[i * cols];
|
||
for (int j = 0; j < cols; j++)
|
||
out[j] += xi * wi[j];
|
||
}
|
||
#endif
|
||
}
|
||
|
||
/* out[rows] = W[rows,cols] @ x[cols] */
|
||
static void t_matvec_nt(const float *W, const float *x, float *out,
|
||
int rows, int cols) {
|
||
#ifdef USE_BLAS
|
||
cblas_sgemv(CblasRowMajor, CblasNoTrans, rows, cols,
|
||
1.0f, W, cols, x, 1, 0.0f, out, 1);
|
||
#else
|
||
for (int i = 0; i < rows; i++) {
|
||
float sum = 0.0f;
|
||
for (int j = 0; j < cols; j++)
|
||
sum += W[i * cols + j] * x[j];
|
||
out[i] = sum;
|
||
}
|
||
#endif
|
||
}
|
||
|
||
/* out[rows] += W[rows,cols] @ x[cols] (accumulate variant) */
|
||
static void t_matvec_nt_acc(const float *W, const float *x, float *out,
|
||
int rows, int cols) {
|
||
#ifdef USE_BLAS
|
||
cblas_sgemv(CblasRowMajor, CblasNoTrans, rows, cols,
|
||
1.0f, W, cols, x, 1, 1.0f, out, 1);
|
||
#else
|
||
for (int i = 0; i < rows; i++) {
|
||
float sum = 0.0f;
|
||
for (int j = 0; j < cols; j++)
|
||
sum += W[i * cols + j] * x[j];
|
||
out[i] += sum;
|
||
}
|
||
#endif
|
||
}
|
||
|
||
/* ================================================================
|
||
* LoRA helpers (identical to rx_train.c patterns)
|
||
* ================================================================ */
|
||
|
||
static void ensure_grad(tm_tensor *t) {
|
||
if (!t) return;
|
||
if (!t->grad) {
|
||
t->grad = tm_zeros(t->shape, t->ndim, t->dtype);
|
||
t->grad->requires_grad = false;
|
||
}
|
||
}
|
||
|
||
/* Forward: out[d_out] += scale * B @ A @ input */
|
||
static void lora_fwd_add(fl_layer *lora, const float *input, float *out) {
|
||
if (!lora) return;
|
||
int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
|
||
float scale = lora->scale;
|
||
float *A = tm_data_f32(lora->lora_A);
|
||
float *B = tm_data_f32(lora->lora_B);
|
||
float *h = (float*)calloc(rank, sizeof(float));
|
||
for (int r = 0; r < rank; r++) {
|
||
float s = 0.0f;
|
||
for (int j = 0; j < d_in; j++) s += A[r * d_in + j] * input[j];
|
||
h[r] = s;
|
||
}
|
||
for (int i = 0; i < d_out; i++) {
|
||
float s = 0.0f;
|
||
for (int r = 0; r < rank; r++) s += B[i * rank + r] * h[r];
|
||
out[i] += scale * s;
|
||
}
|
||
free(h);
|
||
}
|
||
|
||
/* Accumulate LoRA gradients */
|
||
static void lora_grad_acc(fl_layer *lora, const float *input,
|
||
const float *d_output) {
|
||
if (!lora) return;
|
||
int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
|
||
float scale = lora->scale;
|
||
float *A = tm_data_f32(lora->lora_A);
|
||
float *B = tm_data_f32(lora->lora_B);
|
||
ensure_grad(lora->lora_A);
|
||
ensure_grad(lora->lora_B);
|
||
float *dA = tm_data_f32(lora->lora_A->grad);
|
||
float *dB = tm_data_f32(lora->lora_B->grad);
|
||
|
||
float *h = (float*)calloc(rank, sizeof(float));
|
||
for (int r = 0; r < rank; r++) {
|
||
float s = 0.0f;
|
||
for (int j = 0; j < d_in; j++) s += A[r * d_in + j] * input[j];
|
||
h[r] = s;
|
||
}
|
||
for (int i = 0; i < d_out; i++) {
|
||
float ds = scale * d_output[i];
|
||
for (int r = 0; r < rank; r++)
|
||
dB[i * rank + r] += ds * h[r];
|
||
}
|
||
float *dh = (float*)calloc(rank, sizeof(float));
|
||
for (int r = 0; r < rank; r++) {
|
||
float s = 0.0f;
|
||
for (int i = 0; i < d_out; i++) s += B[i * rank + r] * d_output[i];
|
||
dh[r] = scale * s;
|
||
}
|
||
for (int r = 0; r < rank; r++) {
|
||
float dr = dh[r];
|
||
for (int j = 0; j < d_in; j++)
|
||
dA[r * d_in + j] += dr * input[j];
|
||
}
|
||
free(h);
|
||
free(dh);
|
||
}
|
||
|
||
/* Backward: d_input[d_in] += scale * A^T @ B^T @ d_output */
|
||
static void lora_backward_input(fl_layer *lora, const float *d_output,
|
||
float *d_input) {
|
||
if (!lora) return;
|
||
int d_in = lora->d_in, d_out = lora->d_out, rank = lora->rank;
|
||
float scale = lora->scale;
|
||
float *A = tm_data_f32(lora->lora_A);
|
||
float *B = tm_data_f32(lora->lora_B);
|
||
|
||
float *dh = (float*)calloc(rank, sizeof(float));
|
||
for (int r = 0; r < rank; r++) {
|
||
float s = 0.0f;
|
||
for (int i = 0; i < d_out; i++) s += B[i * rank + r] * d_output[i];
|
||
dh[r] = s;
|
||
}
|
||
for (int r = 0; r < rank; r++) {
|
||
float dr = scale * dh[r];
|
||
for (int j = 0; j < d_in; j++)
|
||
d_input[j] += A[r * d_in + j] * dr;
|
||
}
|
||
free(dh);
|
||
}
|
||
|
||
/* ================================================================
|
||
* RMSNorm forward (matching rfa.c convention)
|
||
* ================================================================ */
|
||
|
||
static void rms_norm_fwd(const float *x, const float *w, float *out,
|
||
int n, float eps) {
|
||
float sum_sq = 0.0f;
|
||
for (int i = 0; i < n; i++)
|
||
sum_sq += x[i] * x[i];
|
||
float scale = 1.0f / sqrtf(sum_sq / n + eps);
|
||
for (int i = 0; i < n; i++)
|
||
out[i] = x[i] * scale * w[i];
|
||
}
|
||
|
||
static inline float fast_silu(float x) {
|
||
return x / (1.0f + expf(-x));
|
||
}
|
||
|
||
/* ================================================================
|
||
* Config defaults
|
||
* ================================================================ */
|
||
|
||
rfa_train_config rfa_train_config_defaults(void) {
|
||
rfa_train_config cfg;
|
||
memset(&cfg, 0, sizeof(cfg));
|
||
cfg.lora_rank = 32;
|
||
cfg.lora_alpha = 64.0f;
|
||
cfg.lr = 2e-5f;
|
||
cfg.weight_decay = 0.01f;
|
||
cfg.epochs = 3;
|
||
cfg.max_seq_len = 256;
|
||
cfg.warmup_steps = 100;
|
||
cfg.save_every = 1000;
|
||
cfg.log_every = 10;
|
||
cfg.grad_clip = 1.0f;
|
||
cfg.max_cache = 2048;
|
||
cfg.n_threads = 1;
|
||
cfg.quantize = false;
|
||
return cfg;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training examples
|
||
* ================================================================ */
|
||
|
||
void rfa_train_example_free(rfa_train_example *ex) {
|
||
if (!ex) return;
|
||
free(ex->input_ids);
|
||
free(ex->target_ids);
|
||
}
|
||
|
||
void rfa_train_examples_free(rfa_train_example *examples, int n) {
|
||
if (!examples) return;
|
||
for (int i = 0; i < n; i++) {
|
||
free(examples[i].input_ids);
|
||
free(examples[i].target_ids);
|
||
}
|
||
free(examples);
|
||
}
|
||
|
||
/* ================================================================
|
||
* Binary data loader
|
||
* ================================================================ */
|
||
|
||
rfa_train_example* rfa_load_pretokenized(
|
||
const char *path, int max_seq_len, int *n_examples_out)
|
||
{
|
||
if (!path || !n_examples_out) return NULL;
|
||
*n_examples_out = 0;
|
||
|
||
FILE *f = fopen(path, "rb");
|
||
if (!f) {
|
||
fprintf(stderr, "rfa_train: cannot open %s\n", path);
|
||
return NULL;
|
||
}
|
||
|
||
uint32_t n_total;
|
||
if (fread(&n_total, sizeof(uint32_t), 1, f) != 1) {
|
||
fclose(f);
|
||
return NULL;
|
||
}
|
||
|
||
rfa_train_example *examples = (rfa_train_example*)calloc(n_total, sizeof(rfa_train_example));
|
||
int n_out = 0;
|
||
|
||
for (uint32_t i = 0; i < n_total; i++) {
|
||
uint32_t seq_len;
|
||
if (fread(&seq_len, sizeof(uint32_t), 1, f) != 1) break;
|
||
|
||
uint32_t *tokens = (uint32_t*)malloc(seq_len * sizeof(uint32_t));
|
||
if (fread(tokens, sizeof(uint32_t), seq_len, f) != seq_len) {
|
||
free(tokens);
|
||
break;
|
||
}
|
||
|
||
if (seq_len < 2) { free(tokens); continue; }
|
||
|
||
/* Truncate if needed */
|
||
int n_tok = (int)seq_len;
|
||
if (n_tok > max_seq_len + 1) n_tok = max_seq_len + 1;
|
||
int sl = n_tok - 1;
|
||
|
||
examples[n_out].seq_len = sl;
|
||
examples[n_out].input_ids = (uint32_t*)malloc(sl * sizeof(uint32_t));
|
||
examples[n_out].target_ids = (uint32_t*)malloc(sl * sizeof(uint32_t));
|
||
memcpy(examples[n_out].input_ids, tokens, sl * sizeof(uint32_t));
|
||
memcpy(examples[n_out].target_ids, &tokens[1], sl * sizeof(uint32_t));
|
||
n_out++;
|
||
|
||
free(tokens);
|
||
}
|
||
|
||
fclose(f);
|
||
*n_examples_out = n_out;
|
||
fprintf(stderr, "rfa_train: loaded %d examples from %s\n", n_out, path);
|
||
return examples;
|
||
}
|
||
|
||
/* ================================================================
|
||
* RMSNorm backward
|
||
*
|
||
* y = x * rsqrt(mean(x^2) + eps) * w
|
||
* d_x[i] = inv_rms * (d_out[i] * w[i] - x_hat[i] * dot / n)
|
||
* where x_hat = x * inv_rms, dot = sum(d_out * w * x_hat)
|
||
* ================================================================ */
|
||
|
||
void rfa_rms_norm_backward(
|
||
const float *d_output,
|
||
const float *x,
|
||
const float *weight,
|
||
float *d_x,
|
||
int c, float eps)
|
||
{
|
||
float sum_sq = 0.0f;
|
||
for (int i = 0; i < c; i++)
|
||
sum_sq += x[i] * x[i];
|
||
float rms_sq = sum_sq / c + eps;
|
||
float inv_rms = 1.0f / sqrtf(rms_sq);
|
||
|
||
/* dot = sum(d_output * weight * x_hat) where x_hat = x * inv_rms */
|
||
float dot = 0.0f;
|
||
for (int i = 0; i < c; i++)
|
||
dot += d_output[i] * weight[i] * x[i] * inv_rms;
|
||
|
||
float inv_c = 1.0f / c;
|
||
for (int i = 0; i < c; i++) {
|
||
float dw = d_output[i] * weight[i];
|
||
float x_hat = x[i] * inv_rms;
|
||
d_x[i] = inv_rms * (dw - x_hat * dot * inv_c);
|
||
}
|
||
}
|
||
|
||
/* ================================================================
|
||
* SwiGLU backward
|
||
*
|
||
* Forward: out = down_proj @ (silu(gate_proj @ x) * up_proj @ x)
|
||
* ================================================================ */
|
||
|
||
void rfa_swiglu_backward(
|
||
const float *d_output,
|
||
const float *x_input,
|
||
const rfa_qwen_weights *w,
|
||
fl_layer **lora,
|
||
float *d_input,
|
||
int c, int inter)
|
||
{
|
||
/* Recompute forward intermediates */
|
||
float *gate = (float*)calloc(inter, sizeof(float));
|
||
float *up = (float*)calloc(inter, sizeof(float));
|
||
float *act = (float*)calloc(inter, sizeof(float));
|
||
|
||
/* gate = gate_proj @ x + LoRA[4] */
|
||
t_matvec(x_input, w->gate_proj, gate, c, inter);
|
||
lora_fwd_add(lora[4], x_input, gate);
|
||
|
||
/* up = up_proj @ x + LoRA[5] */
|
||
t_matvec(x_input, w->up_proj, up, c, inter);
|
||
lora_fwd_add(lora[5], x_input, up);
|
||
|
||
/* act = silu(gate) * up */
|
||
for (int i = 0; i < inter; i++)
|
||
act[i] = fast_silu(gate[i]) * up[i];
|
||
|
||
/* --- Backward through down_proj --- */
|
||
/* d_act = down_proj^T @ d_output + LoRA[6] backward */
|
||
float *d_act = (float*)calloc(inter, sizeof(float));
|
||
t_matvec_nt(w->down_proj, d_output, d_act, inter, c);
|
||
lora_grad_acc(lora[6], act, d_output);
|
||
lora_backward_input(lora[6], d_output, d_act);
|
||
|
||
/* --- Backward through SwiGLU activation --- */
|
||
float *d_gate = (float*)calloc(inter, sizeof(float));
|
||
float *d_up = (float*)calloc(inter, sizeof(float));
|
||
|
||
for (int i = 0; i < inter; i++) {
|
||
float sig = 1.0f / (1.0f + expf(-gate[i]));
|
||
float silu_val = gate[i] * sig;
|
||
/* d_gate = d_act * up * silu'(gate) */
|
||
/* silu'(x) = sig(x) + x * sig(x) * (1 - sig(x)) = sig(x) * (1 + x * (1 - sig(x))) */
|
||
float silu_grad = sig * (1.0f + gate[i] * (1.0f - sig));
|
||
d_gate[i] = d_act[i] * up[i] * silu_grad;
|
||
/* d_up = d_act * silu(gate) */
|
||
d_up[i] = d_act[i] * silu_val;
|
||
}
|
||
|
||
/* --- Backward through gate_proj --- */
|
||
memset(d_input, 0, c * sizeof(float));
|
||
t_matvec_nt_acc(w->gate_proj, d_gate, d_input, c, inter);
|
||
lora_grad_acc(lora[4], x_input, d_gate);
|
||
lora_backward_input(lora[4], d_gate, d_input);
|
||
|
||
/* --- Backward through up_proj --- */
|
||
t_matvec_nt_acc(w->up_proj, d_up, d_input, c, inter);
|
||
lora_grad_acc(lora[5], x_input, d_up);
|
||
lora_backward_input(lora[5], d_up, d_input);
|
||
|
||
free(gate);
|
||
free(up);
|
||
free(act);
|
||
free(d_act);
|
||
free(d_gate);
|
||
free(d_up);
|
||
}
|
||
|
||
/* ================================================================
|
||
* GQA Attention backward
|
||
*
|
||
* Forward: for each Q-head qh (with kv_head = qh / group_size):
|
||
* scores[t] = Q[qh] . K[kvh,t] * scale
|
||
* weights[t] = softmax(scores)
|
||
* y[qh] = sum_t weights[t] * V[kvh,t]
|
||
*
|
||
* K/V gradients accumulate across all Q heads sharing the same KV head.
|
||
* ================================================================ */
|
||
|
||
void rfa_gqa_attention_backward(
|
||
const float *d_y,
|
||
const float *q,
|
||
const float *k_cache,
|
||
const float *v_cache,
|
||
const float *attn_weights,
|
||
int cache_len,
|
||
float *d_q,
|
||
float *d_k_cur,
|
||
float *d_v_cur,
|
||
int n_q, int n_kv, int hd)
|
||
{
|
||
int kv_dim = n_kv * hd;
|
||
int group_size = n_q / n_kv;
|
||
float scale = 1.0f / sqrtf((float)hd);
|
||
|
||
memset(d_q, 0, n_q * hd * sizeof(float));
|
||
memset(d_k_cur, 0, kv_dim * sizeof(float));
|
||
memset(d_v_cur, 0, kv_dim * sizeof(float));
|
||
|
||
for (int qh = 0; qh < n_q; qh++) {
|
||
int kvh = qh / group_size;
|
||
const float *w_h = &attn_weights[qh * cache_len]; /* [cache_len] */
|
||
const float *d_y_h = &d_y[qh * hd];
|
||
float *d_q_h = &d_q[qh * hd];
|
||
|
||
/* d_weights[t] = dot(d_y_h, V[kvh, t]) */
|
||
float *d_weights = (float*)calloc(cache_len, sizeof(float));
|
||
for (int t = 0; t < cache_len; t++) {
|
||
float dot = 0.0f;
|
||
const float *vt = &v_cache[t * kv_dim + kvh * hd];
|
||
for (int j = 0; j < hd; j++)
|
||
dot += d_y_h[j] * vt[j];
|
||
d_weights[t] = dot;
|
||
}
|
||
|
||
/* Softmax backward: d_score = w * (d_w - sum(d_w * w)) */
|
||
float wdw_sum = 0.0f;
|
||
for (int t = 0; t < cache_len; t++)
|
||
wdw_sum += d_weights[t] * w_h[t];
|
||
|
||
for (int t = 0; t < cache_len; t++) {
|
||
float d_score = w_h[t] * (d_weights[t] - wdw_sum);
|
||
/* d_Q += scale * d_score * K[kvh, t] */
|
||
const float *kt = &k_cache[t * kv_dim + kvh * hd];
|
||
for (int j = 0; j < hd; j++)
|
||
d_q_h[j] += scale * d_score * kt[j];
|
||
}
|
||
|
||
/* K/V gradients for current token (last entry in cache) */
|
||
int self_idx = cache_len - 1;
|
||
if (self_idx >= 0) {
|
||
float d_score_self = w_h[self_idx] * (d_weights[self_idx] - wdw_sum);
|
||
for (int j = 0; j < hd; j++) {
|
||
/* K gradient accumulates across Q heads sharing this KV head */
|
||
d_k_cur[kvh * hd + j] += scale * d_score_self * q[qh * hd + j];
|
||
d_v_cur[kvh * hd + j] += w_h[self_idx] * d_y_h[j];
|
||
}
|
||
}
|
||
|
||
free(d_weights);
|
||
}
|
||
}
|
||
|
||
/* ================================================================
|
||
* Full Qwen block backward
|
||
*
|
||
* Reverse walk through:
|
||
* 1. FFN residual split
|
||
* 2. down_proj backward + LoRA[6] grads
|
||
* 3. SwiGLU backward (gate/up LoRA[4,5] grads)
|
||
* 4. Pre-FFN RMSNorm backward
|
||
* 5. Attention residual split
|
||
* 6. O projection backward + LoRA[3] grads
|
||
* 7. GQA attention backward
|
||
* 8. V projection backward + LoRA[2] grads
|
||
* 9. K projection backward + LoRA[1] grads
|
||
* 10. Q projection backward + LoRA[0] grads
|
||
* 11. Pre-attention RMSNorm backward
|
||
* ================================================================ */
|
||
|
||
void rfa_qwen_backward(
|
||
const rfa_model *model,
|
||
const rfa_state *state,
|
||
int qwen_idx,
|
||
const rfa_qwen_saved *saved,
|
||
fl_layer **lora,
|
||
const float *d_output,
|
||
float *d_input,
|
||
rfa_config *cfg)
|
||
{
|
||
int c = cfg->n_embd;
|
||
int n_q = cfg->n_q_heads;
|
||
int n_kv = cfg->n_kv_heads;
|
||
int hd = cfg->head_dim;
|
||
int kv_dim = n_kv * hd;
|
||
int inter = cfg->intermediate_size;
|
||
int group_size = n_q / n_kv;
|
||
float eps = cfg->rms_norm_eps;
|
||
|
||
rfa_qwen_weights w;
|
||
rfa_model_get_qwen_weights(model, qwen_idx, &w);
|
||
|
||
/* ---- Recompute forward intermediates from saved x_input ---- */
|
||
|
||
float *ln1_out = (float*)calloc(c, sizeof(float));
|
||
float *q_vec = (float*)calloc(c, sizeof(float));
|
||
float *k_vec = (float*)calloc(kv_dim, sizeof(float));
|
||
float *v_vec = (float*)calloc(kv_dim, sizeof(float));
|
||
float *y_attn = (float*)calloc(c, sizeof(float));
|
||
float *attn_out = (float*)calloc(c, sizeof(float));
|
||
float *x_after_attn = (float*)calloc(c, sizeof(float));
|
||
float *ln2_out = (float*)calloc(c, sizeof(float));
|
||
|
||
/* Pre-attention RMSNorm */
|
||
rms_norm_fwd(saved->x_input, w.ln1_w, ln1_out, c, eps);
|
||
|
||
/* Q/K/V projections + LoRA */
|
||
t_matvec(ln1_out, w.q_proj_w, q_vec, c, c);
|
||
if (w.q_proj_b) for (int i = 0; i < c; i++) q_vec[i] += w.q_proj_b[i];
|
||
lora_fwd_add(lora[0], ln1_out, q_vec);
|
||
|
||
t_matvec(ln1_out, w.k_proj_w, k_vec, c, kv_dim);
|
||
if (w.k_proj_b) for (int i = 0; i < kv_dim; i++) k_vec[i] += w.k_proj_b[i];
|
||
lora_fwd_add(lora[1], ln1_out, k_vec);
|
||
|
||
t_matvec(ln1_out, w.v_proj_w, v_vec, c, kv_dim);
|
||
if (w.v_proj_b) for (int i = 0; i < kv_dim; i++) v_vec[i] += w.v_proj_b[i];
|
||
lora_fwd_add(lora[2], ln1_out, v_vec);
|
||
|
||
/* Apply RoPE to Q and K (needed for correct attention computation) */
|
||
/* Use cache_len - 1 as the position for this token */
|
||
int tok_pos = saved->cache_len > 0 ? saved->cache_len - 1 : 0;
|
||
|
||
/* RoPE application (inline, matching rfa.c apply_rope) */
|
||
{
|
||
int half = hd / 2;
|
||
float theta = cfg->rope_theta;
|
||
for (int h = 0; h < n_q; h++) {
|
||
float *head = &q_vec[h * hd];
|
||
for (int i = 0; i < half; i++) {
|
||
float freq = 1.0f / powf(theta, (float)(2 * i) / hd);
|
||
float angle = tok_pos * freq;
|
||
float cos_a = cosf(angle);
|
||
float sin_a = sinf(angle);
|
||
float x0 = head[2 * i], x1 = head[2 * i + 1];
|
||
head[2 * i] = x0 * cos_a - x1 * sin_a;
|
||
head[2 * i + 1] = x0 * sin_a + x1 * cos_a;
|
||
}
|
||
}
|
||
for (int h = 0; h < n_kv; h++) {
|
||
float *head = &k_vec[h * hd];
|
||
for (int i = 0; i < half; i++) {
|
||
float freq = 1.0f / powf(theta, (float)(2 * i) / hd);
|
||
float angle = tok_pos * freq;
|
||
float cos_a = cosf(angle);
|
||
float sin_a = sinf(angle);
|
||
float x0 = head[2 * i], x1 = head[2 * i + 1];
|
||
head[2 * i] = x0 * cos_a - x1 * sin_a;
|
||
head[2 * i + 1] = x0 * sin_a + x1 * cos_a;
|
||
}
|
||
}
|
||
}
|
||
|
||
/* Recompute attention output using KV cache */
|
||
const float *k_cache, *v_cache;
|
||
int ct = rfa_state_get_kv_cache(state, qwen_idx, &k_cache, &v_cache);
|
||
int total = saved->cache_len > 0 ? saved->cache_len : ct;
|
||
if (total > ct) total = ct;
|
||
|
||
float scale = 1.0f / sqrtf((float)hd);
|
||
|
||
/* Per-head attention + save weights for backward */
|
||
float *attn_weights = (float*)calloc((size_t)n_q * total, sizeof(float));
|
||
memset(y_attn, 0, c * sizeof(float));
|
||
|
||
for (int qh = 0; qh < n_q; qh++) {
|
||
int kvh = qh / group_size;
|
||
float *q_head = &q_vec[qh * hd];
|
||
float *scores = &attn_weights[qh * total];
|
||
|
||
/* Compute scores */
|
||
float max_s = -1e30f;
|
||
for (int t = 0; t < total; t++) {
|
||
const float *kt = &k_cache[t * kv_dim + kvh * hd];
|
||
float dot = 0.0f;
|
||
for (int j = 0; j < hd; j++)
|
||
dot += q_head[j] * kt[j];
|
||
scores[t] = dot * scale;
|
||
if (scores[t] > max_s) max_s = scores[t];
|
||
}
|
||
|
||
/* Softmax */
|
||
float sum_e = 0.0f;
|
||
for (int t = 0; t < total; t++) {
|
||
scores[t] = expf(scores[t] - max_s);
|
||
sum_e += scores[t];
|
||
}
|
||
if (sum_e > 0.0f)
|
||
for (int t = 0; t < total; t++)
|
||
scores[t] /= sum_e;
|
||
|
||
/* Value aggregation */
|
||
float *y_head = &y_attn[qh * hd];
|
||
for (int t = 0; t < total; t++) {
|
||
const float *vt = &v_cache[t * kv_dim + kvh * hd];
|
||
for (int j = 0; j < hd; j++)
|
||
y_head[j] += scores[t] * vt[j];
|
||
}
|
||
}
|
||
|
||
/* O projection + LoRA */
|
||
t_matvec(y_attn, w.o_proj_w, attn_out, c, c);
|
||
lora_fwd_add(lora[3], y_attn, attn_out);
|
||
|
||
/* Attention residual */
|
||
for (int i = 0; i < c; i++)
|
||
x_after_attn[i] = saved->x_input[i] + attn_out[i];
|
||
|
||
/* Pre-FFN RMSNorm */
|
||
rms_norm_fwd(x_after_attn, w.ln2_w, ln2_out, c, eps);
|
||
|
||
/* ---- BACKWARD PASS ---- */
|
||
|
||
float *d_x1 = (float*)calloc(c, sizeof(float)); /* d w.r.t. x_after_attn */
|
||
float *d_ffn_out = (float*)calloc(c, sizeof(float));
|
||
float *d_ln2 = (float*)calloc(c, sizeof(float));
|
||
float *d_att_out = (float*)calloc(c, sizeof(float));
|
||
float *d_y = (float*)calloc(c, sizeof(float));
|
||
float *d_q = (float*)calloc(c, sizeof(float));
|
||
float *d_k_cur = (float*)calloc(kv_dim, sizeof(float));
|
||
float *d_v_cur = (float*)calloc(kv_dim, sizeof(float));
|
||
float *d_ln1 = (float*)calloc(c, sizeof(float));
|
||
|
||
/* --- 1. FFN residual: x_out = x_after_attn + ffn_out --- */
|
||
memcpy(d_x1, d_output, c * sizeof(float));
|
||
memcpy(d_ffn_out, d_output, c * sizeof(float));
|
||
|
||
/* --- 2-3. SwiGLU FFN backward (includes down, gate, up + their LoRAs) --- */
|
||
float *d_ln2_from_ffn = (float*)calloc(c, sizeof(float));
|
||
rfa_swiglu_backward(d_ffn_out, ln2_out, &w, lora, d_ln2_from_ffn, c, inter);
|
||
|
||
/* --- 4. Pre-FFN RMSNorm backward --- */
|
||
float *d_x1_from_ln = (float*)calloc(c, sizeof(float));
|
||
rfa_rms_norm_backward(d_ln2_from_ffn, x_after_attn, w.ln2_w, d_x1_from_ln, c, eps);
|
||
for (int i = 0; i < c; i++)
|
||
d_x1[i] += d_x1_from_ln[i];
|
||
|
||
/* --- 5. Attention residual: x_after_attn = x_input + attn_out --- */
|
||
memcpy(d_input, d_x1, c * sizeof(float));
|
||
memcpy(d_att_out, d_x1, c * sizeof(float));
|
||
|
||
/* --- 6. O projection backward --- */
|
||
t_matvec_nt(w.o_proj_w, d_att_out, d_y, c, c);
|
||
lora_grad_acc(lora[3], y_attn, d_att_out);
|
||
lora_backward_input(lora[3], d_att_out, d_y);
|
||
|
||
/* --- 7. GQA attention backward --- */
|
||
rfa_gqa_attention_backward(d_y, q_vec, k_cache, v_cache,
|
||
attn_weights, total,
|
||
d_q, d_k_cur, d_v_cur,
|
||
n_q, n_kv, hd);
|
||
|
||
/* --- 8-10. Q/K/V LoRA gradients --- */
|
||
/* Q: d_ln1 += Q_proj^T @ d_q (need to undo RoPE for gradient, but since
|
||
* RoPE is positional and orthogonal, we can pass d_q through RoPE inverse.
|
||
* For LoRA gradients, we use the pre-RoPE Q which is what ln1_out was projected to.) */
|
||
/* Approximate: treat RoPE as identity for LoRA gradient computation.
|
||
* This is standard practice — RoPE gradients are small and orthogonal. */
|
||
t_matvec_nt_acc(w.q_proj_w, d_q, d_ln1, c, c);
|
||
lora_grad_acc(lora[0], ln1_out, d_q);
|
||
lora_backward_input(lora[0], d_q, d_ln1);
|
||
|
||
/* K LoRA gradients */
|
||
lora_grad_acc(lora[1], ln1_out, d_k_cur);
|
||
/* V LoRA gradients */
|
||
lora_grad_acc(lora[2], ln1_out, d_v_cur);
|
||
|
||
/* --- 11. Pre-attention RMSNorm backward --- */
|
||
float *d_input_from_ln = (float*)calloc(c, sizeof(float));
|
||
rfa_rms_norm_backward(d_ln1, saved->x_input, w.ln1_w, d_input_from_ln, c, eps);
|
||
for (int i = 0; i < c; i++)
|
||
d_input[i] += d_input_from_ln[i];
|
||
|
||
/* Cleanup */
|
||
free(ln1_out);
|
||
free(q_vec);
|
||
free(k_vec);
|
||
free(v_vec);
|
||
free(y_attn);
|
||
free(attn_out);
|
||
free(x_after_attn);
|
||
free(ln2_out);
|
||
free(attn_weights);
|
||
free(d_x1);
|
||
free(d_ffn_out);
|
||
free(d_ln2);
|
||
free(d_ln2_from_ffn);
|
||
free(d_x1_from_ln);
|
||
free(d_att_out);
|
||
free(d_y);
|
||
free(d_q);
|
||
free(d_k_cur);
|
||
free(d_v_cur);
|
||
free(d_ln1);
|
||
free(d_input_from_ln);
|
||
}
|
||
|
||
/* ================================================================
|
||
* LoRA adapter creation
|
||
* ================================================================ */
|
||
|
||
fl_adapter* rfa_create_qwen_lora(const rfa_model *model, int rank, float alpha) {
|
||
if (!model) return NULL;
|
||
fl_adapter *adapter = fl_adapter_new();
|
||
rfa_config cfg = rfa_model_config(model);
|
||
int c = cfg.n_embd;
|
||
int kv_dim = cfg.n_kv_heads * cfg.head_dim;
|
||
int inter = cfg.intermediate_size;
|
||
int n_qwen = cfg.n_qwen_layers;
|
||
|
||
fl_config lora_cfg = { .rank = rank, .alpha = alpha, .dropout = 0.0f,
|
||
.fan_in_fan_out = false };
|
||
|
||
for (int q = 0; q < n_qwen; q++) {
|
||
char name[128];
|
||
/* Attention projections */
|
||
snprintf(name, sizeof(name), "qwen.%d.attn.q_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
snprintf(name, sizeof(name), "qwen.%d.attn.k_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, kv_dim, &lora_cfg));
|
||
snprintf(name, sizeof(name), "qwen.%d.attn.v_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, kv_dim, &lora_cfg));
|
||
snprintf(name, sizeof(name), "qwen.%d.attn.o_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
/* SwiGLU FFN projections */
|
||
snprintf(name, sizeof(name), "qwen.%d.ffn.gate_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, inter, &lora_cfg));
|
||
snprintf(name, sizeof(name), "qwen.%d.ffn.up_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, inter, &lora_cfg));
|
||
snprintf(name, sizeof(name), "qwen.%d.ffn.down_proj", q);
|
||
fl_adapter_add(adapter, fl_layer_new(name, inter, c, &lora_cfg));
|
||
}
|
||
return adapter;
|
||
}
|
||
|
||
tm_tensor** rfa_collect_params(fl_adapter *adapter, int *n_params_out) {
|
||
if (!adapter || !n_params_out) return NULL;
|
||
int n_layers = fl_adapter_size(adapter);
|
||
int n_params = n_layers * 2;
|
||
*n_params_out = n_params;
|
||
tm_tensor **params = (tm_tensor**)malloc(n_params * sizeof(tm_tensor*));
|
||
for (int i = 0; i < n_layers; i++) {
|
||
fl_layer *l = fl_adapter_get_idx(adapter, i);
|
||
params[i * 2] = l->lora_A;
|
||
params[i * 2 + 1] = l->lora_B;
|
||
}
|
||
return params;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training context for rfa_forward_ex callback
|
||
* ================================================================ */
|
||
|
||
typedef struct {
|
||
fl_adapter *adapter;
|
||
int *token_idx; /* [n_qwen] counter per Qwen block */
|
||
const rfa_model *model;
|
||
const rfa_state *state;
|
||
/* Saved x_input for each Qwen block × token */
|
||
float *x_inputs_flat; /* [n_qwen * max_seq_len * C] */
|
||
int *saved_cache_lens; /* [n_qwen * max_seq_len] */
|
||
int max_seq_len;
|
||
int c;
|
||
} rfa_train_ctx;
|
||
|
||
/* Callback: save x_input before Qwen block runs */
|
||
static int train_block_callback(int phase, int block_idx, int qwen_idx,
|
||
float *x, void *user_data) {
|
||
(void)block_idx;
|
||
rfa_train_ctx *ctx = (rfa_train_ctx*)user_data;
|
||
|
||
if (phase == 0) {
|
||
int t = ctx->token_idx[qwen_idx];
|
||
int idx = qwen_idx * ctx->max_seq_len + t;
|
||
memcpy(&ctx->x_inputs_flat[idx * ctx->c], x, ctx->c * sizeof(float));
|
||
|
||
const float *kc, *vc;
|
||
int cl = rfa_state_get_kv_cache(ctx->state, qwen_idx, &kc, &vc);
|
||
ctx->saved_cache_lens[idx] = cl + 1; /* +1 because forward appends */
|
||
|
||
return 0;
|
||
}
|
||
|
||
if (phase == 1) {
|
||
ctx->token_idx[qwen_idx]++;
|
||
}
|
||
|
||
return 0;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training step
|
||
* ================================================================ */
|
||
|
||
rfa_train_step_result rfa_train_step(
|
||
rfa_model *model,
|
||
rfa_state *state,
|
||
fl_adapter *adapter,
|
||
fo_optimizer *optimizer,
|
||
rfa_train_example *example,
|
||
rfa_train_config *config)
|
||
{
|
||
rfa_train_step_result result;
|
||
memset(&result, 0, sizeof(result));
|
||
|
||
struct timespec t_start, t_end;
|
||
clock_gettime(CLOCK_MONOTONIC, &t_start);
|
||
|
||
rfa_config cfg = rfa_model_config(model);
|
||
int c = cfg.n_embd;
|
||
int vocab = cfg.vocab_size;
|
||
int n_qwen = cfg.n_qwen_layers;
|
||
int seq_len = example->seq_len;
|
||
int lora_per_block = 7;
|
||
|
||
/* Allocate storage */
|
||
float *x_inputs_flat = (float*)calloc((size_t)n_qwen * seq_len * c, sizeof(float));
|
||
int *saved_cache_lens = (int*)calloc((size_t)n_qwen * seq_len, sizeof(int));
|
||
int *token_idx = (int*)calloc(n_qwen, sizeof(int));
|
||
|
||
rfa_train_ctx ctx = {
|
||
.adapter = adapter,
|
||
.token_idx = token_idx,
|
||
.model = model,
|
||
.state = state,
|
||
.x_inputs_flat = x_inputs_flat,
|
||
.saved_cache_lens = saved_cache_lens,
|
||
.max_seq_len = seq_len,
|
||
.c = c,
|
||
};
|
||
|
||
/* Reset state */
|
||
rfa_state_reset(state);
|
||
|
||
/* === FORWARD PASS === */
|
||
float *logits_buf = (float*)malloc((size_t)seq_len * vocab * sizeof(float));
|
||
float total_loss = 0.0f;
|
||
|
||
for (int t = 0; t < seq_len; t++) {
|
||
uint32_t token = example->input_ids[t];
|
||
uint32_t target = example->target_ids[t];
|
||
|
||
tm_tensor *logits = rfa_forward_ex(model, token, state,
|
||
train_block_callback, &ctx);
|
||
if (!logits) continue;
|
||
|
||
float *ldata = tm_data_f32(logits);
|
||
memcpy(&logits_buf[t * vocab], ldata, vocab * sizeof(float));
|
||
|
||
/* Cross-entropy loss */
|
||
float max_l = ldata[0];
|
||
for (int j = 1; j < vocab; j++)
|
||
if (ldata[j] > max_l) max_l = ldata[j];
|
||
float sum_exp = 0.0f;
|
||
for (int j = 0; j < vocab; j++)
|
||
sum_exp += expf(ldata[j] - max_l);
|
||
total_loss += -(ldata[target] - max_l - logf(sum_exp));
|
||
|
||
tm_free(logits);
|
||
}
|
||
|
||
result.loss = seq_len > 0 ? total_loss / seq_len : 0.0f;
|
||
|
||
/* === BACKWARD PASS === */
|
||
|
||
const float *head_w = rfa_model_get_head_weight(model);
|
||
|
||
/* Compute d_logits */
|
||
float *d_logits = (float*)calloc((size_t)seq_len * vocab, sizeof(float));
|
||
for (int t = 0; t < seq_len; t++) {
|
||
float *lgt = &logits_buf[t * vocab];
|
||
float *dlgt = &d_logits[t * vocab];
|
||
uint32_t target = example->target_ids[t];
|
||
|
||
float max_l = lgt[0];
|
||
for (int j = 1; j < vocab; j++)
|
||
if (lgt[j] > max_l) max_l = lgt[j];
|
||
float sum_exp = 0.0f;
|
||
for (int j = 0; j < vocab; j++)
|
||
sum_exp += expf(lgt[j] - max_l);
|
||
float inv_seq = 1.0f / seq_len;
|
||
for (int j = 0; j < vocab; j++)
|
||
dlgt[j] = inv_seq * (expf(lgt[j] - max_l) / sum_exp);
|
||
dlgt[target] -= inv_seq;
|
||
}
|
||
|
||
/* Backward through model for each token */
|
||
float *d_x = (float*)calloc(c, sizeof(float));
|
||
float *d_x_tmp = (float*)calloc(c, sizeof(float));
|
||
|
||
for (int t = seq_len - 1; t >= 0; t--) {
|
||
float *dlgt = &d_logits[t * vocab];
|
||
|
||
/* Head projection backward: d_x = head_w^T @ d_logits */
|
||
/* head_w is [C, vocab] in transposed layout */
|
||
t_matvec_nt(head_w, dlgt, d_x, c, vocab);
|
||
|
||
/* Walk Qwen blocks in reverse order */
|
||
for (int q = n_qwen - 1; q >= 0; q--) {
|
||
if (t >= token_idx[q]) continue;
|
||
|
||
rfa_qwen_saved saved;
|
||
memset(&saved, 0, sizeof(saved));
|
||
int idx = q * seq_len + t;
|
||
saved.x_input = (float*)malloc(c * sizeof(float));
|
||
memcpy(saved.x_input, &x_inputs_flat[idx * c], c * sizeof(float));
|
||
saved.cache_len = saved_cache_lens[idx];
|
||
|
||
fl_layer *lora[7];
|
||
for (int l = 0; l < 7; l++)
|
||
lora[l] = fl_adapter_get_idx(adapter, q * lora_per_block + l);
|
||
|
||
rfa_qwen_backward(model, state, q, &saved, lora,
|
||
d_x, d_x_tmp, &cfg);
|
||
memcpy(d_x, d_x_tmp, c * sizeof(float));
|
||
|
||
free(saved.x_input);
|
||
}
|
||
}
|
||
|
||
/* Gradient clipping + optimizer step */
|
||
if (config->grad_clip > 0.0f)
|
||
fo_clip_grad_norm(optimizer, config->grad_clip);
|
||
fo_step(optimizer);
|
||
fo_zero_grad(optimizer);
|
||
|
||
result.step = fo_get_step(optimizer);
|
||
result.lr = fo_get_lr(optimizer);
|
||
result.tokens_seen = seq_len;
|
||
|
||
clock_gettime(CLOCK_MONOTONIC, &t_end);
|
||
result.elapsed_sec = (t_end.tv_sec - t_start.tv_sec) +
|
||
(t_end.tv_nsec - t_start.tv_nsec) * 1e-9;
|
||
|
||
/* Cleanup */
|
||
free(x_inputs_flat);
|
||
free(saved_cache_lens);
|
||
free(token_idx);
|
||
free(logits_buf);
|
||
free(d_logits);
|
||
free(d_x);
|
||
free(d_x_tmp);
|
||
|
||
return result;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training loop
|
||
* ================================================================ */
|
||
|
||
void rfa_train_loop(
|
||
rfa_model *model,
|
||
rfa_train_config *config,
|
||
rfa_train_example *examples,
|
||
int n_examples)
|
||
{
|
||
if (!model || !config || !examples || n_examples == 0) return;
|
||
|
||
rfa_config cfg = rfa_model_config(model);
|
||
|
||
fprintf(stderr, "rfa_train: %d examples, %d epochs, lr=%.2e, rank=%d\n",
|
||
n_examples, config->epochs, config->lr, config->lora_rank);
|
||
fprintf(stderr, "rfa_train: %d Qwen blocks, 7 LoRA layers each\n",
|
||
cfg.n_qwen_layers);
|
||
|
||
fl_adapter *adapter = rfa_create_qwen_lora(model, config->lora_rank,
|
||
config->lora_alpha);
|
||
fl_adapter_print(adapter);
|
||
|
||
int n_params;
|
||
tm_tensor **params = rfa_collect_params(adapter, &n_params);
|
||
fprintf(stderr, "rfa_train: %d trainable tensors, %lld parameters\n",
|
||
n_params, (long long)fl_adapter_num_params(adapter));
|
||
|
||
fo_adam_config adam_cfg = fo_adam_default(config->lr);
|
||
adam_cfg.weight_decay = config->weight_decay;
|
||
fo_optimizer *opt = fo_adamw_new(params, n_params, &adam_cfg);
|
||
|
||
int total_steps = n_examples * config->epochs;
|
||
fo_scheduler_config sched_cfg;
|
||
sched_cfg.type = FO_SCHED_WARMUP_COSINE;
|
||
sched_cfg.warmup_steps = config->warmup_steps;
|
||
sched_cfg.total_steps = total_steps;
|
||
sched_cfg.min_lr = config->lr * 0.1f;
|
||
sched_cfg.step_size = 1;
|
||
sched_cfg.gamma = 0.1f;
|
||
fo_set_scheduler(opt, &sched_cfg);
|
||
|
||
int max_cache = config->max_cache > 0 ? config->max_cache : 2048;
|
||
rfa_state *state = rfa_state_new(model, max_cache);
|
||
if (!state) {
|
||
fprintf(stderr, "rfa_train: failed to create state\n");
|
||
goto cleanup;
|
||
}
|
||
|
||
if (config->output_dir) mkdir(config->output_dir, 0755);
|
||
|
||
int global_step = 0;
|
||
for (int epoch = 0; epoch < config->epochs; epoch++) {
|
||
fprintf(stderr, "\n=== Epoch %d/%d ===\n", epoch + 1, config->epochs);
|
||
|
||
/* Shuffle */
|
||
int *indices = (int*)malloc(n_examples * sizeof(int));
|
||
for (int i = 0; i < n_examples; i++) indices[i] = i;
|
||
for (int i = n_examples - 1; i > 0; i--) {
|
||
int j = rand() % (i + 1);
|
||
int tmp = indices[i]; indices[i] = indices[j]; indices[j] = tmp;
|
||
}
|
||
|
||
float epoch_loss = 0.0f;
|
||
int epoch_tokens = 0;
|
||
|
||
for (int i = 0; i < n_examples; i++) {
|
||
fo_scheduler_step(opt);
|
||
rfa_train_step_result res = rfa_train_step(
|
||
model, state, adapter, opt, &examples[indices[i]], config);
|
||
|
||
epoch_loss += res.loss * res.tokens_seen;
|
||
epoch_tokens += res.tokens_seen;
|
||
global_step++;
|
||
|
||
if (config->log_every > 0 && global_step % config->log_every == 0) {
|
||
float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
|
||
fprintf(stderr, " step %d | loss %.4f | avg %.4f | lr %.2e | %.1f tok/s\n",
|
||
global_step, res.loss, avg, res.lr,
|
||
res.elapsed_sec > 0 ? res.tokens_seen / res.elapsed_sec : 0.0);
|
||
}
|
||
|
||
if (config->save_every > 0 && config->output_dir &&
|
||
global_step % config->save_every == 0) {
|
||
char path[512];
|
||
snprintf(path, sizeof(path), "%s/checkpoint_%d.safetensors",
|
||
config->output_dir, global_step);
|
||
if (fl_adapter_save(adapter, path))
|
||
fprintf(stderr, " Saved: %s\n", path);
|
||
}
|
||
}
|
||
|
||
float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
|
||
fprintf(stderr, "Epoch %d: avg_loss=%.4f, tokens=%d\n",
|
||
epoch + 1, avg, epoch_tokens);
|
||
free(indices);
|
||
}
|
||
|
||
if (config->output_dir) {
|
||
char path[512];
|
||
snprintf(path, sizeof(path), "%s/adapter_final.safetensors",
|
||
config->output_dir);
|
||
if (fl_adapter_save(adapter, path))
|
||
fprintf(stderr, "Saved: %s\n", path);
|
||
}
|
||
|
||
cleanup:
|
||
if (state) rfa_state_free(state);
|
||
fo_free(opt);
|
||
free(params);
|
||
fl_adapter_free(adapter);
|
||
}
|
||
|
||
/* ================================================================
|
||
* LoRA adapter averaging
|
||
* ================================================================ */
|
||
|
||
fl_adapter* rfa_lora_average(fl_adapter **adapters, int n_adapters) {
|
||
if (!adapters || n_adapters <= 0) return NULL;
|
||
if (n_adapters == 1) {
|
||
/* Clone the single adapter */
|
||
fl_adapter *src = adapters[0];
|
||
int n_layers = fl_adapter_size(src);
|
||
fl_adapter *out = fl_adapter_new();
|
||
for (int i = 0; i < n_layers; i++) {
|
||
fl_layer *sl = fl_adapter_get_idx(src, i);
|
||
tm_tensor *a_clone = tm_clone(sl->lora_A);
|
||
tm_tensor *b_clone = tm_clone(sl->lora_B);
|
||
fl_layer *nl = fl_layer_from_tensors(sl->name, a_clone, b_clone,
|
||
sl->scale * sl->rank);
|
||
fl_adapter_add(out, nl);
|
||
}
|
||
return out;
|
||
}
|
||
|
||
fl_adapter *first = adapters[0];
|
||
int n_layers = fl_adapter_size(first);
|
||
fl_adapter *out = fl_adapter_new();
|
||
float inv_n = 1.0f / n_adapters;
|
||
|
||
for (int i = 0; i < n_layers; i++) {
|
||
fl_layer *ref = fl_adapter_get_idx(first, i);
|
||
int d_in = ref->d_in, d_out = ref->d_out, rank = ref->rank;
|
||
int a_size = rank * d_in;
|
||
int b_size = d_out * rank;
|
||
|
||
/* Allocate averaged tensors */
|
||
int64_t a_shape[] = {rank, d_in};
|
||
int64_t b_shape[] = {d_out, rank};
|
||
tm_tensor *avg_a = tm_zeros(a_shape, 2, TM_F32);
|
||
tm_tensor *avg_b = tm_zeros(b_shape, 2, TM_F32);
|
||
float *ad = tm_data_f32(avg_a);
|
||
float *bd = tm_data_f32(avg_b);
|
||
|
||
/* Sum across all adapters */
|
||
for (int a = 0; a < n_adapters; a++) {
|
||
fl_layer *al = fl_adapter_get_idx(adapters[a], i);
|
||
float *src_a = tm_data_f32(al->lora_A);
|
||
float *src_b = tm_data_f32(al->lora_B);
|
||
for (int j = 0; j < a_size; j++) ad[j] += src_a[j];
|
||
for (int j = 0; j < b_size; j++) bd[j] += src_b[j];
|
||
}
|
||
|
||
/* Average */
|
||
for (int j = 0; j < a_size; j++) ad[j] *= inv_n;
|
||
for (int j = 0; j < b_size; j++) bd[j] *= inv_n;
|
||
|
||
fl_layer *nl = fl_layer_from_tensors(ref->name, avg_a, avg_b,
|
||
ref->scale * ref->rank);
|
||
fl_adapter_add(out, nl);
|
||
}
|
||
|
||
return out;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Parallel training (pthread-based)
|
||
* ================================================================ */
|
||
|
||
typedef struct {
|
||
rfa_model *model; /* shared, read-only */
|
||
rfa_train_config config; /* per-thread copy */
|
||
rfa_train_example *examples; /* per-thread chunk */
|
||
int n_examples;
|
||
int thread_id;
|
||
fl_adapter *adapter; /* output: trained adapter */
|
||
} rfa_train_worker_ctx;
|
||
|
||
static void* rfa_train_worker(void *arg) {
|
||
rfa_train_worker_ctx *ctx = (rfa_train_worker_ctx *)arg;
|
||
rfa_model *model = ctx->model;
|
||
rfa_train_config *config = &ctx->config;
|
||
|
||
fprintf(stderr, "[thread %d] starting: %d examples, %d epochs\n",
|
||
ctx->thread_id, ctx->n_examples, config->epochs);
|
||
|
||
/* Each thread gets its own adapter, state, optimizer */
|
||
fl_adapter *adapter = rfa_create_qwen_lora(model, config->lora_rank,
|
||
config->lora_alpha);
|
||
|
||
int n_params;
|
||
tm_tensor **params = rfa_collect_params(adapter, &n_params);
|
||
|
||
fo_adam_config adam_cfg = fo_adam_default(config->lr);
|
||
adam_cfg.weight_decay = config->weight_decay;
|
||
fo_optimizer *opt = fo_adamw_new(params, n_params, &adam_cfg);
|
||
|
||
int total_steps = ctx->n_examples * config->epochs;
|
||
fo_scheduler_config sched_cfg;
|
||
sched_cfg.type = FO_SCHED_WARMUP_COSINE;
|
||
sched_cfg.warmup_steps = config->warmup_steps;
|
||
sched_cfg.total_steps = total_steps;
|
||
sched_cfg.min_lr = config->lr * 0.1f;
|
||
sched_cfg.step_size = 1;
|
||
sched_cfg.gamma = 0.1f;
|
||
fo_set_scheduler(opt, &sched_cfg);
|
||
|
||
int max_cache = config->max_cache > 0 ? config->max_cache : 2048;
|
||
rfa_state *state = rfa_state_new(model, max_cache);
|
||
if (!state) {
|
||
fprintf(stderr, "[thread %d] failed to create state\n", ctx->thread_id);
|
||
fo_free(opt);
|
||
free(params);
|
||
fl_adapter_free(adapter);
|
||
return NULL;
|
||
}
|
||
|
||
int global_step = 0;
|
||
for (int epoch = 0; epoch < config->epochs; epoch++) {
|
||
/* Shuffle with per-thread seed */
|
||
unsigned int seed = (unsigned int)(ctx->thread_id * 1000 + epoch * 37);
|
||
int *indices = (int*)malloc(ctx->n_examples * sizeof(int));
|
||
for (int i = 0; i < ctx->n_examples; i++) indices[i] = i;
|
||
for (int i = ctx->n_examples - 1; i > 0; i--) {
|
||
int j = (int)(rand_r(&seed) % (unsigned)(i + 1));
|
||
int tmp = indices[i]; indices[i] = indices[j]; indices[j] = tmp;
|
||
}
|
||
|
||
float epoch_loss = 0.0f;
|
||
int epoch_tokens = 0;
|
||
|
||
for (int i = 0; i < ctx->n_examples; i++) {
|
||
fo_scheduler_step(opt);
|
||
rfa_train_step_result res = rfa_train_step(
|
||
model, state, adapter, opt,
|
||
&ctx->examples[indices[i]], config);
|
||
|
||
epoch_loss += res.loss * res.tokens_seen;
|
||
epoch_tokens += res.tokens_seen;
|
||
global_step++;
|
||
|
||
if (config->log_every > 0 && global_step % config->log_every == 0) {
|
||
float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
|
||
fprintf(stderr, " [T%d] step %d | loss %.4f | avg %.4f | lr %.2e\n",
|
||
ctx->thread_id, global_step, res.loss, avg, res.lr);
|
||
}
|
||
}
|
||
|
||
float avg = epoch_tokens > 0 ? epoch_loss / epoch_tokens : 0.0f;
|
||
fprintf(stderr, "[thread %d] epoch %d: avg_loss=%.4f, tokens=%d\n",
|
||
ctx->thread_id, epoch + 1, avg, epoch_tokens);
|
||
free(indices);
|
||
}
|
||
|
||
/* Save per-thread adapter */
|
||
if (config->output_dir) {
|
||
char path[512];
|
||
snprintf(path, sizeof(path), "%s/thread_%d.safetensors",
|
||
config->output_dir, ctx->thread_id);
|
||
if (fl_adapter_save(adapter, path))
|
||
fprintf(stderr, "[thread %d] saved: %s\n", ctx->thread_id, path);
|
||
}
|
||
|
||
rfa_state_free(state);
|
||
fo_free(opt);
|
||
free(params);
|
||
|
||
ctx->adapter = adapter;
|
||
return NULL;
|
||
}
|
||
|
||
void rfa_train_parallel(
|
||
rfa_model *model,
|
||
rfa_train_config *config,
|
||
rfa_train_example *examples,
|
||
int n_examples,
|
||
int n_threads)
|
||
{
|
||
if (!model || !config || !examples || n_examples == 0 || n_threads < 1)
|
||
return;
|
||
|
||
if (n_threads == 1) {
|
||
rfa_train_loop(model, config, examples, n_examples);
|
||
return;
|
||
}
|
||
|
||
fprintf(stderr, "rfa_train_parallel: %d threads, %d examples total\n",
|
||
n_threads, n_examples);
|
||
|
||
if (config->output_dir) mkdir(config->output_dir, 0755);
|
||
|
||
/* Split examples into n_threads chunks */
|
||
int chunk_size = n_examples / n_threads;
|
||
rfa_train_worker_ctx *ctxs = (rfa_train_worker_ctx*)calloc(
|
||
n_threads, sizeof(rfa_train_worker_ctx));
|
||
pthread_t *threads = (pthread_t*)calloc(n_threads, sizeof(pthread_t));
|
||
|
||
for (int t = 0; t < n_threads; t++) {
|
||
int start = t * chunk_size;
|
||
int end = (t < n_threads - 1) ? start + chunk_size : n_examples;
|
||
|
||
ctxs[t].model = model;
|
||
ctxs[t].config = *config; /* copy config per thread */
|
||
ctxs[t].examples = &examples[start];
|
||
ctxs[t].n_examples = end - start;
|
||
ctxs[t].thread_id = t;
|
||
ctxs[t].adapter = NULL;
|
||
|
||
/* Adjust warmup proportionally */
|
||
ctxs[t].config.warmup_steps = config->warmup_steps / n_threads;
|
||
if (ctxs[t].config.warmup_steps < 1 && config->warmup_steps > 0)
|
||
ctxs[t].config.warmup_steps = 1;
|
||
|
||
pthread_create(&threads[t], NULL, rfa_train_worker, &ctxs[t]);
|
||
}
|
||
|
||
/* Wait for all threads */
|
||
for (int t = 0; t < n_threads; t++)
|
||
pthread_join(threads[t], NULL);
|
||
|
||
/* Collect adapters and average */
|
||
fl_adapter **adapters = (fl_adapter**)calloc(n_threads, sizeof(fl_adapter*));
|
||
int n_valid = 0;
|
||
for (int t = 0; t < n_threads; t++) {
|
||
if (ctxs[t].adapter)
|
||
adapters[n_valid++] = ctxs[t].adapter;
|
||
}
|
||
|
||
if (n_valid > 0) {
|
||
fprintf(stderr, "rfa_train_parallel: averaging %d adapters\n", n_valid);
|
||
fl_adapter *merged = rfa_lora_average(adapters, n_valid);
|
||
|
||
if (merged && config->output_dir) {
|
||
char path[512];
|
||
snprintf(path, sizeof(path), "%s/adapter_final.safetensors",
|
||
config->output_dir);
|
||
if (fl_adapter_save(merged, path))
|
||
fprintf(stderr, "Saved merged adapter: %s\n", path);
|
||
}
|
||
|
||
fl_adapter_free(merged);
|
||
}
|
||
|
||
/* Cleanup per-thread adapters */
|
||
for (int t = 0; t < n_threads; t++) {
|
||
if (ctxs[t].adapter)
|
||
fl_adapter_free(ctxs[t].adapter);
|
||
}
|
||
|
||
free(adapters);
|
||
free(ctxs);
|
||
free(threads);
|
||
fprintf(stderr, "rfa_train_parallel: done\n");
|
||
}
|