Accelerate documentation

Использование нескольких моделей с DeepSpeed

Accelerate

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Использование нескольких моделей с DeepSpeed

В этом руководстве предполагается, что вы прочитали и поняли руководство по использованию DeepSpeed.

Запуск нескольких моделей с помощью Accelerate и DeepSpeed полезен для:

Дистилляция знаний
Техники постобучения, такие как RLHF (см. библиотеку TRL для дополнительных примеров)
Обучение нескольких моделей одновременно

В настоящее время в Accelerate есть очень экспериментальный API, который помогает использовать несколько моделей.

В этом руководстве мы сосредоточимся на двух типичных сценариях использования:

Метод дистилляции знаний, при котором более маленькая модель-ученик обучается имитировать более крупную и высокопроизводительную модель-учителя. Если модель-ученик помещается на одном GPU, для обучения можно использовать ZeRO-2, а для шардирования учителя при инференсе — ZeRO-3. Это значительно быстрее, чем использовать ZeRO-3 для обеих моделей.
Обучение нескольких непересекающихся моделей одновременно.

Дистилляция знаний

Knowledge distillation — хороший пример использования нескольких моделей, при этом обучается только одна из них.

Обычно для обеих моделей используется один utils.DeepSpeedPlugin. Однако в данном случае есть две отдельные конфигурации. Accelerate позволяет создавать и использовать несколько плагинов только в том случае, если они находятся в dict, чтобы вы могли обращаться к нужному плагину и активировать его при необходимости.

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}

Файл zero2_config.json следует настроить для полного обучения (то есть указать scheduler и optimizer, если вы не используете собственные), тогда как zero3_config.json нужно настроить только для модели инференса, как показано в примере ниже.

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

Ниже приведен пример конфигурации 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 вызовет ошибку, если train_micro_batch_size_per_gpu не указан, даже если данная модель не обучается.

Отсюда создайте один Accelerator и передайте в него обе конфигурации.

from accelerate import Accelerator

accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)

Теперь посмотрим, как их использовать.

Студент-модель

По умолчанию Accelerate назначает первый элемент в dict в качестве плагина по умолчанию или активного плагина (плагин "student"). Проверить это можно с помощью функции utils.deepspeed.get_active_deepspeed_plugin(), которая показывает, какой плагин активен.

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

AcceleratorState также хранит активный плагин DeepSpeed в state.deepspeed_plugin.

assert active_plugin is accelerator.deepspeed_plugin

Поскольку student — это активный плагин, давайте подготовим модель, оптимизатор и планировщик.

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

Теперь пора заняться моделью-учителем.

Модель-учитель

Сначала необходимо указать в Accelerator, что следует использовать конфигурацию zero3_config.json.

accelerator.state.select_deepspeed_plugin("teacher")

Это отключает плагин "student" и вместо него включает плагин "teacher". Конфигурация DeepSpeed с состоянием внутри Transformers обновляется, и меняется то, какая конфигурация плагина вызывается при использовании deepspeed.initialize(). Это позволяет использовать автоматическую интеграцию контекстного менеджера deepspeed.zero.Init, предоставляемую Transformers.

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

В противном случае вам следует вручную инициализировать модель с помощью deepspeed.zero.Init.

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

Обучение

Отсюда ваш цикл обучения может быть любым, при условии, что teacher_model никогда не обучается.

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

Обучение нескольких непересекающихся моделей

Обучение нескольких моделей представляет собой более сложный сценарий. В текущем состоянии мы предполагаем, что во время обучения каждая модель полностью изолирована от остальных.

В этом сценарии по-прежнему необходимо создать два экземпляра utils.DeepSpeedPlugin. Однако вам также потребуется второй Accelerator, поскольку разные движки deepspeed вызываются в разные моменты времени. Один Accelerator может содержать только один экземпляр за раз.

Однако, поскольку state.AcceleratorState является объектом с состоянием, он уже знает обо всех доступных utils.DeepSpeedPlugin. Вы можете просто создать второй экземпляр Accelerator без дополнительных аргументов.

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

Вы можете вызвать first_accelerator.state.select_deepspeed_plugin(), чтобы включить или отключить конкретный плагин, а затем вызвать 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
)

А теперь можно начать обучение:

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

Ресурсы

Чтобы увидеть больше примеров, ознакомьтесь с соответствующими тестами, которые в настоящее время находятся в [Accelerate].

Update on GitHub

←DeepSpeed DDP Communication Hooks→