6.5 KiB

Raw Blame History

rfa_train

LoRA training CLI for RFA hybrid models (RWKV-7 + Qwen GQA). Supports multi-threaded parallel training and INT8 quantized base weights.

Source: examples/rfa_train.c Build: make USE_BLAS=1 rfa-train Tests: make USE_BLAS=1 test-rfa-train (9 tests), make USE_BLAS=1 test-rfa-parallel (6 tests)

Overview

Trains LoRA adapters on the 7 Qwen GQA blocks in the RWKV-7 + Qwen hybrid model. The 21 RWKV-7 blocks are frozen. Uses pre-tokenized binary data (from pre_tokenize.py).

7 LoRA layers per Qwen block (Q, K, V, O, gate, up, down) = 49 LoRA layers total.

At rank 32: ~9.2M trainable parameters.

Features:

Multi-threaded data-parallel training (--parallel N)
INT8 quantized base weights (--quantize)
Cross-machine LoRA adapter merging (via rfa_lora_merge)
Gradient checkpointing (recompute vs store intermediates)
AdamW optimizer with warmup + cosine annealing

Quick Start

1. Pre-tokenize data

pip install tokenizers  # one-time

python3 rwkv-fullAttention/scripts/pre_tokenize.py \
  --tokenizer models/base/ARWKV-R1-1B5/tokenizer.json \
  --data models/datasets/oasst1_chat.jsonl \
  --output models/datasets/oasst1_qwen_tokens.bin \
  --max-seq-len 256

2. Train (single-threaded)

./examples/rfa_train \
  --model ../models/base/rwkv-qwen-hybrid-1.5B.safetensors \
  --data ../models/datasets/oasst1_qwen_tokens.bin \
  --output ../models/loras/rfa-qwen-lora/ \
  --lora-rank 32 --lora-alpha 64 \
  --lr 2e-5 --epochs 3 --max-seq-len 256 --log-every 10

3. Train (parallel with INT8)

./examples/rfa_train \
  --model ../models/base/rwkv-qwen-hybrid-1.5B.safetensors \
  --data ../models/datasets/oasst1_qwen_tokens.bin \
  --output ../models/loras/rfa-qwen-lora/ \
  --lora-rank 32 --lora-alpha 64 \
  --lr 2e-5 --epochs 3 --max-seq-len 256 \
  --parallel 6 --quantize

CLI Options

Flag	Default	Description
`--model PATH`	(required)	Hybrid model safetensors file
`--data PATH`	(required)	Pre-tokenized binary data file
`--output DIR`	(required)	Output directory for LoRA checkpoints
`--lora-rank N`	32	LoRA rank
`--lora-alpha F`	64	LoRA alpha scaling
`--lr F`	2e-5	Learning rate
`--epochs N`	3	Training epochs
`--max-seq-len N`	256	Max sequence length per example
`--warmup N`	100	LR warmup steps
`--log-every N`	10	Log metrics every N steps
`--save-every N`	1000	Save checkpoint every N steps
`--grad-clip F`	1.0	Gradient clipping norm
`--max-cache N`	2048	Max KV cache per Qwen block
`--parallel N`	1	Number of training threads
`--quantize`	off	Use INT8 quantized base weights

Parallel Training

When --parallel N is set (N > 1), the training data is split into N chunks. Each thread independently trains a LoRA adapter on its chunk with its own state, optimizer, and adapter. After all threads complete, the N adapters are averaged element-wise (FedAvg) and saved as adapter_final.safetensors.

Main Thread
  1. Load model (shared, read-only)
  2. Split dataset into N chunks
  3. Spawn N pthreads
  4. Wait for all threads to finish
  5. Average N adapters -> save merged

  Thread 0          Thread 1          Thread N
  ┌─────────┐      ┌─────────┐      ┌─────────┐
  │ state    │      │ state    │      │ state    │
  │ adapter  │      │ adapter  │      │ adapter  │
  │ optimizer│      │ optimizer│      │ optimizer│
  │ chunk    │      │ chunk    │      │ chunk    │
  └─────────┘      └─────────┘      └─────────┘
       ↑ shared model (read-only) ↑

Output files with parallel training:

output_dir/
├── thread_0.safetensors          # Per-thread adapter
├── thread_1.safetensors
├── ...
└── adapter_final.safetensors     # Averaged adapter

INT8 Quantized Training

The --quantize flag quantizes all base model linear weights from f32 to INT8 (per-row symmetric). Forward-pass matvecs use NEON 4x-unrolled INT8 dot products (~3x faster). LoRA weights remain f32 for gradient accuracy.

Both f32 and INT8 weights are kept in memory — f32 is needed for backward pass gradient computation through base weights.

Pre-Tokenization

The RFA model uses Qwen's vocabulary (151K tokens). Data must be pre-tokenized:

python3 rwkv-fullAttention/scripts/pre_tokenize.py \
  --tokenizer models/base/ARWKV-R1-1B5/tokenizer.json \
  --data models/datasets/oasst1_chat.jsonl \
  --output models/datasets/oasst1_qwen_tokens.bin \
  --max-seq-len 256

Binary format: [n_examples: u32] [seq_len: u32, tokens: u32[seq_len]] x N

Dataset splitting for multi-machine training

# Equal split into N chunks
python3 rwkv-fullAttention/scripts/pre_tokenize.py \
  ... --split 2

# Weighted split (e.g., 75%/25% for M4+M3)
python3 rwkv-fullAttention/scripts/pre_tokenize.py \
  ... --split-weights 6,2

Output: <base>_0.bin, <base>_1.bin, etc.

Cross-Machine Training

Train on separate machines, then merge adapters:

# Machine 1 (M4 24GB):
./examples/rfa_train --model model.safetensors \
  --data chunk_0.bin --output loras_m4/ \
  --parallel 6 --quantize --epochs 3

# Machine 2 (M3 8GB):
./examples/rfa_train --model model.safetensors \
  --data chunk_1.bin --output loras_m3/ \
  --parallel 2 --quantize --epochs 3

# Merge:
scp m3:loras_m3/adapter_final.safetensors ./
./examples/rfa_lora_merge \
  --output merged.safetensors \
  loras_m4/adapter_final.safetensors loras_m3/adapter_final.safetensors

Memory Requirements

Config	Memory (1.5B model)
Single-thread, f32	~8 GB
Single-thread, INT8 (`--quantize`)	~9.5 GB
6 threads, INT8 (`--parallel 6 --quantize`)	~20 GB

Reduce memory with shorter sequences (--max-seq-len 128) or fewer threads.

Output Files

output_dir/
├── checkpoint_1000.safetensors   # Step 1000
├── checkpoint_2000.safetensors   # Step 2000
└── adapter_final.safetensors     # End of training (or merged from threads)

6.5 KiB Raw Blame History