Обучение на TPU

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

Обучение в Notebook

Главный момент при обучении на TPU связан с notebook_launcher(). Как упоминалось в руководстве по ноутбукам, вам нужно перестроить код обучения в функцию, которую можно передать в функцию notebook_launcher(), и быть осторожным, чтобы не объявлять тензоры на GPU.

Хотя на TPU эта часть не так важна, критически важно понимать, что при запуске кода из ноутбука вы делаете это через процесс, называемый forking. При запуске из командной строки вы выполняете spawning, когда процесс python в данный момент не запущен, и вы spawn новый процесс. Поскольку ваш Jupyter notebook уже использует процесс python, вам нужно fork новый процесс из него, чтобы запустить ваш код.

Это становится важным при объявлении модели. В разветвлённых TPU-процессах рекомендуется создавать экземпляр модели один раз и передавать его в функцию обучения. Это отличается от обучения на GPU, где создаётся n моделей, градиенты которых синхронизируются и обратная передача выполняется в определённые моменты. Вместо этого один экземпляр модели разделяется между всеми узлами и передаётся между ними. Это особенно важно при обучении на TPU с ограниченными ресурсами, например, в ядрах Kaggle или в Google Colaboratory.

Ниже приведен пример функции обучения, передаваемой в notebook_launcher() при обучении на CPU или GPU:

Этот фрагмент кода основан на примере из блокнота simple_nlp_example, который можно найти здесь, с небольшими изменениями ради простоты

def training_function():
    # Initialize accelerator
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
    )

    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters["num_epochs"]
    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()

from accelerate import notebook_launcher

notebook_launcher(training_function)

По умолчанию notebook_launcher запускает 8 процессов, если Accelerate настроен для TPU

Если вы используете этот пример и объявляете модель внутри цикла обучения, то на системе с ограниченными ресурсами вы, возможно, столкнётесь с ошибкой вида:

ProcessExitedException: process 0 terminated with signal SIGSEGV

Эта ошибка крайне непонятна, но суть в том, что у вас закончилась системная оперативная память. Вы можете полностью избежать этой проблемы, переопределив функцию обучения так, чтобы она принимала один аргумент model, и объявив его в отдельной ячейке:

# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

+ def training_function(model):
      # Initialize accelerator
      accelerator = Accelerator()
-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
      )
  ...

И наконец, вызов функции обучения с параметрами:

  from accelerate import notebook_launcher
- notebook_launcher(training_function)
+ notebook_launcher(training_function, (model,))

Описанный выше обходной путь необходим только при запуске TPU-инстанса из Jupyter Notebook на сервере с ограниченными ресурсами, например Google Colaboratory или Kaggle. Если вы используете скрипт или запускаете задачу на более мощном сервере, предварительное объявление модели не требуется.

Смешанная точность и глобальные переменные

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

При использовании bf16 и Accelerate на TPU есть два «слоя»: базовый уровень и уровень операций.

На базовом уровне это включается при передаче mixed_precision="bf16" в Accelerator, например:

accelerator = Accelerator(mixed_precision="bf16")

По умолчанию на TPU torch.float и torch.double приводятся к bfloat16. Настройка, которая устанавливается, — это переменная окружения XLA_USE_BF16, значение которой равно 1.

Также можно выполнить дополнительную настройку, задав переменную окружения XLA_DOWNCAST_BF16. Если установить её значение в 1, то torch.float станет bfloat16, а torch.double — float32.

Это выполняется в объекте Accelerator при передаче downcast_bf16=True:

accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)

Использование downcasting вместо bf16 во всех местах полезно, когда вы пытаетесь вычислять метрики, записывать значения в лог и т. д., где «сырые» тензоры bf16 были бы непригодны.

Время обучения на TPU

При запуске скрипта вы можете заметить, что обучение сначала идёт исключительно медленно. Это происходит потому, что TPU сначала обрабатывают несколько батчей данных, чтобы определить, сколько памяти необходимо выделить, а затем используют эту настроенную выделенную память чрезвычайно эффективно.

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

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