465 lines
23 KiB
C
465 lines
23 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later AND Apache-2.0 */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* Portions derived from wkv7_cuda_fp32.cu in RWKV-LM
|
|
* (https://github.com/BlinkDL/RWKV-LM), by Bo Peng (BlinkDL),
|
|
* licensed under the Apache License, Version 2.0. Modified by
|
|
* David (Corey) Herrera, 2026. See NOTICE and LICENSES/Apache-2.0.txt.
|
|
*/
|
|
/*
|
|
* rx_pretrain.h — RWKV-7 pretraining kernels (backward-through-time)
|
|
*
|
|
* Prefix: rx_ (RWKV-X). This is the start of the pretrain module; for now it
|
|
* contains ONLY the WKV7 backward-through-time (the counterpart to the forward
|
|
* `rx_wkv7_one` in rwkvx.c) plus a full-sequence forward used to drive it and
|
|
* the gradient checks. The CE loop / base-unfreeze / full pretrain loop come
|
|
* later and are intentionally NOT here.
|
|
*
|
|
* WKV7 inner recurrence (per head, per timestep; identical to rx_wkv7_one):
|
|
* sa[i] = sum_j a[j] * S_prev[i][j] (a = nkk = -kk)
|
|
* S[i][j] = S_prev[i][j]*w[j] + v[i]*k[j] + sa[i]*b[j] (b = kka = kk*iclr)
|
|
* y[i] = sum_j S[i][j] * r[j]
|
|
* State element convention matches foundry `s_d[(h*N+i)*N + j]`: row i is the
|
|
* value/output dim, col j is the key dim.
|
|
*
|
|
* Backward equations transcribed from BlinkDL's Apache-2.0 CUDA
|
|
* wkv7_cuda_fp32.cu backward_kernel via the verified fp64 PyTorch oracle
|
|
* (Ternary_scale_1-10m/reference/wkv7/wkv7_reference.py).
|
|
*
|
|
* IMPORTANT chain-rule hand-offs (see wkv7_backward_reference.md §5):
|
|
* 1. `dw` here is the grad w.r.t. the EFFECTIVE decay w_vec (already in (0,1)),
|
|
* NOT the log-space pre-activation. We emit bare dw; the surrounding
|
|
* time-mix backward chains it through exp(-0.606531*sigmoid(...)). We do
|
|
* NOT apply the CUDA kernel's `dw*wi*wi_fac` factor.
|
|
* 2. `da`/`db` are grads w.r.t. the composites nkk(=a) and kka(=b). The
|
|
* recombination onto shared kk and a_vec (and a_vec's second path into
|
|
* k_mod) is the time-mix backward's job, out of scope here.
|
|
* 3. `dS` after the reverse scan is the grad w.r.t. the initial state carry
|
|
* S_{-1}. Zero it for full-sequence-from-zero; it is the true exerciser
|
|
* of the reverse scan and is checked via finite differences.
|
|
*/
|
|
#ifndef RWKVX_PRETRAIN_H
|
|
#define RWKVX_PRETRAIN_H
|
|
|
|
#include <stdint.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Full-sequence WKV7 forward. Layout for the per-timestep vectors is
|
|
* [B][T][H][N] flattened row-major: index = (((b*T + t)*H + h)*N + n).
|
|
* State layout is [B][H][N][N]: index = (((b*H + h)*N + i)*N + j).
|
|
*
|
|
* s_init : initial state carry [B*H*N*N], or NULL for a zero start.
|
|
* y_out : forward output [B*T*H*N], or NULL if not needed.
|
|
* s_final: final state after T steps [B*H*N*N], or NULL if not needed.
|
|
*
|
|
* Returns 0 on success, -1 on allocation failure.
|
|
*/
|
|
int rx_wkv7_forward_seq(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *s_init,
|
|
float *y_out, float *s_final,
|
|
int B, int T, int n_head, int head_size);
|
|
|
|
/*
|
|
* WKV7 backward-through-time (reverse scan). Same layouts as above.
|
|
*
|
|
* dy : upstream grad dL/dy [B*T*H*N].
|
|
* s_init : initial state carry used in the forward [B*H*N*N], or NULL (zero).
|
|
* dr..db : output grads [B*T*H*N] each (dw is w.r.t. effective decay, see §5).
|
|
* dS_out : output grad w.r.t. the initial state carry S_{-1} [B*H*N*N], or
|
|
* NULL if training full-sequence-from-zero (grad discarded).
|
|
*
|
|
* The forward is recomputed internally and every pre-step state S_prev is
|
|
* cached (correctness-first; the CUDA kernel reconstructs to save memory).
|
|
* Returns 0 on success, -1 on allocation failure.
|
|
*/
|
|
int rx_wkv7_backward(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *dy, const float *s_init,
|
|
float *dr, float *dw, float *dk, float *dv, float *da, float *db,
|
|
float *dS_out,
|
|
int B, int T, int n_head, int head_size);
|
|
|
|
/* ================================================================
|
|
* RWKV-7 TIME-MIX BLOCK backward (wraps the WKV7 backward above).
|
|
*
|
|
* This is the counterpart of rwkvx.c `rwkv7_time_mix_fast`: it covers the
|
|
* WHOLE time-mix block (token-shift/lerp, the R/K/V/output projections, the
|
|
* low-rank w/a/g/v LoRA gates, key L2-norm + key-modification, value-residual
|
|
* gating, the WKV7 recurrence, group-norm, the r*k*r_k bonus and the g gate),
|
|
* EXCLUDING the pre-block ln1 and the residual add. The block input is the
|
|
* (already LN'd) hidden state; the output feeds the residual.
|
|
*
|
|
* It computes dL/d(block input) and grads for every learnable time-mix param,
|
|
* internally recomputing the fp32 forward (caching activations) and delegating
|
|
* the recurrence to rx_wkv7_backward. The two §5 hand-offs live here:
|
|
* - `dw` (effective decay) is chained through exp(-0.606531*sigmoid(...)).
|
|
* - `da`/`db` (grads w.r.t. nkk=-kk and kka=kk*a_vec) are recombined onto the
|
|
* shared kk and a_vec (which also feeds k_mod).
|
|
*
|
|
* Layouts: block input x, output, dy, dx are [B][T][C] row-major (C=n_head*
|
|
* head_size). Projection weights are row-major [in,out] exactly as
|
|
* rwkvx.c's fast_matvec consumes them. v_first is [B][T][C] (per-token; for
|
|
* layer_id==0 it is an OUTPUT copy of v and its grad is ignored).
|
|
*/
|
|
typedef struct {
|
|
int c; /* n_head*head_size */
|
|
int n_head, head_size;
|
|
int Dw, Da, Dg, Dv; /* LoRA inner dims (Dv ignored when layer_id==0) */
|
|
int layer_id;
|
|
const float *x_r, *x_k, *x_v, *x_w, *x_a, *x_g; /* [c] lerps */
|
|
const float *receptance_w, *key_w, *value_w, *output_w; /* [c*c] [in,out] */
|
|
const float *w0, *w1, *w2; /* w0[c], w1[c*Dw], w2[Dw*c] */
|
|
const float *a0, *a1, *a2; /* a0[c], a1[c*Da], a2[Da*c] */
|
|
const float *g1, *g2; /* g1[c*Dg], g2[Dg*c] */
|
|
const float *v0, *v1, *v2; /* v0[c], v1[c*Dv], v2[Dv*c] (layer_id>0) */
|
|
const float *k_k, *k_a; /* [c] */
|
|
const float *r_k; /* [c] */
|
|
const float *ln_x_w, *ln_x_b;/* [c] group-norm affine */
|
|
} rx_tm_params;
|
|
|
|
/* Grad accumulators — same shapes as rx_tm_params, non-const. Every pointer
|
|
* must be caller-allocated; rx_timemix_backward ZEROES them at entry. Pointers
|
|
* for params absent in the configuration (v0/v1/v2 at layer 0) may be NULL. */
|
|
typedef struct {
|
|
float *x_r, *x_k, *x_v, *x_w, *x_a, *x_g;
|
|
float *receptance_w, *key_w, *value_w, *output_w;
|
|
float *w0, *w1, *w2;
|
|
float *a0, *a1, *a2;
|
|
float *g1, *g2;
|
|
float *v0, *v1, *v2;
|
|
float *k_k, *k_a;
|
|
float *r_k;
|
|
float *ln_x_w, *ln_x_b;
|
|
} rx_tm_grads;
|
|
|
|
/* fp32 time-mix forward over a full sequence (per-batch, T tokens).
|
|
* x : block input [B*T*c] (LN'd hidden state)
|
|
* x_prev_init : token-shift carry at t=0 [B*c], or NULL for zero.
|
|
* v_first : [B*T*c]. layer_id==0: WRITTEN (copy of v). layer_id>0: READ.
|
|
* out : block output [B*T*c] (pre-residual).
|
|
* Returns 0 on success, -1 on allocation failure. */
|
|
int rx_timemix_forward_seq(
|
|
const rx_tm_params *P, const float *x, const float *x_prev_init,
|
|
float *v_first, float *out, int B, int T);
|
|
|
|
/* Time-mix backward. dy is dL/d(out) [B*T*c]. Produces dx [B*T*c], all param
|
|
* grads in G, and (optional) d_xprev_init [B*c] / d_v_first [B*T*c]. NULLs for
|
|
* unused outputs. Returns 0 on success, -1 on allocation failure. */
|
|
int rx_timemix_backward(
|
|
const rx_tm_params *P, const float *x, const float *x_prev_init,
|
|
const float *v_first, const float *dy,
|
|
rx_tm_grads *G, float *dx, float *d_xprev_init, float *d_v_first,
|
|
int B, int T);
|
|
|
|
/* ================================================================
|
|
* RWKV-7 CHANNEL-MIX (ReLU^2 FFN) backward.
|
|
*
|
|
* Counterpart of rwkvx.c `rwkv7_channel_mix_fast`: token-shift/lerp, the key
|
|
* projection [C,ffn_dim], the squared-ReLU nonlinearity (k = relu(fk)^2), and
|
|
* the value projection [ffn_dim,C]. EXCLUDING the pre-block ln2 and the
|
|
* residual add (those live in the block). Input x is the (already LN'd)
|
|
* hidden state [B][T][c]; output feeds the residual.
|
|
*
|
|
* Layouts match the time-mix block: x/out/dy/dx are [B][T][c] row-major.
|
|
* ffn_key_w is row-major [c,ffn_dim]; ffn_value_w is [ffn_dim,c] (both as
|
|
* fast_matvec consumes them). The token-shift carry (ffn_x_prev) at t=0 is
|
|
* x_prev_init [B][c] or NULL for a zero start; its grad is d_xprev_init.
|
|
*/
|
|
typedef struct {
|
|
int c;
|
|
int ffn_dim;
|
|
const float *ffn_x_k; /* [c] token-shift lerp */
|
|
const float *ffn_key_w; /* [c*ffn_dim] [in,out] */
|
|
const float *ffn_value_w; /* [ffn_dim*c] [in,out] */
|
|
} rx_cm_params;
|
|
|
|
typedef struct {
|
|
float *ffn_x_k;
|
|
float *ffn_key_w;
|
|
float *ffn_value_w;
|
|
} rx_cm_grads;
|
|
|
|
/* fp32 channel-mix forward over a full sequence. Returns 0 / -1(alloc). */
|
|
int rx_channelmix_forward_seq(
|
|
const rx_cm_params *P, const float *x, const float *x_prev_init,
|
|
float *out, int B, int T);
|
|
|
|
/* Channel-mix backward. dy is dL/d(out) [B*T*c]. Produces dx [B*T*c], all
|
|
* param grads in G (ZEROED at entry), and optional d_xprev_init [B*c].
|
|
* Returns 0 / -1(alloc). */
|
|
int rx_channelmix_backward(
|
|
const rx_cm_params *P, const float *x, const float *x_prev_init,
|
|
const float *dy, rx_cm_grads *G, float *dx, float *d_xprev_init,
|
|
int B, int T);
|
|
|
|
/* ================================================================
|
|
* LayerNorm backward (the block's ln1/ln2; foundry `fast_layer_norm`,
|
|
* eps = 1e-5, normalised over the full C channel). Standard LN backward.
|
|
*
|
|
* ⚠ The mean/variance reductions cancel in fp32 exactly as the L2-norm did;
|
|
* they are accumulated in DOUBLE locally so gamma/beta/dx stay clean.
|
|
*
|
|
* x/dy/dx are [B][T][C] row-major. dw/db are [C] and are ZEROED at entry
|
|
* then accumulated across all B*T tokens. dx is written (=), not accumulated.
|
|
* Returns 0 / -1(alloc).
|
|
* ================================================================ */
|
|
int rx_layernorm_forward_seq(
|
|
const float *x, const float *w, const float *b, float *out,
|
|
int B, int T, int C);
|
|
|
|
int rx_layernorm_backward(
|
|
const float *x, const float *w, const float *b, const float *dy,
|
|
float *dx, float *dw, float *db, int B, int T, int C);
|
|
|
|
/* ================================================================
|
|
* RWKV-7 BLOCK forward/backward — the full residual block.
|
|
* Mirrors rwkvx.c `rwkv7_block_forward_fast`:
|
|
* x -> ln1 -> time-mix -> +x (residual) -> ln2 -> channel-mix -> +x1
|
|
* Backward walks it in reverse, threading residual grads to BOTH the block
|
|
* path and the skip. Produces dL/d(block input) + every block param grad
|
|
* (time-mix + channel-mix + ln1 + ln2).
|
|
* ================================================================ */
|
|
typedef struct {
|
|
rx_tm_params tm;
|
|
rx_cm_params cm;
|
|
const float *ln1_w, *ln1_b; /* [c] */
|
|
const float *ln2_w, *ln2_b; /* [c] */
|
|
} rx_block_params;
|
|
|
|
typedef struct {
|
|
rx_tm_grads tm;
|
|
rx_cm_grads cm;
|
|
float *ln1_w, *ln1_b, *ln2_w, *ln2_b;
|
|
} rx_block_grads;
|
|
|
|
/* Block forward. v_first [B*T*c]: WRITTEN at layer 0 (copy of v), READ at
|
|
* layer>0. tm_xprev_init / cm_xprev_init are the two token-shift carries
|
|
* [B*c] (over ln1-out and ln2-out respectively) or NULL for zero. Returns
|
|
* 0 / -1(alloc). */
|
|
int rx_block_forward_seq(
|
|
const rx_block_params *P, const float *x,
|
|
const float *tm_xprev_init, const float *cm_xprev_init,
|
|
float *v_first, float *out, int B, int T);
|
|
|
|
/* Block backward. dy is dL/d(out) [B*T*c]. Produces dx [B*T*c], all grads in
|
|
* G (sub-grads ZEROED by their sub-backwards; ln grads zeroed here), and the
|
|
* optional carry / v_first grads. Returns 0 / -1(alloc).
|
|
*
|
|
* d_v_first is dual-purpose (see rx_timemix_backward): for a layer>0 block it is
|
|
* an OUTPUT (grad w.r.t. the v_first input, zeroed then written); for the
|
|
* layer-0 block (which PRODUCES v_first) it is an optional INPUT carrying the
|
|
* accumulated downstream grad on v_first, injected into the value path. Pass
|
|
* NULL when there is no v_first coupling (single RWKV layer / legacy). */
|
|
int rx_block_backward(
|
|
const rx_block_params *P, const float *x,
|
|
const float *tm_xprev_init, const float *cm_xprev_init,
|
|
const float *v_first, const float *dy,
|
|
rx_block_grads *G, float *dx,
|
|
float *d_tm_xprev_init, float *d_cm_xprev_init, float *d_v_first,
|
|
int B, int T);
|
|
|
|
/* ================================================================
|
|
* SPARSE (MoBA-style) ATTENTION BLOCK forward/backward — BASE weights.
|
|
*
|
|
* Counterpart of rwkvx.c `attn_block_forward_fast` for the full residual
|
|
* attention block (RWKV-x interleaves these with the RWKV-7 blocks above):
|
|
* h1 = ln1(x)
|
|
* q,k,v = h1 @ {Wq,Wk,Wv} (base projections, [c,c])
|
|
* att = MoBA-sparse causal-softmax attention(q,k,v) scale = 1/sqrt(head_size)
|
|
* att_out = att @ Wo ; x1 = x + att_out (residual)
|
|
* h2 = ln2(x1)
|
|
* cm = channel-mix(h2) ; out = x1 + cm (residual)
|
|
*
|
|
* This computes dL/d(block input) and grads for EVERY learnable attention-block
|
|
* base parameter (Q/K/V/O projections, ln1/ln2 affine, and the channel-mix
|
|
* params) — the full base, not just the LoRA of rx_train.c. The softmax scale
|
|
* is the fixed constant 1/sqrt(head_size) (not learnable); the forward has no
|
|
* RoPE, so there is nothing to invert there.
|
|
*
|
|
* MoBA top-k CHUNK SELECTION IS HARD / NON-DIFFERENTIABLE. The selected chunk
|
|
* set is treated as FIXED and the backward runs the standard sparse-attention
|
|
* gradient over exactly the selected keys (keys of unselected chunks get
|
|
* attention weight 0, so they drop out of every reduction automatically). For
|
|
* a finite-difference check, perturbations must not flip the selection — test
|
|
* in a regime where it is stable (e.g. moba_top_k >= #candidate chunks so all
|
|
* are selected => plain causal attention). See tests/test_attn_bwd.c.
|
|
*
|
|
* Selection rule (per query t, per head): keys are partitioned into chunks of
|
|
* `chunk_size`. The chunk containing t is always attended (causally, keys
|
|
* s<=t). Each earlier complete chunk gets a gate = <q_t, mean_key(chunk)>; the
|
|
* top `moba_top_k` earlier chunks are selected. moba_top_k <= 0 or
|
|
* >= #candidate chunks both mean "keep all" (full causal).
|
|
* ================================================================ */
|
|
typedef struct {
|
|
int c; /* n_head*head_size */
|
|
int n_head, head_size;
|
|
int chunk_size; /* MoBA chunk size (>=1) */
|
|
int moba_top_k; /* # earlier chunks kept; <=0 => keep all */
|
|
const float *ln1_w, *ln1_b; /* [c] pre-attention LN affine */
|
|
const float *receptance_w; /* [c*c] Q proj [in,out] */
|
|
const float *key_w; /* [c*c] K proj */
|
|
const float *value_w; /* [c*c] V proj */
|
|
const float *output_w; /* [c*c] O proj */
|
|
const float *ln2_w, *ln2_b; /* [c] pre-FFN LN affine */
|
|
rx_cm_params cm; /* channel-mix (ffn_x_k, ffn_key_w, ffn_value_w) */
|
|
} rx_attn_params;
|
|
|
|
typedef struct {
|
|
float *ln1_w, *ln1_b;
|
|
float *receptance_w, *key_w, *value_w, *output_w;
|
|
float *ln2_w, *ln2_b;
|
|
rx_cm_grads cm;
|
|
} rx_attn_grads;
|
|
|
|
/* fp32 forward over a full sequence (per batch, causal, MoBA-sparse).
|
|
* x : block input [B*T*c]
|
|
* cm_xprev_init : channel-mix token-shift carry at t=0 [B*c] or NULL(zero)
|
|
* out : block output [B*T*c]
|
|
* Returns 0 / -1(alloc). */
|
|
int rx_attn_forward_seq(
|
|
const rx_attn_params *P, const float *x, const float *cm_xprev_init,
|
|
float *out, int B, int T);
|
|
|
|
/* Attention-block backward with BASE-weight gradients. dy is dL/d(out)
|
|
* [B*T*c]. Produces dx [B*T*c] and every param grad in G (all ZEROED at
|
|
* entry), plus optional d_cm_xprev_init [B*c]. The forward is recomputed
|
|
* internally; the MoBA selection is recomputed from the base weights and held
|
|
* fixed through the attention backward. Returns 0 / -1(alloc). */
|
|
int rx_attn_backward_full(
|
|
const rx_attn_params *P, const float *x, const float *cm_xprev_init,
|
|
const float *dy, rx_attn_grads *G, float *dx, float *d_cm_xprev_init,
|
|
int B, int T);
|
|
|
|
/* Exported test hooks: the ISOLATED MoBA attention core (scores->softmax->AV),
|
|
* mirrored on GPU in cuda_kernels/moba_attn.cu and parity-checked by
|
|
* cuda_kernels/tests/test_moba_attn_cuda.cu. Inputs are the already-projected
|
|
* per-token q,k,v [B*T*c]; forward fills softmax weights p [B*n_head*T*T] and
|
|
* attention output y [B*T*c]; backward recomputes the MoBA selection + p from
|
|
* q,k,v and accumulates dq,dk,dv [B*T*c] (all zeroed at entry) from dy=d(y). */
|
|
int rx_attn_score_av_forward(const rx_attn_params *P, const float *q, const float *k,
|
|
const float *v, int B, int T, float *p, float *y);
|
|
int rx_attn_score_av_backward(const rx_attn_params *P, const float *q, const float *k,
|
|
const float *v, const float *dy, int B, int T,
|
|
float *dq, float *dk, float *dv);
|
|
|
|
/* ================================================================
|
|
* FULL RWKV-x MODEL: embedding, head/unembedding, cross-entropy, from-scratch
|
|
* init, and the whole-model forward/backward assembly. Single sequence
|
|
* (B=1); the underlying block kernels are called per-sequence.
|
|
*
|
|
* Architecture mirrors rwkvx.c rx_forward:
|
|
* x0 = emb[token] (raw embedding; foundry fuses ln0 into
|
|
* emb at load, so a from-scratch random
|
|
* emb absorbs it — no separate ln0 here)
|
|
* for each block (RWKV-7 and sparse-attention, interleaved per the RWKV-x
|
|
* pattern — interval = n_rwkv/n_attn, attention after each RWKV group):
|
|
* x = block(x) (residual stream)
|
|
* h = ln_out(x) (final LayerNorm)
|
|
* logits = h @ head_w ([C,vocab]; head is NOT weight-tied to
|
|
* emb — foundry stores rwkv.head.weight
|
|
* separately)
|
|
* loss = mean_t CE(logits[t], target[t])
|
|
*
|
|
* v_first coupling: the RWKV-7 block at rwkv-array index 0 produces v_first;
|
|
* later RWKV-7 blocks consume it. The reverse pass accumulates the grad on
|
|
* v_first from the consumers and injects it into producer block 0 via the
|
|
* rx_timemix_backward d_v_first-as-input path (see above).
|
|
* ================================================================ */
|
|
|
|
/* ---- Embedding (gather / scatter-add) ---- */
|
|
void rx_embedding_forward(const float *emb, const int *tokens,
|
|
int T, int C, float *out); /* out[T*C] */
|
|
void rx_embedding_backward(const int *tokens, int T, int C, int vocab,
|
|
const float *d_out, float *d_emb); /* d_emb ZEROED then scatter-add */
|
|
|
|
/* ---- Head / unembedding. head is [C,vocab] row-major (as fast_matvec eats). ---- */
|
|
void rx_head_forward(const float *x, const float *head,
|
|
int T, int C, int vocab, float *logits); /* logits[T*vocab] */
|
|
void rx_head_backward(const float *x, const float *head, const float *dlogits,
|
|
int T, int C, int vocab,
|
|
float *dx, float *d_head); /* dx written(=); d_head ZEROED then accumulate */
|
|
|
|
/* ---- Fused log-softmax cross-entropy over the vocab. Mean over valid tokens
|
|
* (targets != ignore_index; pass ignore_index<0 to disable). log-sum-exp is
|
|
* double-accumulated for numerical cleanliness. ---- */
|
|
double rx_cross_entropy_forward(const float *logits, const int *targets,
|
|
int T, int vocab, int ignore_index);
|
|
/* Returns loss and writes dlogits = (softmax - onehot)/n_valid (ignored token
|
|
* rows are zero). */
|
|
double rx_cross_entropy_backward(const float *logits, const int *targets,
|
|
int T, int vocab, int ignore_index,
|
|
float *dlogits); /* dlogits[T*vocab] */
|
|
|
|
/* ---- From-scratch model ---- */
|
|
typedef struct {
|
|
int n_embd, n_head, head_size; /* head_size*n_head == n_embd */
|
|
int n_rwkv_layer, n_attn_layer;
|
|
int vocab_size;
|
|
int Dw, Da, Dg, Dv; /* time-mix LoRA inner dims */
|
|
int ffn_dim; /* channel-mix hidden dim */
|
|
int chunk_size, moba_top_k; /* attention MoBA config (>=candidates => full causal) */
|
|
} rx_pt_config;
|
|
|
|
typedef struct rx_ptmodel rx_ptmodel; /* opaque; owns all params + grads */
|
|
|
|
/* Deterministic seeded init of ALL params (splitmix64 PRNG). Init scheme:
|
|
* linear projections & LoRA matrices : uniform(-0.5,0.5) * 0.3
|
|
* token-shift / gate lerps, w0/a0/v0/k_k/k_a/r_k : uniform(-0.5,0.5) * 0.5
|
|
* LayerNorm / group-norm gamma : 1 + 0.2*uniform(-0.5,0.5)
|
|
* LayerNorm / group-norm beta : 0.2*uniform(-0.5,0.5)
|
|
* embedding : uniform(-0.5,0.5) * 0.5
|
|
* head : uniform(-0.5,0.5) * 0.3
|
|
* Small-scale, non-degenerate, keeps CE finite; a production trainer would
|
|
* refine these. Returns NULL on bad config / alloc failure. */
|
|
rx_ptmodel* rx_model_new_random(const rx_pt_config *pc, uint64_t seed);
|
|
void rx_ptmodel_free(rx_ptmodel *m);
|
|
|
|
/* Introspection for the gradient check / reporting. */
|
|
int rx_ptmodel_num_params(const rx_ptmodel *m);
|
|
/* Per-tensor accessor: name, mutable data, matching analytic grad, element count. */
|
|
int rx_ptmodel_param(rx_ptmodel *m, int i, const char **name,
|
|
float **data, float **grad, int *n);
|
|
rx_pt_config rx_ptmodel_config(const rx_ptmodel *m);
|
|
/* Direct views onto a block's live (possibly-perturbed) params — for an fp64
|
|
* mirror / gradient check. idx is the index within the RWKV / attn array. */
|
|
const rx_block_params* rx_ptmodel_rwkv(rx_ptmodel *m, int idx);
|
|
const rx_attn_params* rx_ptmodel_attn(rx_ptmodel *m, int idx);
|
|
|
|
/* ---- Activation cache for backward ---- */
|
|
typedef struct {
|
|
int T, C, vocab, n_blocks;
|
|
int *tokens; /* [T] */
|
|
float *block_in; /* [(n_blocks+1)*T*C]: entry residual per block; [n_blocks]=final */
|
|
float *v_first; /* [T*C] */
|
|
float *lnout; /* [T*C] ln_out output (head input) */
|
|
float *logits; /* [T*vocab] */
|
|
} rx_pt_acts;
|
|
|
|
rx_pt_acts* rx_pt_acts_new(const rx_ptmodel *m, int T);
|
|
void rx_pt_acts_free(rx_pt_acts *a);
|
|
|
|
/* Full forward: gather emb -> interleaved blocks -> ln_out -> head -> logits.
|
|
* Fills acts (incl. logits). Returns 0 / -1(alloc). */
|
|
int rx_pretrain_forward(rx_ptmodel *m, const int *tokens, int T,
|
|
rx_pt_acts *acts);
|
|
|
|
/* Full backward: CE -> head -> ln_out -> reverse blocks (interleaved) ->
|
|
* embedding scatter-add. Accumulates every param grad into the model's grad
|
|
* buffers (zeroed by the sub-backwards / here) and returns the scalar loss. */
|
|
double rx_pretrain_backward(rx_ptmodel *m, const rx_pt_acts *acts,
|
|
const int *targets, int T, int ignore_index);
|
|
|
|
/* Zero all analytic grad buffers (call before a backward if reusing). */
|
|
void rx_ptmodel_zero_grad(rx_ptmodel *m);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* RWKVX_PRETRAIN_H */
|