MegascaleXLA — это компилятор и среда выполнения, обеспечивающие крупномасштабное обучение на TPU. Он реализует примитивы коллективной связи, позволяющие нескольким срезам TPU взаимодействовать, что дает возможность запускать задачи обучения, выходящие за пределы одной области ICI.
В руководстве по отладке рассматривается, как выявлять и диагностировать источники проблем с производительностью, таких как замедление работы, зависания или ошибки в многокомпонентном задании, выполняемом с помощью Megascale.
Терминология
Ломтик
- Срез представляет собой набор чипов, расположенных внутри одного и того же модуля TPU и соединенных высокоскоростными межчиповыми соединениями (ICI). Срезы описываются с помощью чипов или TensorCores, в зависимости от версии TPU.
- Multislice — это группа срезов, расширяющая возможности подключения TPU за пределы межчиповых соединений (ICI) и использующая сеть центров обработки данных (DCN) для передачи данных за пределы среза. Данные внутри каждого среза по-прежнему передаются по ICI. Благодаря этой гибридной связи Multislice обеспечивает параллелизм между срезами и позволяет использовать большее количество ядер TPU для одной задачи, чем может вместить один срез.
- TPU можно использовать для выполнения задания как на одном, так и на нескольких срезах.
RapidEye
- RapidEye — это система, цель которой — предоставить глобальную инфраструктуру отладки машинного обучения для быстрого выявления и определения первопричин проблем, вызванных неисправным оборудованием или программными ошибками. Она отслеживает задания MegaScale для автоматического обнаружения, анализа и классификации событий зависания. Процесс включает сбор данных от всех обработчиков заданий, координацию действий при возникновении зависания и генерацию сводного файла-дайджеста для каждого события.
- RapidEye включен по умолчанию для всех многослойных рабочих нагрузок. Диагностику можно найти в диспетчере ресурсов Pathways или координаторе MXLA (slice0 task0 для многоконтроллерной рабочей нагрузки JAX). RapidEye также используется для автоматического удаления неисправных TPU и сетевых адаптеров на основе данных RapidEye, полученных в масштабах всего парка устройств.
Мегамасштабный Коллектив
- Коллективные операции XLA поддерживаются через примитивы Megascale XLA (MXLA), которые не могут быть непосредственно использованы конечным пользователем. На момент написания этой статьи примитивы MXLA включают в себя коллективные операции AllGather, AllReduce, AllToAll, ReduceScatter и OneToOne. В настоящее время поддерживаются следующие операции редукции: суммирование, вычисление максимума/минимума и умножение.