TRL documentation

Treinamento com Jobs

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Treinamento com Jobs

model badge

Hugging Face Jobs executa seu treinamento em GPUs da Hugging Face. Você escolhe o hardware para cada execução e paga apenas pelos segundos em que ela roda. Com --push_to_hub, o modelo treinado é enviado ao Hub ao final.

Neste guia, você aprenderá como:

  • Inicie um script de treinamento do TRL com um único comando
  • Acompanhe uma execução e veja suas curvas de perda
  • Execute seu próprio script, com as configurações de hardware no arquivo
  • Mantenha checkpoints e retome uma execução interrompida
  • Treinar em várias GPUs

Para saber como os Jobs funcionam em geral, incluindo hardware e preços, consulte a documentação de Jobs.

Requisitos

  • Uma conta da Hugging Face com um saldo de crédito positivo. Os jobs são pagos conforme o uso: você paga apenas pelos segundos utilizados.
  • Conectado ao Hugging Face Hub (hf auth login)

Uma primeira execução

Este comando faz o fine-tuning de Qwen2-0.5B-Instruct no conjunto de dados de chat Capybara com o script SFT do TRL e, em seguida, envia o modelo para a sua conta:

bash
python
hf jobs uv run \
    --flavor a10g-small \
    --timeout 30m \
    --secrets HF_TOKEN \
    -- \
    https://raw.githubusercontent.com/huggingface/trl/refs/heads/main/trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/Capybara \
    --max_steps 100 \
    --output_dir Qwen2-0.5B-SFT \
    --push_to_hub

Leva alguns minutos e custa bem menos de US$ 1 em a10g-small. As opções antes de -- são para Jobs: o hardware, um limite de tempo e o token usado para enviar o modelo. --secrets HF_TOKEN envia o token com o qual você fez login. Em Python, passe o próprio token. Os argumentos após a URL do script vão para o script. hf jobs uv run executa um script Python com uv. sft.py lista suas dependências em um cabeçalho no topo do arquivo, para que o Jobs instale o TRL antes de o script começar. O Job para quando o script termina, e a cobrança para junto com ele.

--max_steps 100 mantém essa primeira execução curta. Remova-o para a execução completa: três épocas, o padrão do script, levam algumas horas em a10g-small, então aumente --timeout para corresponder, por exemplo --timeout 3h. O Jobs interrompe uma execução quando ela atinge o tempo limite, que é de 30 minutos por padrão. Uma GPU maior, como a100-large, termina mais rápido. Consulte Hardware para ver os sabores e seus preços.

Acompanhe a execução

hf jobs uv run transmite os logs e mantém seu terminal aberto até o término da execução. Ctrl+C interrompe apenas o fluxo de logs: o Job continua em execução. Adicione --detach para obter o ID do Job imediatamente e, em seguida:

hf jobs logs -f <job_id>   # stream the logs
hf jobs ps                 # list your running Jobs
hf jobs inspect <job_id>   # status, and the error message if it failed
hf jobs cancel <job_id>    # stop it

Veja Gerenciar Jobs para mais informações.

Para as curvas de perda, o TRL registra no Trackio. Adicione --env TRACKIO_SPACE_ID=<your-username>/trackio antes de -- para nomear um Space para o painel e --report_to trackio após a URL do script. O Trackio cria o Space e um bucket para as métricas, caso não existam. Ambos são públicos por padrão. Outros rastreadores suportados pelo Transformers também funcionam. Para o Weights & Biases, adicione --with wandb --secrets WANDB_API_KEY antes de -- e --report_to wandb após a URL do script. --secrets WANDB_API_KEY lê o valor de WANDB_API_KEY no seu ambiente local. Consulte variáveis de ambiente e segredos para outras formas de passá-los.

Execute seu próprio script

Escreva seu código de treinamento em um arquivo Python (por exemplo, train.py). Liste suas dependências no topo do arquivo, em um cabeçalho de script, da mesma forma que os scripts do TRL fazem. Os Jobs os instalam antes de o script começar:

# /// script
# dependencies = [
#     "trl",
#     "peft",
# ]
# ///

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=LoraConfig(),
    args=SFTConfig(output_dir="Qwen2.5-0.5B-SFT", max_steps=100),
)
trainer.train()
trainer.push_to_hub()

Em seguida, inicie-o com a CLI hf jobs ou a API Python:

bash
python
hf jobs uv run \
    --flavor a10g-small \
    --timeout 30m \
    --secrets HF_TOKEN \
    train.py

Para um script sem cabeçalho, liste as dependências na inicialização com --with, uma vez por pacote: --with trl --with peft (dependencies=["trl", "peft"] em Python). O script também pode ser uma URL, como um link raw do GitHub, um Gist ou um arquivo em um repositório público do Hub.

As configurações de lançamento também podem ficar no script. Uma tabela [tool.hf-jobs] no cabeçalho define o hardware, o tempo limite e os segredos:

# /// script
# dependencies = [
#     "trl",
#     "peft",
# ]
#
# [tool.hf-jobs]
# flavor = "a10g-small"
# timeout = "30m"
# secrets = ["HF_TOKEN"]
# ///

hf jobs uv run train.py então não requer opções, e uma opção que você passa ainda substitui a do script. O script carrega tudo o que precisa para executar, então você pode compartilhá-lo, ou voltar a ele mais tarde, sem precisar lembrar qual GPU e timeout ele requer. hf jobs uv run --dry-run train.py exibe as configurações resolvidas e marca os valores que vêm do script. A tabela é lida apenas pelo CLI hf: run_uv_job() a ignora. Veja Definir a configuração de lançamento no script.

Execute qualquer script do TRL

O mesmo padrão é usado em todos os scripts de treinador do TRL (SFT, DPO, GRPO, modelagem de recompensa e outros) e nos exemplos independentes .py no Índice de Exemplos. Cada um declara suas dependências no cabeçalho do script. Um exemplo que lê outros arquivos adjacentes precisa que sua pasta seja montada no Job. Veja Locais. Exemplos exclusivos de notebook não podem ser submetidos dessa forma. Os argumentos do script são os mesmos de quando você executa o script localmente.

Manter checkpoints

O disco de um Job é excluído quando o Job termina, incluindo quando ele atinge o tempo limite ou falha. Para uma execução longa, salve checkpoints fora do Job para que eles sobrevivam a ele.

A maneira mais simples é enviar cada checkpoint para o repositório do modelo de saída conforme ele é salvo:

bash
python
hf jobs uv run \
    --flavor a10g-small \
    --timeout 3h \
    --secrets HF_TOKEN \
    -- \
    https://raw.githubusercontent.com/huggingface/trl/refs/heads/main/trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/Capybara \
    --output_dir Qwen2-0.5B-SFT \
    --push_to_hub \
    --save_steps 500 \
    --hub_strategy checkpoint

Com --hub_strategy checkpoint, o checkpoint mais recente, incluindo o estado do otimizador, é mantido em uma pasta last-checkpoint do repositório. Para continuar uma execução interrompida, monte o repositório em um novo Job e retome a partir dessa pasta: inicie o mesmo comando com --volume hf://<your-username>/Qwen2-0.5B-SFT:/previous antes de -- e --resume_from_checkpoint /previous/last-checkpoint após a URL do script.

Para manter vários checkpoints fora do repositório do modelo, escreva-os em um Storage Bucket em vez disso. Crie um com hf buckets create <your-username>/checkpoints --private, monte-o no Job e aponte --output_dir para ele:

bash
python
hf jobs uv run \
    --flavor a10g-small \
    --timeout 3h \
    --secrets HF_TOKEN \
    --volume hf://buckets/<your-username>/checkpoints:/checkpoints \
    -- \
    https://raw.githubusercontent.com/huggingface/trl/refs/heads/main/trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/Capybara \
    --output_dir /checkpoints/Qwen2-0.5B-SFT \
    --save_steps 500

Para continuar uma execução interrompida, execute novamente o mesmo comando com --resume_from_checkpoint /checkpoints/Qwen2-0.5B-SFT/checkpoint-<step>. Consulte Escrever em um bucket conforme avança para mais informações e Volumes para as opções de montagem.

Cada checkpoint contém o estado do modelo e do otimizador, ocupando várias vezes o tamanho do modelo. Adicione --save_total_limit 2 para manter apenas os mais recentes. Se você também passar --push_to_hub, defina --hub_model_id, ou o repositório será nomeado com base na última parte do caminho de saída.

Várias GPUs

Um sabor com várias GPUs encurta uma execução. Para isso, use hf jobs run, que executa um comando em uma imagem Docker em vez de um script, com a imagem huggingface/trl, que tem TRL e Accelerate instalados. Em seguida, inicie um processo por GPU. Isso funciona para os scripts do TRL e para o seu próprio código de treinamento.

Acelere um script do TRL

A execução de SFT do início deste guia, em quatro GPUs:

bash
python
hf jobs run \
    --flavor l4x4 \
    --timeout 30m \
    --secrets HF_TOKEN \
    huggingface/trl \
    -- \
    trl sft --num_processes 4 \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/Capybara \
    --max_steps 100 \
    --output_dir Qwen2-0.5B-SFT \
    --push_to_hub

O trl CLI passa --num_processes para accelerate launch. Defina-o para o número de GPUs no flavor, aqui quatro L4s. Cada processo assume seu próprio lote, portanto o tamanho efetivo do lote é quatro vezes o de uma execução com GPU única. Esta execução leva alguns minutos. Sem --max_steps, os três epochs completos levam aproximadamente uma hora, então aumente --timeout, por exemplo para 2h.

Seu próprio script

Você não está limitado aos scripts do TRL. O código de treinamento que é executado com accelerate launch localmente funciona da mesma forma nos Jobs, desde que a imagem tenha os pacotes necessários. Por exemplo, salve um script GRPO com sua própria função de recompensa como train_grpo.py em uma pasta e monte a pasta no Job. Os Jobs enviam a pasta para um bucket privado e a montam somente para leitura:

bash
python
hf jobs run \
    --flavor l4x4 \
    --timeout 30m \
    --secrets HF_TOKEN \
    --volume ./my-project:/code \
    huggingface/trl \
    -- \
    accelerate launch --num_processes 4 /code/train_grpo.py

O script é executado com o TRL instalado na imagem, portanto não precisa de cabeçalho de script. GRPOTrainer gera completamentos com transformers por padrão, então isso é executado em um único Job, sem servidor vLLM. Para geração mais rápida, consulte vLLM Integration. A próxima seção explica como hf jobs run usa a imagem.

Imagens Docker

Jobs executa seu script com uv, que instala as dependências declaradas no cabeçalho # /// script em um ambiente novo. O TRL que seu script importa, portanto, vem desse cabeçalho, não da imagem, e os exemplos de hf jobs uv run acima não precisam de --image de forma alguma.

Uma imagem Docker com o TRL pré-instalado está disponível em huggingface/trl. Passá-la para hf jobs uv run fornece ao job a camada de sistema da imagem, como sua toolchain CUDA, o que é importante para dependências que compilam contra ela:

bash
python
hf jobs uv run \
    --flavor a100-large \
    --secrets HF_TOKEN \
    --image huggingface/trl \
    train.py

Para executar o TRL que está instalado na imagem, use hf jobs run no lugar. Ele executa um comando diretamente na imagem, sem um cabeçalho de script para resolver, de modo que o próprio TRL da imagem é o que executa. O -- separa o comando das opções do Jobs, o que é necessário sempre que o próprio comando aceita opções:

bash
python
hf jobs run \
    --flavor a100-large \
    --secrets HF_TOKEN \
    huggingface/trl \
    -- \
    trl sft --model_name_or_path Qwen/Qwen2-0.5B-Instruct --dataset_name trl-lib/Capybara --output_dir Qwen2-0.5B-SFT

A imagem é publicada com três tipos de tag:

TagConteúdos
X.Y.ZA liberação do TRL dessa versão, compilada quando essa versão é lançada
latestA versão mais recente
devO ramo main, reconstruído a cada mesclagem

Use dev para executar a versão de desenvolvimento, o que é útil para testar uma correção que já foi incorporada ao main, mas ainda não foi lançada:

bash
python
hf jobs run \
    --flavor a100-large \
    --secrets HF_TOKEN \
    huggingface/trl:dev \
    -- \
    trl sft --model_name_or_path Qwen/Qwen2-0.5B-Instruct --dataset_name trl-lib/Capybara --output_dir Qwen2-0.5B-SFT

As tags apenas selecionam o que está instalado na imagem, portanto, são relevantes para hf jobs run. Com hf jobs uv run, a versão vem do cabeçalho do script, e a tag não faz diferença para qual TRL é importado.

A combinação dos dois, de modo que uv resolva o cabeçalho do script enquanto alguns imports ainda vêm da imagem, requer flags extras. Consulte Reutilize os pacotes da imagem e adicione dependências com UV para esse formato e os caminhos que ele exige.

Jobs é executado em uma imagem Docker do Hugging Face Spaces ou do Docker Hub, então você também pode especificar qualquer imagem personalizada:

bash
python
hf jobs uv run \
    --flavor a100-large \
    --secrets HF_TOKEN \
    --image <docker-image> \
    train.py

O wrapper TRL Jobs inicia os scripts do TRL em Jobs com configurações predefinidas para alguns modelos, por exemplo trl-jobs sft --model_name Qwen/Qwen3-0.6B --dataset_name trl-lib/Capybara.

Update on GitHub