Tensor parallelism — метод распределения крупных тензоров между несколькими устройствами для параллельного выполнения матричных операций в больших моделях.
Определение
Tensor parallelism — это форма model parallelism, при которой отдельные тензоры и матрицы внутри модели разбиваются на части и обрабатываются параллельно несколькими GPU. В отличие от pipeline parallelism, который делит модель по слоям, tensor parallelism делит сами матричные операции, распределяя вычисления по устройствам.
Такой подход используется при обучении моделей, где матричные умножения слишком велики, чтобы уместиться в память одного GPU или выполняются слишком медленно при последовательном вычислении.
Как работает
Основная идея tensor parallelism — разделить большие матрицы на шардированные блоки и выполнять операции над ними параллельно. Например, матрица весов или входной тензор может быть разрезан вдоль одной из осей, и каждая часть обрабатывается на своём GPU.
Типичные стратегии:
- Column parallelism — разделение колонок матрицы весов между GPU.
- Row parallelism — разделение строк матрицы.
- Attention parallelism — разбиение матриц Q, K, V между устройствами.
- Projected parallelism — шардирование FFN-блоков.
После частичных вычислений GPU выполняют коллективные коммуникации (all-reduce, all-gather), чтобы собрать результат или передать данные на следующую операцию.
Эффективность tensor parallelism зависит от скорости межсоединений (NVLink, InfiniBand), так как коммуникационная стоимость может стать узким местом.
Где применяется
- Обучение больших Transformer-моделей.
- Обучение архитектур с широкими FFN-блоками.
- Распределение attention-операций между несколькими GPU.
- Комбинация в гибридных схемах с pipeline parallelism.
- Модели с большим количеством параметров, не помещающихся в один GPU.
- Суперкомпьютерные тренинги в Megatron-LM, DeepSpeed и Colossal-AI.
Практические примеры использования
В Megatron-LM tensor parallelism используется для шардирования больших матриц в FFN и attention. Например, матрица размером 12k × 48k делится на равные части между 8 GPU, уменьшая объём памяти на каждом устройстве и ускоряя вычисления.
В больших LLM (GPT-3, GPT-4-подобные архитектуры) tensor parallelism позволяет использовать очень широкие скрытые слои, которые невозможно разместить в одном GPU.
В системах с MoE tensor parallelism применяется внутри экспертов или для оптимизации роутеров, если экспертов много и они распределены по разным устройствам.
При распределённом обучении tensor parallelism часто комбинируется с data и pipeline parallelism для максимальной масштабируемости.
Ключевые свойства tensor parallelism
- Шардирование матриц — ключевой принцип распределения вычислений.
- Коммуникационные операции — обязательный этап для синхронизации частичных результатов.
- Масштабируемость — позволяет тренировать модели с экстремально широкими слоями.
- Гибридность — легко комбинируется с другими видами параллелизма.
Проблемы и ограничения
- Высокие коммуникационные накладные расходы — all-reduce может стать узким местом.
- Требование симметричности — матрицы должны делиться равномерно для оптимальной загрузки GPU.
- Чувствительность к топологии сети — NVLink или InfiniBand критичны.
- Сложность профилирования — трудно определить, где именно происходит дисбаланс.
Преимущества и ограничения
- Плюс: позволяет обучать модели, слишком большие для одного GPU.
- Плюс: снижает потребление памяти.
- Плюс: увеличивает производительность матричных операций.
- Плюс: хорошо сочетается с pipeline и data parallelism.
- Минус: требует быстрых межсоединений.
- Минус: вызывает overhead из-за all-gather и all-reduce.
- Минус: сложнее в реализации, чем data parallelism.
- Минус: высокие требования к балансировке размеров шарда.
Связанные термины
- Model parallelism
- Pipeline parallelism
- Expert parallelism
- Data parallelism
- Distributed training
- All-reduce
- Megatron-LM
- DeepSpeed