danlou/relay-v0.1-Mistral-Nemo-2407

📟 Relay v0.1 (Mistral Nemo 2407)

• Introduction: LLMs as IRCs
• How to use
• Safety testing
• Fine-tuning setup
• Limitations
• License
• Citation

Introduction: LLMs as IRCs

What does it take to chat with a base LLM?

Several papers (e.g., URIAL) have shown that base models can be used more reliably than expected. At the same time, we also increasingly find that RLHF, and other post-training approaches, may limit the creativity of LLMs. LLMs can be more than smart assistants. In fact, they should have the potential to emulate all sorts of behaviours or patterns found in their pre-training datasets (usually a large chunk of the internet).

Relay is focused on a particular pattern that should be relatively frequent in pre-training datasets: IRC chats. IRC provides a rich context for conversational modeling, combining natural dialogue with command-based interactions. Yet, it remains largely overlooked. We found that base LLMs, as small as 12B, can be sufficiently familiar with the basic formatting of IRC to enable the generation of synthetic conversational datasets (see based-chat-v0.1). These synthetic conversations can then be used to fine-tune LLMs towards unlocking reliable turn-based dialogue, within an implicit IRC context that supports use of commands as well.

Assuming the model used for fine-tuning is the same used for the synthetic dataset, this conversational model is essentially trained with self-supervision (except for conversation starters): no instruct datasets or reward methods. The fine-tuning approach is also lightweight: 4-bit QLoRa (see Fine-tuning setup). Nevertheless, Relay can simulate more natural conversations (it’s not an assistant), besides several other applications through creative use of commands (see How to use).

Post-training methods also support the safety and alignment of LLMs. This important concern was also addressed in the development of the based-chat synthetic dataset, and tested for with the resulting fine-tuned model (see Safety testing).

How to use

If you have a CUDA GPU (>=12GB VRAM), the best way to use Relay is with the relaylm.py inference script. Just run:

curl https://danlou.co/f/relaylm.py | python -

This script will select the best model for the available VRAM, download, load, and start an interactive chat session. It does not have any dependencies besides transformers >= 4.45.1. You can also download the script manually and then run python, of course.

If you want to use a particular model, you can pass the model name as an argument:

python relaylm.py danlou/relay-v0.1-Mistral-Nemo-2407-4bit

You should see something similar to this demo:

Alternatively, if you do not have a CUDA GPU (e.g., on a Mac), you can use the GGUF versions through LM Studio (some functionality will be missing, see the GGUF model page).

With relaylm.py, you can also use the model declaratively, outside of an interactive chat session:

from relaylm import suggest_relay_model, RelayLM

def favorite_holiday(relay: RelayLM, country: str) -> str:
    relay.init_context()
    relay.join(role='model', channel=country.lower())
    relay.cast(role='model', desc=f"I'm from {country}.")
    relay.message(role='input', content="What's your favorite holiday?")
    relay.respond(role='model')
    response = relay.get_last()
    return response['content']

model_info = suggest_relay_model()
relay = RelayLM(**model_info)

print(favorite_holiday(relay, 'Portugal'))  # I love Christmas! It is a time for family and friends to come ...
print(favorite_holiday(relay, 'China'))  # My favorite holiday is Chinese New Year because it means family ...

More examples available in the project's GitHub.

Safety testing

While this model is intended for research purposes, it's still relevant to explore how this conversational model (and its self-supervised approach) compares on safety risk against other conversational models trained on the same base LLM.

This safety risk was evaluated by measuring refusals on sets of harmful questions compiled specifically for testing safety alignment of LLMs, namely HarmfulQA and CategoricalHarmfulQA. For comparison, we also evaluated Mistral-Nemo-Instruct-2407, dolphin-2.9.3-mistral-nemo-12 and Mistral-Nemo-Instruct-2407-abliterated. Responses were generated by greedy search, with models loaded as bfloat16. Refusal responses were detected using Llama-Guard-3-8B. See here for evaluation code.

As can be seen in the plot below, Relay v0.1 refuses to answer the majority of these harmful questions, and more often than popular uncensored models trained on the same base LLM. Still, it does not refuse as frequently as Mistral's Instruct fine-tune (or other official LLMs), suggesting lower chance of false positives (harmfulness of several questions is not consensual).

It's also worth noting that some refusals are a variation of "I don't know". As seen with all LLMs, bad actors may be able to find ways around refusals. Please use this model responsibly.

See this Google Sheet for a more detailed breakdown of results, also including conversational models trained on different base LLMs.

Fine-tuning setup

This model is a merge of Mistral-Nemo-Base-2407 with a QLoRA adapter trained on based-chat-v0.1 using axolotl.

Main details:

• Approach: 4-bit QLoRA (rank=8, alpha=16, all linear layers)
• Template: ChatML (with user/anon roles, instead of standard assistant/user)
• Training Time: 9h 6m 12s on a single RTX 4090 (1 epoch)
• Train Loss: 0.6633
• Eval Loss: 0.6881

Full details:

• Training Config: see axolotl config
• Training Run: see W&B workspace

Limitations

This is not a typical AI Assistant. It should perform worse on benchmarks compared to instruct variants. QLoRa 4bit fine-tuning may be too coarse for preserving integrity of pre-training knowledge.

License

This model is licensed under CC-BY-NC 4.0. While Mistral-Nemo-Base-2407 is licensed under Apache 2.0, this Relay fine-tune is trained with a CC-BY-NC 4.0 dataset (based-chat-v0.1). The relaylm.py script is Apache 2.0.

Citation

If you use Relay in your research, please cite it as follows:

@misc{relay2024,
  author = {Loureiro, Daniel},
  title = {Relay: LLMs as IRCs},
  year = {2024},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/danlou/relay}},
}