217 lines
7.7 KiB
C
217 lines
7.7 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* rfa_train.h — LoRA training for RFA hybrid model (Qwen blocks only)
|
|
*
|
|
* Trains Qwen GQA attention blocks to work with frozen RWKV-7 blocks.
|
|
* 7 LoRA layers per Qwen block: Q, K, V, O, gate, up, down.
|
|
* Gradient flows through RWKV blocks via residual connections.
|
|
*
|
|
* Prefix: rfa_ (for RFA training)
|
|
*/
|
|
|
|
#ifndef RFA_TRAIN_H
|
|
#define RFA_TRAIN_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include "rwkvx/rfa.h"
|
|
#include "lora/lora.h"
|
|
#include "optim/optim.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Training example (pre-tokenized binary format)
|
|
* ================================================================ */
|
|
|
|
typedef struct {
|
|
uint32_t *input_ids; /* [seq_len] */
|
|
uint32_t *target_ids; /* [seq_len] */
|
|
int seq_len;
|
|
} rfa_train_example;
|
|
|
|
/* Free a single training example */
|
|
void rfa_train_example_free(rfa_train_example *ex);
|
|
|
|
/* Free array of training examples */
|
|
void rfa_train_examples_free(rfa_train_example *examples, int n);
|
|
|
|
/* Load pre-tokenized binary file (from pre_tokenize.py).
|
|
* Format: [n_examples: u32] [seq_len: u32, tokens: u32[seq_len]] x N
|
|
* Creates input/target pairs shifted by 1.
|
|
* Returns malloc'd array, sets *n_examples_out. */
|
|
rfa_train_example* rfa_load_pretokenized(
|
|
const char *path, int max_seq_len, int *n_examples_out
|
|
);
|
|
|
|
/* ================================================================
|
|
* Saved activations (for backward pass through Qwen blocks)
|
|
* ================================================================ */
|
|
|
|
typedef struct {
|
|
float *x_input; /* [C] block input (saved during forward) */
|
|
int cache_len; /* KV cache length at this token */
|
|
} rfa_qwen_saved;
|
|
|
|
/* ================================================================
|
|
* Training config
|
|
* ================================================================ */
|
|
|
|
typedef struct {
|
|
/* LoRA */
|
|
int lora_rank; /* LoRA rank (default: 32) */
|
|
float lora_alpha; /* LoRA alpha (default: 64) */
|
|
|
|
/* Training */
|
|
float lr; /* Learning rate (default: 2e-5) */
|
|
float weight_decay; /* AdamW weight decay (default: 0.01) */
|
|
int epochs; /* Number of epochs (default: 3) */
|
|
int max_seq_len; /* Max sequence length (default: 256) */
|
|
int warmup_steps; /* LR warmup steps (default: 100) */
|
|
int save_every; /* Checkpoint every N steps (default: 1000) */
|
|
int log_every; /* Log every N steps (default: 10) */
|
|
float grad_clip; /* Max gradient norm (default: 1.0) */
|
|
int max_cache; /* Max KV cache per Qwen block (default: 2048) */
|
|
int n_threads; /* Parallel training threads (default: 1) */
|
|
bool quantize; /* Use INT8 base weights (default: false) */
|
|
|
|
/* Paths */
|
|
const char *model_path;
|
|
const char *data_path; /* Pre-tokenized binary file */
|
|
const char *output_dir; /* Checkpoint output directory */
|
|
} rfa_train_config;
|
|
|
|
/* Return default training config */
|
|
rfa_train_config rfa_train_config_defaults(void);
|
|
|
|
/* ================================================================
|
|
* Backward pass primitives
|
|
* ================================================================ */
|
|
|
|
/* RMSNorm backward: given d_output, original input x, weight
|
|
* Computes d_x (gradient to propagate through) */
|
|
void rfa_rms_norm_backward(
|
|
const float *d_output, /* [C] gradient from above */
|
|
const float *x, /* [C] original input to RMSNorm */
|
|
const float *weight, /* [C] RMSNorm weight */
|
|
float *d_x, /* [C] gradient to propagate (output) */
|
|
int c, float eps
|
|
);
|
|
|
|
/* SwiGLU backward: given d_output [C], computes d_input [C].
|
|
* Requires recomputed gate/up activations. */
|
|
void rfa_swiglu_backward(
|
|
const float *d_output, /* [C] gradient from above */
|
|
const float *x_input, /* [C] input to SwiGLU */
|
|
const rfa_qwen_weights *w,
|
|
fl_layer **lora, /* [7] LoRA layers: Q,K,V,O,gate,up,down */
|
|
float *d_input, /* [C] gradient to propagate (output) */
|
|
int c, int inter
|
|
);
|
|
|
|
/* GQA attention backward: given d_y [C], computes d_q, d_k, d_v.
|
|
* Uses saved attention weights from forward. */
|
|
void rfa_gqa_attention_backward(
|
|
const float *d_y, /* [C] gradient into attention output */
|
|
const float *q, /* [C] query vector */
|
|
const float *k_cache, /* [cache_len, kv_dim] */
|
|
const float *v_cache, /* [cache_len, kv_dim] */
|
|
const float *attn_weights, /* [n_q_heads, cache_len] saved softmax */
|
|
int cache_len,
|
|
float *d_q, /* [C] output */
|
|
float *d_k_cur, /* [kv_dim] output: gradient for current token's K */
|
|
float *d_v_cur, /* [kv_dim] output: gradient for current token's V */
|
|
int n_q, int n_kv, int hd
|
|
);
|
|
|
|
/* ================================================================
|
|
* Full Qwen block backward
|
|
* ================================================================ */
|
|
|
|
/* Backward through one Qwen block. Computes LoRA gradients.
|
|
* d_output: [C] gradient from above
|
|
* d_input: [C] gradient to propagate down (output) */
|
|
void rfa_qwen_backward(
|
|
const rfa_model *model,
|
|
const rfa_state *state,
|
|
int qwen_idx,
|
|
const rfa_qwen_saved *saved,
|
|
fl_layer **lora, /* [7] LoRA layers for this block */
|
|
const float *d_output, /* [C] gradient from above */
|
|
float *d_input, /* [C] gradient to propagate down */
|
|
rfa_config *cfg
|
|
);
|
|
|
|
/* ================================================================
|
|
* Training step result
|
|
* ================================================================ */
|
|
|
|
typedef struct {
|
|
int step;
|
|
float loss;
|
|
float lr;
|
|
int tokens_seen;
|
|
double elapsed_sec;
|
|
} rfa_train_step_result;
|
|
|
|
/* ================================================================
|
|
* LoRA adapter + training
|
|
* ================================================================ */
|
|
|
|
/* Create LoRA adapter for Qwen blocks.
|
|
* Creates 7 LoRA layers per Qwen block: Q, K, V, O, gate, up, down.
|
|
* Returns adapter with all layers. */
|
|
fl_adapter* rfa_create_qwen_lora(const rfa_model *model, int rank, float alpha);
|
|
|
|
/* Collect all trainable parameters from adapter.
|
|
* Returns malloc'd array of tm_tensor* pointers (A and B per layer).
|
|
* Sets *n_params_out. */
|
|
tm_tensor** rfa_collect_params(fl_adapter *adapter, int *n_params_out);
|
|
|
|
/* Single training step: forward + loss + backward + optimizer 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
|
|
);
|
|
|
|
/* Full training loop over all examples for all epochs. */
|
|
void rfa_train_loop(
|
|
rfa_model *model,
|
|
rfa_train_config *config,
|
|
rfa_train_example *examples,
|
|
int n_examples
|
|
);
|
|
|
|
/* ================================================================
|
|
* Parallel training + LoRA averaging
|
|
* ================================================================ */
|
|
|
|
/* Average N LoRA adapters element-wise.
|
|
* All adapters must have the same structure (same layers, same dimensions).
|
|
* Returns a new adapter with averaged A and B matrices. */
|
|
fl_adapter* rfa_lora_average(fl_adapter **adapters, int n_adapters);
|
|
|
|
/* Parallel training: spawns n_threads workers, each training on a chunk
|
|
* of the dataset. Merges adapters via averaging when all complete.
|
|
* Equivalent to calling rfa_train_loop on each chunk, then averaging. */
|
|
void rfa_train_parallel(
|
|
rfa_model *model,
|
|
rfa_train_config *config,
|
|
rfa_train_example *examples,
|
|
int n_examples,
|
|
int n_threads
|
|
);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* RFA_TRAIN_H */
|