Это продолжение цикла статей о масштабировании тренировки и инференса LLM. Предыдущая глава находится по этой ссылке.
Итак, с основами разобрались, давайте теперь разбираться с тем, как распихать матрицы по нескольким чипам, перемножить, а затем собрать это все в удобоваримый результат. По-умному это называется шардинг.
Для начала давайте определимся, зачем этот шардинг вообще нужен. А нужен он потому что, как я уже писал в предыдущей статье, при работе с действительно большими нейронками матрицы и вектора практически никогда целиком не влезают в память одного GPU/TPU, поэтому их приходится разделять или шардировать. От того, насколько грамотно произведен шардинг, зависит то, насколько эффективно используется наш массив ускорителей, а следовательно и скорость тренировки, эффективность расхода вычислительных ресурсов и т.д.
Возьмем для примера матрицу A размера [I, J] и распределим ее на 4 ускорителя: