В последние годы большие языковые модели (LLM) на базе архитектуры Transformer привлекли широкое внимание благодаря выдающейся производительности. Промышленные LLM должны обрабатывать длинные последовательности для предоставления качественных услуг. Однако потребление памяти возрастает квадратично с длиной последовательности, что ограничивает масштабируемость обучения длинных последовательностей. Существующие параллельные методы создают избыточные тензоры во время выполнения, что оставляет пространство для оптимизации памяти; кроме того, тензорный параллелизм (TP) не позволяет эффективно перекрывать вычисления и коммуникации. Для решения этих проблем в данной работе предлагается универсальный параллельный метод — память-эффективный тензорный параллелизм (METP), специально разработанный для ключевых вычислительных блоков обучения Transformer (то есть два последовательных матричных умножения и возможные функции между ними O=f(AB)C). METP распределяет подзадачи вычисления O между несколькими устройствами, используя коммуникацию «точка-точка» (send/recv) вместо коллективной передачи для обмена подматрицами, что позволяет избежать создания избыточных тензоров. С помощью техники двойной буферизации достигается глубокое перекрытие вычислений и передачи данных, а также предложены теоретические условия полного перекрытия для обучения длинных последовательностей Transformer. Теоретический анализ показывает, что при степени параллелизма p нагрузка по памяти METP без использования FlashAttention для вычисления внимания составляет O(1/p³); при использовании FlashAttention для вычисления многошагового само-внимания METP по сравнению с TP экономит не менее 41.7% памяти. Эксперименты демонстрируют, что на конфигурации из 8 GPU A100 METP может увеличить длину последовательности в 2.38–2.99 раза по сравнению с другими методами.

распределенное обучение;крупномасштабные языковые модели;длинные последовательности;системы машинного обучения;эффективность памяти;тензорный параллелизм