diff --git a/CMakeLists.txt b/CMakeLists.txt index b687c1c1..6776da45 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -308,6 +308,7 @@ endif() add_library(jllama SHARED src/main/cpp/jllama.cpp src/main/cpp/tts_engine.cpp + src/main/cpp/train_engine.cpp ${JLLAMA_TTS_GEN_CPP} src/main/cpp/utils.hpp ${llama.cpp_SOURCE_DIR}/tools/server/server-common.cpp diff --git a/src/main/cpp/train_engine.cpp b/src/main/cpp/train_engine.cpp new file mode 100644 index 00000000..eade0556 --- /dev/null +++ b/src/main/cpp/train_engine.cpp @@ -0,0 +1,165 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +#include "train_engine.h" + +#include "common.h" +#include "ggml-opt.h" +#include "llama.h" + +#include + +#include + +#include +#include +#include +#include +#include + +using json = nlohmann::json; + +namespace jllama_train { + +bool finetune(const finetune_config &cfg, std::string &err) { + common_params params; + params.escape = false; + params.model.path = cfg.model_path; + params.out_file = cfg.output_path; + params.n_ctx = cfg.n_ctx; + params.n_gpu_layers = cfg.n_gpu_layers; + params.val_split = cfg.val_split; + if (cfg.n_batch > 0) { + params.n_batch = cfg.n_batch; + } + if (cfg.n_ubatch > 0) { + params.n_ubatch = cfg.n_ubatch; + } + + params.optimizer = + cfg.optimizer == 1 ? GGML_OPT_OPTIMIZER_TYPE_SGD : GGML_OPT_OPTIMIZER_TYPE_ADAMW; + params.lr.lr0 = cfg.learning_rate; + params.lr.lr_min = cfg.lr_min; + params.lr.decay_epochs = cfg.decay_epochs; + params.lr.wd = cfg.weight_decay; + params.lr.epochs = static_cast(cfg.epochs > 0 ? cfg.epochs : 1); + params.lr.init(); // required after setting lr fields, before the optimizer reads get_lr() + + // The corpus is either read from a file or supplied inline. + if (!cfg.training_file.empty()) { + std::ifstream in(cfg.training_file, std::ios::binary); + if (!in) { + err = "cannot open training file: " + cfg.training_file; + return false; + } + params.prompt.assign(std::istreambuf_iterator(in), std::istreambuf_iterator()); + } else { + params.prompt = cfg.training_text; + } + + // Training needs writable weights (mmap yields read-only pointers) and an f32 KV cache + // (OUT_PROD has no f16 support) — same forced settings as upstream finetune.cpp. + params.use_mmap = false; + params.cache_type_k = GGML_TYPE_F32; + params.cache_type_v = GGML_TYPE_F32; + + llama_backend_init(); + llama_numa_init(params.numa); + + common_init_result_ptr llama_init = common_init_from_params(params); + llama_model *model = llama_init->model(); + llama_context *ctx = llama_init->context(); + if (model == nullptr || ctx == nullptr) { + err = "failed to load model for training: " + cfg.model_path; + return false; + } + + std::vector tokens = common_tokenize(ctx, params.prompt, true); + if (tokens.size() < 2) { + err = "training corpus produced too few tokens (need at least 2)"; + return false; + } + + ggml_opt_dataset_t dataset = common_opt_dataset_init(ctx, tokens, llama_n_ctx(ctx) / 2); + + llama_opt_params lopt_params = { + /*n_ctx_train =*/0, + /*param_filter =*/llama_opt_param_filter_all, + /*param_filter_ud =*/nullptr, + /*get_opt_pars =*/common_opt_lr_pars, + /*get_opt_pars_ud =*/¶ms.lr, + /*optimizer_type =*/params.optimizer, + }; + llama_opt_init(ctx, model, lopt_params); + + const int64_t idata_split = ggml_opt_dataset_ndata(dataset) * (1.0f - params.val_split); + + ggml_opt_result_t result_train = ggml_opt_result_init(); + ggml_opt_result_t result_eval = ggml_opt_result_init(); + + for (params.lr.epoch = 0; params.lr.epoch < params.lr.epochs; ++params.lr.epoch) { + llama_opt_epoch(ctx, dataset, result_train, result_eval, idata_split, + ggml_opt_epoch_callback_progress_bar, ggml_opt_epoch_callback_progress_bar); + ggml_opt_result_reset(result_train); + ggml_opt_result_reset(result_eval); + } + + ggml_opt_result_free(result_train); + ggml_opt_result_free(result_eval); + ggml_opt_dataset_free(dataset); + + llama_model_save_to_file(model, params.out_file.c_str()); + + // Deliberately NOT calling llama_backend_free(): other live llama contexts in this JVM + // (e.g. an inference LlamaModel) may still depend on the initialized backend. + return true; +} + +} // namespace jllama_train + +extern "C" JNIEXPORT jstring JNICALL +Java_net_ladenthin_llama_LlamaTrainer_finetuneNative(JNIEnv *env, jclass, jstring jconfig) { + std::string config_json; + if (jconfig != nullptr) { + const char *c = env->GetStringUTFChars(jconfig, nullptr); + if (c != nullptr) { + config_json = c; + env->ReleaseStringUTFChars(jconfig, c); + } + } + + jllama_train::finetune_config cfg; + try { + const json j = json::parse(config_json); + cfg.model_path = j.value("model_path", std::string()); + cfg.training_text = j.value("training_text", std::string()); + cfg.training_file = j.value("training_file", std::string()); + cfg.output_path = j.value("output_path", std::string()); + cfg.epochs = j.value("epochs", 2); + cfg.learning_rate = j.value("learning_rate", 1e-5f); + cfg.lr_min = j.value("lr_min", -1.0f); + cfg.decay_epochs = j.value("decay_epochs", -1.0f); + cfg.weight_decay = j.value("weight_decay", 0.0f); + cfg.optimizer = j.value("optimizer", 0); + cfg.n_ctx = j.value("n_ctx", 0); + cfg.n_gpu_layers = j.value("n_gpu_layers", -1); + cfg.val_split = j.value("val_split", 0.05f); + cfg.n_batch = j.value("n_batch", 0); + cfg.n_ubatch = j.value("n_ubatch", 0); + } catch (const std::exception &e) { + return env->NewStringUTF((std::string("invalid training config: ") + e.what()).c_str()); + } + + std::string err; + try { + if (jllama_train::finetune(cfg, err)) { + return env->NewStringUTF(""); // empty == success + } + } catch (const std::exception &e) { + err = e.what(); + } catch (...) { + err = "unknown C++ exception during fine-tuning"; + } + return env->NewStringUTF(err.c_str()); +} diff --git a/src/main/cpp/train_engine.h b/src/main/cpp/train_engine.h new file mode 100644 index 00000000..ba29beef --- /dev/null +++ b/src/main/cpp/train_engine.h @@ -0,0 +1,44 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT +// +// Native fine-tuning engine (proof-of-concept): a self-contained wrapper over llama.cpp's +// ggml-opt training path (llama_opt_init / llama_opt_epoch), mirroring upstream +// examples/training/finetune.cpp. Loads its own model + context (independent of the inference +// server_context in jllama.cpp), fine-tunes on a text corpus, and writes a new GGUF via +// llama_model_save_to_file. Kept out of jllama.cpp so the JNI layer stays thin. + +#ifndef JLLAMA_TRAIN_ENGINE_H +#define JLLAMA_TRAIN_ENGINE_H + +#include + +namespace jllama_train { + +// One fine-tuning run's inputs. +struct finetune_config { + std::string model_path; // base GGUF to fine-tune + std::string training_text; // corpus supplied inline (used when training_file is empty) + std::string training_file; // corpus read from this path instead of training_text + std::string output_path; // where the fine-tuned GGUF is written + int epochs; // number of passes over the corpus (>= 1) + float learning_rate; // lr at the first epoch + float lr_min; // minimum lr for decay; < 0 = no decay + float decay_epochs; // decay lr0 -> lr_min over this many epochs; <= 0 = disabled + float weight_decay; // weight decay; 0 = disabled + int optimizer; // ggml_opt_optimizer_type: 0 = AdamW, 1 = SGD + int n_ctx; // context size; 0 = the model's trained context + int n_gpu_layers; // layers offloaded to the GPU; -1 = auto + float val_split; // fraction of the corpus held out for validation + int n_batch; // logical batch size; 0 = native default + int n_ubatch; // physical (micro) batch size; 0 = native default +}; + +// Run one fine-tuning job end to end. Returns true on success; on failure returns false and sets +// `err`. Not re-entrant; intended to be called off the JVM's critical threads (it blocks for the +// full training run). +bool finetune(const finetune_config &cfg, std::string &err); + +} // namespace jllama_train + +#endif // JLLAMA_TRAIN_ENGINE_H diff --git a/src/main/java/net/ladenthin/llama/LlamaTrainer.java b/src/main/java/net/ladenthin/llama/LlamaTrainer.java new file mode 100644 index 00000000..7ba9db41 --- /dev/null +++ b/src/main/java/net/ladenthin/llama/LlamaTrainer.java @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +package net.ladenthin.llama; + +import java.nio.file.Path; +import net.ladenthin.llama.exception.LlamaException; +import net.ladenthin.llama.loader.LlamaLoader; +import net.ladenthin.llama.parameters.TrainingParameters; + +/** + * In-process fine-tuning entry point, wrapping llama.cpp's ggml-opt training path + * ({@code llama_opt_init} / {@code llama_opt_epoch}) the same way the upstream + * {@code examples/training/finetune.cpp} tool does. Loads its own model and context (independent of + * {@link LlamaModel}), fine-tunes on a text corpus, and writes a new GGUF. + * + *

Configure a run with {@link TrainingParameters} and pass it to {@link #finetune(TrainingParameters)}. + * Full-model fine-tuning is compute- and memory-intensive and blocks for the whole run; upstream + * training support is itself experimental. + */ +public final class LlamaTrainer { + + static { + LlamaLoader.initialize(); + } + + private LlamaTrainer() {} + + /** + * Run one fine-tuning job to completion. + * + * @param parameters the training configuration (model, corpus, output, optimizer, schedule, ...) + * @throws LlamaException if the model cannot be loaded or training fails + */ + public static void finetune(TrainingParameters parameters) { + String error = finetuneNative(parameters.toJson()); + if (error != null && !error.isEmpty()) { + throw new LlamaException(error); + } + } + + /** + * Convenience fine-tune with inline text and otherwise-default settings. + * + * @param model the base GGUF model to fine-tune + * @param trainingText the training corpus (tokenized in-process) + * @param output the path the fine-tuned GGUF is written to + * @param epochs number of passes over the corpus (at least 1) + * @param learningRate the AdamW learning rate at the first epoch (e.g. {@code 1e-5f}) + * @throws LlamaException if the model cannot be loaded or training fails + */ + public static void finetune(Path model, String trainingText, Path output, int epochs, float learningRate) { + finetune( + TrainingParameters.builder() + .modelPath(model) + .trainingText(trainingText) + .outputPath(output) + .epochs(epochs) + .learningRate(learningRate) + .build()); + } + + private static native String finetuneNative(String configJson); +} diff --git a/src/main/java/net/ladenthin/llama/args/Optimizer.java b/src/main/java/net/ladenthin/llama/args/Optimizer.java new file mode 100644 index 00000000..5d3e2a17 --- /dev/null +++ b/src/main/java/net/ladenthin/llama/args/Optimizer.java @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +package net.ladenthin.llama.args; + +/** + * Optimizer used by {@link net.ladenthin.llama.LlamaTrainer} fine-tuning, mapping to llama.cpp's + * {@code ggml_opt_optimizer_type}. + */ +public enum Optimizer { + + /** Adam with decoupled weight decay ({@code GGML_OPT_OPTIMIZER_TYPE_ADAMW}). The default. */ + ADAMW(0), + + /** Stochastic gradient descent ({@code GGML_OPT_OPTIMIZER_TYPE_SGD}). */ + SGD(1); + + private final int nativeValue; + + Optimizer(int nativeValue) { + this.nativeValue = nativeValue; + } + + /** + * The integer value passed to the native layer (matches the {@code ggml_opt_optimizer_type} enum). + * + * @return the native optimizer-type ordinal + */ + public int getNativeValue() { + return nativeValue; + } +} diff --git a/src/main/java/net/ladenthin/llama/parameters/TrainingParameters.java b/src/main/java/net/ladenthin/llama/parameters/TrainingParameters.java new file mode 100644 index 00000000..71640b12 --- /dev/null +++ b/src/main/java/net/ladenthin/llama/parameters/TrainingParameters.java @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +package net.ladenthin.llama.parameters; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.nio.file.Path; +import lombok.Builder; +import lombok.Getter; +import net.ladenthin.llama.args.Optimizer; +import org.jspecify.annotations.Nullable; + +/** + * Immutable configuration for a {@link net.ladenthin.llama.LlamaTrainer} fine-tuning run. + * + *

Build with {@link #builder()}; only {@code modelPath} and {@code outputPath} are required, and + * exactly one of {@code trainingText} / {@code trainingFile} should be set. All other fields default + * to values that mirror upstream llama.cpp's fine-tuning defaults. The configuration is serialized to + * JSON via {@link #toJson()} and parsed by the native layer, the same way {@link ModelParameters} and + * {@link InferenceParameters} cross the JNI boundary. + */ +@Builder +@Getter +public final class TrainingParameters { + + // Base GGUF model to fine-tune. + private final Path modelPath; + + // Training corpus supplied inline; mutually exclusive with trainingFile. + private final @Nullable String trainingText; + + // Training corpus read from a file by the native layer; mutually exclusive with trainingText. + private final @Nullable Path trainingFile; + + // Destination path for the fine-tuned GGUF. + private final Path outputPath; + + // Number of passes over the corpus (at least 1). + @Builder.Default private final int epochs = 2; + + // Learning rate at the first epoch. + @Builder.Default private final float learningRate = 1e-5f; + + // Minimum learning rate for decay, or -1 to disable decay. + @Builder.Default private final float lrMin = -1f; + + // If > 0, decay the learning rate from learningRate to lrMin over this many epochs. + @Builder.Default private final float decayEpochs = -1f; + + // Weight decay (0 disables it). + @Builder.Default private final float weightDecay = 0f; + + // Optimizer algorithm. + @Builder.Default private final Optimizer optimizer = Optimizer.ADAMW; + + // Context size in tokens, or 0 to use the model's trained context. + @Builder.Default private final int nCtx = 0; + + // Layers to offload to the GPU, or -1 for automatic. + @Builder.Default private final int nGpuLayers = -1; + + // Fraction of the corpus held out for validation. + @Builder.Default private final float valSplit = 0.05f; + + // Logical batch size, or 0 to use the native default. + @Builder.Default private final int nBatch = 0; + + // Physical (micro) batch size, or 0 to use the native default. + @Builder.Default private final int nUbatch = 0; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** + * Serialize this configuration to the JSON object the native fine-tuning layer expects. + * + * @return a compact JSON string + */ + public String toJson() { + ObjectNode node = MAPPER.createObjectNode(); + node.put("model_path", modelPath.toString()); + if (trainingText != null) { + node.put("training_text", trainingText); + } + if (trainingFile != null) { + node.put("training_file", trainingFile.toString()); + } + node.put("output_path", outputPath.toString()); + node.put("epochs", epochs); + node.put("learning_rate", learningRate); + node.put("lr_min", lrMin); + node.put("decay_epochs", decayEpochs); + node.put("weight_decay", weightDecay); + node.put("optimizer", optimizer.getNativeValue()); + node.put("n_ctx", nCtx); + node.put("n_gpu_layers", nGpuLayers); + node.put("val_split", valSplit); + node.put("n_batch", nBatch); + node.put("n_ubatch", nUbatch); + return node.toString(); + } +} diff --git a/src/test/java/net/ladenthin/llama/LlamaTrainerIntegrationTest.java b/src/test/java/net/ladenthin/llama/LlamaTrainerIntegrationTest.java new file mode 100644 index 00000000..c7027d22 --- /dev/null +++ b/src/test/java/net/ladenthin/llama/LlamaTrainerIntegrationTest.java @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +package net.ladenthin.llama; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.is; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import net.ladenthin.llama.args.Optimizer; +import net.ladenthin.llama.parameters.TrainingParameters; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * End-to-end fine-tuning smoke over a real model. Self-skips unless a (small) GGUF is provided via + * {@code -Dnet.ladenthin.llama.train.model=/abs/path/to/model.gguf}. Full-model fine-tuning is + * compute- and memory-intensive, so this is opt-in and never runs in a default build. + */ +class LlamaTrainerIntegrationTest { + + @Test + void finetuneWritesAnOutputModel(@TempDir Path tmp) throws Exception { + String modelPath = System.getProperty("net.ladenthin.llama.train.model"); + Assumptions.assumeTrue( + modelPath != null && !modelPath.isEmpty() && Files.exists(Paths.get(modelPath)), + "set -Dnet.ladenthin.llama.train.model=/path/to/small.gguf to run the fine-tune smoke"); + + StringBuilder corpus = new StringBuilder(); + for (int i = 0; i < 64; i++) { + corpus.append("The quick brown fox jumps over the lazy dog. "); + } + + Path output = tmp.resolve("finetuned.gguf"); + LlamaTrainer.finetune( + TrainingParameters.builder() + .modelPath(Paths.get(modelPath)) + .trainingText(corpus.toString()) + .outputPath(output) + .epochs(1) + .learningRate(1e-5f) + .optimizer(Optimizer.ADAMW) + .build()); + + assertThat(Files.exists(output), is(true)); + assertThat(Files.size(output), greaterThan(0L)); + } +} diff --git a/src/test/java/net/ladenthin/llama/parameters/TrainingParametersTest.java b/src/test/java/net/ladenthin/llama/parameters/TrainingParametersTest.java new file mode 100644 index 00000000..f5dfe0d3 --- /dev/null +++ b/src/test/java/net/ladenthin/llama/parameters/TrainingParametersTest.java @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: 2026 Bernard Ladenthin +// +// SPDX-License-Identifier: MIT + +package net.ladenthin.llama.parameters; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.nio.file.Paths; +import net.ladenthin.llama.args.Optimizer; +import org.junit.jupiter.api.Test; + +/** Model-free tests for {@link TrainingParameters} builder defaults and JSON serialization. */ +class TrainingParametersTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private JsonNode json(TrainingParameters parameters) throws Exception { + return MAPPER.readTree(parameters.toJson()); + } + + @Test + void defaultsSerializeWithExpectedValues() throws Exception { + TrainingParameters parameters = + TrainingParameters.builder() + .modelPath(Paths.get("base.gguf")) + .trainingText("hello world") + .outputPath(Paths.get("tuned.gguf")) + .build(); + + JsonNode node = json(parameters); + assertThat(node.get("model_path").asText(), is("base.gguf")); + assertThat(node.get("training_text").asText(), is("hello world")); + assertThat(node.get("output_path").asText(), is("tuned.gguf")); + assertThat(node.get("epochs").asInt(), is(2)); + assertThat(node.get("optimizer").asInt(), is(0)); // ADAMW + assertThat(node.get("learning_rate").floatValue(), is(1e-5f)); + assertThat(node.get("n_gpu_layers").asInt(), is(-1)); + assertThat(node.get("n_batch").asInt(), is(0)); + // training_file is omitted when only inline text is given. + assertThat(node.has("training_file"), is(false)); + } + + @Test + void customValuesSerialize() throws Exception { + TrainingParameters parameters = + TrainingParameters.builder() + .modelPath(Paths.get("base.gguf")) + .trainingFile(Paths.get("corpus.txt")) + .outputPath(Paths.get("tuned.gguf")) + .epochs(5) + .learningRate(3e-4f) + .optimizer(Optimizer.SGD) + .nCtx(512) + .nGpuLayers(0) + .valSplit(0.1f) + .nBatch(256) + .nUbatch(64) + .build(); + + JsonNode node = json(parameters); + assertThat(node.get("epochs").asInt(), is(5)); + assertThat(node.get("optimizer").asInt(), is(1)); // SGD + assertThat(node.get("n_ctx").asInt(), is(512)); + assertThat(node.get("n_gpu_layers").asInt(), is(0)); + assertThat(node.get("n_batch").asInt(), is(256)); + assertThat(node.get("n_ubatch").asInt(), is(64)); + assertThat(node.get("training_file").asText(), is("corpus.txt")); + // training_text is omitted when a corpus file is given. + assertThat(node.has("training_text"), is(false)); + } + + @Test + void optimizerNativeValuesMatchGgml() { + assertThat(Optimizer.ADAMW.getNativeValue(), is(0)); + assertThat(Optimizer.SGD.getNativeValue(), is(1)); + } +}