Методы обучения с низкой точностью

Accelerate предоставляет интеграции для обучения с использованием методов пониженной точности на поддерживаемом оборудовании через пакеты TransformersEngine и torchao. В этой документации описано, какое оборудование поддерживается, как настроить ваш Accelerator для использования методов пониженной точности, а также чего ожидать в процессе обучения.

Что означает обучение в FP8

Чтобы разобраться в тонкостях обучения в FP8 с PyTorch и Accelerate, ознакомьтесь с concept_guide о том, почему это может быть сложно. По сути, вместо обучения в BF16 некоторые (или все) аспекты обучения модели можно выполнять с использованием 8 бит вместо 16. Сложность заключается в том, чтобы сделать это без ухудшения итоговой производительности.

Это поддерживается только на определённом оборудовании NVIDIA, а именно:

• Всё, что выше потребительских графических карт серии 3000 (например, 4090)
• Архитектуры GPU на базе Hopper (например, H100 и H200)

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

Настройка Accelerator

В настоящее время поддерживаются два активно поддерживаемых бэкенда для FP8 (TransformersEngine и torchao), каждый из которых имеет свои возможности и конфигурации.

Для использования любого из них применяется один и тот же основной API. Просто передайте mixed_precision="fp8" в Accelerator, при выполнении accelerate config в ответ на запрос о смешанной точности, или в качестве части вашего файла config.yaml в ключе mixed_precision:

from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")

Чтобы указать бэкенд (и настроить другие аспекты конфигурации смешанной точности FP8), можно использовать один из датаклассов RecipeKwargs, например utils.AORecipeKwargs или utils.TERecipeKwargs; также это можно сделать в конфигурационном файле yaml или при запуске accelerate launch. Для новых проектов мы рекомендуем использовать TransformersEngine или torchao:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs, AORecipeKwargs
# Use TransformersEngine
kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false

Настройка TransformersEngine

TransformersEngine предлагает множество вариантов настройки того, как и какие вычисления в формате FP8 выполняются. Полный список поддерживаемых аргументов и их значения доступен в документации NVIDIA, однако для вашего удобства они также приведены в docstring класса FP8KwargsHandler.

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

Для его использования укажите backend="te" и измените любые аргументы, которые вы хотите, в рамках вашего обработчика kwargs:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs
kwargs = [TERecipeKwargs(...)]
accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs)

Или во время accelerate launch с помощью --fp8_backend=te .... Используйте accelerate launch --fp8_backend=te -h, чтобы просмотреть соответствующие аргументы.

Аналогично, это можно задать в вашем config.yaml:

mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false

Настройка torchao

torchao — это созданный на базе PyTorch модифицируемый FP8-бэкенд, который призван быть более доступным, чем два предыдущих движка. Одно из ключевых отличий ao от двух предыдущих состоит в том, что для численной стабильности, как правило, лучше сохранять первые и последние слои модели в обычной точности (FP32 или BF16), а остальные слои квантовать до FP8. В результате конфигурация для ao выглядит немного иначе:

Примечание: этот API является экспериментальным и может быть изменён

from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin
from torchao.float8 import Float8LinearConfig

fsdp2_plugin = FullyShardedDataParallelPlugin(
  fsdp_version=2,
  cpu_ram_efficient_loading=False, # CPU RAM efficient loading CANNOT work with fp8 torchao
  fsdp_auto_wrap_policy="TRANSFORMER_BASED_WRAP",
)
dynamo_plugin = TorchDynamoPlugin(
  backend="inductor",
  use_regional_compilation=True,
)
fp8_config = Float8LinearConfig(
  enable_fsdp_float8_all_gather=True, # Use FP8 all_gather in FSDP2
  pad_inner_dim=True,
)
kwargs = [AORecipeKwargs(
  config=fp8_config
)]
accelerator = Accelerator(
  mixed_precision="fp8",
  fsdp_plugin=fsdp2_plugin,
  dynamo_plugin=dynamo_plugin,
  kwarg_handlers=kwargs,
)

Или во время accelerate launch с помощью --fp8_backend=ao .... Используйте accelerate launch --fp8_backend=ao -h, чтобы просмотреть соответствующие аргументы.

Аналогично, это можно задать в config.yaml:

mixed_precision: fp8
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_version: 2
fp8_config:
  backend: AO
  pad_inner_dim: true
  enable_fsdp_float8_all_gather: true
dynamo_config:
  dynamo_backend: INDUCTOR
  dynamo_use_regional_compilation: true

Чтобы узнать больше о конкретных параметрах, которые следует использовать, см. официальный репозиторий torchao.

Примеры Zoo

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

• Один GPU
• Распределённая параллельная обработка данных (Multi-GPU)
• Полностью шардированная параллельность данных
• DeepSpeed ZeRO 1–3

Подробнее — здесь

Дополнительные материалы

Чтобы узнать больше о обучении в FP8, ознакомьтесь со следующими ресурсами:

• Наш концептуальный гид с подробной информацией о TransformersEngine и torchao
• Документация transformers-engine
• Документация torchao