213 lines
6.5 KiB
Markdown
213 lines
6.5 KiB
Markdown
# 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
|
|
|
|
```bash
|
|
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)
|
|
|
|
```bash
|
|
./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)
|
|
|
|
```bash
|
|
./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:
|
|
|
|
```bash
|
|
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
|
|
|
|
```bash
|
|
# 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:
|
|
|
|
```bash
|
|
# 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)
|
|
```
|
|
|
|
---
|
|
|
|
## See Also
|
|
|
|
- [rfa-generate](rfa-generate.md) - Inference with trained adapters
|
|
- [rfa-lora-merge](rfa-lora-merge.md) - Merge adapters from multiple machines
|
|
- [rwkvx-train](rwkvx-train.md) - Training for RWKV-X models
|
|
- [lora](lora.md) - LoRA adapter module
|