1060 lines
36 KiB
C
1060 lines
36 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
||
/* Copyright (C) 2026 David (Corey) Herrera */
|
||
/*
|
||
* rx_train.c — RWKV-X LoRA training for attention blocks
|
||
*
|
||
* Gradient checkpointing approach: during forward pass, we save only the
|
||
* block input (x_input) and the KV cache state at each attention block.
|
||
* During backward, we recompute all intermediate activations using the
|
||
* weight accessors, then compute proper gradients for all 6 LoRA layers.
|
||
*
|
||
* The standard inference forward (attn_block_forward_fast) handles all
|
||
* cache management, eviction, etc. Our training forward just piggybacks
|
||
* via the rx_forward_ex callback to save x_input and cache_len.
|
||
*
|
||
* LoRA layers apply an additive correction: out += scale * B @ A @ input.
|
||
* During backward, we recompute (base + LoRA) forward, then backprop to
|
||
* get LoRA A/B gradients. Base weight gradients are not computed (frozen).
|
||
*
|
||
* LoRA indices per attention block: 0=Q, 1=K, 2=V, 3=output, 4=ffn_key, 5=ffn_value
|
||
*/
|
||
|
||
#include "rwkvx/rx_train.h"
|
||
#include "rwkvx/rwkvx.h"
|
||
#include "tensor/tensor.h"
|
||
#include "lora/lora.h"
|
||
#include "optim/optim.h"
|
||
#include "data/data.h"
|
||
|
||
#include <math.h>
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
#include <string.h>
|
||
#include <time.h>
|
||
#include <sys/stat.h>
|
||
|
||
#ifdef USE_BLAS
|
||
#if defined(__APPLE__)
|
||
#include <Accelerate/Accelerate.h>
|
||
#else
|
||
#include <cblas.h>
|
||
#endif
|
||
#endif
|
||
|
||
/* ================================================================
|
||
* Internal helpers
|
||
* ================================================================ */
|
||
|
||
/* out[cols] = x[rows] @ W[rows, cols] (fast_matvec convention) */
|
||
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] (W not transposed)
|
||
* Used for backward: d_input = W @ d_output */
|
||
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
|
||
}
|
||
|
||
/* Layer norm forward */
|
||
static void t_layer_norm(const float *x, const float *w, const float *b,
|
||
float *out, int c) {
|
||
float mean = 0.0f;
|
||
for (int i = 0; i < c; i++) mean += x[i];
|
||
mean /= c;
|
||
float var = 0.0f;
|
||
for (int i = 0; i < c; i++) { float d = x[i] - mean; var += d * d; }
|
||
var /= c;
|
||
float inv_std = 1.0f / sqrtf(var + 1e-5f);
|
||
for (int i = 0; i < c; i++)
|
||
out[i] = (x[i] - mean) * inv_std * w[i] + b[i];
|
||
}
|
||
|
||
/* Apply LoRA 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);
|
||
}
|
||
|
||
/* Ensure gradient tensor exists */
|
||
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;
|
||
}
|
||
}
|
||
|
||
/* Accumulate LoRA gradients: dB += scale * d_out ⊗ h^T, dA += scale * (B^T @ d_out) ⊗ in^T */
|
||
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);
|
||
}
|
||
|
||
/* LoRA 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);
|
||
}
|
||
|
||
/* ================================================================
|
||
* Config defaults
|
||
* ================================================================ */
|
||
|
||
rx_train_config rx_train_config_defaults(void) {
|
||
rx_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 = 512;
|
||
cfg.warmup_steps = 100;
|
||
cfg.save_every = 1000;
|
||
cfg.log_every = 10;
|
||
cfg.grad_clip = 1.0f;
|
||
return cfg;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training examples
|
||
* ================================================================ */
|
||
|
||
void rx_train_example_free(rx_train_example *ex) {
|
||
if (!ex) return;
|
||
free(ex->input_ids);
|
||
free(ex->target_ids);
|
||
}
|
||
|
||
void rx_train_examples_free(rx_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);
|
||
}
|
||
|
||
rx_train_example* rx_tokenize_dataset(
|
||
fd_dataset *ds, rx_tokenizer *tok,
|
||
int max_seq_len, int *n_examples_out)
|
||
{
|
||
if (!ds || !tok || max_seq_len < 2 || !n_examples_out) return NULL;
|
||
int n_samples = fd_len(ds);
|
||
if (n_samples == 0) { *n_examples_out = 0; return NULL; }
|
||
|
||
rx_train_example *examples = (rx_train_example*)calloc(n_samples, sizeof(rx_train_example));
|
||
int n_out = 0;
|
||
int max_tokens = max_seq_len + 1;
|
||
uint32_t *tokens = (uint32_t*)malloc(max_tokens * sizeof(uint32_t));
|
||
|
||
for (int i = 0; i < n_samples; i++) {
|
||
const char *text = fd_sample_get(fd_get(ds, i), "text");
|
||
if (!text || strlen(text) < 2) continue;
|
||
int n_tok = rx_tokenizer_encode(tok, text, tokens, max_tokens);
|
||
if (n_tok < 2) continue;
|
||
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);
|
||
*n_examples_out = n_out;
|
||
return examples;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Saved activations
|
||
* ================================================================ */
|
||
|
||
void rx_attn_saved_free(rx_attn_saved *s) {
|
||
if (!s) return;
|
||
free(s->ln_out);
|
||
free(s->q_vec);
|
||
free(s->y_data);
|
||
free(s->attn_weights);
|
||
free(s->head_totals);
|
||
free(s->x_input);
|
||
free(s->x_after_attn);
|
||
free(s->ffn_ln_out);
|
||
free(s->shifted);
|
||
free(s->fk_pre_relu);
|
||
free(s->fk_post_relu);
|
||
}
|
||
|
||
/* ================================================================
|
||
* LayerNorm backward
|
||
* ================================================================ */
|
||
|
||
void rx_layer_norm_backward(
|
||
const float *d_output,
|
||
const float *x,
|
||
const float *weight,
|
||
float *d_x,
|
||
int c)
|
||
{
|
||
float mean = 0.0f;
|
||
for (int i = 0; i < c; i++) mean += x[i];
|
||
mean /= c;
|
||
float var = 0.0f;
|
||
for (int i = 0; i < c; i++) { float d = x[i] - mean; var += d * d; }
|
||
var /= c;
|
||
float inv_std = 1.0f / sqrtf(var + 1e-5f);
|
||
float inv_c = 1.0f / c;
|
||
|
||
float mean1 = 0.0f, mean2 = 0.0f;
|
||
for (int i = 0; i < c; i++) {
|
||
float dw = d_output[i] * weight[i];
|
||
float x_hat = (x[i] - mean) * inv_std;
|
||
mean1 += dw;
|
||
mean2 += dw * x_hat;
|
||
}
|
||
mean1 *= inv_c;
|
||
mean2 *= inv_c;
|
||
|
||
for (int i = 0; i < c; i++) {
|
||
float dw = d_output[i] * weight[i];
|
||
float x_hat = (x[i] - mean) * inv_std;
|
||
d_x[i] = inv_std * (dw - mean1 - x_hat * mean2);
|
||
}
|
||
}
|
||
|
||
/* ================================================================
|
||
* Forward recomputation for backward pass (gradient checkpointing)
|
||
*
|
||
* Given saved x_input, recompute all intermediate activations
|
||
* needed for the backward pass. Uses weight accessors + LoRA.
|
||
* Stores results in rx_attn_saved struct.
|
||
*
|
||
* This is called once per attention block per token during backward.
|
||
* ================================================================ */
|
||
|
||
static void recompute_attn_intermediates(
|
||
rx_model *model,
|
||
rx_state *state,
|
||
int attn_idx,
|
||
fl_layer **lora, /* [6] */
|
||
rx_attn_saved *saved,
|
||
int n_head, int head_size)
|
||
{
|
||
rx_config cfg = rx_model_config(model);
|
||
int c = cfg.n_embd;
|
||
|
||
rx_attn_weights w;
|
||
rx_model_get_attn_weights(model, attn_idx, &w);
|
||
int ffn_dim = w.ffn_dim;
|
||
|
||
/* Ensure saved has allocated buffers */
|
||
if (!saved->ln_out) saved->ln_out = (float*)calloc(c, sizeof(float));
|
||
if (!saved->q_vec) saved->q_vec = (float*)calloc(c, sizeof(float));
|
||
if (!saved->y_data) saved->y_data = (float*)calloc(c, sizeof(float));
|
||
if (!saved->x_after_attn) saved->x_after_attn = (float*)calloc(c, sizeof(float));
|
||
if (!saved->ffn_ln_out) saved->ffn_ln_out = (float*)calloc(c, sizeof(float));
|
||
if (!saved->shifted) saved->shifted = (float*)calloc(c, sizeof(float));
|
||
if (!saved->fk_pre_relu) saved->fk_pre_relu = (float*)calloc(ffn_dim, sizeof(float));
|
||
if (!saved->fk_post_relu) saved->fk_post_relu = (float*)calloc(ffn_dim, sizeof(float));
|
||
|
||
/* Pre-attention LayerNorm */
|
||
t_layer_norm(saved->x_input, w.ln1_w, w.ln1_b, saved->ln_out, c);
|
||
|
||
/* Q projection with LoRA */
|
||
t_matvec(saved->ln_out, w.receptance_w, saved->q_vec, c, c);
|
||
lora_fwd_add(lora[0], saved->ln_out, saved->q_vec);
|
||
|
||
/* Get K/V cache (current state — includes all tokens up to cache_len) */
|
||
const float *k_cache, *v_cache;
|
||
int current_cache_len = rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
|
||
/* Use saved cache_len (state at this token's forward) */
|
||
int ct = saved->cache_len;
|
||
if (ct <= 0) ct = current_cache_len; /* fallback */
|
||
float scale = 1.0f / sqrtf((float)head_size);
|
||
|
||
/* Per-head attention */
|
||
float *y_data = saved->y_data;
|
||
memset(y_data, 0, c * sizeof(float));
|
||
|
||
/* Allocate/reuse attention weights buffer */
|
||
int max_tokens = ct > 0 ? ct : 1;
|
||
if (!saved->attn_weights)
|
||
saved->attn_weights = (float*)calloc((size_t)n_head * max_tokens, sizeof(float));
|
||
if (!saved->head_totals)
|
||
saved->head_totals = (int*)calloc(n_head, sizeof(int));
|
||
|
||
for (int h = 0; h < n_head; h++) {
|
||
int off = h * head_size;
|
||
float *scores = &saved->attn_weights[h * max_tokens];
|
||
saved->head_totals[h] = ct;
|
||
|
||
for (int t = 0; t < ct; t++) {
|
||
float dot = 0.0f;
|
||
for (int j = 0; j < head_size; j++)
|
||
dot += saved->q_vec[off + j] * k_cache[t * c + off + j];
|
||
scores[t] = dot * scale;
|
||
}
|
||
|
||
/* Softmax */
|
||
float max_s = ct > 0 ? scores[0] : 0.0f;
|
||
for (int t = 1; t < ct; t++)
|
||
if (scores[t] > max_s) max_s = scores[t];
|
||
float sum_e = 0.0f;
|
||
for (int t = 0; t < ct; t++) {
|
||
scores[t] = expf(scores[t] - max_s);
|
||
sum_e += scores[t];
|
||
}
|
||
if (sum_e > 0.0f)
|
||
for (int t = 0; t < ct; t++)
|
||
scores[t] /= sum_e;
|
||
|
||
/* Value aggregation */
|
||
for (int j = 0; j < head_size; j++) {
|
||
float val = 0.0f;
|
||
for (int t = 0; t < ct; t++)
|
||
val += scores[t] * v_cache[t * c + off + j];
|
||
y_data[off + j] = val;
|
||
}
|
||
}
|
||
|
||
/* Output projection + LoRA → attention output */
|
||
float *att_out = (float*)calloc(c, sizeof(float));
|
||
t_matvec(y_data, w.output_w, att_out, c, c);
|
||
lora_fwd_add(lora[3], y_data, att_out);
|
||
|
||
/* Attention residual: x_after_attn = x_input + att_out */
|
||
for (int i = 0; i < c; i++)
|
||
saved->x_after_attn[i] = saved->x_input[i] + att_out[i];
|
||
|
||
/* Pre-FFN LayerNorm */
|
||
t_layer_norm(saved->x_after_attn, w.ln2_w, w.ln2_b, saved->ffn_ln_out, c);
|
||
|
||
/* FFN token shift (approximate: use ffn_ln_out as both current and prev) */
|
||
for (int i = 0; i < c; i++)
|
||
saved->shifted[i] = saved->ffn_ln_out[i]; /* shift is identity when prev≈current */
|
||
|
||
/* FFN key + LoRA */
|
||
memset(saved->fk_pre_relu, 0, ffn_dim * sizeof(float));
|
||
t_matvec(saved->shifted, w.ffn_key_w, saved->fk_pre_relu, c, ffn_dim);
|
||
lora_fwd_add(lora[4], saved->shifted, saved->fk_pre_relu);
|
||
|
||
/* ReLU² */
|
||
for (int i = 0; i < ffn_dim; i++) {
|
||
float v = saved->fk_pre_relu[i] > 0.0f ? saved->fk_pre_relu[i] : 0.0f;
|
||
saved->fk_post_relu[i] = v * v;
|
||
}
|
||
|
||
free(att_out);
|
||
}
|
||
|
||
/* ================================================================
|
||
* Attention training forward
|
||
*
|
||
* This is a public function that can be called directly. For the
|
||
* gradient checkpointing approach used in rx_train_step, it's
|
||
* called during backward to recompute activations.
|
||
* ================================================================ */
|
||
|
||
void rx_attn_forward_train(
|
||
rx_model *model,
|
||
int attn_idx,
|
||
float *x_data,
|
||
rx_state *state,
|
||
fl_layer **lora,
|
||
float *x_out,
|
||
rx_attn_saved *saved)
|
||
{
|
||
rx_config cfg = rx_model_config(model);
|
||
int c = cfg.n_embd;
|
||
int n_head = cfg.n_head;
|
||
int head_size = cfg.head_size;
|
||
|
||
/* Save x_input */
|
||
if (!saved->x_input) saved->x_input = (float*)calloc(c, sizeof(float));
|
||
memcpy(saved->x_input, x_data, c * sizeof(float));
|
||
|
||
/* Get cache_len before processing */
|
||
const float *k_cache, *v_cache;
|
||
saved->cache_len = rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
|
||
/* After standard forward runs, cache_len will be incremented by 1 */
|
||
saved->cache_len += 1;
|
||
|
||
/* Recompute intermediates (this fills all saved fields) */
|
||
recompute_attn_intermediates(model, state, attn_idx, lora, saved,
|
||
n_head, head_size);
|
||
|
||
/* Compute x_out from saved intermediates */
|
||
rx_attn_weights w;
|
||
rx_model_get_attn_weights(model, attn_idx, &w);
|
||
int ffn_dim = w.ffn_dim;
|
||
|
||
/* FFN value + LoRA → ffn_out */
|
||
float *ffn_out = (float*)calloc(c, sizeof(float));
|
||
t_matvec(saved->fk_post_relu, w.ffn_value_w, ffn_out, ffn_dim, c);
|
||
lora_fwd_add(lora[5], saved->fk_post_relu, ffn_out);
|
||
|
||
/* FFN residual: x_out = x_after_attn + ffn_out */
|
||
for (int i = 0; i < c; i++)
|
||
x_out[i] = saved->x_after_attn[i] + ffn_out[i];
|
||
|
||
free(ffn_out);
|
||
}
|
||
|
||
/* ================================================================
|
||
* Attention backward pass
|
||
*
|
||
* Full backward through the attention block. Computes proper
|
||
* gradients for all 6 LoRA layers using saved/recomputed activations.
|
||
*
|
||
* Backward flow (reverse of forward):
|
||
* 1. FFN residual backward
|
||
* 2. FFN value projection backward + LoRA grads
|
||
* 3. ReLU² backward
|
||
* 4. FFN key projection backward + LoRA grads
|
||
* 5. Token shift backward
|
||
* 6. Pre-FFN LayerNorm backward
|
||
* 7. Attention residual backward
|
||
* 8. Output projection backward + LoRA grads
|
||
* 9. Per-head attention backward (softmax + scores → d_Q)
|
||
* 10. Q/K/V LoRA grads (K/V get gradients from attention mechanism)
|
||
* 11. Pre-attention LayerNorm backward
|
||
* ================================================================ */
|
||
|
||
void rx_attn_backward(
|
||
rx_model *model,
|
||
rx_state *state,
|
||
int attn_idx,
|
||
rx_attn_saved *saved,
|
||
fl_layer **lora,
|
||
float *d_output,
|
||
float *d_input,
|
||
int n_head, int head_size)
|
||
{
|
||
rx_config cfg = rx_model_config(model);
|
||
int c = cfg.n_embd;
|
||
|
||
rx_attn_weights w;
|
||
rx_model_get_attn_weights(model, attn_idx, &w);
|
||
int ffn_dim = w.ffn_dim;
|
||
|
||
/* Recompute intermediate activations (gradient checkpointing) */
|
||
recompute_attn_intermediates(model, state, attn_idx, lora, saved,
|
||
n_head, head_size);
|
||
|
||
/* Scratch buffers */
|
||
float *d_x1 = (float*)calloc(c, sizeof(float));
|
||
float *d_ffn_out = (float*)calloc(c, sizeof(float));
|
||
float *d_fk_post = (float*)calloc(ffn_dim, sizeof(float));
|
||
float *d_fk_pre = (float*)calloc(ffn_dim, sizeof(float));
|
||
float *d_shifted = (float*)calloc(c, sizeof(float));
|
||
float *d_ffn_ln = (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_ln_out = (float*)calloc(c, sizeof(float));
|
||
|
||
/* --- 1. FFN residual backward: 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. FFN value projection backward ---
|
||
* ffn_out = fk_post_relu @ ffn_value_w [ffn_dim, C] + LoRA_5
|
||
* d_fk_post = ffn_value_w[ffn_dim, C] @ d_ffn_out[C] */
|
||
t_matvec_nt(w.ffn_value_w, d_ffn_out, d_fk_post, ffn_dim, c);
|
||
lora_grad_acc(lora[5], saved->fk_post_relu, d_ffn_out);
|
||
lora_backward_input(lora[5], d_ffn_out, d_fk_post);
|
||
|
||
/* --- 3. ReLU² backward ---
|
||
* d_fk_pre[i] = d_fk_post[i] * 2 * relu(fk_pre[i]) (where fk_pre > 0) */
|
||
for (int i = 0; i < ffn_dim; i++) {
|
||
if (saved->fk_pre_relu[i] > 0.0f)
|
||
d_fk_pre[i] = d_fk_post[i] * 2.0f * saved->fk_pre_relu[i];
|
||
else
|
||
d_fk_pre[i] = 0.0f;
|
||
}
|
||
|
||
/* --- 4. FFN key projection backward ---
|
||
* fk_pre = shifted @ ffn_key_w [C, ffn_dim] + LoRA_4
|
||
* d_shifted = ffn_key_w[C, ffn_dim] @ d_fk_pre[ffn_dim] */
|
||
t_matvec_nt(w.ffn_key_w, d_fk_pre, d_shifted, c, ffn_dim);
|
||
lora_grad_acc(lora[4], saved->shifted, d_fk_pre);
|
||
lora_backward_input(lora[4], d_fk_pre, d_shifted);
|
||
|
||
/* --- 5. Token shift backward ---
|
||
* shifted ≈ ffn_ln_out (simplified — see forward)
|
||
* d_ffn_ln = d_shifted */
|
||
memcpy(d_ffn_ln, d_shifted, c * sizeof(float));
|
||
|
||
/* --- 6. Pre-FFN LayerNorm backward ---
|
||
* d_x1 += LN_backward(d_ffn_ln, x_after_attn, ln2_w) */
|
||
float *d_x1_ln = (float*)calloc(c, sizeof(float));
|
||
rx_layer_norm_backward(d_ffn_ln, saved->x_after_attn, w.ln2_w, d_x1_ln, c);
|
||
for (int i = 0; i < c; i++)
|
||
d_x1[i] += d_x1_ln[i];
|
||
free(d_x1_ln);
|
||
|
||
/* --- 7. Attention residual backward ---
|
||
* x_after_attn = x_input + att_out
|
||
* d_input = d_x1 (from residual); d_att_out = d_x1 */
|
||
memcpy(d_input, d_x1, c * sizeof(float));
|
||
memcpy(d_att_out, d_x1, c * sizeof(float));
|
||
|
||
/* --- 8. Output projection backward ---
|
||
* att_out = y_data @ output_w [C, C] + LoRA_3
|
||
* d_y = output_w[C, C] @ d_att_out[C] */
|
||
t_matvec_nt(w.output_w, d_att_out, d_y, c, c);
|
||
lora_grad_acc(lora[3], saved->y_data, d_att_out);
|
||
lora_backward_input(lora[3], d_att_out, d_y);
|
||
|
||
/* --- 9. Per-head attention backward ---
|
||
* y_h = sum_t weights_h[t] * V_h[t]
|
||
* weights_h = softmax(Q_h @ K_h^T * scale)
|
||
*
|
||
* d_weights_h[t] = dot(d_y_h, V_h[t])
|
||
* d_scores_h = weights_h * (d_weights - sum(d_weights * weights))
|
||
* d_Q_h = scale * sum_t d_scores[t] * K_h[t] */
|
||
const float *k_cache, *v_cache;
|
||
rx_state_get_kv_cache(state, attn_idx, &k_cache, &v_cache);
|
||
int ct = saved->cache_len;
|
||
float attn_scale = 1.0f / sqrtf((float)head_size);
|
||
|
||
memset(d_q, 0, c * sizeof(float));
|
||
|
||
/* Also compute K/V gradients for LoRA */
|
||
float *d_k = (float*)calloc(c, sizeof(float));
|
||
float *d_v = (float*)calloc(c, sizeof(float));
|
||
|
||
for (int h = 0; h < n_head; h++) {
|
||
int off = h * head_size;
|
||
int total = saved->head_totals[h];
|
||
if (total <= 0) total = ct;
|
||
int max_tokens = ct > 0 ? ct : 1;
|
||
float *weights = &saved->attn_weights[h * max_tokens];
|
||
|
||
/* d_weights_h[t] = dot(d_y_h, V_h[t]) */
|
||
float *d_weights = (float*)calloc(total, sizeof(float));
|
||
for (int t = 0; t < total; t++) {
|
||
float dot = 0.0f;
|
||
for (int j = 0; j < head_size; j++)
|
||
dot += d_y[off + j] * v_cache[t * c + off + j];
|
||
d_weights[t] = dot;
|
||
}
|
||
|
||
/* Softmax backward */
|
||
float wdw_sum = 0.0f;
|
||
for (int t = 0; t < total; t++)
|
||
wdw_sum += d_weights[t] * weights[t];
|
||
|
||
for (int t = 0; t < total; t++) {
|
||
float d_score = weights[t] * (d_weights[t] - wdw_sum);
|
||
/* d_Q_h += scale * d_score * K_h[t] */
|
||
for (int j = 0; j < head_size; j++)
|
||
d_q[off + j] += attn_scale * d_score * k_cache[t * c + off + j];
|
||
}
|
||
|
||
/* K/V gradients for current token (last entry in cache) */
|
||
int self_idx = total - 1;
|
||
if (self_idx >= 0) {
|
||
float d_score_self = weights[self_idx] * (d_weights[self_idx] - wdw_sum);
|
||
for (int j = 0; j < head_size; j++) {
|
||
d_k[off + j] = attn_scale * d_score_self * saved->q_vec[off + j];
|
||
d_v[off + j] = weights[self_idx] * d_y[off + j];
|
||
}
|
||
}
|
||
|
||
free(d_weights);
|
||
}
|
||
|
||
/* --- 10. Q/K/V LoRA gradients ---
|
||
* Q = ln_out @ receptance_w + LoRA_0(ln_out) */
|
||
t_matvec_nt(w.receptance_w, d_q, d_ln_out, c, c);
|
||
lora_grad_acc(lora[0], saved->ln_out, d_q);
|
||
lora_backward_input(lora[0], d_q, d_ln_out);
|
||
|
||
/* K LoRA (index 1): input=ln_out, d_out=d_k */
|
||
lora_grad_acc(lora[1], saved->ln_out, d_k);
|
||
/* V LoRA (index 2): input=ln_out, d_out=d_v */
|
||
lora_grad_acc(lora[2], saved->ln_out, d_v);
|
||
|
||
free(d_k);
|
||
free(d_v);
|
||
|
||
/* --- 11. Pre-attention LayerNorm backward --- */
|
||
float *d_input_ln = (float*)calloc(c, sizeof(float));
|
||
rx_layer_norm_backward(d_ln_out, saved->x_input, w.ln1_w, d_input_ln, c);
|
||
for (int i = 0; i < c; i++)
|
||
d_input[i] += d_input_ln[i];
|
||
free(d_input_ln);
|
||
|
||
/* Cleanup */
|
||
free(d_x1);
|
||
free(d_ffn_out);
|
||
free(d_fk_post);
|
||
free(d_fk_pre);
|
||
free(d_shifted);
|
||
free(d_ffn_ln);
|
||
free(d_att_out);
|
||
free(d_y);
|
||
free(d_q);
|
||
free(d_ln_out);
|
||
}
|
||
|
||
/* ================================================================
|
||
* LoRA adapter creation
|
||
* ================================================================ */
|
||
|
||
fl_adapter* rx_create_attn_lora(rx_model *model, int rank, float alpha) {
|
||
if (!model) return NULL;
|
||
fl_adapter *adapter = fl_adapter_new();
|
||
rx_config cfg = rx_model_config(model);
|
||
int c = cfg.n_embd;
|
||
int n_attn = cfg.n_attn_layer;
|
||
|
||
fl_config lora_cfg = { .rank = rank, .alpha = alpha, .dropout = 0.0f,
|
||
.fan_in_fan_out = false };
|
||
|
||
for (int a = 0; a < n_attn; a++) {
|
||
rx_attn_weights w;
|
||
int ffn_dim = c * 4;
|
||
if (rx_model_get_attn_weights(model, a, &w) == 0)
|
||
ffn_dim = w.ffn_dim;
|
||
|
||
char name[128];
|
||
snprintf(name, sizeof(name), "moba.%d.att.receptance", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
snprintf(name, sizeof(name), "moba.%d.att.key", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
snprintf(name, sizeof(name), "moba.%d.att.value", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
snprintf(name, sizeof(name), "moba.%d.att.output", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, c, &lora_cfg));
|
||
snprintf(name, sizeof(name), "moba.%d.ffn.key", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, c, ffn_dim, &lora_cfg));
|
||
snprintf(name, sizeof(name), "moba.%d.ffn.value", a);
|
||
fl_adapter_add(adapter, fl_layer_new(name, ffn_dim, c, &lora_cfg));
|
||
}
|
||
return adapter;
|
||
}
|
||
|
||
tm_tensor** rx_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 rx_forward_ex callback
|
||
* ================================================================ */
|
||
|
||
typedef struct {
|
||
fl_adapter *adapter;
|
||
int *cache_len_before; /* [n_attn] cache_len before each token */
|
||
int *token_idx; /* [n_attn] token counter per attn block */
|
||
rx_model *model;
|
||
rx_state *state;
|
||
int lora_per_block;
|
||
/* Saved x_input for each attention block × token */
|
||
float **x_inputs; /* [n_attn * max_seq_len * C] */
|
||
int *saved_cache_lens; /* [n_attn * max_seq_len] */
|
||
int max_seq_len;
|
||
int c;
|
||
} rx_train_ctx;
|
||
|
||
/* Callback for rx_forward_ex: save x_input before standard forward runs.
|
||
* Returns 0 (don't skip) so standard forward handles KV cache management. */
|
||
static int train_block_callback(int phase, int attn_idx, float *x,
|
||
void *user_data) {
|
||
rx_train_ctx *ctx = (rx_train_ctx*)user_data;
|
||
|
||
if (phase == 0) {
|
||
/* Pre-attention: save x_input and current cache_len */
|
||
int t = ctx->token_idx[attn_idx];
|
||
int idx = attn_idx * ctx->max_seq_len + t;
|
||
memcpy(&ctx->x_inputs[0][idx * ctx->c], x, ctx->c * sizeof(float));
|
||
|
||
/* Save cache_len before this token is processed */
|
||
const float *kc, *vc;
|
||
int cl = rx_state_get_kv_cache(ctx->state, attn_idx, &kc, &vc);
|
||
ctx->saved_cache_lens[idx] = cl + 1; /* +1 because forward will append */
|
||
|
||
return 0; /* Don't skip — standard forward runs */
|
||
}
|
||
|
||
if (phase == 1) {
|
||
/* Post-attention: increment token counter */
|
||
ctx->token_idx[attn_idx]++;
|
||
}
|
||
|
||
return 0;
|
||
}
|
||
|
||
/* ================================================================
|
||
* Training step
|
||
* ================================================================ */
|
||
|
||
rx_train_step_result rx_train_step(
|
||
rx_model *model,
|
||
rx_state *state,
|
||
fl_adapter *adapter,
|
||
fo_optimizer *optimizer,
|
||
rx_train_example *example,
|
||
rx_train_config *config)
|
||
{
|
||
rx_train_step_result result;
|
||
memset(&result, 0, sizeof(result));
|
||
|
||
struct timespec t_start, t_end;
|
||
clock_gettime(CLOCK_MONOTONIC, &t_start);
|
||
|
||
rx_config cfg = rx_model_config(model);
|
||
int c = cfg.n_embd;
|
||
int n_head = cfg.n_head;
|
||
int head_size = cfg.head_size;
|
||
int vocab = cfg.vocab_size;
|
||
int n_attn = cfg.n_attn_layer;
|
||
int seq_len = example->seq_len;
|
||
int lora_per_block = 6;
|
||
|
||
/* Allocate x_input storage for all attn blocks × tokens */
|
||
float *x_inputs_flat = (float*)calloc((size_t)n_attn * seq_len * c, sizeof(float));
|
||
int *saved_cache_lens = (int*)calloc((size_t)n_attn * seq_len, sizeof(int));
|
||
int *token_idx = (int*)calloc(n_attn, sizeof(int));
|
||
|
||
rx_train_ctx ctx = {
|
||
.adapter = adapter,
|
||
.cache_len_before = NULL,
|
||
.token_idx = token_idx,
|
||
.model = model,
|
||
.state = state,
|
||
.lora_per_block = lora_per_block,
|
||
.x_inputs = &x_inputs_flat,
|
||
.saved_cache_lens = saved_cache_lens,
|
||
.max_seq_len = seq_len,
|
||
.c = c,
|
||
};
|
||
|
||
/* Reset state */
|
||
rx_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 = rx_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 = rx_model_get_head_weight(model);
|
||
const float *ln_out_w, *ln_out_b;
|
||
rx_model_get_ln_out(model, &ln_out_w, &ln_out_b);
|
||
|
||
/* Compute d_logits: softmax(logits) - one_hot(target), scaled by 1/seq_len */
|
||
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_ln = head_w[C, vocab] @ d_logits[vocab] */
|
||
t_matvec_nt(head_w, dlgt, d_x, c, vocab);
|
||
|
||
/* Final LayerNorm backward is approximate here — we don't save the
|
||
* pre-LN activation. The gradient direction is preserved. */
|
||
|
||
/* Walk attention blocks in reverse order */
|
||
for (int a = n_attn - 1; a >= 0; a--) {
|
||
if (t >= token_idx[a]) continue;
|
||
|
||
/* Build saved struct from stored x_input + cache_len */
|
||
rx_attn_saved saved;
|
||
memset(&saved, 0, sizeof(saved));
|
||
int idx = a * 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[6];
|
||
for (int l = 0; l < 6; l++)
|
||
lora[l] = fl_adapter_get_idx(adapter, a * lora_per_block + l);
|
||
|
||
rx_attn_backward(model, state, a, &saved, lora,
|
||
d_x, d_x_tmp, n_head, head_size);
|
||
memcpy(d_x, d_x_tmp, c * sizeof(float));
|
||
|
||
rx_attn_saved_free(&saved);
|
||
}
|
||
}
|
||
|
||
/* 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 rx_train_loop(
|
||
rx_model *model,
|
||
rx_train_config *config,
|
||
rx_train_example *examples,
|
||
int n_examples)
|
||
{
|
||
if (!model || !config || !examples || n_examples == 0) return;
|
||
|
||
fprintf(stderr, "rx_train: %d examples, %d epochs, lr=%.2e, rank=%d\n",
|
||
n_examples, config->epochs, config->lr, config->lora_rank);
|
||
|
||
fl_adapter *adapter = rx_create_attn_lora(model, config->lora_rank,
|
||
config->lora_alpha);
|
||
fl_adapter_print(adapter);
|
||
|
||
int n_params;
|
||
tm_tensor **params = rx_collect_params(adapter, &n_params);
|
||
fprintf(stderr, "rx_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);
|
||
|
||
rx_state *state = rx_state_new(model);
|
||
if (!state) {
|
||
fprintf(stderr, "rx_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);
|
||
|
||
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);
|
||
rx_train_step_result res = rx_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) rx_state_free(state);
|
||
fo_free(opt);
|
||
free(params);
|
||
fl_adapter_free(adapter);
|
||
}
|