Ребята из Thinking Machines Labs эту всю проблему решили, придумав batch-invariant computation. Суть в чем, давайте сделаем так, чтобы ядра для матмула, рмснорма и атеншена давали одинаковый результат вне зависимости от размера батча путём того, что мы фиксируем порядок операций внутри ядра. Конкретно мысль какая: матмул по дефолту при малых батчах использует сплит-к, это когда K-размерность (по которой идёт скалярное произведение) разбивается на куски, считается параллельно, потом складывается. Batch-invariant matmul от сплит-к отказывается вообще, всегда использует один и тот же data-parallel тайлинг. С рмснормом ситуация аналогичная, фиксированное разбиение по hidden dim независимо от количества активных SM. Атеншен — фиксированный размер KV-splits, никакой адаптации под длину последовательности или загрузку GPU. Всё это впихнули в sglang под флаг —enable-deterministic-inference и получили детерминизм ценой того, что производительность пупупу. Больше чем вполовину теряем по скорости, это грусть-печаль.
Я вот это всё к чему. Я для себя недавно открыл LLM-42, Майкрософт рисерч придумал. Ребята увидели, что дивергенция, то есть расхождение, это не такое уж частое событие, большинство токенов генерируются с такой уверенностью, что даже при разных схемах редукции аргмакс возвращает одно и то же, а проблемными остаются те токены, у которых два-три кандидата с примерно одинаковыми логитами. Но в силу того, что даже один токен уже меняет контекст, последовательности начинают расходиться сильнее и все это yet again аккумулируется.
Собсна, идея простая. Генерируем последовательности обычными ядрами, проверяем раз в N времени не случилась ли дивергенция, если случилась - откатываемся, исправляем. Decode-Verify-Rollback. Реализация вот так выглядит: есть два пути, fast-path с обычным динамическим батчингом и верификатор, форвард с фиксированным батчсайзом, который дает ground truth. Фастпасс дает кандидатные токены, верификатор параллельно проверяет их, беря за точку отсчёта последний известно-корректный токен, первый токен из префилла, детерминированный по конструкции потому что префилл всегда обрабатывается одним запросом с фиксированным батчем. При этом помним что у фастпасса кв нестабилен, он был посчитан при одном батчсайзе, если его оставить, то форварды будут наследовать численный дрейф, поэтому при каждой верификации мы перезаписываем кв фастпасса консистентными значениями от верификатора.
Для того, чтобы верификатор не боттлнэчил, вместо того чтобы накапливать W (где W - батч-сайз верификатора и количество токенов, что верификатор проверяет) токенов одного запроса, мы набираем K токенов от каждого из N запросов, получая W = N * K. Таким образом верификатор работает над нормальным батчем, мы нормально утилим гпу и выигрываем в несколько (!) раз по производительности от стандартной реализации детерминизма в сгланге.
Вывод простой, люди опять изобрели спекдек и он опять великолепен.
Я вот это всё к чему. Я для себя недавно открыл LLM-42, Майкрософт рисерч придумал. Ребята увидели, что дивергенция, то есть расхождение, это не такое уж частое событие, большинство токенов генерируются с такой уверенностью, что даже при разных схемах редукции аргмакс возвращает одно и то же, а проблемными остаются те токены, у которых два-три кандидата с примерно одинаковыми логитами. Но в силу того, что даже один токен уже меняет контекст, последовательности начинают расходиться сильнее и все это yet again аккумулируется.
Собсна, идея простая. Генерируем последовательности обычными ядрами, проверяем раз в N времени не случилась ли дивергенция, если случилась - откатываемся, исправляем. Decode-Verify-Rollback. Реализация вот так выглядит: есть два пути, fast-path с обычным динамическим батчингом и верификатор, форвард с фиксированным батчсайзом, который дает ground truth. Фастпасс дает кандидатные токены, верификатор параллельно проверяет их, беря за точку отсчёта последний известно-корректный токен, первый токен из префилла, детерминированный по конструкции потому что префилл всегда обрабатывается одним запросом с фиксированным батчем. При этом помним что у фастпасса кв нестабилен, он был посчитан при одном батчсайзе, если его оставить, то форварды будут наследовать численный дрейф, поэтому при каждой верификации мы перезаписываем кв фастпасса консистентными значениями от верификатора.
Для того, чтобы верификатор не боттлнэчил, вместо того чтобы накапливать W (где W - батч-сайз верификатора и количество токенов, что верификатор проверяет) токенов одного запроса, мы набираем K токенов от каждого из N запросов, получая W = N * K. Таким образом верификатор работает над нормальным батчем, мы нормально утилим гпу и выигрываем в несколько (!) раз по производительности от стандартной реализации детерминизма в сгланге.
Вывод простой, люди опять изобрели спекдек и он опять великолепен.