Балансировка параллелизма конвейера с параллелизмом словаря

Параллелизм конвейера широко используется для масштабирования обучения крупных языковых моделей на базе трансформеров; различные работы были проведены для улучшения его пропускной способности и объема памяти. В данной статье мы рассматриваем часто упускаемую из виду проблему: слои словаря могут вызывать несбалансированное использование вычислительных ресурсов и памяти на различных этапах конвейера, что усугубляет проблемы с "пузырями" конвейера и узкими местами памяти. Чтобы решить эту проблему, мы равномерно распределяем слои словаря по устройствам конвейера и группируем вычисления в проходы конвейера. Для уменьшения накладных расходов на память активации мы предлагаем несколько алгоритмов для снижения барьеров коммуникации внутри слоев словаря. Кроме того, мы используем обобщенный метод для интеграции параллелизма словаря с существующими расписаниями конвейера. Сочетая эти методы, наши подходы эффективно балансируют вычисления и память параметров с небольшим постоянным накладным расходом на память активации. Особо следует отметить, что при сочетании с расписаниями, балансирующими память активаций, такими как V-Half, наш подход достигает идеального баланса как по памяти, так и по вычислениям. Обширные оценки показывают, что наш метод обеспечивает баланс вычислений и памяти независимо от размера словаря, что приводит к улучшению пропускной способности на 5% до 51% по сравнению с наивными подходами, одновременно значительно снижая пиковое использование памяти, особенно в сценариях с большим словарём. Наша реализация доступна в открытом доступе по адресу https://github.com/sail-sg/VocabularyParallelism.