Using multiple models with DeepSpeed

This guide assumes that you have read and understood the DeepSpeed usage guide.

Running multiple models with Accelerate and DeepSpeed is useful for:

Knowledge distillation
Post-training techniques like RLHF (see the TRL library for more examples)
Training multiple models at once

Currently, Accelerate has a very experimental API to help you use multiple models.

This tutorial will focus on two common use cases:

Knowledge distillation, where a smaller student model is trained to mimic a larger, better-performing teacher. If the student model fits on a single GPU, we can use ZeRO-2 for training and ZeRO-3 to shard the teacher for inference. This is significantly faster than using ZeRO-3 for both models.
Training multiple disjoint models at once.

Knowledge distillation

Knowledge distillation is a good example of using multiple models, but only training one of them.

Normally, you would use a single utils.DeepSpeedPlugin for both models. However, in this case, there are two separate configurations. Accelerate allows you to create and use multiple plugins if and only if they are in a dict so that you can reference and enable the proper plugin when needed.

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}

The zero2_config.json should be configured for full training (so specify scheduler and optimizer if you are not utilizing your own), while zero3_config.json should only be configured for the inference model, as shown in the example below.

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

An example zero2_config.json configuration is shown below.

{
    "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 will raise an error if train_micro_batch_size_per_gpu isn’t specified, even if this particular model isn’t being trained.

From here, create a single Accelerator and pass in both configurations.

from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)

Now let’s see how to use them.

Student model

By default, Accelerate sets the first item in the dict as the default or enabled plugin ("student" plugin). Verify this by using the utils.deepspeed.get_active_deepspeed_plugin() function to see which plugin is enabled.

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

AcceleratorState also keeps the active DeepSpeed plugin saved in state.deepspeed_plugin.

assert active_plugin is accelerator.deepspeed_plugin

Since student is the currently active plugin, let’s go ahead and prepare the model, optimizer, and scheduler.

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

Now it’s time to deal with the teacher model.

Teacher model

First, you need to specify in Accelerator that the zero3_config.json configuration should be used.

accelerator.state.select_deepspeed_plugin("teacher")

This disables the "student" plugin and enables the "teacher" plugin instead. The DeepSpeed stateful config inside of Transformers is updated, and it changes which plugin configuration gets called when using deepspeed.initialize(). This allows you to use the automatic deepspeed.zero.Init context manager integration Transformers provides.

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

Otherwise, you should manually initialize the model with deepspeed.zero.Init.

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

Training

From here, your training loop can be whatever you like, as long as teacher_model is never being trained on.

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

Train multiple disjoint models

Training multiple models is a more complicated scenario. In its current state, we assume each model is completely disjointed from the other during training.

This scenario still requires two utils.DeepSpeedPlugin’s to be made. However, you also need a second Accelerator, since different deepspeed engines are being called at different times. A single Accelerator can only carry one instance at a time.

Since the state.AcceleratorState is a stateful object though, it is already aware of both utils.DeepSpeedPlugin’s available. You can just instantiate a second Accelerator with no extra arguments.

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

You can call either first_accelerator.state.select_deepspeed_plugin() to enable or disable a particular plugin, and then call 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
)

And now you can train:

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

Resources

To see more examples, please check out the related tests currently in [Accelerate].

< > Update on GitHub