Shardy — это тензорная система секционирования на основе MLIR для всех диалектов. Созданная в результате сотрудничества команд GSPMD и PartIR , она объединяет лучшее из обеих систем, а также общий опыт команд и пользователей.
Преимущества
- Больше контроля и предсказуемости для пользователей за счет объединения распространения GSPMD с инкрементным секционированием PartIR.
- Новые функции, основанные на общем опыте, например новая поддержка изменений, которые, как известно, генерируют дополнительное общение, если пользователи не знают, как их обойти.
- Повышенное удобство использования и возможность отладки для повышения скорости работы конечного пользователя, например, за счет использования представления сегментирования на основе осей.
- Простая база кода с открытым исходным кодом, использующая MLIR, с более широким набором активных участников (внутренних, внешних и в разных часовых поясах) для поддержки пользователей.
Компоненты
- Представление сегментирования: представление сегментирования на основе осей, которое привязано к определенной логической сетке (из потенциально нескольких сеток) и поддерживает ограничение сегментирования измерений и осей, разделение осей для таких операций, как изменение формы, приоритеты для инкрементного разделения и многое другое.
- API-интерфейсы компилятора: набор компонентов компилятора, которые можно использовать вместе с представлением сегментирования, чтобы влиять на распространение сегментирования.
- Шардинги ввода/вывода — прикрепите сегментацию к входу или выходу основной функции, чтобы указать, что именно так должен сегментироваться тензор ввода/вывода при передаче в функцию или возврате из функции.
- Ограничение сегментирования — прикрепите сегментирование к промежуточному тензору (например, результату matmul), чтобы указать, что именно так должен быть сегментирован этот тензор или подмножество его применений.
- Shard As/Like — сгруппируйте несколько тензоров по идентификатору, чтобы указать, что они должны быть сегментированы одинаковым образом.
- Ручное вычисление — включает в себя подвычисление, которое разделено вручную с использованием подмножества осей сетки, где сегменты вдоль этих ручных осей указаны для всех входов и выходов, а внутри подвычислений типы тензоров являются локальными относительно этих сегментов.
- Распространение сегментирования: алгоритм распространения, который сочетает в себе приоритеты пользователя и ограничения сегментирования с моделями затрат компилятора и эвристикой:
- Определяемые пользователем приоритеты, например, пакетный параллелизм, затем ZeRO
- Приоритеты на основе операций, например, сначала поэлементные операции, затем матмулы и т. д.
- Более детальная эвристика, например, предпочтение размерам партии.
- SPMD Partitioner: компонент, который снижает необходимость принятия решений о распространении сегментов за счет разделения программы на программу SPMD, добавляя в процесс необходимое перемещение/форматирование данных и коллективные операции.
- В краткосрочной перспективе первоначальная реализация будет использовать текущий разделитель GSPMD SPMD.
- В долгосрочной перспективе мы планируем создать новый разделитель SPMD на основе MLIR.
Репозиторий кода
Проект Shardy находится в активной разработке, и мы ждем отзывов от сообщества открытого исходного кода. Код Shardy доступен по адресу https://github.com/openxla/shardy.