DeepSeek R1 Training Guide

# Generate response function
def generate_response(prompt_text):
    messages = [
        {"role": "system", "content": "You are a helpful assistant that provides step-by-step solutions."},
        {"role": "user", "content": prompt_text}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("<|im_start|>assistant\n")[-1].strip()

# Few-shot examples
few_shot_prompt = """
Problem: What's the square root of 9 plus 5?
Solution: <|special_token|> First, find the square root of 9, which is 3. Then, add 5 to 3. 3 + 5 equals 8. <|special_token|> Summary: The answer is 8.

Problem: Train travels at 60 mph for 2 hours, how far?
Solution: <|special_token|> Use the formula: Distance = Speed times Time. Speed is 60 mph, Time is 2 hours. Distance = 60 * 2 = 120 miles. <|special_token|> Summary: Train travels 120 miles.

Problem: What is 2 + 3 * 4?
Solution:
"""

# Generate structured response
target_problem = "What is 2 + 3 * 4?"
model_response = generate_response(few_shot_prompt + target_problem)
print("Few-shot CoT Response:")
print(model_response)

# Direct prompting approach
direct_prompt = """
Problem: Solve this, show reasoning step-by-step, and verify:
What is 2 + 3 * 4?
"""

direct_response = generate_response(direct_prompt)
print("Direct Prompting Response:")
print(direct_response)

# Refine messy R1 Zero outputs
def refine_output(messy_text):
    """Refine messy reasoning output into structured format."""
    try:
        think_content = messy_text.split("")[1].split("")[0].strip()
        answer_content = messy_text.split("")[1].split("")[0].strip()
        
        # Clean up the reasoning
        cleaned_reasoning = think_content.replace('ummm...', '').replace('...', '').strip()
        
        refined_text = f"""<|special_token|> Reasoning: {cleaned_reasoning.capitalize()}.
<|special_token|> Summary: The answer is {answer_content}."""
        return refined_text
    except:
        return messy_text

# Example refinement
messy_output = "  ummm... multiply 3 and 4... get 12... then add 2...\n 14 "
refined = refine_output(messy_output)

print("Before refinement:")
print(messy_output)
print("\nAfter refinement:")
print(refined)

# SFT Training Configuration
OUTPUT_DIR_SFT = "data/Qwen-SFT-training"
os.makedirs(OUTPUT_DIR_SFT, exist_ok=True)

sft_training_args = TrainingArguments(
    output_dir=OUTPUT_DIR_SFT,
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,
    gradient_accumulation_steps=2,
    learning_rate=2e-5,  # Lower learning rate for SFT
    warmup_ratio=0.1,
    weight_decay=0.01,
    logging_steps=10,
    evaluation_strategy="no",
    save_strategy="steps",
    save_steps=50,
    save_total_limit=2,
    dataloader_num_workers=2,
    seed=42,
    bf16=True,
    push_to_hub=False,
    gradient_checkpointing=True,
    report_to="none",
)

# Load high-quality reasoning dataset
dataset_sft = load_dataset("HuggingFaceH4/Bespoke-Stratos-17k", split='train')

# Initialize model for SFT
model_sft = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# Initialize SFT Trainer
sft_trainer = SFTTrainer(
    model=model_sft,
    train_dataset=dataset_sft,
    tokenizer=tokenizer,
    args=sft_training_args,
)

# Start SFT training
print("Starting SFT training...")
sft_result = sft_trainer.train()

# Save SFT model
TRAINED_SFT_PATH = "data/Qwen-SFT-training"
tokenizer.save_pretrained(TRAINED_SFT_PATH)
sft_trainer.save_model(TRAINED_SFT_PATH)
print(f"SFT trained model saved to {TRAINED_SFT_PATH}")

# Conceptual implementation of language consistency reward
def language_consistency_reward(completions, questions, **kwargs):
    """Reward function to ensure consistent language usage."""
    rewards = []
    
    for completion, question in zip(completions, questions):
        content = completion[0]["content"]
        
        # Detect languages (simplified)
        question_lang = detect_language(question)
        response_lang = detect_language(content)
        
        # Reward consistency
        if question_lang == response_lang:
            rewards.append(1.0)
        else:
            rewards.append(0.0)  # Penalty for language mixing
    
    return rewards

# Enhanced GRPO training with language consistency
enhanced_reward_functions = [
    accuracy_reward,
    format_reward,
    reasoning_steps_reward,
    language_consistency_reward,
    get_cosine_scaled_reward(),
    get_repetition_penalty_reward()
]

# Conceptual rejection sampling implementation
def rejection_sampling(model, tokenizer, problems, num_samples=10, quality_threshold=0.8):
    """Generate high-quality reasoning examples through rejection sampling."""
    high_quality_examples = []
    
    for problem in problems:
        best_response = None
        best_score = 0
        
        # Generate multiple responses
        for _ in range(num_samples):
            response = generate_response(problem)
            
            # Evaluate quality (simplified)
            score = evaluate_response_quality(response, problem)
            
            if score > best_score and score >= quality_threshold:
                best_response = response
                best_score = score
        
        if best_response:
            high_quality_examples.append({
                'problem': problem,
                'response': best_response,
                'quality_score': best_score
            })
    
    return high_quality_examples

# Use rejection sampling to create refined dataset
refined_data = rejection_sampling(model, tokenizer, sample_problems)

# Final stage training configuration
final_training_args = TrainingArguments(
    output_dir="data/Qwen-R1-final",
    num_train_epochs=1,
    per_device_train_batch_size=4,  # Smaller batch for diverse data
    learning_rate=1e-5,  # Very low learning rate for fine-tuning
    warmup_ratio=0.05,
    weight_decay=0.01,
    logging_steps=5,
    save_strategy="steps",
    save_steps=25,
    evaluation_strategy="steps",
    eval_steps=25,
    seed=42,
    bf16=True,
    gradient_checkpointing=True,
    report_to="none",
)

# Multi-objective reward function for final stage
def helpfulness_harmlessness_reward(completions, **kwargs):
    """Reward function balancing helpfulness and harmlessness."""
    rewards = []
    
    for completion in completions:
        content = completion[0]["content"]
        
        # Evaluate helpfulness (simplified)
        helpfulness_score = evaluate_helpfulness(content)
        
        # Evaluate harmlessness (simplified)
        harmlessness_score = evaluate_harmlessness(content)
        
        # Combine scores
        combined_reward = 0.6 * helpfulness_score + 0.4 * harmlessness_score
        rewards.append(combined_reward)
    
    return rewards

# Knowledge distillation setup
def distillation_training(teacher_model, student_model, dataset, temperature=3.0):
    """Distill knowledge from large teacher to smaller student model."""
    
    class DistillationTrainer(SFTTrainer):
        def compute_loss(self, model, inputs, return_outputs=False):
            # Student forward pass
            student_outputs = model(**inputs)
            student_logits = student_outputs.logits
            
            # Teacher forward pass (no gradient)
            with torch.no_grad():
                teacher_outputs = teacher_model(**inputs)
                teacher_logits = teacher_outputs.logits
            
            # Distillation loss
            distill_loss = F.kl_div(
                F.log_softmax(student_logits / temperature, dim=-1),
                F.softmax(teacher_logits / temperature, dim=-1),
                reduction='batchmean'
            ) * (temperature ** 2)
            
            # Standard cross-entropy loss
            ce_loss = F.cross_entropy(
                student_logits.view(-1, student_logits.size(-1)),
                inputs['labels'].view(-1)
            )
            
            # Combined loss
            total_loss = 0.7 * distill_loss + 0.3 * ce_loss
            
            return (total_loss, student_outputs) if return_outputs else total_loss
    
    # Initialize distillation trainer
    distill_trainer = DistillationTrainer(
        model=student_model,
        train_dataset=dataset,
        tokenizer=tokenizer,
        args=final_training_args,
    )
    
    return distill_trainer

# Example usage
# teacher_model = load_trained_r1_model()
# student_model = load_smaller_base_model()
# distill_trainer = distillation_training(teacher_model, student_model, dataset)
# distill_trainer.train()

# Comprehensive evaluation function
def evaluate_r1_model(model, tokenizer, test_dataset):
    """Evaluate the trained R1 model on multiple metrics."""
    
    results = {
        'accuracy': 0,
        'format_compliance': 0,
        'reasoning_quality': 0,
        'language_consistency': 0,
        'response_length': [],
        'reasoning_steps': []
    }
    
    for example in test_dataset:
        # Generate response
        response = generate_response(example['problem'])
        
        # Evaluate accuracy
        accuracy = evaluate_mathematical_accuracy(response, example['solution'])
        results['accuracy'] += accuracy
        
        # Evaluate format compliance
        format_score = evaluate_format_compliance(response)
        results['format_compliance'] += format_score
        
        # Evaluate reasoning quality
        reasoning_score = evaluate_reasoning_quality(response)
        results['reasoning_quality'] += reasoning_score
        
        # Track response metrics
        results['response_length'].append(len(response))
        results['reasoning_steps'].append(count_reasoning_steps(response))
    
    # Calculate averages
    n_examples = len(test_dataset)
    results['accuracy'] /= n_examples
    results['format_compliance'] /= n_examples
    results['reasoning_quality'] /= n_examples
    
    return results

# Example evaluation
# evaluation_results = evaluate_r1_model(trained_model, tokenizer, test_dataset)
# print("Evaluation Results:", evaluation_results)