Accelerate documentation
Uso de varios modelos con DeepSpeed
Uso de varios modelos con DeepSpeed
Esta guía asume que has leído y comprendido la guía de uso de DeepSpeed.
Ejecutar varios modelos con Accelerate y DeepSpeed es útil para:
- Destilación de conocimiento
- Técnicas de posentrenamiento como RLHF (consulte la biblioteca TRL para más ejemplos)
- Entrenar varios modelos a la vez
Actualmente, Accelerate cuenta con una API muy experimental para ayudarte a utilizar varios modelos.
Este tutorial se centrará en dos casos de uso comunes:
- Destilación de conocimiento, donde un modelo estudiante más pequeño se entrena para imitar a un modelo maestro más grande y con mejor rendimiento. Si el modelo estudiante cabe en una sola GPU, podemos usar ZeRO-2 para el entrenamiento y ZeRO-3 para dividir el maestro para la inferencia. Esto es significativamente más rápido que usar ZeRO-3 para ambos modelos.
- Entrenar varios modelos disjuntos a la vez.
Destilación de conocimiento
La destilación de conocimiento es un buen ejemplo del uso de varios modelos, pero entrenando solo uno de ellos.
Normalmente, se utilizaría un único utils.DeepSpeedPlugin para ambos modelos. Sin embargo, en este caso, hay dos configuraciones separadas. Accelerate permite crear y usar varios plugins si y solo si están en un dict para que puedas referenciar y habilitar el plugin adecuado cuando sea necesario.
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}El zero2_config.json debe configurarse para el entrenamiento completo (es decir, especifica scheduler y optimizer si no estás utilizando los tuyos), mientras que zero3_config.json solo debe configurarse para el modelo de inferencia, como se muestra en el ejemplo siguiente.
{
"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
}A continuación se muestra un ejemplo de configuración zero2_config.json.
{
"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",
}DeepSpeed generará un error si no se especifica
train_micro_batch_size_per_gpu, incluso si este modelo en particular no se está entrenando.
A partir de aquí, cree una única instancia de Accelerator y pase ambas configuraciones.
from accelerate import Accelerator
accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)Ahora veamos cómo usarlos.
Modelo estudiante
Por defecto, Accelerate establece el primer elemento del dict como el plugin predeterminado o habilitado (plugin "student"). Verifique esto usando la función utils.deepspeed.get_active_deepspeed_plugin() para ver qué plugin está habilitado.
active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]AcceleratorState también guarda el plugin de DeepSpeed activo en state.deepspeed_plugin.
assert active_plugin is accelerator.deepspeed_pluginDado que student es el plugin activo actualmente, preparemos el modelo, el optimizador y el programador.
student_model, optimizer, scheduler = ... student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)
Ahora es el momento de ocuparse del modelo maestro.
Modelo maestro
Primero, debe especificar en Accelerator que se debe usar la configuración zero3_config.json.
accelerator.state.select_deepspeed_plugin("teacher")Esto deshabilita el plugin "student" y habilita el plugin "teacher" en su lugar. La configuración con estado de DeepSpeed dentro de Transformers se actualiza, y cambia la configuración de plugin que se llama al usar deepspeed.initialize(). Esto permite usar la integración automática del gestor de contexto deepspeed.zero.Init que proporciona Transformers.
teacher_model = AutoModel.from_pretrained(...) teacher_model = accelerator.prepare(teacher_model)
En caso contrario, debe inicializar manualmente el modelo con deepspeed.zero.Init.
with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
model = MyModel(...)Entrenamiento
A partir de aquí, tu bucle de entrenamiento puede ser el que prefieras, siempre y cuando teacher_model no se esté entrenando.
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()Entrenar varios modelos disjuntos
Entrenar varios modelos es un escenario más complicado. En su estado actual, asumimos que cada modelo es completamente independiente de los demás durante el entrenamiento.
Este escenario sigue requiriendo la creación de dos utils.DeepSpeedPlugin. Sin embargo, también se necesita un segundo Accelerator, ya que se invocan motores deepspeed diferentes en momentos distintos. Un único Accelerator solo puede gestionar una instancia a la vez.
Sin embargo, dado que state.AcceleratorState es un objeto con estado, ya conoce los utils.DeepSpeedPlugin disponibles. Solo tienes que instanciar un segundo Accelerator sin argumentos adicionales.
first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins) second_accelerator = Accelerator()
Puede llamar a first_accelerator.state.select_deepspeed_plugin() para habilitar o deshabilitar un plugin en particular y, a continuación, llamar a 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
)Y ahora puedes entrenar:
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 más ejemplos, consulte las pruebas relacionadas que actualmente se encuentran en [Accelerate].
Update on GitHub