Выполнение и откладывание заданий

При запуске обычного скрипта инструкции выполняются последовательно. Использование Accelerate для развёртывания скрипта на нескольких GPU одновременно усложняет ситуацию: хотя каждый процесс выполняет все инструкции по порядку, некоторые из них могут выполняться быстрее других.

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

accelerator.wait_for_everyone()

Эта инструкция заблокирует все процессы, пришедшие первыми, до тех пор, пока все остальные процессы не достигнут этой точки (если вы запускаете скрипт на одном GPU или CPU, это не будет иметь никакого эффекта).

Ниже приведены несколько примеров случаев, когда следует использовать эту утилиту:

Some of these are utilized with the main_process_first() context manager, which utilizes wait_for_everyone() to run a particular set of code on the main process beforehand before triggering and launching the other processes

Загрузка набора данных

При загрузке набора данных сначала скачайте его в основном процессе, а затем загрузите кэшированный набор данных

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

with accelerator.main_process_first():
    datasets = load_dataset("glue", "mrpc")

Под капотом это то же самое, что и вызов:

# First do something on the main process
if accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

# And then send it to the rest of them
if not accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

Сохранение state_dict

При сохранении state_dict модели, поскольку обычно один файл сохраняется только на основном процессе, это следует указать:

if accelerator.is_main_process:
    model = accelerator.unwrap_model(model)
    torch.save(model.state_dict(), "weights.pth")

Загрузка в state_dict

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

with accelerator.main_process_first():
    state = torch.load("weights.pth")
    model.load_state_dict(state)

Применение многопоточной CPU-операции

Применение операции map() на нескольких рабочих процессах, например токенизация, должно выполняться сначала в основном процессе, а затем распространяться на каждый из них.

datasets = load_dataset("glue", "mrpc")

with accelerator.main_process_first():
    tokenized_datasets = datasets.map(
        tokenize_function,
        batched=True,
        remove_columns=["idx", "sentence1", "sentence2"],
    )

Применение проверок, таких как Early Stopping

Чтобы проверка срабатывала на флаге, установленном конкретным процессом, следует использовать API set_trigger и check_trigger. Полезными примерами такого применения могут быть ситуации, когда используется ранняя остановка и отслеживается loss (поскольку значение loss немного различается на каждом процессе).

Вызывайте Accelerator.set_trigger(), когда ваше условие выполнено, и Accelerator.check_trigger(), чтобы проверить, выполнено ли это условие в каком-либо процессе:

for (x,y) in data_loader:
    logits = model(x)
    loss = loss_func(logits, y)
    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
    if should_do_early_stopping(loss):
        accelerator.set_trigger()

    # Later in the training script when we need to check for the breakpoint
    if accelerator.check_trigger():
        break