“Thinking” LLMs on a Shoestring (S1): Fine-Tuning and Budget Forcing for Reasoning

from datasets import load_dataset
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          BitsAndBytesConfig, TrainingArguments)
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
import torch

# Define model and dataset identifiers
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
DATASET_ID = "simplescaling/s1K_tokenized"

# Initialize tokenizer, handling the pad token carefully.
# Qwen's official pad token is <|extra_0|>. It's safer to use a model's
# designated pad_token than to default to the EOS token.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
if tokenizer.pad_token is None:
    if tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token
    else:
        # Fallback for models without a default pad or eos token.
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})

dataset = load_dataset(DATASET_ID, split="train")

# Define a formatting function to structure data into the ChatML format Qwen expects.
def format_example(example):
    # This function converts prompt/chosen pairs into a single text field
    # adhering to the model's required conversation structure.
    return {
        "text": f"<|im_start|>user\n{example['prompt']}<|im_end|>\n<|im_start|>assistant\n{example['chosen']}<|im_end|>"
    }

formatted_dataset = dataset.map(format_example)

# Configure 4-bit quantization (QLoRA) to reduce memory footprint.
# bfloat16 is the compute dtype of choice for modern GPUs (Ampere architecture or newer).
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    attn_implementation="flash_attention_2", # Use Flash Attention 2 for efficiency
    device_map="auto" # Automatically map model layers across available devices
)
model = prepare_model_for_kbit_training(model)

# LoRA configuration: defines which layers to adapt and with what rank.
# Target modules are specific to the Qwen2 architecture.
lora_config = LoraConfig(
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "up_proj", "down_proj", "gate_proj"],
    bias="none",
    task_type="CAUSAL_LM"
)

# Use a completion-only data collator to ensure loss is calculated only on the
# assistant's response, not the user prompt. This is crucial for effective tuning.
data_collator = DataCollatorForCompletionOnlyLM(
    instruction_template="<|im_start|>user",
    response_template="<|im_start|>assistant\n",
    tokenizer=tokenizer,
    mlm=False # This is a Causal LM, not Masked LM.
)

# Training arguments
OUTPUT_DIR = "qwen2.5-7b-s1k-lora-checkpoints"
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=32, # Achieves an effective batch size of 32
    num_train_epochs=5,
    learning_rate=1e-6, # A conservative LR, often appropriate for fine-tuning
    bf16=True,
    logging_dir=f"{OUTPUT_DIR}/logs",
    logging_steps=10,
    save_strategy="epoch",
)

# Initialize the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=formatted_dataset,
    peft_config=lora_config,
    dataset_text_field="text",
    max_seq_length=model.config.max_position_embeddings if hasattr(model.config, 'max_position_embeddings') else 2048,
    data_collator=data_collator,
    args=training_args
)

trainer.train()
trainer.model.save_pretrained(f"{OUTPUT_DIR}/final_adapter") # Save only the trained adapter

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest
import torch

# Initialize the vLLM engine with LoRA enabled.
# The base model is loaded, and LoRA adapters are applied on-the-fly per request.
llm_engine = LLM("Qwen/Qwen2.5-7B-Instruct",
                  enable_lora=True,
                  max_lora_rank=16,
                  tensor_parallel_size=1,
                 )

# Path to the trained LoRA adapter directory.
LORA_ADAPTER_PATH = "qwen2.5-7b-s1k-lora-checkpoints/final_adapter"

# Create a LoRA request object to identify the adapter for inference.
lora_request = LoRARequest(
    lora_name="s1_reasoning_adapter",
    lora_int_id=1, # A unique integer ID for this adapter
    lora_local_path=LORA_ADAPTER_PATH
)

tokenizer = llm_engine.get_tokenizer()
im_end_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
if im_end_token_id == tokenizer.unk_token_id:
    raise ValueError("<|im_end|> token not found in tokenizer. Check model configuration.")

def budget_force(question: str, max_total_think_tokens: int = 1000, num_forced_waits: int = 2) -> str:
    system_prompt = "You are a helpful AI assistant. Please think step by step to answer the question."
    
    # The prompt buffer starts with system instructions and the user's question,
    # then cues the model to begin its internal monologue.
    current_prompt_str = (
        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
        f"<|im_start|>user\n{question}<|im_end|>\n"
        f"<|im_start|>assistant\n<|im_start|>think\n"
    )
    
    remaining_think_budget = max_total_think_tokens
    
    for i in range(num_forced_waits):
        if remaining_think_budget <= 0:
            break

        # Cap each thought-generation step to a reasonable length.
        tokens_for_this_step = min(remaining_think_budget, 512)

        # Generate a chunk of thought.
        # This implementation re-sends the entire prompt buffer. While simple,
        # a production system would optimize this by managing the KV cache directly.
        outputs = llm_engine.generate(
            [current_prompt_str],
            SamplingParams(
                max_new_tokens=tokens_for_this_step,
                stop_token_ids=[im_end_token_id], # Stop if the model concludes its turn prematurely
                temperature=0.0, # Greedy decoding for deterministic reasoning
            ),
            lora_request=lora_request
        )
        
        generated_output = outputs[0].outputs[0]
        thought_chunk_text = generated_output.text
        num_tokens_in_chunk = len(generated_output.token_ids)

        current_prompt_str += thought_chunk_text
        remaining_think_budget -= num_tokens_in_chunk

        # If more thinking is forced, add a continuation cue.
        if i < num_forced_waits - 1:
            current_prompt_str += " Wait. Let me continue thinking.\n<|im_start|>think\n"
        else:
            # Transition to generating the final, user-facing answer.
            current_prompt_str += "\n<|im_start|>assistant\n"

    # Generate the final answer based on the accumulated thought process.
    final_outputs = llm_engine.generate(
        [current_prompt_str],
        SamplingParams(
            max_new_tokens=512,
            stop_token_ids=[im_end_token_id],
            temperature=0.0
        ),
        lora_request=lora_request
    )
    
    final_answer = final_outputs[0].outputs[0].text
    return final_answer.strip()