Cognition Lab (разработчики Devin) выпустили модель для написания эффективных ядер CUDA

Модель называется Kevin 32В, буквально K(ernel D)evin. Релиз, на самом деле, был бы не самый примечательный, если бы не техника: для дообучения инженеры завели RL с мультиступенчатым фидбэком от среды.

Разбираемся:

⚪️ В качестве базовой LM взяли QwQ-32B, это ризонинг модель. Ее дообучали всего на 180 задачах из датасета KernelBench. Это задачи на PyTorch, в которых надо заменить базовые операции (например, матричные операции или свертки) оптимизированными CUDA-ядрами.

⚪️ Для дообучения использовали только RL с GRPO (наш разбор метода). Но GRPO не обычный, а многоэтапный. Модель получает фидбэк не один раз за всю генерацию, а генерирует ядро -> получает ревард -> исправляет -> получает ревард -> ….

⚪️ Чтобы не было проблем с перегрузкой контекста, вместо полного CoT на следующие итерации передают только резюме. Плюс для стабилизации всего этого каждый шаг получает не обычное вознаграждение, а дисконтирование, то есть равное дисконтированной сумме скорингов всех последующих ядер.

По метрикам видно, что относительно привычного однопроходного RL такой подход сильно бустит обучение. Интересно было бы посмотреть, сработает ли такое с обычным кодингом.

Блогпост | Модель на HF