Accelerate documentation
Usando múltiplos modelos com DeepSpeed
Usando múltiplos modelos com DeepSpeed
Este guia pressupõe que você tenha lido e compreendido o guia de uso do DeepSpeed.
Executar vários modelos com Accelerate e DeepSpeed é útil para:
- Destilação de conhecimento
- Técnicas de pós-treinamento como RLHF (veja a biblioteca TRL para mais exemplos)
- Treinando vários modelos ao mesmo tempo
Atualmente, o Accelerate possui uma API muito experimental para ajudar você a usar vários modelos.
Este tutorial focará em dois casos de uso comuns:
- Destilação de conhecimento, em que um modelo aluno menor é treinado para imitar um modelo professor maior e com melhor desempenho. Se o modelo aluno couber em uma única GPU, podemos usar ZeRO-2 para o treinamento e ZeRO-3 para dividir o professor para inferência. Isso é significativamente mais rápido do que usar ZeRO-3 para ambos os modelos.
- Treinando vários modelos disjuntos ao mesmo tempo.
Destilação de conhecimento
A destilação de conhecimento é um bom exemplo do uso de múltiplos modelos, treinando apenas um deles.
Normalmente, você usaria um único utils.DeepSpeedPlugin para ambos os modelos. No entanto, neste caso, há duas configurações separadas. O Accelerate permite que você crie e use múltiplos plugins se e somente se eles estiverem em um dict, para que você possa referenciar e habilitar o plugin adequado quando necessário.
from accelerate.utils import DeepSpeedPlugin
zero2_plugin = DeepSpeedPlugin(hf_ds_config="zero2_config.json")
zero3_plugin = DeepSpeedPlugin(hf_ds_config="zero3_config.json")
deepspeed_plugins = {"student": zero2_plugin, "teacher": zero3_plugin}O zero2_config.json deve ser configurado para treinamento completo (portanto, especifique scheduler e optimizer se você não estiver utilizando os seus próprios), enquanto o zero3_config.json deve ser configurado apenas para o modelo de inferência, conforme mostrado no exemplo abaixo.
{
"bf16": {
"enabled": "auto"
},
"zero_optimization": {
"stage": 3,
"overlap_comm": true,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": "auto",
"stage3_max_reuse_distance": "auto",
},
"train_micro_batch_size_per_gpu": 1
}Um exemplo de configuração zero2_config.json é mostrado abaixo.
{
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"weight_decay": "auto",
"torch_adam": true,
"adam_w_mode": true
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
},
"gradient_accumulation_steps": 1,
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
}O DeepSpeed lançará um erro se
train_micro_batch_size_per_gpunão for especificado, mesmo que esse modelo específico não esteja sendo treinado.
A partir daqui, crie um único Accelerator e passe ambas as configurações.
from accelerate import Accelerator
accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)Agora, vamos ver como usá-los.
Modelo de estudante
Por padrão, o Accelerate define o primeiro item do dict como o plugin padrão ou habilitado (plugin "student"). Verifique isso usando a função utils.deepspeed.get_active_deepspeed_plugin() para ver qual plugin está habilitado.
active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]AcceleratorState também mantém o plugin DeepSpeed ativo salvo em state.deepspeed_plugin.
assert active_plugin is accelerator.deepspeed_pluginComo student é o plugin ativo no momento, vamos preparar o modelo, o otimizador e o agendador.
student_model, optimizer, scheduler = ... student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)
Agora é hora de lidar com o modelo professor.
Modelo professor
Primeiro, você precisa especificar no Accelerator que a configuração zero3_config.json deve ser usada.
accelerator.state.select_deepspeed_plugin("teacher")Isso desativa o plugin "student" e ativa o plugin "teacher" em vez disso. A configuração com estado do DeepSpeed dentro do Transformers é atualizada, e isso altera qual configuração de plugin é chamada ao usar deepspeed.initialize(). Isso permite que você use a integração automática com o gerenciador de contexto deepspeed.zero.Init fornecida pelo Transformers.
teacher_model = AutoModel.from_pretrained(...) teacher_model = accelerator.prepare(teacher_model)
Caso contrário, você deve inicializar manualmente o modelo com deepspeed.zero.Init.
with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
model = MyModel(...)Treinamento
A partir daqui, seu loop de treinamento pode ser o que você quiser, desde que teacher_model nunca esteja sendo treinado.
teacher_model.eval()
student_model.train()
for batch in train_dataloader:
with torch.no_grad():
output_teacher = teacher_model(**batch)
output_student = student_model(**batch)
# Combine the losses or modify it in some way
loss = output_teacher.loss + output_student.loss
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()Treinar múltiplos modelos disjuntos
Treinar múltiplos modelos é um cenário mais complicado. No estado atual, assumimos que cada modelo é completamente independente dos outros durante o treinamento.
Esse cenário ainda exige que dois utils.DeepSpeedPlugin sejam criados. No entanto, você também precisa de um segundo Accelerator, pois diferentes motores deepspeed são chamados em momentos distintos. Um único Accelerator só pode carregar uma instância por vez.
Como o state.AcceleratorState é um objeto com estado, ele já está ciente dos dois utils.DeepSpeedPlugin disponíveis. Você pode simplesmente instanciar um segundo Accelerator sem argumentos extras.
first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins) second_accelerator = Accelerator()
Você pode chamar first_accelerator.state.select_deepspeed_plugin() para habilitar ou desabilitar um plugin específico e, em seguida, chamar prepare.
# can be `accelerator_0`, `accelerator_1`, or by calling `AcceleratorState().select_deepspeed_plugin(...)`
first_accelerator.state.select_deepspeed_plugin("first_model")
first_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
first_optimizer, first_scheduler, train_dl, eval_dl = get_training_items(model1)
first_model, first_optimizer, first_scheduler, train_dl, eval_dl = accelerator.prepare(
first_model, first_optimizer, first_scheduler, train_dl, eval_dl
)
second_accelerator.state.select_deepspeed_plugin("second_model")
second_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
second_optimizer, second_scheduler, _, _ = get_training_items(model2)
second_model, second_optimizer, second_scheduler = accelerator.prepare(
second_model, second_optimizer, second_scheduler
)E agora você pode treinar:
for batch in dl:
outputs1 = first_model(**batch)
first_accelerator.backward(outputs1.loss)
first_optimizer.step()
first_scheduler.step()
first_optimizer.zero_grad()
outputs2 = model2(**batch)
second_accelerator.backward(outputs2.loss)
second_optimizer.step()
second_scheduler.step()
second_optimizer.zero_grad()Recursos
Para ver mais exemplos, consulte os testes relacionados atualmente em [Accelerate].
Update on GitHub