138 lines
3.4 KiB
Markdown
138 lines
3.4 KiB
Markdown
# rwkvx_train
|
|
|
|
LoRA training CLI for RWKV-X sparse attention (moba) blocks. Pure C, no Python required.
|
|
|
|
**Source:** `examples/rwkvx_train.c`
|
|
**Build:** `make USE_BLAS=1 rwkvx-train`
|
|
**Tests:** `make USE_BLAS=1 test-rx-train` (8 tests)
|
|
|
|
---
|
|
|
|
## Overview
|
|
|
|
Trains LoRA adapters on the 4 sparse attention (moba) blocks in RWKV-X. The 12 RWKV-7 recurrent blocks are frozen. Supports JSONL and SQLite chat history as data sources.
|
|
|
|
6 LoRA layers per attention block (Q, K, V, output, ffn_key, ffn_value) = **24 LoRA layers** total.
|
|
|
|
At rank 32: ~4.7M trainable parameters.
|
|
|
|
---
|
|
|
|
## Usage
|
|
|
|
### From JSONL data
|
|
|
|
```bash
|
|
./examples/rwkvx_train \
|
|
--model ../models/base/RWKV-X-3.6B-64k-Base.safetensors \
|
|
--vocab rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
--data ../models/datasets/oasst1_chat.jsonl \
|
|
--output ../models/loras/attn-lora/ \
|
|
--lora-rank 32 --lora-alpha 64 \
|
|
--lr 2e-5 --epochs 3 --max-seq-len 512
|
|
```
|
|
|
|
### From SQLite chat history
|
|
|
|
```bash
|
|
# Build with sqlite support
|
|
make clean && make USE_BLAS=1 HAS_SQLITE=1 rwkvx-train
|
|
|
|
./examples/rwkvx_train \
|
|
--model ../models/base/RWKV-X-3.6B-64k-Base.safetensors \
|
|
--vocab rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
--data ../path/to/conversations.db \
|
|
--output ../models/loras/chat-lora/ \
|
|
--max-conv 1000
|
|
```
|
|
|
|
### With INT8 quantized base weights
|
|
|
|
```bash
|
|
./examples/rwkvx_train \
|
|
--model ../models/base/RWKV-X-3.6B-64k-Base.safetensors \
|
|
--vocab rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
--data ../models/datasets/oasst1_chat.jsonl \
|
|
--output ../models/loras/attn-lora/ \
|
|
--quantize
|
|
```
|
|
|
|
---
|
|
|
|
## CLI Options
|
|
|
|
### Model
|
|
|
|
| Flag | Default | Description |
|
|
|------|---------|-------------|
|
|
| `--model PATH` | (required) | Path to .safetensors model file |
|
|
| `--vocab PATH` | (required) | Path to RWKV vocab file |
|
|
| `--quantize` | off | Load model with INT8 quantization |
|
|
|
|
### Data
|
|
|
|
| Flag | Default | Description |
|
|
|------|---------|-------------|
|
|
| `--data PATH` | (required) | Training data: .jsonl or .db (sqlite) |
|
|
| `--max-conv N` | 0 (all) | Max conversations to load from sqlite |
|
|
|
|
### LoRA
|
|
|
|
| Flag | Default | Description |
|
|
|------|---------|-------------|
|
|
| `--lora-rank N` | 32 | LoRA rank |
|
|
| `--lora-alpha F` | 64 | LoRA alpha scaling |
|
|
|
|
### Training
|
|
|
|
| Flag | Default | Description |
|
|
|------|---------|-------------|
|
|
| `--lr F` | 2e-5 | Learning rate |
|
|
| `--weight-decay F` | 0.01 | AdamW weight decay |
|
|
| `--epochs N` | 3 | Number of training epochs |
|
|
| `--max-seq-len N` | 512 | Max sequence length |
|
|
| `--warmup-steps N` | 100 | LR warmup steps |
|
|
| `--grad-clip F` | 1.0 | Max gradient norm |
|
|
| `--log-every N` | 10 | Log metrics every N steps |
|
|
| `--save-every N` | 1000 | Checkpoint every N steps |
|
|
|
|
### Output
|
|
|
|
| Flag | Default | Description |
|
|
|------|---------|-------------|
|
|
| `--output DIR` | (required) | Checkpoint output directory |
|
|
|
|
---
|
|
|
|
## Output Files
|
|
|
|
```
|
|
output_dir/
|
|
├── checkpoint_1000.safetensors # Step 1000
|
|
├── checkpoint_2000.safetensors # Step 2000
|
|
└── adapter_final.safetensors # End of training
|
|
```
|
|
|
|
---
|
|
|
|
## Build Variants
|
|
|
|
```bash
|
|
# Standard (BLAS only)
|
|
make clean && make USE_BLAS=1 rwkvx-train
|
|
|
|
# With sqlite data loading
|
|
make clean && make USE_BLAS=1 HAS_SQLITE=1 rwkvx-train
|
|
|
|
# With Metal GPU acceleration
|
|
make clean && make USE_BLAS=1 USE_METAL=1 rwkvx-train
|
|
```
|
|
|
|
---
|
|
|
|
## See Also
|
|
|
|
- [rwkvx-generate](rwkvx-generate.md) - Inference with trained adapters
|
|
- [rfa-train](rfa-train.md) - Training for RFA hybrid models
|
|
- [lora](lora.md) - LoRA adapter module
|