/* 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 #include #include 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; }