tgoop.com/quant_prune_distill/102
Last Update:
Memory-Efficient Selective Fine-Tuning
[Статья][Кода нет]
Введение
С ростом размера моделей, используемых в разных приложениях, все острее встает проблема о том, как обучать нейросети в условиях ограниченной памяти GPU.
На GPU надо как-то разместить модель + состояния оптимизатора, и промежуточные активации, необходимые для большинства операций на обратном проходе. Для экономии памяти на состояниях оптимизатора существуют различные методы PEFT. Но память, расходуемая на хранение активаций, тоже бывает весьма существенной, особенно при обучении на больших батчах и длинных последовательностей. Есть давно известный способ - gradient checkpointing - позволяющий экономить расход по памяти за счет дополнильных вычислений, но в идеале хотелось бы обойтись без этого.
Метод
В рассматриваемой статье предлагается делать backprogation только через часть выбранных случайным образом токенов (одинаковых для всех блоков трансформера), а для остальных можно высвобождать активации после выполнения операции. При длине последовательности L и количестве выбранных токенов k экономия по памяти L/k раз.
Предложенный подход применяется к задачам классификации и в качестве конечного состояния трансформера, подающегося в классификатор, используют усредненный эмбеддинг выбранных токенов.
Результаты
Авторы проверяют свой метод на бенчмарках из GLUE, используя маленькую по современным меркам BERT-Large модель. При одинаковых гиперпараметрах процедуры обучения (за исключением подбираемого индивидуально под каждый метод learning rate) Selective finetuning c k=16 не сильно уступает full-finetuning и слегка опережает LoRA.
Далее авторы показывают, что их метод гораздо лучше масштабируется с размером батча по сравнению с альтернативами. Кроме того, Selective finetuning можно применять в связке с LoRA для экономии на активациях и состояниях оптимизатора одновременно.
Потом приводится ablation качества с количеством выбранных токенов. Первоначально качество модели сильно улучшается с ростом k, но затем выходит почти на плато.
Вывод
Идея проста в реализации и довольно интуитивна. Интересно, насколько хорошо полученные в статье результаты будут обобщаться на другие задачи, например, пресловутый instructing finetuning. Да и GLUE содержит преимущественно последовальности длиной в несколько десятков токенов. Идея об экономии памяти за счет использования части токенов не нова, например, работы по token pruning и merging (DynamicViT, Differentiable Patch Selection, ToMe и др) предлагают оставлять только часть или объединять в один несколько токенов по ходу прогонки последовательности через трансформер. Здесь же длина последовательности фиксирована и экономия достигается за счет backward pass.
Для полноты картины нехватает сравнения с gradient checkpointing по скорости шага обучения и расходу памяти.
BY КПД
Share with your friend now:
tgoop.com/quant_prune_distill/102