HelixNet-LMoE

HelixNet-LMoE is a simple LoRA based Mixture of Experts version of the HelixNet 3-model system by Migel Tissera.
It is a 6bpw multi-lora exl2 model for use with ExLlamaV2.

For each HelixNet model, a separate LoRA adapter was extracted:

• lora-actor
• lora-critic
• lora-regenerator

These are then loaded together with the base Mistral 7b model which has been quantized to 6bpw-exl2 to give the combined LMoE model.

As HelixNet processes its inputs using the actor, critic and regenerator actions, the corresponding LoRA adapter is dynamically enabled as required.

It is similar in approach to Airoboro's MoE implementation allowing GPU memory requirements in this (6bpw quantized) instance to be reduced from 20GB for the 3 separate 6bpw models to 8GB for the 6bpw multi lora model. The LoRAs were extracted based on the process given in https://github.com/uukuguy/multi_loras, with a Rank of 64 and an Alpha of 128.

Performance:

Testing on an RTX-4090 to compare with the separate 6bpw exl2 models from https://huggingface.co/LoneStriker?search_models=helixnet gives:

3 separate models: 120 tokens / second, using 20GB GPU
LMoE combined model: 91 tokens / second, using 8GB GPU

Prompt format:

SYSTEM: Elaborate on the topic using a Tree of Thoughts and backtrack when necessary to construct a clear, cohesive Chain of Thought reasoning. Always answer without hesitation.
USER: What is the relationship between Earth's atmosphere, magnetic field and gravity?
ASSISTANT:

Example Usage

The following is a code example on how to use HelixNet-LMoE. No special system-context messages are needed for the critic and the regenerator.
At the You: prompt, enter a question such as What is the relationship between Earth's atmosphere, magnetic field and gravity?

import time
import sys, os
import dataclasses
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from exllamav2 import(
    ExLlamaV2,
    ExLlamaV2Config,
    ExLlamaV2Cache,
    ExLlamaV2Tokenizer,
    ExLlamaV2Lora,    
)

from exllamav2.generator import (
    ExLlamaV2BaseGenerator,
    ExLlamaV2Sampler
)


class ModelClass:
    def __init__(self, generator, tokenizer, model):
        self.generator = generator
        self.tokenizer = tokenizer
        self.model = model

DEBUG = os.environ.get("DEBUG") and True or False

# Initialize model and cache
def load_model(model_directory, max_seq_len=8192):
    """
    Loads a model from a directory and return the generator and tokenizer
    """
    config = ExLlamaV2Config()
    config.model_dir = model_directory
    config.max_seq_len = max_seq_len
    config.prepare()

    model = ExLlamaV2(config)
    print("Loading model: " + model_directory)

    cache = ExLlamaV2Cache(model, lazy = True, max_seq_len=max_seq_len)
    model.load_autosplit(cache)

    tokenizer = ExLlamaV2Tokenizer(config)
    generator = ExLlamaV2BaseGenerator(model, cache, tokenizer)
    model = ModelClass(generator=generator, tokenizer=tokenizer, model=model)
    generator.warmup()
    return model

def generate_text(prompt, lora, settings, max_new_tokens):
    time_begin = time.time()
    response = base_model.generator.generate_simple(prompt, settings, max_new_tokens, loras=lora)
    response = response[len(prompt):]
    time_end = time.time()
    time_total = time_end - time_begin
    tokens = base_model.tokenizer.encode(response)
    count = tokens.shape[-1]
    print(f"Response generated in {time_total:.2f} seconds, {count} tokens, {count / time_total:.2f} tokens/second, character len: {len(response)}")
    return response

base_model = load_model("models/HelixNet-LMoE-6.0bpw-h6-exl2")
lora_actor = ExLlamaV2Lora.from_directory(base_model.model, "models/HelixNet-LMoE-6.0bpw-h6-exl2/lora-actor")
lora_critic = ExLlamaV2Lora.from_directory(base_model.model, "models/HelixNet-LMoE-6.0bpw-h6-exl2/lora-critic")
lora_regenerator = ExLlamaV2Lora.from_directory(base_model.model, "models/HelixNet-LMoE-6.0bpw-h6-exl2/lora-regenerator")

settings = ExLlamaV2Sampler.Settings()
settings.temperature = 0.75
settings.top_k = 50
settings.top_p = 1.0
max_new_tokens = 2000

system_prompt = "You are HelixNet. Elaborate on the topic using a Tree of Thoughts and backtrack when necessary to construct a clear, cohesive Chain of Thought reasoning. Always answer without hesitation."

while True:
    user_input = input("You: ")

    prompt_actor = f"SYSTEM: {system_prompt}\nUSER: {user_input}\nASSISTANT: "
    if DEBUG: print(f"{prompt_actor}\n\n")
    print("ACTOR:")
    response_actor = generate_text(prompt_actor, lora_actor, settings, max_new_tokens)
    if DEBUG: print(f"{response_actor}\n\n")
    print("="*132)

    prompt_critic = f"SYSTEM: {system_prompt}\nUSER: {user_input}\nRESPONSE: {response_actor}\nCRITIQUE: "
    if DEBUG: print(f"{prompt_critic}\n\n")
    print("CRITIQUE:")
    response_critic = generate_text(prompt_critic, lora_critic, settings, max_new_tokens)
    if DEBUG: print(f"{response_critic}\n\n")
    print("="*132)

    prompt_regenerator = f"SYSTEM: {system_prompt}\nUSER: {user_input}\nRESPONSE: {response_actor}\nCRITIQUE: {response_critic}\nREGENERATOR: "
    if DEBUG: print(f"{prompt_regenerator}\n\n")
    print("REGENERATION:")
    response_regenerator = generate_text(prompt_regenerator, lora_regenerator, settings, max_new_tokens)
    print("="*132)
    conversation = f"SYSTEM: {system_prompt}\nUSER: {user_input}\nASSISTANT: {response_regenerator}"
    print(conversation)

LLM Evaluation

Evaluation on a merged version of each base+lora models has yet to be done on the HuggingFaceH4 Open LLM Leaderboard to see how it compares to the equivalent full HelixNet model.

HelixNet Details

HelixNet is a Deep Learning architecture consisting of 3 x Mistral-7B LLMs. It has an actor, a critic, and a regenerator. The actor LLM produces an initial response to a given system-context and a question. The critic then takes in as input, a tuple of (system-context, question, response) and provides a critique based on the provided answer to the given system-context and the question. Its job is not to criticize, but to provide an intelligent critique so that the answer can be modified/regenerated to address the question better. Finally, the regenerator takes in a tuple of (system-context, question, response, critique) and regenerates the answer.

HelixNet is insprired from an actor-critic architecture most prominent in Reinforcement Learning algorithms. The name derives from Helix, referring to the spiral structure of a DNA molecule. It symbolizes the intertwined nature of the three networks, working in tandem, much like the strands of a DNA molecule.

HelixNet regenerates very pleasing and accurate responses, due to the entropy preservation of the regenerator. The regenerator was only trained on a dataset of 1000 samples, similar to Meta's LIMA. The actor network here was trained on about 250K very high-quality samples, and the critic network was trained on further 10K samples.

Full details on how HelixNet was trained and evaluated is located at https://huggingface.co/migtissera/HelixNet \

The 6bpw separate models for HelixNet are available at https://huggingface.co/LoneStriker?search_models=helixnet