Microajuste Eficiente de DeepSeek-R1 con Unsloth: Guía Práctica en Entornos con Recursos Limitados

Introducción a la Optimización del Microajuste con Unsloth

Unsloth es una biblioteca avanzada diseñada para acelerar el proceso de fine-tuning en modelos basados en arquitecturas Llama y sus variantes derivadas, como los modelos DeepSeek-R1. Destaca por su capacidad para reducir significativamente el uso de memoria GPU (hasta un 80% menos que las implementaciones estándar de Hugging Face), permitiendo entrenamientos incluso en tarjetas gráficas con solo 7 GB de VRAM. Esta eficiencia se logra mediante optimizaciones internas, incluyendo una versión mejorada del algoritmo GRPO (Generalized Reward-Policy Optimization), clave en los procesos de aprendizaje por refuerzo utilizados durante el entrenamiento original de DeepSeek.

Dado que los modelos distilados de DeepSeek están construidos sobre estructuras compatibles con Llama o Qwen, son plenamente soportados por Unsloth sin necesidad de modificaciones. A continuación, se detalla un flujo completo de ajuste fino utilizando un modelo médico especializado.

Configuración del Entorno

  • Modelo base: unsloth/DeepSeek-R1-Distill-Llama-8B
  • Biblioteca de microajuste: Unsloth + integración con TRL (Transformer Reinforcement Learning)
  • Conjunto de datos: FreedomIntelligence/medical-o1-reasoning-SFT (versión en chino, subconjunto curado)

Instalación e Inicialización del Modelo

Primero, instale las versiones más recientes de las herramientas necesarias:

pip install --upgrade pip
pip uninstall unsloth -y && pip install --upgrade --no-cache-dir git+https://github.com/unslothai/unsloth.git
pip install bitsandbytes trl peft accelerate datasets

Ahora cargue el modelo cuantificado en 4 bits para minimizar el consumo de memoria:

from unsloth import FastLanguageModel
import torch

max_length = 2048
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/DeepSeek-R1-Distill-Llama-8B",
    max_seq_length=max_length,
    load_in_4bit=True,
    dtype=None,
)

Evaluación Inicial del Modelo

Antes del entrenamiento, verifiquemos el comportamiento actual del modelo ante una consulta médica compleja:

prompt_template = """Below is an instruction that describes a task, paired with an input providing context.
Write a response completing the request. Think step-by-step.

### Instruction:
You are a medical expert with advanced knowledge in clinical reasoning and diagnostics.
Answer the following question.

### Question:
{}

### Response:
<think>{}"""

question = "Un paciente con apendicitis aguda desde hace 5 días presenta alivio parcial del dolor abdominal, pero persiste fiebre y tiene una masa palpable y dolorosa en fosa ilíaca derecha. ¿Cuál es el manejo indicado?"

FastLanguageModel.for_inference(model)
inputs = tokenizer([prompt_template.format(question, "")], return_tensors="pt").to("cuda")

outputs = model.generate(**inputs, max_new_tokens=1024, use_cache=True)
result = tokenizer.decode(outputs[0], skip_special_tokens=False)
print(result.split("### Response:")[-1])

Preparación del Dataset para SFT

El formato del conjunto de datos debe seguir un estilo de prompt coherente, incluyendo razonamiento paso a paso (Chain-of-Thought) y terminando cada entrada con el token de fin de secuencia (EOS). Este diseño mejora la capacidad del modelo para generar respuestas estructuradas.

def generate_training_prompt(data):
    instructions = data["Question"]
    chains = data["Complex_CoT"]
    responses = data["Response"]
    texts = []
    eos_token = tokenizer.eos_token
    
    for inst, chain, resp in zip(instructions, chains, responses):
        full_text = prompt_template.format(inst, chain) + "\n" + resp + eos_token
        texts.append(full_text)
    
    return {"text": texts}

from datasets import load_dataset
dataset = load_dataset("FreedomIntelligence/medical-o1-reasoning-SFT", "zh", split="train[:500]")
dataset = dataset.map(generate_training_prompt, batched=True, remove_columns=dataset.column_names)

Configuración del Entrenamiento con LoRA

Utilizamos adaptadores LoRA (Low-Rank Adaptation) para ajustar selectivamente capas clave del modelo, manteniendo bajo el costo computacional:

model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", 
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth"
)

Ejecución del Ajuste Fino

Configuramos el entrenador SFTTrainer con parámetros óptimos para recursos limitados:

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_length,
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        output_dir="fine_tuned_medical_model",
        report_to="none"
    ),
    packing=False
)

# Iniciar entrenamiento
trainer.train()

Validación Post-Microajuste

Después del entrenamiento, volvemos a evaluar la misma pregunta médica para observar mejoras en calidad, coherencia y profundidad diagnóstica. El modelo ahora genera respuestas más clínicas, con énfasis en pruebas de imagen, manejo antibiótico y decisiones quirúrgicas guiadas por evidencia.

Exportación del Modelo Ajustado

Guarde el adaptador LoRA localmente o súbelo a Huggging Face Hub:

model.save_pretrained("med_adapter_lora")
tokenizer.save_pretrained("med_adapter_lora")

# Opcional: publicar en línea
# model.push_to_hub("tu_usuario/med-adapter-deepseek", private=True)
# tokenizer.push_to_hub("tu_usuario/med-adapter-deepseek", private=True)

Etiquetas: unsloth deepseek-r1 LoRA sft-trainer medical-llm

Publicado el 8-15 20:00