Как BERT использует память TPU?

Question

README в репозитории Google BERT гласит, что даже одно предложение длиной 512 не может вместиться в 12 ГБ Titan X для модели BERT-Large.

Но в BERTВ документе говорится, что для обучения BERT-Large используются 64 чипа TPU с максимальной длиной 512 и размером пакета 256. Как они могут вместить партию, превышающую в 256 раз большую, только в 171 раз больше памяти?

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

• TPU: Предположим, что TPUv3 используется в предварительном обучении,общий объем памяти ТПУ составляет 32 ГБ / чип * 64 чипа = 2048 ГБ.Согласно документу, размер пакета 256 с максимальной длиной 512 хорошо работает в этой конфигурации, что означает, что 8 ГБ памяти может содержать одну выборку. Кроме того, использование памяти на выборку сократится только до 4ГБ, если используется GPUv2.
• GPU: Титан X объемом 12 ГБ не может вместить даже один образец длиной 512 .

Почему используется памятьГПУ намного больше?Означает ли это, что потребление памяти на TPU оптимизировано лучше, чем на GPU?

Answer 1

Вероятно, это связано с продвинутым компилятором, который поставляется с TPU и оптимизирован для операций с тензорным потоком. Как говорит readme - проблемы нехватки памяти в BERT,

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

Однако при компиляции TPU используется специальная XLA (предметно-ориентированный компилятор для линейной алгебры, которая оптимизирует вычисления TensorFlow), которая называется fusion

может объединить несколько команд из разных операций TensorFlow в одно вычисление. Операция TensorFlow, соответствующая корневой инструкции в объединении, используется в качестве пространства имен операции объединения.

С другой стороны, работа на GPU с ванильным TF практически не имеет (или очень ограничена) оптимизаций.