Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
В PyTorch удалось выполнять цепочки операций над пакетами с несколькими осями переменного размера без расчётов на пустом дополнении и ручной передачи смещений. В нерецензированном препринте DanLing Team, где все числа получили сами авторы, типичное ускорение в скомпилированном режиме составило 3,39 раза по четырём масштабам BERT. Такой тензор может заменить собственный слой упаковки там, где модель создаёт, преобразует и сворачивает многомерные структуры разного размера.
Структура данных остаётся внутри тензора
Обычный пакет выравнивает все примеры по самому крупному. Для последовательностей это добавляет пустые токены, а для парного представления размером N × N лишняя работа растёт по двум осям: объект вдвое короче максимального занимает лишь четверть выделенного блока.
Упаковка решает первую часть задачи: полезные значения из разных примеров можно склеить в один буфер. Но после этого операции не знают, где проходят границы примеров, какие строки относятся к одной матрице и в каком порядке располагались логические оси. Обычно модель передаёт индексы и смещения вручную либо распаковывает данные перед очередным преобразованием.
DanLing NestedTensor хранит вместе с буфером размеры каждого примера, иерархические границы сегментов и порядок осей. Эти сведения меняются вместе с данными. Расширение размерности создаёт новую переменную ось, линейный слой или нормализация сохраняют структуру, а свёртка по оси удаляет использованную часть структуры.
Например, модель может взять представления отдельных элементов, развернуть их в парную матрицу, изменить признаки каждой пары и усреднить матрицу обратно до исходной последовательности. Код модели выглядит как обычные операции PyTorch: ни одна строка не строит индексы упаковки. Обработчики NestedTensor охватывают 404 операции верхнего уровня и низкоуровневые операции ATen.
Та же структура проходит через автоматическое дифференцирование. При компиляции схема в каждой точке программы остаётся фиксированной: известны число осей, их назначение и расположение признаков. Конкретные длины и границы сегментов приходят как тензоры во время выполнения, поэтому программа может повторно использовать скомпилированный граф для разных наборов длин.
Экономия растёт вместе с разбросом размеров
На четырёхблочной нагрузке в стиле Pairformer DanLing работал в 2,40–4,32 раза быстрее варианта с дополнением. Верхняя граница пришлась на пакет с сильным разбросом длин: обычный вариант продолжал обрабатывать квадрат по максимальному объекту, а упакованный — только существующие пары.
В том же режиме пиковое выделение памяти сократилось с 38,08 до 5,41 ГиБ. Это особенно существенно для моделей белков и других систем, где парное состояние сохраняется между блоками, а не существует несколько операций внутри специализированного ядра.
Ускорение даёт не только отсутствие пустых ячеек. Когда явная упаковка и NestedTensor вызывали ядра отдельно для каждого примера, накладные расходы интерфейса не превышали 1,60%. Основной выигрыш появлялся, когда операции получали весь пакет и обрабатывали его сегментами за один вызов.
Замеры проводили на A100 с BERT, GPT-2, энкодер-декодерным Transformer, ViT, DETR, FCN и нагрузками в стиле Pairformer. Сравнивали дополнение, DanLing, нативный torch.jagged и ручную упаковку; модели инициализировали случайно, поэтому работа измеряет время и память выполнения, а не качество обучения. Основные тесты динамической компиляции сохраняли размер пакета постоянным, меняя длины примеров.
Когда NestedTensor меняет архитектурный план
Работа меняет планы для систем, где переменные размеры проходят через значительную часть модели. Упаковать только attention недостаточно, если проекции, нормализация, остаточные связи и парные состояния по-прежнему считают дополненные позиции. NestedTensor позволяет сохранить упаковку между этими операциями и не превращать управление смещениями в отдельный слой прикладного кода.
Наиболее прямой сценарий — модель с несколькими переменными осями: изображения разного разрешения, прямоугольные или квадратные парные состояния, переходы от одиночных признаков к парным и обратно. Нативный torch.jagged рассчитан на одну такую ось, тогда как DanLing сохраняет несколько осей и различает структуры, которые после упаковки случайно получили одинаковое число строк.
Если пакетирование уже собирает почти одинаковые примеры, причина миграции слабее. В опыте с IMDB группировка близких длин подняла заполнение пакета до 97,7%, после чего преимущество DanLing в обычном режиме сократилось до 1,01 раза. Упаковка не заменяет группировку по размеру: первая сохраняет выбранный состав пакета, а вторая уменьшает объём пустой работы заранее.
Перед внедрением придётся сопоставить граф модели со списком поддержанных операций. Представление компонуется только там, где обработчики умеют корректно обновлять границы и порядок осей; собственные операции потребуют такой же структурной логики. Практический повод для прототипа появляется, если профилирование показывает большие дополненные области, а данные нельзя удобно сгруппировать без изменения порядка или состава пакетов.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
