177 lines
6.6 KiB
C
177 lines
6.6 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* rwkvx_train.c — CLI LoRA training for RWKV-X attention blocks
|
|
*
|
|
* Usage:
|
|
* ./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
|
|
*/
|
|
|
|
#include "rwkvx/rwkvx.h"
|
|
#include "rwkvx/rx_train.h"
|
|
#include "data/data.h"
|
|
#include "tensor/tensor.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
static void print_usage(const char *prog) {
|
|
fprintf(stderr,
|
|
"Usage: %s [options]\n\n"
|
|
"Model:\n"
|
|
" --model PATH Path to .safetensors model file (required)\n"
|
|
" --vocab PATH Path to vocab file (required)\n"
|
|
" --quantize Load model with INT8 quantization\n\n"
|
|
"Data:\n"
|
|
" --data PATH Training data: .jsonl or .db (sqlite) (required)\n"
|
|
" --max-conv N Max conversations to load from sqlite (0=all)\n\n"
|
|
"LoRA:\n"
|
|
" --lora-rank N LoRA rank (default: 32)\n"
|
|
" --lora-alpha F LoRA alpha (default: 64)\n\n"
|
|
"Training:\n"
|
|
" --lr F Learning rate (default: 2e-5)\n"
|
|
" --weight-decay F AdamW weight decay (default: 0.01)\n"
|
|
" --epochs N Number of epochs (default: 3)\n"
|
|
" --max-seq-len N Max sequence length (default: 512)\n"
|
|
" --warmup-steps N LR warmup steps (default: 100)\n"
|
|
" --grad-clip F Max gradient norm (default: 1.0)\n"
|
|
" --log-every N Log every N steps (default: 10)\n"
|
|
" --save-every N Checkpoint every N steps (default: 1000)\n\n"
|
|
"Output:\n"
|
|
" --output DIR Checkpoint output directory (required)\n",
|
|
prog);
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
rx_train_config cfg = rx_train_config_defaults();
|
|
const char *model_path = NULL;
|
|
const char *vocab_path = NULL;
|
|
const char *data_path = NULL;
|
|
int max_conv = 0;
|
|
int quantize = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "--model") == 0 && i + 1 < argc)
|
|
model_path = argv[++i];
|
|
else if (strcmp(argv[i], "--vocab") == 0 && i + 1 < argc)
|
|
vocab_path = argv[++i];
|
|
else if (strcmp(argv[i], "--data") == 0 && i + 1 < argc)
|
|
data_path = argv[++i];
|
|
else if (strcmp(argv[i], "--output") == 0 && i + 1 < argc)
|
|
cfg.output_dir = argv[++i];
|
|
else if (strcmp(argv[i], "--lora-rank") == 0 && i + 1 < argc)
|
|
cfg.lora_rank = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--lora-alpha") == 0 && i + 1 < argc)
|
|
cfg.lora_alpha = (float)atof(argv[++i]);
|
|
else if (strcmp(argv[i], "--lr") == 0 && i + 1 < argc)
|
|
cfg.lr = (float)atof(argv[++i]);
|
|
else if (strcmp(argv[i], "--weight-decay") == 0 && i + 1 < argc)
|
|
cfg.weight_decay = (float)atof(argv[++i]);
|
|
else if (strcmp(argv[i], "--epochs") == 0 && i + 1 < argc)
|
|
cfg.epochs = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--max-seq-len") == 0 && i + 1 < argc)
|
|
cfg.max_seq_len = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--warmup-steps") == 0 && i + 1 < argc)
|
|
cfg.warmup_steps = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--grad-clip") == 0 && i + 1 < argc)
|
|
cfg.grad_clip = (float)atof(argv[++i]);
|
|
else if (strcmp(argv[i], "--log-every") == 0 && i + 1 < argc)
|
|
cfg.log_every = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--save-every") == 0 && i + 1 < argc)
|
|
cfg.save_every = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--max-conv") == 0 && i + 1 < argc)
|
|
max_conv = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--quantize") == 0)
|
|
quantize = 1;
|
|
else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
|
|
print_usage(argv[0]);
|
|
return 0;
|
|
} else {
|
|
fprintf(stderr, "Unknown option: %s\n", argv[i]);
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
if (!model_path || !vocab_path || !data_path || !cfg.output_dir) {
|
|
fprintf(stderr, "Error: --model, --vocab, --data, and --output are required\n\n");
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
cfg.quantize = quantize;
|
|
cfg.model_path = model_path;
|
|
cfg.vocab_path = vocab_path;
|
|
cfg.data_path = data_path;
|
|
|
|
/* Load model */
|
|
fprintf(stderr, "Loading model: %s%s\n", model_path, quantize ? " (INT8)" : "");
|
|
rx_model *model = quantize ? rx_model_load_q8(model_path) : rx_model_load(model_path);
|
|
if (!model) {
|
|
fprintf(stderr, "Failed to load model\n");
|
|
return 1;
|
|
}
|
|
rx_config mcfg = rx_model_config(model);
|
|
fprintf(stderr, "Model: n_embd=%d, n_head=%d, n_rwkv=%d, n_attn=%d, vocab=%d\n",
|
|
mcfg.n_embd, mcfg.n_head, mcfg.n_rwkv_layer, mcfg.n_attn_layer, mcfg.vocab_size);
|
|
|
|
/* Load tokenizer */
|
|
fprintf(stderr, "Loading tokenizer: %s\n", vocab_path);
|
|
rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
|
|
if (!tok) {
|
|
fprintf(stderr, "Failed to load tokenizer\n");
|
|
rx_model_free(model);
|
|
return 1;
|
|
}
|
|
|
|
/* Load dataset */
|
|
fprintf(stderr, "Loading data: %s\n", data_path);
|
|
fd_dataset *ds = NULL;
|
|
|
|
size_t path_len = strlen(data_path);
|
|
if (path_len > 3 && strcmp(data_path + path_len - 3, ".db") == 0) {
|
|
ds = fd_load_chat_sqlite(data_path, max_conv);
|
|
} else {
|
|
ds = fd_load_jsonl(data_path);
|
|
}
|
|
|
|
if (!ds || fd_len(ds) == 0) {
|
|
fprintf(stderr, "Failed to load dataset or dataset is empty\n");
|
|
rx_tokenizer_free(tok);
|
|
rx_model_free(model);
|
|
return 1;
|
|
}
|
|
fprintf(stderr, "Loaded %d samples\n", fd_len(ds));
|
|
|
|
/* Tokenize */
|
|
fprintf(stderr, "Tokenizing (max_seq_len=%d)...\n", cfg.max_seq_len);
|
|
int n_examples = 0;
|
|
rx_train_example *examples = rx_tokenize_dataset(ds, tok, cfg.max_seq_len, &n_examples);
|
|
fd_free(ds);
|
|
rx_tokenizer_free(tok);
|
|
|
|
if (!examples || n_examples == 0) {
|
|
fprintf(stderr, "No valid training examples after tokenization\n");
|
|
rx_model_free(model);
|
|
return 1;
|
|
}
|
|
fprintf(stderr, "Created %d training examples\n", n_examples);
|
|
|
|
/* Train */
|
|
rx_train_loop(model, &cfg, examples, n_examples);
|
|
|
|
/* Cleanup */
|
|
rx_train_examples_free(examples, n_examples);
|
|
rx_model_free(model);
|
|
|
|
fprintf(stderr, "Training complete.\n");
|
|
return 0;
|
|
}
|