Асинхронный пайплайн для обучения больших моделей можно ускорить без лишних пауз

Авторы взяли проблему, знакомую всем, кто обучает очень большие модели кусками: одни этапы пайплайна простаивают, пока другие заняты, и из-за этого обучение теряет время. Вместо того чтобы лечить эту проблему привычными способами — хранить копии весов, делать предсказания или постоянно синхронизировать все части модели — они предложили схему PACI, где версия параметров «стареет» под контролем. Проще говоря, они специально сделали так, чтобы отдельные микробатчи не убегали слишком далеко по версии весов, и добились этого локальным накоплением градиентов, то есть накапливали небольшие порции обновлений перед шагом оптимизатора.

Главная идея тут не в том, чтобы полностью убрать расхождение между прямым и обратным проходом, а в том, чтобы держать его в разумных границах. Это важно, потому что в асинхронном пайплайне именно этот разъезд версий обычно и ломает стабильность обучения. Авторы показывают, что если этот разъезд явно ограничить, можно сохранить почти ту же аккуратность обучения, что и у более осторожной схемы 1F1B-flush, но без лишних пауз в работе пайплайна.

На GPT-подобных языковых моделях PACI не потерял в итоговом качестве и устойчивости обучения, сохранил тот же пик памяти, что и сравниваемая схема, и при этом полностью загрузил пайплайн. Самый заметный практический эффект — ускорение времени до нужного качества: в лучшем случае обучение шло до цели до 1,69 раза быстрее, чем у самого быстрого flush-базлайна. То есть выигрыш получился не за счет дополнительной памяти или сложной синхронизации, а за счет более умного управления тем, насколько «несвежими» могут быть веса для разных микробатчей.