Fine-Tuning DeepSeek R1 with Unsloth & LoRA: Memory-Efficient Reasoning Models

User: <prompt>
Assistant: <think>Thought process leading to the answer.</think><answer>Final answer.</answer>

[
  {
    "messages": [
      {
        "role": "user",
        "content": "What is the sum of the first 10 prime numbers?"
      },
      {
        "role": "assistant",
        "content": "<think>The first 10 prime numbers are 2, 3, 5, 7, 11, 13, 17, 19, 23, 29. Summing them: 2+3+5+7+11+13+17+19+23+29 = 129.</think><answer>129</answer>"
      }
    ]
  },
  {
    "messages": [
      {
        "role": "user",
        "content": "If a car travels at 60 mph for 2 hours, how far does it travel?"
      },
      {
        "role": "assistant",
        "content": "<think>Distance = Speed × Time. Speed = 60 mph, Time = 2 hours. Distance = 60 * 2 = 120 miles.</think><answer>120 miles</answer>"
      }
    ]
  }
]

pip install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git"
pip install transformers peft accelerate bitsandbytes trl datasets torch

import torch
from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments, AutoTokenizer
from datasets import load_dataset
import os

# 1. Configuration
max_seq_length = 2048 # Max sequence length for DeepSeek-R1-Distill
model_name = "deepseek-ai/deepseek-r1-3b-base" # Or deepseek-ai/deepseek-r1-7b-base
dataset_path = "your_synthetic_reasoning_data.json" # Path to your JSON dataset

# 2. Load Model and Tokenizer with Unsloth
# Unsloth automatically handles 4-bit quantization (QLoRA)
# and loads the model with optimized kernels.
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_name,
    max_seq_length = max_seq_length,
    dtype = None, # None for auto detection (bfloat16 if supported, else float16)
    load_in_4bit = True, # Enable QLoRA
)

# 3. Configure LoRA Adapters
# Target all linear layers for optimal performance.
model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # LoRA rank
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0.05,
    bias = "none",
    use_gradient_checkpointing = "unsloth", # Use Unsloth's optimized gradient checkpointing
    random_state = 3407,
    max_seq_length = max_seq_length,
)

# 4. Load and Format Dataset
# The dataset should be a JSON file with the structure described above.
# We use `apply_chat_template` to format messages into the model's expected input format.
# DeepSeek models typically use a specific chat template.
# Ensure the tokenizer has a chat template or define one.
# For DeepSeek, it's often similar to:
# {% for message in messages %}{% if message['role'] == 'user' %}{{ '<|User|>
# ' + message['content'] + '<|EOT|>
# ' }}{% elif message['role'] == 'assistant' %}{{ '<|Bot|>
# ' + message['content'] + '<|EOT|>
# ' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|Bot|>' }}{% endif %}

# If your tokenizer doesn't have a default, you might need to set it:
# tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ '<|User|>\n' + message['content'] + '<|EOT|>\n' }}{% elif message['role'] == 'assistant' %}{{ '<|Bot|>\n' + message['content'] + '<|EOT|>\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|Bot|>' }}{% endif %}"

def formatting_prompts_func(examples):
    """
    Formats the dataset examples into the model's expected chat template.
    Ensures the <think> and <answer> tokens are preserved within the assistant's response.
    """
    texts = []
    for i in range(len(examples["messages"])):
        # Apply chat template to each conversation
        # `add_generation_prompt=False` is crucial for fine-tuning to ensure
        # the assistant's response is fully included, not just the prompt for generation.
        formatted_text = tokenizer.apply_chat_template(
            examples["messages"][i],
            tokenize=False,
            add_generation_prompt=False
        )
        texts.append(formatted_text)
    return { "text" : texts }

dataset = load_dataset("json", data_files=dataset_path, split="train")
dataset = dataset.map(
    formatting_prompts_func,
    batched = True,
)

# 5. Configure Training Arguments
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text", # Field containing the formatted text
    max_seq_length = max_seq_length,
    dataset_num_proc = os.cpu_count(), # Use all CPU cores for data processing
    packing = False, # Set to True for more efficient packing of short sequences
    args = TrainingArguments(
        per_device_train_batch_size = 2, # Adjust based on VRAM
        gradient_accumulation_steps = 4, # Accumulate gradients over 4 steps
        warmup_steps = 5,
        num_train_epochs = 3,
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(), # Use fp16 if bfloat16 not supported
        bf16 = torch.cuda.is_bf16_supported(), # Use bf16 if supported
        logging_steps = 1,
        optim = "adamw_8bit", # Use 8-bit AdamW optimizer
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
        report_to = "none", # Disable reporting to W&B etc. for simplicity
    ),
)

# 6. Train the Model
trainer.train()

# 7. Save the LoRA Adapters
model.save_pretrained("deepseek_r1_lora_adapters")
tokenizer.save_pretrained("deepseek_r1_lora_adapters")

print("Fine-tuning complete. LoRA adapters saved to deepseek_r1_lora_adapters.")

import torch
from unsloth import FastLanguageModel
from transformers import AutoTokenizer

model_name = "deepseek-ai/deepseek-r1-3b-base"
lora_adapters_path = "deepseek_r1_lora_adapters"
output_merged_path = "deepseek_r1_merged_model"

# Load the base model and tokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_name,
    max_seq_length = 2048, # Must match training max_seq_length
    dtype = None, # Use the same dtype as during training or float16/bfloat16
    load_in_4bit = False, # Load in full precision for merging
)

# Load the LoRA adapters
model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Must match training r
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0.05,
    bias = "none",
    use_gradient_checkpointing = False, # Not needed for inference/merging
    random_state = 3407,
    max_seq_length = 2048, # Must match training max_seq_length
)
model.load_adapter(lora_adapters_path)

# Merge LoRA adapters and save the full model
model.save_pretrained_merged(output_merged_path, tokenizer, save_method = "merged_16bit")
print(f"Merged model saved to {output_merged_path}")

import os
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import HfApi, create_repo
from pathlib import Path

merged_model_path = "deepseek_r1_merged_model"
output_gguf_path = "deepseek_r1_merged_model_gguf"
quantization_type = "q4_k_m" # Example: q4_k_m, q5_k_m, q8_0

# Ensure llama.cpp is installed and converted.py is available
# You might need to clone llama.cpp and build it:
# git clone https://github.com/ggerganov/llama.cpp.git
# cd llama.cpp && make

# Load the merged model
model = AutoModelForCausalLM.from_pretrained(
    merged_model_path,
    torch_dtype=torch.float16, # Or bfloat16
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(merged_model_path)

# Save in a format compatible with llama.cpp's convert.py
# This typically means saving as a standard Hugging Face model
model.save_pretrained(output_gguf_path)
tokenizer.save_pretrained(output_gguf_path)

# Convert to GGUF using llama.cpp's convert.py
# Make sure you have llama.cpp cloned and built, and its `convert.py` is accessible.
# Adjust `llama_cpp_dir` to your llama.cpp installation path.
llama_cpp_dir = "/path/to/llama.cpp" # IMPORTANT: Set this path
convert_script = os.path.join(llama_cpp_dir, "convert.py")

if not os.path.exists(convert_script):
    raise FileNotFoundError(f"llama.cpp convert.py not found at {convert_script}. "
                            "Please clone and build llama.cpp.")

# Step 1: Convert PyTorch weights to ggml format (intermediate step for convert.py)
# This step is often implicitly handled by newer convert.py versions,
# but historically involved a separate script or specific arguments.
# For DeepSeek, ensure convert.py supports its architecture.
print(f"Converting {output_gguf_path} to GGUF...")
os.system(f"python {convert_script} {output_gguf_path} --outfile {output_gguf_path}/model.gguf")

# Step 2: Quantize the GGUF model
quantize_script = os.path.join(llama_cpp_dir, "quantize")
if not os.path.exists(quantize_script):
    # For older llama.cpp versions, quantize might be part of convert.py or a separate script
    # For newer versions, `quantize` is a compiled binary.
    # If not found, try `python {convert_script} --quantize {quantization_type} ...`
    raise FileNotFoundError(f"llama.cpp quantize binary not found at {quantize_script}. "
                            "Please build llama.cpp or check its documentation for quantization.")

os.system(f"{quantize_script} {output_gguf_path}/model.gguf {output_gguf_path}/model-{quantization_type}.gguf {quantization_type}")

print(f"GGUF model ({quantization_type}) saved to {output_gguf_path}/model-{quantization_type}.gguf")

# Optional: Upload to Hugging Face Hub
# api = HfApi()
# repo_id = "your_hf_username/deepseek-r1-3b-reasoning-gguf"
# create_repo(repo_id, repo_type="model", exist_ok=True)
# api.upload_file(
#     path_or_fileobj=f"{output_gguf_path}/model-{quantization_type}.gguf",
#     path_in_repo=f"deepseek-r1-3b-reasoning-{quantization_type}.gguf",
#     repo_id=repo_id,
# )
# print(f"GGUF model uploaded to Hugging Face Hub: {repo_id}")

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

merged_model_path = "deepseek_r1_merged_model"
output_vllm_path = "deepseek_r1_vllm_model"

# Load the merged model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    merged_model_path,
    torch_dtype=torch.bfloat16, # Use bfloat16 for vLLM if supported, else float16
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(merged_model_path)

# Save the model in a format vLLM can directly load
# vLLM expects a standard Hugging Face model directory.
model.save_pretrained(output_vllm_path)
tokenizer.save_pretrained(output_vllm_path)

print(f"vLLM-compatible model saved to {output_vllm_path}")

# Example vLLM inference (requires vLLM to be installed)
# from vllm import LLM, SamplingParams
#
# llm = LLM(model=output_vllm_path, dtype=torch.bfloat16)
#
# prompts = [
#     "<|User|>\nWhat is the capital of France?<|EOT|>\n<|Bot|>",
#     "<|User|>\nIf x = 5 and y = 3, what is x + y?<|EOT|>\n<|Bot|>"
# ]
#
# sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=128)
#
# outputs = llm.generate(prompts, sampling_params)
#
# for output in outputs:
#     prompt = output.prompt
#     generated_text = output.outputs[0].text
#     print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")