В последние годы огромная языковая модель (LLM), основанная на архитектуре Transformer, вызвала широкий интерес благодаря своим выдающимся характеристикам. Промышленная LLM должна обрабатывать длинные входные последовательности для обеспечения высококачественного обслуживания. Однако потребление памяти растет квадратично с увеличением длины последовательности, что ограничивает расширяемость тренировки длинных последовательностей. Нынешние параллельные методы создают избыточные тензоры в процессе выполнения, что предоставляет пространство для оптимизации памяти; в то же время, тензорный параллелизм (TP) не может обеспечить эффективное перекрытие вычислений и коммуникаций. Для решения вышеупомянутых проблем в данной статье предлагается общий параллельный метод - память-эффективный тензорный параллелизм (METP), специально разработанный для основной вычислительной единицы Transformer (то есть двух последовательных умножений матриц и возможных операций функции между ними O=f (AB)C). METP распределяет подзадачи вычисления O на несколько устройств, используя точка-точка связь (отправка/прием) вместо обмена коллективной связью для завершения вычислений, избегая создания избыточных тензоров. Через технику двойной буферизации достигается глубокое перекрытие вычислений и коммуникаций, и предложены полностью перекрывающиеся теоретические условия для тренировки длинных последовательностей Transformer. Теоретический анализ показывает: когда параллелизм равен p, METP в затратах памяти в случае отсутствия вычисления внимания FlashAttention составляет O(1/p3); при использовании FlashAttention для вычисления многоголового самовнимания, по сравнению с TP может быть сэкономлено как минимум 41,7% памяти. Эксперименты подтверждают, что настройка на 8 блоках GPU A100 METP может увеличить длину последовательности в 2,38–2,99 раза по сравнению с другими методами.
Keywords
Распределенное обучение; Большая языковая модель (LLM); Длинная последовательность; Система машинного обучения; Эффективность памяти; Тензорный параллелизм