group-telegram.com/kitty_bytes/27
Last Update:
Тут Meta релизнули новую либу для обучения моделей - optimizers и я решил рассказать вам про основные особенности чуть более подробно, чем это сделали @data_secrets
Для чего она нужна? Optimizers сильно расширяют функционал контроля обучения моделей для ускорения их сходимости. По видимому, в бигтехах устали от самописно-костыльных решений и Meta решила помочь всем и сразу
На данный момент в либе единственным оптимизатором является Distributed Shampoo, с помощью которого реализованы остальные методы оптимизации: SGD, Adagrad, RMSProp и Adam. К нему добавили новые гиперпараметры max_preconditioner_dim и precondition_frequency , которые позволяют регулировать размерность и частоту обновления матрицы preconditioner - оператора, который используется для преобразования задачи оптимизации, чтобы ускорить сходимость, изменяя масштаб его градиентов. Ограничивая его размер, мы регулируем объем вычисления и затраты на память, что может быть важно при обучении больших моделей. Частота обновления preconditioner влияет на точность сходимости, однако при высоких значениях требует больше памяти.
Порадовало, что Distributed Shampoo поддерживает разные варианты распределенного обучения, включая DDP и FSDP. Для DDP предусмотрена поддержка ZeRO-1, что снижает требования к памяти и времени выполнения. Однако для FSDP требуется использование дополнительной памяти для восстановления исходных параметров. Кстати в DDP можно обучать квантизованные модельки (BF16, FP16, FP32)
Либа поддерживает чекпоинты, но в типе DTensor, который не поддерживается торчом, поэтому нужно применять методы distributed_state_dict и load_distributed_state_dict
Теперь основной вопрос - как эту всю красоту использовать?
Имея основной метод оптимизации, мы заменяем его на конфиг метода из либы (i.e. SGDGraftingConfig ), оставляя при этом старые значения некоторых гиперпараметров. Другие гиперпараметры тюнятся для повышения эффективности сходимости и этому делу посвящена целая глава в ридми проекта. Все это дело оборачивается в DistributedShampoo , который и является основным методом оптимизации.
Но что именно происходит, когда мы заменяем наш метод оптимизации на Shampoo? Понятно, что имея такой инструментарий, мы хотим применять этот оптимизатор для уточнения сходимости ранее обученных моделей. Но как подобрать верные параметры Shampoo, чтобы он не дестабилизировал модель, а действительно улучшил ее? Зная гиперпараметры заменяемого оптимизатора, мы можем вычислить последовательность и размер шагов обучения и приблизить Shampoo к этому процессу. На языке линала это происходит так:
1) Мы берём норму шага обучения из прошлого оптимизатора
2) Нормализуем шаг Shampoo, чтобы он имел ту же норму
3) Масштабируем нормализованный шаг Shampoo, умножая его на произведение нормы шага нашего оптимизатора и коэффициента обучения
В результате мы получили новый оптимизатор, настроенный на обучение нашей модели с любого этапа, который может обеспечить нам лучшую сходимость модели