Integração com o DeepSpeed

O TRL suporta treinamento com DeepSpeed, uma biblioteca que implementa técnicas avançadas de otimização de treinamento. Isso inclui particionamento de estado do otimizador, offloading, particionamento de gradientes e muito mais.

O DeepSpeed integra o Zero Redundancy Optimizer (ZeRO), que permite dimensionar o tamanho do modelo proporcionalmente ao número de dispositivos com alta eficiência sustentada.

Instalação

Para usar o DeepSpeed com o TRL, instale-o usando o seguinte comando:

pip install deepspeed

Executando scripts de treinamento com DeepSpeed

Nenhuma modificação no seu script de treinamento é necessária. O DeepSpeed é habilitado por meio dos argumentos de treinamento, como para qualquer Trainer: defina deepspeed como o caminho para um arquivo de configuração do DeepSpeed e inicie com torchrun.

training_args = SFTConfig(..., deepspeed="<DEEPSPEED_CONFIG_FILE.json>")

torchrun --nproc_per_node 8 train.py

Scripts que analisam seus argumentos com TrlParser, como os que estão por trás da CLI trl, obtêm-no da linha de comando:

trl sft ... --deepspeed <DEEPSPEED_CONFIG_FILE.json>

Fornecemos arquivos de configuração do DeepSpeed prontos para uso no diretório examples/deepspeed_configs. Por exemplo, para executar o treinamento com ZeRO Stage 2:

trl sft ... --deepspeed examples/deepspeed_configs/zero2.json

Os valores definidos como "auto" nesses arquivos (tamanho do lote, acúmulo de gradiente, precisão, …) são preenchidos a partir dos argumentos de treinamento.

Recursos adicionais

Consulte a documentação do DeepSpeed do 🤗 Transformers para mais informações sobre a integração com o DeepSpeed.