Обучение на TPU

TPU (Tensor Processing Unit) — это тип аппаратного обеспечения, специально предназначенного для эффективной обучения моделей. Accelerate поддерживает обучение на TPU, однако есть несколько моментов, о которых следует знать, в частности, компиляция графов. В этом руководстве кратко рассматривается компиляция, а для получения подробной информации обратитесь к гайду Training on TPUs with Accelerate.

Компиляция

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

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

• все тензоры в ваших батчах должны иметь одинаковую длину (например, без динамического выравнивания для задач NLP)
• ваш код должен быть статическим (например, без слоёв с циклами for, длина которых зависит от входных данных, как в случае с LSTM)

Связывание весов

Распространённый подход в проектировании языковых моделей — связать веса слоёв эмбеддингов и softmax. Однако при переносе модели на TPU (как при самостоятельном переносе, так и при передаче модели в метод prepare()) связь весов нарушается, и их необходимо связать заново.

Чтобы добавить особое поведение (например, связывание весов) в вашем скрипте для TPU, сначала установите distributed_type в DistributedType.TPU. После этого вы можете использовать метод tie_weights для связывания весов.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()