Недавно удалось чуть-чуть погрузиться в Tinker (спасибо коллеге, который сделал обзор). Помимо общего положительного впечатления, одна вещь в дизайне платформы мне особенно понравилась – имплементация поддержки кастомных лосс-функций.

Пару слов о Tinker – это API, которым вы пользуетесь для обучения LLM. Вы пишете скрипт с подгрузкой данных и логикой обучения (включая лосс и эвалы), но весь инференс и обучение (sample, forward, backward, save_model) происходят на серверах Thinking Machines. То есть вы можете запустить скрипт на локальном компьютере с CPU и хорошим интернетом и на нем тюнить DeepSeek. Точнее, не весь DeepSeek, а только лоры. На это есть любопытная причина: для высокой утилизации GPU нужны большие батчи, особенно для MoE, а с лорами можно эффективно инференсить все еще одну LLM для пользователей с разными тюнами. Небольшой тред от одного из разработчиков Tinker в эту же тему. Вот пример скрипта, как может выглядеть обучение SFT.

Так вот по умолчанию Tinker дает доступ к трем лоссам: cross_entropy, importance_sampling и ppo, но вы можете заимплементировать любой свой, который будет принимать на вход (data: tensor, logprobs: tensor). Первое, что ожидаешь увидеть в таком случае – пользовательский код будет сериализовываться и отправляться по сети исполняться на сервере. Но здесь появляется очень элегантное, на мой взгляд, решение: forward_backward_custom. Forward_pass с сервера возвращает вам логпробы, по которым вы локально считаете лосс и производные, но только dLoss/dLogprobs (весов-то у вас нет). Далее, при вызове backward, сервер еще раз делает forward, считает новый лосс sum(logprobs * dLoss/dLogprobs) и по нему апдейтит веса модели. Цена за это – два forward pass’а и, как следствие, 1.5x FLOPS на шаг. Но зато Тинкеру не нужно вообще никак связываться со сторонним кодом.

Другое интересное архитектурное решение – это Clock Cycles, но об этом возможно напишу в другой раз.

PyTorch Profiling 101 with Modded-NanoGPT https://blog.underfit.ai/profiling-101-nanogpt