Gradient checkpointing и локальный SGD

Эффективное обучение моделей ИИ с PyTorch

Dennis Lee

Data Engineer, Amazon

Повышение эффективности обучения

Иконки, представляющие эффективность памяти, коммуникаций и вычислений.

Gradient checkpointing повышает эффективность памяти

Иконки, представляющие эффективность памяти, коммуникаций и вычислений.

Локальный SGD решает проблему эффективности коммуникаций

Иконки, представляющие эффективность памяти, коммуникаций и вычислений.

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C
- A, B не нужны для остального прямого прохода
Сохранить или удалить A и B?

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C
- A, B не нужны для остального прямого прохода
Сохранить или удалить A и B?
- Без gradient checkpointing: сохранить A, B

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C
- A, B не нужны для остального прямого прохода
Сохранить или удалить A и B?
- Без gradient checkpointing: сохранить A, B
- Gradient checkpointing: удалить A, B

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C
- A, B не нужны для остального прямого прохода
Сохранить или удалить A и B?
- Без gradient checkpointing: сохранить A, B
- Gradient checkpointing: удалить A, B
- Пересчитать A, B в ходе обратного прохода

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Что такое gradient checkpointing?

Gradient checkpointing: снижает память, выбирая сохраняемые активации
Пример: вычислить A + B = C
- Сначала вычислить A, B, затем C
- A, B не нужны для остального прямого прохода
Сохранить или удалить A и B?
- Без gradient checkpointing: сохранить A, B
- Gradient checkpointing: удалить A, B
- Пересчитать A, B в ходе обратного прохода
- Если пересчёт B дорог — сохранить его

График, иллюстрирующий gradient checkpointing с узлами и рёбрами

Trainer и Accelerator

График, сравнивающий простоту использования и гибкость настройки для Accelerator и Trainer.

Trainer и Accelerator

График, сравнивающий простоту использования и гибкость настройки для Accelerator и Trainer.

Gradient checkpointing с Trainer

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=4)

Gradient checkpointing с Trainer

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=4,
                                  gradient_checkpointing=True)

trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=dataset["train"],
                  eval_dataset=dataset["validation"],
                  compute_metrics=compute_metrics)

trainer.train()

{'epoch': 1.0, 'eval_loss': 0.73, 'eval_accuracy': 0.03, 'eval_f1': 0.05}

От Trainer к Accelerator

График, сравнивающий простоту использования и гибкость настройки для Accelerator и Trainer.

Gradient checkpointing с Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)


for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Gradient checkpointing с Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)
model.gradient_checkpointing_enable()

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Локальный SGD повышает эффективность коммуникаций

Иконки, представляющие эффективность памяти, коммуникаций и вычислений.

Что такое локальный SGD?

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

Каждое устройство вычисляет градиенты параллельно

Что такое локальный SGD?

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

Каждое устройство вычисляет градиенты параллельно
Синхронизация градиентов: главный узел обновляет параметры модели на каждом устройстве
Локальный SGD: снижает частоту синхронизации градиентов

Локальный SGD с Accelerator





for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Локальный SGD с Accelerator

from accelerate.local_sgd import LocalSGD

with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, 
              enabled=True) as local_sgd:
    for index, batch in enumerate(dataloader):
        with accelerator.accumulate(model):
            inputs, targets = batch["input_ids"], batch["labels"]
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

Локальный SGD с Accelerator

from accelerate.local_sgd import LocalSGD

with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, 
              enabled=True) as local_sgd:
    for index, batch in enumerate(dataloader):
        with accelerator.accumulate(model):
            inputs, targets = batch["input_ids"], batch["labels"]
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            local_sgd.step()

Lass uns üben!

Эффективное обучение моделей ИИ с PyTorch

Preparing Video For Download...