Módulo de Entrenamiento RL

AITraining incluye un módulo completo de aprendizaje por refuerzo para escenarios avanzados de entrenamiento de LLM.

Los comandos CLI (--trainer ppo, --trainer dpo, --trainer reward) usan implementaciones de la biblioteca TRL para estabilidad. El módulo autotrain.trainers.rl documentado aquí proporciona bloques de construcción de nivel inferior para pipelines de entrenamiento RL personalizados.

Visión General

El módulo RL proporciona:

PPO Trainer - Proximal Policy Optimization con penalización KL y GAE
DPO Trainer - Optimización Directa de Preferencias desde datos de preferencia
Modelos de Recompensa - Modelos de recompensa estándar, pareados y multi-objetivo
Entornos RL - Entornos de generación de texto, matemáticas, código y comparación de preferencias
Pipeline Asíncrono - Entrenamiento forward-backward con acumulación de gradiente

Entrenamiento PPO

Configuración

from autotrain.trainers.rl import PPOConfig, PPOTrainer

config = PPOConfig(
    model_name="google/gemma-2-2b",
    learning_rate=1e-5,
    batch_size=16,
    mini_batch_size=4,
    gradient_accumulation_steps=1,

    # PPO hyperparameters
    ppo_epochs=4,
    gamma=0.99,           # Discount factor
    lam=0.95,             # GAE lambda
    clip_ratio=0.2,       # PPO clip ratio
    value_clip=0.2,       # Value function clip
    max_grad_norm=1.0,    # Gradient clipping

    # KL penalty
    kl_penalty_coef=0.01,
    kl_target=0.01,
    kl_horizon=10000,     # Horizon for adaptive KL

    # Coefficients
    entropy_coef=0.01,    # Entropy regularization
    value_coef=0.5,       # Value function coefficient

    # Generation
    max_new_tokens=128,
    temperature=0.7,
    top_p=0.9,

    # Training loop
    num_iterations=100,
    save_every=10,
    eval_every=5,
    device=None,          # Auto-detected
)

Arquitectura PPO

La implementación PPO usa un wrapper PPOModel que añade una cabeza de valor a cualquier LM causal:

# PPOModel wraps base model with ValueHead
class PPOModel(nn.Module):
    def __init__(self, base_model):
        self.base_model = base_model
        self.value_head = ValueHead(hidden_size)

# ValueHead architecture
class ValueHead(nn.Module):
    # hidden -> ReLU -> output (scalar value)

Controlador KL Adaptativo

El AdaptiveKLController ajusta automáticamente el coeficiente de penalización KL para mantener la divergencia KL cerca del objetivo:

# Automatically managed by PPOTrainer
# Adjusts kl_penalty_coef based on current KL vs target

Bucle de Entrenamiento

# Initialize trainer with custom reward function
def my_reward_fn(prompts, responses, metadata=None):
    rewards = []
    for response in responses:
        score = evaluate_response(response)
        rewards.append(score)
    return rewards

trainer = PPOTrainer(
    config=config,
    tokenizer=tokenizer,  # Optional, loaded from model if not provided
    reward_fn=my_reward_fn,
)

# Train on prompts
prompts = ["Write a poem about...", "Explain quantum..."]
metrics = trainer.train(prompts, num_iterations=100)

Características Clave

Característica	Descripción
Adaptive KL Controller	Ajusta automáticamente el coeficiente de penalización KL basado en KL actual vs objetivo
GAE Advantage Estimation	Estimación de Ventaja Generalizada para entrenamiento estable
Value Head	Función de valor separada para crítico (PPOModel envuelve modelo base)
Reference Model	Copia congelada para prevenir deriva
Async Training	Usa AsyncTrainingClient para forward-backward eficiente

Entrenamiento DPO

Entrena directamente desde datos de preferencia sin un modelo de recompensa separado.

Configuración

from autotrain.trainers.rl import DPOConfig, DPOTrainer

config = DPOConfig(
    model_name="google/gemma-2-2b",
    learning_rate=1e-6,
    batch_size=8,
    gradient_accumulation_steps=2,

    # DPO hyperparameters
    beta=0.1,              # Temperature parameter
    label_smoothing=0.0,   # For robustness
    reference_free=False,  # Use reference model

    # Training
    num_epochs=1,
    max_grad_norm=1.0,
    warmup_ratio=0.1,

    # Sequence lengths
    max_length=512,
    max_prompt_length=256,

    # Checkpointing
    eval_every=100,
    save_every=500,
    device=None,           # Auto-detected
)

Dataset de Preferencia

from autotrain.trainers.rl.dpo import PreferenceDataset

# Create dataset from preference pairs
dataset = PreferenceDataset(
    prompts=["What is AI?", "Explain gravity"],
    chosen=["AI is artificial intelligence...", "Gravity is a force..."],
    rejected=["idk lol", "its like magnets"],
    tokenizer=tokenizer,
    max_length=512,
    max_prompt_length=256,
)

# Train
trainer = DPOTrainer(config=config, tokenizer=tokenizer)
metrics = trainer.train(dataset, eval_dataset=eval_dataset)

PreferenceDataset debe ser importado directamente de autotrain.trainers.rl.dpo ya que no está exportado en el __init__.py principal.

DPO Sin Referencia

Para entrenamiento sin modelo de referencia:

config = DPOConfig(
    model_name="google/gemma-2-2b",
    reference_free=True,  # No reference model needed
    beta=0.1,
)

Modelos de Recompensa

Modelo de Recompensa Estándar

from autotrain.trainers.rl import RewardModel, RewardModelConfig, RewardModelTrainer

config = RewardModelConfig(
    model_name="bert-base-uncased",
    num_labels=1,
    pooling_strategy="last",  # "mean", "last", or "cls"
    dropout_prob=0.1,
    temperature=1.0,          # Temperature scaling for rewards

    # LoRA settings
    use_lora=True,
    lora_rank=8,
    lora_alpha=16,
    lora_dropout=0.1,

    # Training
    learning_rate=1e-4,
    warmup_steps=100,
    gradient_accumulation_steps=1,
)

model = RewardModel(config)

Entrenamiento en Preferencias

trainer = RewardModelTrainer(
    model=model,
    tokenizer=tokenizer,
    config=config,
    device=None,  # Auto-detected
)

trainer.train_on_preferences(
    chosen_texts=["Good response 1", "Good response 2"],
    rejected_texts=["Bad response 1", "Bad response 2"],
    num_epochs=3,
    batch_size=8,
)

# Save/load
trainer.save_model("reward_model.pt")
trainer.load_model("reward_model.pt")

Modelo de Recompensa Pareado

Para comparación directa de preferencias usando modelo Bradley-Terry:

from autotrain.trainers.rl import PairwiseRewardModel

model = PairwiseRewardModel(config)

# Forward pass compares two inputs
preference_score = model.forward_pair(
    input_ids_a, attention_mask_a,
    input_ids_b, attention_mask_b,
)

# Bradley-Terry loss for training
loss = model.compute_bradley_terry_loss(
    input_ids_a, attention_mask_a,
    input_ids_b, attention_mask_b,
    labels,  # 1 if A preferred, 0 if B preferred
)

Modelo de Recompensa Multi-Objetivo

Combina múltiples señales de recompensa:

from autotrain.trainers.rl import MultiObjectiveRewardModel

model = MultiObjectiveRewardModel(
    config=config,
    num_objectives=3,
    objective_weights=[0.5, 0.3, 0.2],  # Helpfulness, safety, honesty
)

# Get all objectives
outputs = model(input_ids, attention_mask, return_all_objectives=True)
# outputs["rewards"] shape: (batch_size, 3)
# outputs["combined_reward"] shape: (batch_size, 1)

# Multi-objective loss
loss, per_objective_losses = model.compute_multi_objective_loss(
    input_ids, attention_mask,
    target_rewards,      # Shape: (batch_size, num_objectives)
    objective_mask=None, # Optional: which objectives to train
)

Entornos RL

Dataclasses de Entorno

from autotrain.trainers.rl.environments import Observation, StepResult, Trajectory

# Observation from environment
@dataclass
class Observation:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    prompt: str
    metadata: Dict[str, Any]

# Result from env.step()
@dataclass
class StepResult:
    reward: float
    done: bool
    next_observation: Optional[Observation]
    info: Dict[str, Any]
    metrics: Dict[str, float]

# Full episode trajectory
@dataclass
class Trajectory:
    observations: List[Observation]
    actions: List[torch.Tensor]
    rewards: List[float]
    logprobs: List[torch.Tensor]
    done: bool
    total_reward: float
    metrics: Dict[str, Any]

Entorno de Generación de Texto

from autotrain.trainers.rl import TextGenerationEnv

env = TextGenerationEnv(
    tokenizer=tokenizer,
    prompts=["Write a story about...", "Explain how..."],
    max_length=512,
    reward_fn=my_reward_function,  # Optional, default is length-based
    stop_sequences=["</s>", "\n\n"],
    temperature=1.0,
)

# Reset and step
observation = env.reset()
result = env.step(action_token)
# result.reward, result.done, result.next_observation

# Render current state
print(env.render())

Entorno Multi-Objetivo

from autotrain.trainers.rl import MultiObjectiveRewardEnv

def correctness_reward(prompt, generated, full_text):
    return 1.0 if is_correct(generated) else 0.0

def formatting_reward(prompt, generated, full_text):
    return 0.5 if properly_formatted(generated) else 0.0

env = MultiObjectiveRewardEnv(
    tokenizer=tokenizer,
    prompts=prompts,
    reward_components={
        "correctness": correctness_reward,
        "formatting": formatting_reward,
    },
    reward_weights={
        "correctness": 1.0,
        "formatting": 0.1,
    },
)

# Step returns component rewards in metrics
result = env.step(action)
# result.metrics["reward_correctness"], result.metrics["reward_formatting"]

Entorno de Comparación de Preferencias

Para recolección de datos RLHF y DPO:

from autotrain.trainers.rl import PreferenceComparisonEnv

env = PreferenceComparisonEnv(
    tokenizer=tokenizer,
    prompts=prompts,
    preference_model=preference_model,  # Optional
    human_feedback_fn=feedback_fn,      # Optional callback
    max_length=512,
)

# Generates pairs of responses and computes preference
observation = env.reset()
result1 = env.step(response1_tokens)  # First response
result2 = env.step(response2_tokens)  # Second response, computes preference

Entornos Integrados

from autotrain.trainers.rl import create_math_problem_env, create_code_generation_env

# Math problem solving (correctness + formatting rewards)
math_env = create_math_problem_env(tokenizer)

# Code generation (syntax + style rewards)
code_env = create_code_generation_env(tokenizer)

Pipeline Forward-Backward

Entrenamiento asíncrono con acumulación de gradiente:

from autotrain.trainers.rl import ForwardBackwardPipeline

# Low-level pipeline
pipeline = ForwardBackwardPipeline(
    model=model,
    device="cuda",
    max_workers=2,                    # Thread pool size
    gradient_accumulation_steps=4,
)

# Queue forward-backward pass
future = pipeline.forward_backward(
    input_ids=input_ids,
    attention_mask=attention_mask,
    labels=labels,
    loss_fn="cross_entropy",
)

# Get result (blocks until complete)
result = future.result()
print(f"Loss: {result.loss}")

# Queue optimizer step
optim_future = pipeline.optim_step(
    optimizer=optimizer,
    scheduler=scheduler,  # Optional
    max_grad_norm=1.0,
)
optim_result = optim_future.result()

Funciones de Pérdida Integradas

El pipeline soporta varias funciones de pérdida integradas:

Función de Pérdida	Descripción	kwargs Requeridos
`"cross_entropy"`	Pérdida estándar de modelado de lenguaje	None
`"importance_sampling"`	RL con muestreo de importancia	`old_logprobs`, `advantages`
`"ppo"`	Pérdida PPO completa	`old_log_probs`, `advantages`, opcionalmente `values`, `returns`

Funciones de Pérdida Personalizadas

def custom_loss_fn(model, inputs, outputs, **kwargs):
    # Your custom loss computation
    logits = outputs.logits
    # ... compute loss ...
    return loss_tensor  # Must be scalar

future = pipeline.forward_backward_custom(
    input_ids=input_ids,
    custom_loss_fn=custom_loss_fn,
    attention_mask=attention_mask,  # Optional
    my_param=42,  # Passed to loss function via kwargs
)

Cliente de Alto Nivel

from autotrain.trainers.rl.forward_backward import AsyncTrainingClient

client = AsyncTrainingClient(
    model=model,
    reference_model=reference_model,  # For PPO/DPO
    device="cuda",
    gradient_accumulation_steps=4,
)

# Training step
fwd_future = client.forward_backward(batch, loss_fn="cross_entropy")
optim_future = client.optim_step(optimizer, max_grad_norm=1.0)

# Forward only (for reference model)
ref_future = client.forward(batch, use_reference=True)

# Clean up
client.shutdown()

AsyncTrainingClient debe ser importado directamente de autotrain.trainers.rl.forward_backward ya que no está exportado en el __init__.py principal.

Checkpointing

# Save checkpoint
checkpoint_info = pipeline.save_state("checkpoint_1000")
# Returns: {"path": ..., "model_path": ..., "optimizer_path": ..., "state_path": ...}

# Load checkpoint
pipeline.load_state("checkpoints/checkpoint_1000")

Muestreo

Genera muestras durante el entrenamiento:

samples = pipeline.sample(
    prompt=prompt_tokens,  # List[int] or Tensor
    max_tokens=100,
    temperature=0.7,
    top_k=50,
    top_p=0.9,
    stop=[tokenizer.eos_token_id],
)

print(f"Generated: {samples['tokens']}")
print(f"Logprobs: {samples['logprobs']}")
print(f"Prompt: {samples['prompt']}")

Mejores Prácticas

Entrenamiento PPO

Comienza con coeficiente KL pequeño - Deja que el controlador adaptativo ajuste
Usa acumulación de gradiente - Tamaños de batch efectivos más grandes son más estables
Monitorea divergencia KL - Debe permanecer cerca del objetivo
Calienta la función de valor - Entrena el crítico antes de PPO completo

Entrenamiento DPO

Datos de preferencia de alta calidad - La calidad importa más que la cantidad
Tasa de aprendizaje baja - 1e-6 a 1e-5 recomendado
Suavizado de etiquetas - 0.1 puede mejorar robustez
Evalúa frecuentemente - Rastrea precisión y margen de recompensa

Modelado de Recompensas

Datos balanceados - Ejemplos elegido/rechazado iguales
Prompts diversos - Cubre casos de uso esperados
LoRA para eficiencia - Fine-tuning eficiente de modelos grandes
Multi-objetivo - Separa señales de seguridad y utilidad

Integración CLI

Para uso en producción, el CLI proporciona interfaces más simples usando implementaciones TRL:

# PPO training (uses TRL PPOTrainer)
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer ppo \
  --reward-model ./my-reward-model

# DPO training (uses TRL DPOTrainer)
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer dpo \
  --dpo-beta 0.1

# Reward model training
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer reward \
  --data-path ./preference_data.jsonl

Próximos Pasos

PPO Training

Guía CLI para PPO

DPO Training

Guía CLI para DPO

Reward Modeling

Entrenar modelos de recompensa

ORPO Training

Optimización de Preferencias por Razón de Probabilidades

​Módulo de Entrenamiento RL

​Visión General

​Entrenamiento PPO

​Configuración

​Arquitectura PPO

​Controlador KL Adaptativo

​Bucle de Entrenamiento

​Características Clave

​Entrenamiento DPO

​Configuración

​Dataset de Preferencia

​DPO Sin Referencia

​Modelos de Recompensa

​Modelo de Recompensa Estándar

​Entrenamiento en Preferencias

​Modelo de Recompensa Pareado

​Modelo de Recompensa Multi-Objetivo

​Entornos RL

​Dataclasses de Entorno

​Entorno de Generación de Texto

​Entorno Multi-Objetivo

​Entorno de Comparación de Preferencias

​Entornos Integrados

​Pipeline Forward-Backward

​Funciones de Pérdida Integradas

​Funciones de Pérdida Personalizadas

​Cliente de Alto Nivel

​Checkpointing

​Muestreo

​Mejores Prácticas

​Entrenamiento PPO

​Entrenamiento DPO

​Modelado de Recompensas

​Integración CLI

​Próximos Pasos

PPO Training

DPO Training

Reward Modeling

ORPO Training

Módulo de Entrenamiento RL

Visión General

Entrenamiento PPO

Configuración

Arquitectura PPO

Controlador KL Adaptativo

Bucle de Entrenamiento

Características Clave

Entrenamiento DPO

Configuración

Dataset de Preferencia

DPO Sin Referencia

Modelos de Recompensa

Modelo de Recompensa Estándar

Entrenamiento en Preferencias

Modelo de Recompensa Pareado

Modelo de Recompensa Multi-Objetivo

Entornos RL

Dataclasses de Entorno

Entorno de Generación de Texto

Entorno Multi-Objetivo

Entorno de Comparación de Preferencias

Entornos Integrados

Pipeline Forward-Backward

Funciones de Pérdida Integradas

Funciones de Pérdida Personalizadas

Cliente de Alto Nivel

Checkpointing

Muestreo

Mejores Prácticas

Entrenamiento PPO

Entrenamiento DPO

Modelado de Recompensas

Integración CLI

Próximos Pasos