Промпт, цель и запуск в Colab с Codex!

Разработчик из OpenAI проверил, может ли Codex по одному промпту создать и обучить 10-миллионный трансформер на задаче сложения -- от написания кода до запуска в Google Colab и аудита результатов через субагентов.

Это AI-перевод статьи, сделанный каналом Про AI: Лучшие Статьи и Исследования.

Промпт, цель и запуск в Colab с Codex!

A Prompt, a Goal, and a Colab Run with Codex! Автор: Vaibhav (VB) Srivastav Оригинальный текст:

Несколько дней назад @swyx опубликовал челлендж: написать трансформер примерно на 10 миллионов параметров на JAX, Flax и Optax, запустить его в бесплатном Colab и обучить на простом сложении чисел.

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

Я скопировал промпт почти дословно и добавил одну инструкцию: /goal. В Codex команда goal привязывает цель к потоку работы и не даёт отклониться от задачи.

TL;DR: Codex создал ноутбук в Colab с decoder-only трансформером на 10 652 557 параметров. Обучение заняло 4000 шагов за 1139,5 секунды -- около 19 минут. Модель достигла точности 99/100 на случайных примерах и 199/200 на верификационном проходе.

Результат 199/200 на верификации был получен из мониторинга вывода Colab. Исходный ноутбук также включал проверку на 100 точных примерах.

Промпт

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

Для работы с Colab требовалась авторизованная сессия Google пользователя. Codex использовал для этого расширение Chrome для Codex. Расширение позволяет Codex управлять браузером Chrome пользователя, включая страницы, зависящие от cookies, состояния авторизации или интерфейса, доступного только в браузере.

Первая неудача

Первая проблема возникла с редактированием ячеек Colab. Редактор ноутбуков хранит состояние, имеет автосохранение, диалоги выбора среды выполнения, фокус клавиатуры и поведение редактора Monaco.

Сначала Codex попытался переписать ячейку напрямую. Некоторые правки, казалось, применились, но затем возвращалось устаревшее содержимое. Практическое решение оказалось простым: скопировать всю ячейку, установить фокус на редактор, выделить всё, очистить, вставить и убедиться, что активная ячейка содержит ожидаемые маркеры.

Инспекция через браузер также стала ненадёжной во время выполнения ячеек. Codex переключился на использование сохранённого вывода ноутбука и метаданных Drive как более надёжного способа отслеживания.

Настройка обучения

Словарь состоял ровно из 13 символов: цифры 0-9, пробел, + и =. Самый длинный возможный пример -- 999+999=1998, поэтому каждую сгенерированную последовательность можно было дополнить до фиксированной длины 12. Модель представляла собой decoder-only трансформер на Flax с шестью слоями, шестью головами внимания, шириной 384 и шириной MLP 1536.

После инициализации ноутбук вывел:

Это попало в запрошенный диапазон ~10M без последующей корректировки архитектуры.

Генератор данных семплировал операнды из диапазона 0..999 и форматировал строки вида 123+456=579. Использовалось сочетание равномерных и сбалансированных по длине примеров, что помогало более коротким числам появляться достаточно часто во время обучения.

Функция потерь применялась только к ответу: после a+b= модель училась генерировать цифры ответа, а пробел в конце использовался как стоп-токен (stop token), когда оставалось место.

Ноутбук давал Codex достаточно обратной связи, чтобы определить, работает ли обучение. Он выводил видимые устройства JAX, генерировал примеры, выполнял три пробных шага, а затем проводил 4000 полных шагов обучения, печатая значение функции потерь, точность по токенам (token accuracy) и прошедшее время каждые 400 шагов.

TPU, затем T4

Более ранний запуск в сессии, по-видимому, завершился на бесплатной среде TPU. Во время чистого мониторируемого повторного запуска бэкенд TPU в Colab был недоступен.

Исходный промпт допускал T4 как запасной вариант. Codex переключил среду выполнения, перезапустил ноутбук и отслеживал обучение до завершения.

На T4 всё получилось:

Значение token acc -- это точность по токенам (token accuracy) ответа в текущем батче во время обучения. Проверки на случайных примерах давали сигнал на уровне целых последовательностей.

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

Эти проверки показали, что ноутбук успешно обучил модель, которая обобщает на небольших случайных выборках. Для более сильного утверждения потребовались бы более крупные оценочные группы или исчерпывающий проход по всему домену из 1 000 000 упорядоченных пар операндов от 0..999 + 0..999.

Аудит

После запуска я попросил Codex использовать субагентов (subagents) для аудита результата с разных сторон: соответствие требованиям, механика модели, генерация данных, качество оценки, воспроизводимость в Colab, удобство ноутбука и скептическое принятие.

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

Аудит показал, что дизайн модели был правильным для данной задачи. Каузальная маска была корректной. Сдвинутая языково-модельная целевая функция соответствовала жадному предсказанию. Количество параметров было в диапазоне. Словарь и фиксированное дополнение до нужной длины соответствовали промпту.

Оценки требовали доработки. Финальные точные проверки покрывали 100 случайных примеров и 200 верификационных примеров. Оба набора брались из тех же примеров, что использовались при обучении. Точность по токенам (token accuracy) при обучении бралась из текущего батча. Пробел служил одновременно и символом дополнения, и стоп-токеном. Ноутбуку также требовались зафиксированные версии пакетов, чекпоинт и более чёткий путь быстрого отказа при работе на CPU.

P.S. Ещё многое предстоит улучшить, прежде чем это станет по-настоящему полезным, но направление верное!

Что использовал Codex

• /goal: привязывал финишную черту к потоку работы.
• Chrome: давал Codex доступ к авторизованной сессии Colab.
• Локальные файлы и инспекция через шелл: давали Codex стабильный источник для сравнения с живым ноутбуком и сохранёнными результатами.
• Субагенты: разделяли финальный ревью на фокусные проходы аудита.

Как сделать надёжнее

Если бы я запускал это снова, я бы сделал ноутбук надёжнее пятью способами.

Во-первых, сделать среду выполнения явной. Напечатать бэкенд, тип ускорителя и версии пакетов в начале. Быстро завершать работу на CPU. Обрабатывать недоступность TPU как чёткое ветвление, затем переключаться на T4, когда он доступен.

Во-вторых, разделить режим пробного запуска и полный режим. Пробный режим должен выполнить несколько шагов и доказать, что код работает. Полный режим должен запустить реальный цикл обучения и печатать прогресс через предсказуемые интервалы.

В-третьих, усилить оценку. Сохранить маленькую случайную проверку для быстрой обратной связи, затем добавить группы по длине операнда, поведению при переносе разряда, длине результата и граничным случаям вроде 999+999. Добавить опциональный исчерпывающий проход по всем 1 000 000 парам операндов.

В-четвёртых, сохранять компактный артефакт результата. Включить конфигурацию, среду выполнения, финальные метрики, курируемые примеры, результаты случайных проверок и, в идеале, чекпоинт. Вывод ноутбука полезен, но артефакт результата упрощает сравнение запусков.

В-пятых, использовать возможности Codex как часть настройки запуска. Установить /goal с явными критериями приёмки. Использовать Chrome для авторизованной работы в Colab. Просить Codex проверять сохранённые результаты, когда браузер становится ненадёжным. Использовать субагентов для скептического финального аудита.

Для будущих запусков Codex: установить цель с чёткими критериями приёмки, начать с пробного теста, выполнить полную задачу в реальной среде, сохранить отчёт и запросить скептический аудит.

Кстати: этот пост в блоге написал Codex, он же создал скриншоты, твит и всё вокруг -- всё, что я сделал, это отформатировал и нажал «Опубликовать»!

Подпишитесь на канал и каждый день читайте лучшие материалы про AI переведенные на русский!

Нашли интересную статью для перевода? Пришлите нашему боту: @ailongreadsbot