Utiliser plusieurs modèles avec DeepSpeed

Ce guide suppose que vous avez lu et compris le guide d'utilisation de DeepSpeed.

L'exécution de plusieurs modèles avec Accelerate et DeepSpeed est utile pour :

• Distillation de connaissances
• Les techniques de post-entraînement telles que le RLHF (voir la bibliothèque TRL pour plus d'exemples)
• Entraîner plusieurs modèles à la fois

Actuellement, Accelerate dispose d'une API très expérimentale pour vous aider à utiliser plusieurs modèles.

Ce tutoriel se concentrera sur deux cas d'usage courants :

1. La distillation de connaissances, où un modèle étudiant plus petit est entraîné à imiter un modèle enseignant plus grand et plus performant. Si le modèle étudiant tient sur un seul GPU, nous pouvons utiliser ZeRO-2 pour l'entraînement et ZeRO-3 pour fractionner l'enseignant lors de l'inférence. C'est nettement plus rapide que d'utiliser ZeRO-3 pour les deux modèles.
2. Entraîner plusieurs modèles disjoints simultanément.

Distillation de connaissances

La distillation de connaissances est un bon exemple d'utilisation de plusieurs modèles, mais avec l'entraînement d'un seul d'entre eux.

Normalement, vous utiliseriez un seul utils.DeepSpeedPlugin pour les deux modèles. Cependant, dans ce cas, il existe deux configurations distinctes. Accelerate vous permet de créer et d’utiliser plusieurs plugins si et seulement si ils sont dans un dict, afin que vous puissiez référencer et activer le bon plugin au moment opportun.

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}

Le fichier zero2_config.json doit être configuré pour l’entraînement complet (il faut donc spécifier scheduler et optimizer si vous n’utilisez pas les vôtres), tandis que zero3_config.json ne doit être configuré que pour le modèle d’inférence, comme le montre l’exemple ci-dessous.

{
    "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
}

Un exemple de configuration zero2_config.json est présenté ci-dessous.

{
    "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 lèvera une erreur si train_micro_batch_size_per_gpu n’est pas spécifié, même si ce modèle particulier n’est pas entraîné.

À partir de là, créez un seul Accelerator et passez-y les deux configurations.

from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)

Voyons maintenant comment les utiliser.

Modèle étudiant

Par défaut, Accelerate définit le premier élément du dict comme plugin par défaut ou activé (plugin "student"). Vous pouvez vérifier cela en utilisant la fonction utils.deepspeed.get_active_deepspeed_plugin() pour voir quel plugin est activé.

active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]

AcceleratorState conserve également le plugin DeepSpeed actif dans state.deepspeed_plugin.

assert active_plugin is accelerator.deepspeed_plugin

Puisque student est le plugin actuellement actif, préparons le modèle, l’optimiseur et le planificateur.

student_model, optimizer, scheduler = ...
student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)

Il est maintenant temps de s’occuper du modèle enseignant.

Modèle enseignant

Tout d'abord, vous devez spécifier dans Accelerator que la configuration zero3_config.json doit être utilisée.

accelerator.state.select_deepspeed_plugin("teacher")

Cela désactive le plugin "student" et active à la place le plugin "teacher". La configuration stateful de DeepSpeed au sein de Transformers est mise à jour, et cela modifie la configuration de plugin appelée lors de l'utilisation de deepspeed.initialize(). Cela vous permet d'utiliser l'intégration automatique du gestionnaire de contexte deepspeed.zero.Init fournie par Transformers.

teacher_model = AutoModel.from_pretrained(...)
teacher_model = accelerator.prepare(teacher_model)

Sinon, vous devez initialiser manuellement le modèle avec deepspeed.zero.Init.

with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
    model = MyModel(...)

Entraînement

À partir de là, votre boucle d'entraînement peut être ce que vous voulez, à condition que teacher_model ne soit jamais entraîné.

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()

Entraîner plusieurs modèles disjoints

L'entraînement de plusieurs modèles est un scénario plus complexe. À l'heure actuelle, nous supposons que chaque modèle est totalement indépendant des autres pendant l'entraînement.

Ce scénario nécessite toujours de créer deux utils.DeepSpeedPlugin. Cependant, vous avez également besoin d’un second Accelerator, car différents moteurs deepspeed sont appelés à des moments différents. Un seul Accelerator ne peut porter qu’une seule instance à la fois.

Puisque state.AcceleratorState est un objet avec état, il est déjà au courant des utils.DeepSpeedPlugin disponibles. Vous pouvez simplement instancier un second Accelerator sans arguments supplémentaires.

first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)
second_accelerator = Accelerator()

Vous pouvez appeler first_accelerator.state.select_deepspeed_plugin() pour activer ou désactiver un plugin particulier, puis appeler 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
)

Et maintenant, vous pouvez entraîner :

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()

Ressources

Pour voir plus d'exemples, veuillez consulter les tests associés actuellement dans [Accelerate].