Mistral-Small-3-Reasoner-s1

A simple Mistral-Small-24B-Instruct-2501 finetune on the LIMA-like s1K reasoning dataset by Muennighoff et al. to give the original model basic reasoning capabilities within <think> tags, like DeepSeek R1. Surprisingly, the model can reason even outside math/STEM subjects.

Usage notes

Prepend the assistant response with <think> to make the model engage in a chain-of-thought. This should happen automatically with math questions on an empty context, but it needs to be forced in longer conversations. When done thinking, it will generate </think> and then the final response.

Make sure that the model's output length is long enough; be prepared to make the it continue its response if it stops prematurely.

Low-depth instructions (perhaps at depth-0, just before the assistant's response) can be beneficial in steering how the model should think. An additional [SYSTEM_PROMPT] could be used there.

From tests it seems beneficial to keep at least one chain-of-thought in the context in addition to the one being generated. More experimentation required here.

Prompting format

[INST]User question.[/INST]<think>
Chain of thought.
</think>

Model response.</s>

Observed quirks and issues

• Assistant responses may lack the final punctuation mark as a probable result of how the source training data was formatted.
• Not really a true issue, but information in the system prompt that contradicts that in the chat history or that is meant to be only temporarily valid may cause coherency issues with this model, since it will follow instructions very precisely.
• Without user control on chain-of-thought length, the model can ramble for several thousands of tokens.
• Besides multi-turn capabilities, other non-reasoning capabilities of the original Mistral-Small-24B-Instruct-2501 model might have degraded.
• Most default guardrails apparently still work, but can be very easily bypassed with a suitable prompt as with the original model.

What's in this repository

• Checkpoints for epochs 1~5
• LoRA adapter for the final model
• Static GGUF quantizations
• HF FP16 weights

Dataset

Almost the entirety of the s1K dataset was used with minimal modifications to make it properly work with Mistral-Small-3-Instruct, except 4 rows that didn't fit within the training sequence length of 8192 tokens and 16 of the shortest ones that have been used as the test set instead. No samples have been clipped and no system prompt was added. All were single-turn.

Training hyperparameters

I tried to roughly follow indications from appendix C in the paper with the notable exception of using 4-bit LoRA finetuning and tuning the learning rate accordingly. The loss was not computed on questions within [INST]...[/INST] tags (including the tags), just on reasoning traces and solutions. The training sequence length was close to the maximum I could use on one NVidia RTX3090 24GB GPU. The total training time was about 18 hours.

Overfitting (increasing eval loss) occurred after 1 epoch, but the training loss behavior was similar to that observed in the paper.

lora_r = 32
lora_alpha = 64
lora_dropout = 0.05
load_in_4bit = True
max_seq_length = 8192
gradient_accumulation_steps = 16
eval_accumulation_steps = 16
learning_rate = 0.000400
weight_decay = 1e-4
num_train_epochs = 5
warmup_ratio = 0.05
lr_scheduler_type = "cosine"
optim = "adamw_8bit"

Training Graphs

Appears to more or less correspond to the general trend and values seen in Appendix C in the paper.

Some hyperparameters including the learning rate were changed before the end of the first epoch to better match those of the paper.