121 lines
4.2 KiB
C
121 lines
4.2 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* rfa_lora_merge — CLI tool for averaging multiple LoRA adapters
|
|
*
|
|
* Loads N adapter safetensors files, averages A/B weights element-wise,
|
|
* saves the merged adapter. Used for cross-machine parallel training.
|
|
*
|
|
* Usage:
|
|
* ./examples/rfa_lora_merge \
|
|
* --output merged_adapter.safetensors \
|
|
* adapter_0.safetensors adapter_1.safetensors ... adapter_N.safetensors
|
|
*/
|
|
|
|
#include "rwkvx/rfa_train.h"
|
|
#include "lora/lora.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#define MAX_ADAPTERS 64
|
|
|
|
static void print_usage(const char *prog) {
|
|
fprintf(stderr, "Usage: %s --output OUTPUT ADAPTER1 ADAPTER2 [... ADAPTER_N]\n\n", prog);
|
|
fprintf(stderr, "Averages multiple LoRA adapter files element-wise.\n");
|
|
fprintf(stderr, "All adapters must have the same structure (same layers/dimensions).\n\n");
|
|
fprintf(stderr, "Options:\n");
|
|
fprintf(stderr, " --output PATH Output merged adapter path (required)\n");
|
|
fprintf(stderr, " --help Show this help message\n");
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *output_path = NULL;
|
|
const char *adapter_paths[MAX_ADAPTERS];
|
|
int n_adapters = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "--output") == 0 && i + 1 < argc) {
|
|
output_path = argv[++i];
|
|
} else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
|
|
print_usage(argv[0]);
|
|
return 0;
|
|
} else if (argv[i][0] == '-') {
|
|
fprintf(stderr, "Unknown option: %s\n", argv[i]);
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
} else {
|
|
if (n_adapters >= MAX_ADAPTERS) {
|
|
fprintf(stderr, "Error: max %d adapters supported\n", MAX_ADAPTERS);
|
|
return 1;
|
|
}
|
|
adapter_paths[n_adapters++] = argv[i];
|
|
}
|
|
}
|
|
|
|
if (!output_path || n_adapters < 2) {
|
|
fprintf(stderr, "Error: need --output and at least 2 adapter files\n\n");
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
fprintf(stderr, "Loading %d adapters...\n", n_adapters);
|
|
|
|
/* Load all adapters */
|
|
fl_config lora_cfg = fl_config_default();
|
|
fl_adapter **adapters = (fl_adapter**)calloc(n_adapters, sizeof(fl_adapter*));
|
|
for (int i = 0; i < n_adapters; i++) {
|
|
adapters[i] = fl_adapter_load(adapter_paths[i], &lora_cfg);
|
|
if (!adapters[i]) {
|
|
fprintf(stderr, "Error: failed to load %s\n", adapter_paths[i]);
|
|
for (int j = 0; j < i; j++) fl_adapter_free(adapters[j]);
|
|
free(adapters);
|
|
return 1;
|
|
}
|
|
fprintf(stderr, " [%d] %s: %d layers, %lld params\n",
|
|
i, adapter_paths[i],
|
|
fl_adapter_size(adapters[i]),
|
|
(long long)fl_adapter_num_params(adapters[i]));
|
|
}
|
|
|
|
/* Verify all have same structure */
|
|
int n_layers = fl_adapter_size(adapters[0]);
|
|
for (int i = 1; i < n_adapters; i++) {
|
|
if (fl_adapter_size(adapters[i]) != n_layers) {
|
|
fprintf(stderr, "Error: adapter %d has %d layers (expected %d)\n",
|
|
i, fl_adapter_size(adapters[i]), n_layers);
|
|
for (int j = 0; j < n_adapters; j++) fl_adapter_free(adapters[j]);
|
|
free(adapters);
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
/* Average */
|
|
fprintf(stderr, "Averaging %d adapters (%d layers each)...\n",
|
|
n_adapters, n_layers);
|
|
fl_adapter *merged = rfa_lora_average(adapters, n_adapters);
|
|
if (!merged) {
|
|
fprintf(stderr, "Error: averaging failed\n");
|
|
for (int j = 0; j < n_adapters; j++) fl_adapter_free(adapters[j]);
|
|
free(adapters);
|
|
return 1;
|
|
}
|
|
|
|
/* Save */
|
|
if (fl_adapter_save(merged, output_path)) {
|
|
fprintf(stderr, "Saved: %s (%d layers, %lld params)\n",
|
|
output_path, fl_adapter_size(merged),
|
|
(long long)fl_adapter_num_params(merged));
|
|
} else {
|
|
fprintf(stderr, "Error: failed to save %s\n", output_path);
|
|
}
|
|
|
|
/* Cleanup */
|
|
fl_adapter_free(merged);
|
|
for (int i = 0; i < n_adapters; i++)
|
|
fl_adapter_free(adapters[i]);
|
|
free(adapters);
|
|
|
|
return 0;
|
|
}
|