Добавление Accelerate в ваш код

Каждая фреймворк распределённого обучения имеет собственный подход, который может потребовать написания большого количества собственного кода для адаптации к вашему коду обучения на PyTorch и среде обучения. Accelerate предлагает удобный способ взаимодействия с этими фреймворками распределённого обучения без необходимости изучать специфические детали каждого из них. Accelerate берёт на себя заботу об этих деталях, поэтому вы можете сосредоточиться на коде обучения и масштабировать его в любую среду распределённого обучения.

В этом руководстве вы узнаете, как адаптировать существующий код PyTorch с помощью Accelerate и начать легко обучать модели на распределённых системах! Вы начнёте с базового цикла обучения PyTorch (предполагается, что все объекты обучения, такие как model и optimizer, уже настроены) и постепенно интегрируете в него Accelerate.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Ускоритель

Класс Accelerator является основным для адаптации вашего кода к работе с Accelerate. Он знает о конфигурации распределённой системы, которую вы используете, например, о количестве различных процессов и типе вашего оборудования. Этот класс также предоставляет доступ ко многим необходимым методам для обеспечения работы вашего кода на PyTorch в любой среде распределённого обучения, а также для управления и выполнения процессов на разных устройствах.

Поэтому всегда начинайте с импорта и создания экземпляра Accelerator в вашем скрипте.

from accelerate import Accelerator

accelerator = Accelerator()

Объект Accelerator также знает, на какое устройство следует переместить ваши объекты PyTorch, поэтому рекомендуется доверить это Accelerate.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Подготовка объектов PyTorch

Далее вам необходимо подготовить объекты PyTorch (модель, оптимизатор, планировщик и т. д.) для распределённого обучения. Метод prepare() обеспечивает размещение вашей модели в подходящем контейнере (например, один GPU или несколько GPU) для вашей конфигурации обучения, адаптацию оптимизатора и планировщика для использования AcceleratedOptimizer и AcceleratedScheduler из Accelerate, а также создание нового загрузчика данных, который может быть разбит на части между процессами.

Accelerate подготавливает только объекты, наследующиеся от соответствующих классов PyTorch, таких как torch.optim.Optimizer.

Объекты PyTorch возвращаются в том же порядке, в котором они были отправлены.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

Цикл обучения

Наконец, удалите вызовы to(device) для входных данных и целевых значений в цикле обучения, так как классы DataLoader от Accelerate автоматически размещают их на нужном устройстве. Также следует заменить обычный проход backward() на метод backward() от Accelerate, который масштабирует градиенты за вас и использует подходящий метод backward() в зависимости от вашей распределённой конфигурации (например, DeepSpeed или Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Соберите всё вместе, и ваш новый цикл обучения с Accelerate должен выглядеть так!

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

Возможности обучения

Accelerate предлагает дополнительные функции — такие как накопление градиентов, обрезание градиентов, обучение с смешанной точностью и другие — которые можно добавить в ваш скрипт для улучшения процесса обучения. Давайте рассмотрим эти три функции.

Накопление градиентов

Накопление градиентов позволяет обучать модели на более крупных размерах батчей, накапливая градиенты в течение нескольких батчей перед обновлением весов. Это может быть полезно для обхода ограничений по памяти. Чтобы включить эту функцию в Accelerate, укажите параметр gradient_accumulation_steps в классе Accelerator и добавьте контекстный менеджер accumulate() в ваш скрипт.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Обрезка градиентов

Обрезка градиентов — это техника, предотвращающая «взрыв градиентов», и Accelerate предлагает:

• clipgrad_value() для ограничения градиентов минимальным и максимальным значениями
• clipgrad_norm() для нормализации градиентов до заданного значения

Смешанная точность

Смешанная точность ускоряет обучение, используя тип данных с меньшей точностью, например fp16 (half-precision), для вычисления градиентов. Для достижения наилучшей производительности с Accelerate потерю следует вычислять внутри вашей модели (как в моделях Transformers), поскольку вычисления вне модели выполняются с полной точностью.

Установите тип смешанной точности, используемый в Accelerator, а затем воспользуйтесь контекстным менеджером autocast(), чтобы автоматически приводить значения к указанному типу данных.

Accelerate обеспечивает автоматическое смешанное разрешение, поэтому autocast() требуется только в том случае, если присутствуют другие операции со смешанным разрешением, помимо тех, что выполняются над потерей в backward(), который уже обрабатывает масштабирование.

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Сохранение и загрузка

Accelerate также позволяет сохранять и загружать модель после завершения обучения, а также сохранять состояние модели и оптимизатора, что может быть полезно для возобновления обучения.

Модель

После завершения всех процессов распакуйте модель с помощью метода unwrap_model() перед сохранением, так как метод prepare() обернул вашу модель в соответствующий интерфейс для распределённого обучения. Если не распаковать модель, при сохранении словаря состояния модели будут сохранены и потенциальные дополнительные слои из более крупной модели, и вы не сможете загрузить веса обратно в базовую модель.

Для распаковки и сохранения словаря состояния модели следует использовать метод save_model(). Этот метод также позволяет сохранять модель в разбитые контрольные точки или в формате safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

Состояние

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

Чтобы дополнительно настроить, где и как сохраняются состояния с помощью save_state(), используйте класс ProjectConfiguration. Например, если включён automatic_checkpoint_naming, каждая сохранённая контрольная точка хранится в Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Все остальные объекты с состоянием, которые необходимо сохранить, следует регистрировать с помощью метода register_for_checkpointing(), чтобы их можно было сохранить и загрузить. Каждый объект, передаваемый в этот метод для сохранения, должен иметь функции load_state_dict и state_dict.

Если у вас установлен torchdata>=0.8.0, вы дополнительно можете передать use_stateful_dataloader=True в ваш DataLoaderConfiguration. Это расширяет классы DataLoader в Accelerate функциями load_state_dict и state_dict и обеспечивает отслеживание того, насколько далеко прочитан обучающий набор данных, при сохранении модели с помощью Accelerator.save_state и Accelerator.load_state.