Цель: минимизация суммарных FLOPS на обучение и весь...

Цель: минимизация суммарных FLOPS на обучение и весь инференс. Для трансформера с N параметров используется стандартная аппроксимация для FLOPS: 6N на токен обучения и 2N на токен в инференсе, соответственно нужно найти argmin 6NDtr + 2NDinf, где Dtr и Dinf -- число токенов в обучении и инференсе соответственно.

Задача в работе про Шиншиллу ставилась как “при заданном вычислительном бюджете найти оптимальные N (параметры) и Dtr (токены), минимизирующие лосс предобучения”.

В текущей работе задача иная “при заданном лоссе найти N и Dtr, минимизирующие стоимость компьюта”. Стоимость компьюта зависит от будущего инференса, так что его надо оценить до обучения.

Полученные графики показывают отношение предлагаемых работой FLOPs, параметров и токенов предобучения к значениям оптимальным для Шиншиллы. Это представлено в виде кривых для разных значений лосса и потребностей в инференсе. Когда потребности в инференсе невелики, то Шиншилла оптимальна.

Например, для модели уровня качества Chinchilla-7B при потребности в инференсе 10^11 токенов оптимальнее обучить 6B модель на 1.18x оригинальных данных. Этот выбор приведёт к уменьшению итоговых FLOPs.

Но интереснее конечно не FLOPs уменьшать, а деньги. На инференсе и на обучении загрузка (Model FLOPs Utilization, MFU) железа очень разная, последовательная генерация приводит к особенно низкой загрузке. К тому же на инференсе можно сделать квантизацию и на том же железе более эффективно это считать. А ещё можно и инференс на другом более дешёвом железе делать…

В обновлённую формулу добавляются MFU обучения и инференса, а также цена за FLOPs, получаются другие кривые (подразумевая обучение на A100-80GB, инференс на A100-40GB с квантизацией INT8 и по ценам Lambda Labs). Короче, можно экономить. В некоторых условиях до 50%+.

Тренд на обучение меньших моделей подольше вовсю цветёт и разворачивается!