TACO уместил полное дообучение 32B-модели на одной H100

Полное дообучение LLM с 32B параметров удалось провести на одной H100 с 80 ГБ памяти. Работа University of Central Florida и Mohammed VI Polytechnic University не рецензирована, и числа получили сами авторы; на OPT-13B оптимизатор TACO снизил пик памяти почти втрое относительно AdamW8bit при сопоставимых качестве и времени обучения. Обновление всех весов становится вариантом для одной GPU там, где прежде приходилось распределять обучение или переходить на адаптеры.

Почему одного элемента на столбец хватает для обновления

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

TACO сохраняет обучение первого порядка: модель по-прежнему вычисляет градиенты обратным распространением ошибки. Экономия начинается после этого. Для каждой колонки двумерной матрицы оптимизатор находит компонент градиента с наибольшим модулем, оставляет только его знак, а остальные компоненты обновления обнуляет.

Получается матрица, где в каждой колонке ненулевым остаётся не более одного элемента. Он принимает одно из трёх значений: отрицательное, нулевое или положительное. Это не дополнительное прореживание готового шага, а точное решение локальной задачи наискорейшего спуска для выбранной авторами геометрии, которая ограничивает сумму модулей отдельно в каждой колонке.

Выбор максимального компонента может меняться от пакета к пакету из-за шума градиентов. Поэтому практическая версия TACO хранит для каждой колонки небольшой набор доминирующих компонентов в FP8 и обновляет их через экспоненциальное среднее. Полный градиент всё ещё возникает во время обратного прохода, но оптимизатор сразу обрабатывает его и не сохраняет плотную историю между шагами.

Теоретическая часть объясняет, почему TACO может вести себя ближе к Adam, чем Muon. В упрощённой задаче с фиксированным входом TACO и модель Adam проходят по одному семейству обновлений и сходятся к одному решению, тогда как Muon выбирает другой путь. Это узкий теоретический сценарий, а не доказательство одинакового поведения на произвольной LLM.

Состояние оптимизатора перестало определять размер модели

На OPT-13B постоянное состояние TACO оказалось в 174 раза меньше, чем у AdamW8bit, а общий пик памяти снизился в 2,9 раза. Разница между этими отношениями показывает границу подхода: TACO почти убирает собственное состояние, но не отменяет память для весов, активаций и временного градиента.

Освободившаяся память позволила полностью дообучить Qwen3-32B на одной H100 с 80 ГБ. На проверенных задачах TACO показал сопоставимые с другими оптимизаторами качество и время работы, но не выигрывал у них по каждой метрике и на каждой модели.

Эксперименты охватывают семейства OPT, Qwen3, Pythia, Llama и Mistral, а также классификацию, выбор ответа и извлечение ответа из текста. Работа проверяет перенос готовых моделей на прикладные задачи через обновление всех параметров; выводы относятся именно к этому режиму, а результаты заметно зависят от сочетания архитектуры и задачи.

Когда TACO меняет план обучения

TACO влияет на выбор инфраструктуры, если продукту нужно менять все веса модели, а проект упирается в память состояния оптимизатора. В таком случае одна GPU может заменить распределённую схему не за счёт выгрузки состояний в оперативную память, а потому что этих состояний остаётся мало. Это упрощает экспериментальную установку и убирает обмен состояниями между ускорителями.

Если задачу уже решают LoRA или другие адаптеры, работа не даёт основания автоматически переходить на полное дообучение. TACO экономит память оптимизатора, но полный обратный проход и хранение всех весов остаются. Выбор по-прежнему зависит от того, даёт ли обновление всех параметров нужный прирост качества на конкретных данных.

Для производственного решения TACO пока разумнее рассматривать как кандидата на проверку, а не как замену AdamW по умолчанию. Реализация опирается на разреженную историю градиентов, низкую точность и специальную обработку матриц; перед сменой стека стоит воспроизвести качество, пик памяти и скорость на целевой модели. Авторы приложили код и описали настройки экспериментов, поэтому такую проверку можно провести без восстановления алгоритма по формулам.

Источники

• TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning

Иллюстрация: рисунок из статьи «TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning», Jichao Jiang (University of Central Florida), Cristian McGee (University of Central Florida), El Houcine Bergou (Mohammed VI Polytechnic University) и др., CC BY 4.0