uzaLEAT LowEndAITrainer

uzaLEAT is a framework to train AI models on low-end hardware. The framework uses Vulkan for GPU acceleration. It includes the GUTR interpreter. You use GUTR to build custom AI architectures. You can build basic RNN or RWKV models. You can also build new models.

You can write C++ plugins for maximum speed. The core program loads shared libraries. A shared library must export C functions. The core program looks for specific function names.

Features

• Train AI models on low-end hardware.
• Use Vulkan for GPU acceleration.
• Build models with the GUTR language.
• Use compiled C++ plugins for maximum speed.
• Read JSONL and Parquet datasets.
• Train BPE tokenizers automatically.
• Save and load model states.

Build Instructions

You use the UMK build system to compile the program.

Run umk build to make the main program.

Usage

To train a model, use the --train option. You must specify the data path and the model path.

./uzaLEAT --train --data ./data --model-so ./libslesa.so --gpu

To chat with a model, use the --chat option. You must specify the model file.

./uzaLEAT --chat --model-so ./libslesa.so --model ./model.uleat

Configuration Options

• --train: Start training mode.
• --chat: Start chat mode.
• --data: Set the path to the training data.
• --model-so: Set the compiled C++ model file.
• --plugin: Set the GUTR plugin file.
• --model: Set the model save and load path.
• --tokenizer: Set the tokenizer file path.
• --gpu: Enable Vulkan GPU acceleration.
• --epochs: Set the number of training epochs.
• --lr: Set the learning rate.
• --batch: Set the batch size.
• --threads: Set the number of CPU threads.
• --hidden: Set the hidden size.
• --layers: Set the number of layers.
• --context: Set the context size.

C++ Plugin API

You can write a C++ plugin to run models faster. The core program loads your plugin as a shared library. The plugin must export standard C functions. The core program calls these functions during training and generation.

Required Functions

Your plugin must implement these functions:

• model_init: Initialize the model.
• model_train_step: Train the model on one batch.
• model_generate: Generate text tokens.
• model_save: Save the model weights to a file.
• model_load: Load the model weights from a file.

Optional Functions

Your plugin can implement these functions:

• model_cleanup: Free memory when the program exits.
• model_get_tokenizer_size: Get the size of the tokenizer data.
• model_get_tokenizer_data: Get the tokenizer data.
• model_load_holdout: Load validation data.

C++ Plugin Example

Create a file named mymodel.cpp. Write the code below.

#include <cstdint>
#include <cstring>
#include <vector>

struct ModelMetrics {
    uint64_t step = 0;
    float loss = 0.0f;
    float ppl = 0.0f;
    uint64_t fl_updates = 0;
    uint64_t gpu_calls = 0;
    uint64_t cpu_fallbacks = 0;
    char custom_info[256] = {0};
};

extern "C" {

void model_init(int hidden_size, int num_layers, int vocab_size, int context_size, int tt_rank, int num_experts, int window_size, int update_interval) {
    // Allocate memory for your model.
    // Set the initial weights.
}

float model_train_step(const int* input, const int* target, int seq_len, float learning_rate, ModelMetrics* metrics) {
    // Do the forward pass.
    // Calculate the loss.
    // Do the backward pass.
    // Update the weights.
    
    metrics->loss = 0.5f;
    metrics->step += 1;
    return metrics->loss;
}

int model_generate(const int* prompt, int prompt_len, int* output, int max_tokens, float temperature, float top_p) {
    // Generate tokens.
    // Copy the generated tokens to the output array.
    int generated_count = 1;
    output[0] = 10;
    return generated_count;
}

void model_save(const char* path) {
    // Open the file.
    // Write the weights to the file.
    // Close the file.
}

void model_load(const char* path) {
    // Open the file.
    // Read the weights from the file.
    // Close the file.
}

void model_cleanup() {
    // Free the allocated memory.
}

} // extern "C"

Build the Plugin

Compile the C++ file into a shared library. Use a C++ compiler.

g++ -O3 -march=native -fPIC -shared mymodel.cpp -o libmymodel.so

Run the Plugin

Use the --model-so option to load your plugin.

./uzaLEAT --train --data ./data --model-so ./libmymodel.so --model ./mymodel.uleat