Otimização de Raciocínio Matemático em LLMs com SimpleRL-reason

O SimpleRL-reason é um framework de código aberto desenvolvido para potencializar a capacidade de raciocínio lógico-matemático em modelos de linguagem (LLMs) de pequeno porte. Inspirado nas metodologias de treinamento do DeepSeek-R1-Zero e DeepSeek-R1, este projeto permite que desenvolvedores repliquem comportamentos de "Chain of Thought" (Cadeia de Pensamento) utilizando Reinforcement Learning (RL) em ambientes com recursos computacionais limitados.

Arquitetura e Funcionamento

O framework utiliza uma arquitetura de treinamento distribuído baseada no Ray. O fluxo de trabalho coordena múltiplos componentes para otimizar a política do modelo:

Modelo Actor: A rede neural que está sendo treinada para resolver problemas.
Modelo Reference: Uma versão estática do modelo para evitar o colapso da política (KL divergence).
Modelo de Recompensa (Reward Model): Avalia a precisão da resposta final e, em alguns casos, a qualidade do processo de raciocínio.

Principais Funcionalidades

Algoritmo PPO (Proximal Policy Optimization): Implementação robusta para estabilizar o aprendizado por reforço em tarefas de texto.
Compatibilidade com Datasets Matemáticos: Suporte nativo para benchmarks como GSM8K (matemática básica) e MATH (nível avançado/olimpíadas).
Treinamento Distribuído: Integração com Ray e DeepSpeed (ZeRO-2/3) para suportar treinamento em múltiplos nós e GPUs.
Mecanismo de Verificação: Ferramentas integradas para validar se a resposta gerada pelo modelo coincide com o gabarito matemático via regras determinísticas.

Configuração do Ambiente

Para iniciar o desenvolvimento, é necessário configurar um ambiente Python 3.8+ com suporte a CUDA. A instalação das dependências pode ser feita da seguinte forma:

# Clonar o repositório principal
git clone https://gitcode.com/gh_mirrors/si/simpleRL-reason
cd simpleRL-reason

# Instalar pacotes necessários para o módulo de treino
python3 -m pip install -r train/requirements.txt

Para ambientes isolados, o projeto disponibiliza um Dockerfile pronto para uso, facilitando a portabilidade em clusters de alta performence:

cd train/dockerfile
docker build -t rl-reasoning-env .
docker run --gpus all -it rl-reasoning-env

Processo de Treinamento

O treinamento foca na capacidade do modelo de gerar passos intermediários de raciocínio antes de entregar a resposta final. O script abaixo ilustra como iniciar um experimento PPO utilizando um modelo base como o Qwen:

#!/bin/bash
# Exemplo de configuração para treinamento em um único nó
NODE_RANK=0
MASTER_ADDR="localhost"

python3 -m train.ppo_trainer \
    --model_name_or_path "qwen-base-model" \
    --dataset_name "math_lv3_5" \
    --output_dir "./models/math_reasoning_v1" \
    --deepspeed_config "./configs/ds_zero3_offload.json" \
    --learning_rate 1e-6 \
    --batch_size 8

Avaliação e Métricass

Após a conclusão do ciclo de RL, é fundamental validar se o modelo não apenas decorou as respostas, mas aprendeu a lógica subjacente. O framework inclui um pipeline de avaliação que gera relatórios detalhados:

# Execução do script de validação
cd eval/sh
bash run_evaluation.sh --model_path ../../models/math_reasoning_v1 --test_set gsm8k

Os resultados são armazenados em formato JSONL, permitindo analisar a taxa de acerto (accuracy) e o tempo médio de inferência por token de raciocínio.

Otimização de Memória e Performance

Trabalhar com Reinforcement Learning exige uma gestão eficiente de VRAM. O SimpleRL-reason suporta as seguintes técnicas de otimização:

DeepSpeed ZeRO-3: Fragmenta os estados do otimizador, gradientes e parâmetros entre as GPUs disponíveis.
Gradient Checkpointing: Reduz o consumo de memória durante o forward pass ao custo de um leve aumento no tempo de computação.
Quantização: Suporte para carregar modelos em baixa precisão (4-bit ou 8-bit) para acelerar a iteração em hradware doméstico.

Customização de Dados

Para treinar o modelo em domínios específicos, você deve formatar seus dados no padrão esperado pelo carregador interno. O formato recomendado é o JSONL, onde cada entrada contém uma pergunta e um caminho de solução estruturado:

{
  "instruction": "Se x + 5 = 12, qual o valor de x?",
  "expected_output": "7",
  "reasoning_steps": "Subtraia 5 de ambos os lados: x = 12 - 5; x = 7."
}

Ao adicionar novos dados, certifique-se de atualizar o arquivo data_loader.py para incluir os novos caminhos de diretório.

Tags: Reinforcement-Learning LLM PPO DeepSpeed Ray

Publicado em 7-1 21:58

Doido Dev