Развертывание Marco o1 на локальном PC. Языковая модель рассуждений

Недавно я запускал и тестировал Marco o1. Это одна из первых опенсорсных языковых моделей с многоступенчатой логикой, эта модель использует Chain-of-Thoughts и некоторые другие алгоритмы, которые помогают с решением задач на математику, логику и кодинг. Marco-o1 названа по аналогии с OpenAI o1, благодаря которой Chain-of-Thoughts промптинг и файнтюнинг получил особую популярность в GenAI индустрии.

В последнее время разные компании, в основном из Китая, стремятся повторить возможности o1. Самые впечатляющие результаты — у DeepSeek-R1-Lite-Preview, но веса этой модели не были опубликованы на момент проведения моих тестов. Однако разработчики DeepSeek R1 Lite обещали открыть доступ в свое время, и это будет очень интересно для нас.

А пока я решил поиграть с весами Marco-o1, модели хотя и легковесной, но реализующей те продвинутые алгоритмы, которые стоят за удивительными возможностями оригинальной o1. Как видно из карточки модели на HuggingFace, она создана путем файнтюнинга Qwen 2 7B на Chain-of-Thoughts датасете. Это комбинация датасетов Open-O1 и двух дополнительных наборов данных, которые разработчики из Alibaba Cloud сгенерировали, используя разные стратегии промптинга — Chain of Thoughts и обычные инструкции. Опубликована, к сожалению, только часть данных, но по ним ясно видно, какой формат использовали для файнтюнинга Chain-of-Thoughts:

Сам по себе Chain-of-Thoughts — это формат промпта, который заставляет модель строить цепочки мыслей вроде этой. Но как в случае с моделью Marco, чтобы нейросеть могла эффективно работать с таким форматом, ее нужно файнтюнить на Chain-of-Thoughts датасете. Точно так же Instruct модель требует файнтюнинга на данных с определенной структурой промптов и ответов.

Marco o1 также использует алгоритм поиска по дереву Монте-Карло (MCTS), который, согласно статье разработчиков, позволяет исследовать несколько путей рассуждения, используя показатель достоверности, полученный из логарифмических вероятностей топ-K альтернативных токенов. К этому значению вероятности для каждого токена и к пяти альтернативным значениям применяют функцию softmax и получают показатель достоверности C(i) для i-того токена. Потом вычисляют среднее значение для всех пяти альтернатив. Более высокое среднее значение показывает большую уверенность в том, что выбранный путь рассуждения является более оптимальным.

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

Дальше есть еще пара интересных идей, например, Action Selection. Чтобы использовать только что описанный алгоритм MCTS, нужно определиться, что использовать в качестве единицы в пространстве решений, или одного шага рассуждения, и для этого разработчики применяли разные подходы — как целый шаг или действие в качестве единицы, так и мини-шаг — 32 или 63 токена.

Наконец, разработчики добавили механизм саморефлексии, в виде промпта в конце каждой цепочки рассуждений, который побуждает модель проверить свои решения: «Подожди, может быть я сделал какие-то ошибки, мне нужно обдумать все сначала». Это позволяет модели критически оценивать собственные рассуждения и находить в них ошибки.

А теперь немного практической части. Веса модели Marco-o1 имеют всего 7 миллиардов параметров. Я запускал ее и локально, и в облаке, в обоих случаях особо мощная видеокарта не требуется, особенно если применять квантизацию. Локально я запустил модель на RTX 4060 c 4-битной квантизацией bitsandbytes.

У меня была идея задеплоить полноценный LLM-server, совместимый с openai или gradio клиентом, и я выбрал два решения:

Первое — Text Generation Inference: на этот сервер уже был обзор на моем YouTube-канале, он высокопроизводительный, поддерживается платформой Huggingface, предлагает квантизацию и другие полезные фичи из коробки.

Запустить TGI локально не проблема, простой и правильный путь — в докере. Иначе придется устанавливать Rust и прочие зависимости, это ни к чему. Команда для запуска:

docker run --gpus all --shm-size 1g -p 8080:80 -v $path_to_volume:/data \   ghcr.io/huggingface/text-generation-inference:2.4.1 --model-id AIDC-AI/Marco-o1 --quantize bitsandbytes-nf4

Докер должен поддерживать GPU, что вообще не проблема на Linux или на Windows с WSL2. Флаг —quantize опциональный — можете выбрать опцию bitsandbytes (по умолчанию 8 бит) или вообще без сжатия, если есть достаточно видеопамяти.

В итоге на 4060 инференс получился очень быстрый. Я использовал следующий код для фронтенда на gradio:

import gradio as gr from huggingface_hub import InferenceClient  client = InferenceClient("http://127.0.0.1:8080")  def respond(     message,     history: list[tuple[str, str]],     system_message,     max_tokens,     temperature,     top_p, ):  messages = [{"role": "system", "content": system_message}]      for val in history:         if val[0]:             messages.append({"role": "user", "content": val[0]})         if val[1]:             messages.append({"role": "assistant", "content": val[1]})      messages.append({"role": "user", "content": message})      response = ""      for message in client.chat_completion(         messages,         max_tokens=max_tokens,         stream=True,         temperature=temperature,         top_p=top_p,     ):         token = message.choices[0].delta.content          response += token         yield response system_prompt = """You are a well-trained AI assistant, your name is Marco-o1. Created by AI Business of Alibaba International Digital Business Group.  ## IMPORTANT!!!!!! When you answer questions, your thinking should be done in <Thought>, and your results should be output in <Output>. <Thought> should be in English as much as possible, but there are 2 exceptions, one is the reference to the original text, and the other is that mathematics should use markdown format, and the output in <Output> needs to follow the language of the user input. """  demo = gr.ChatInterface(     respond,     additional_inputs=[         gr.Textbox(value=system_prompt, label="System message"),         gr.Slider(minimum=1, maximum=2048, value=1, step=1, label="Max new tokens"),         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),  gr.Slider(             minimum=0.1,             maximum=1.0,             value=0.95,             step=0.05,             label="Top-p (nucleus sampling)",         ),     ], )   if __name__ == "__main__":     demo.launch()

Однако я столкнулся с одной проблемой — системный промпт. Он взят из ollama, которая с таким промптом прекрасно работает. А вот TGI выдает ошибку, когда промпт превышает определенную длину. Так что для тех, кто не хочет возиться, предпочтительнее будет ollama, с ней все просто:

ollama pull marco-o1 ollama serve

А о том, как запустить ollama как сервис, чтобы создать таким образом свой LLM-сервер, я расскажу в следующей статье.