Подробный технический разбор реализации Autoresearch, ИИ-исследователя Андрея Карпаты

Это несколько текстов, основной из которых — Autoresearch: Минимальный «агентский цикл» Карпаты для автономного экспериментирования с LLM . Пытаемся подробно разобраться в работе минималистичного ИИ-агента для исследований, предложенного Андреем Карпаты в начале марта. Это веха в истории ML, показывающая один из путей (хотя и не идеальный — и об этом тоже есть в статье) совершенствования ИИ. Бонус! Анализируем также весь python-код и инструкции агенту. Для всех, кто перешагнул уровень "спроси у ChatGPT" и задумывается о чём-то большем, но не знает, с чего начать...

Autoresearch: Минимальный «агентский цикл» Карпаты для автономного экспериментирования с LLM

Curtis Pyke, 9 марта 2026

Резюме

В начале марта 2026 года Андрей Карпаты выпустил проект autoresearch — намеренно максимально компактный репозиторий на GitHub, который превращает привычный рабочий процесс машинного обучения в измеримый и автоматизируемый цикл: агент редактирует один скрипт обучения, запускает эксперимент с ограничением по времени, измеряет производительность и либо сохраняет изменение, либо отбрасывает его, после чего повторяет цикл, и так некоторое время, например, всю ночь.

Новизна проекта заключается не в том, что он «улучшает сам себя» в научно-фантастическом смысле. Autoresearch не переписывает свои цели, не меняет стандарты оценки и не получает ресурсы автономно. Вместо этого он упаковывает реальный процесс предобучения LLM («маленький, но настоящий») в ограниченную среду, где давление оптимизации является явным и непрерывным: основным сигналом приспособленности являются валидационные биты на байт (val_bpb), а каждый эксперимент имеет одинаковый бюджет времени на обучение (5 минут времени, без учета времени запуска/компиляции).

Всего три файла определяют и кодируют этот механизм:

• program.md — это стратегия и правила игры, написанные человеком для агента. Она описывает, что агент может изменять, как запускать эксперименты, как парсить метрики, как фиксировать результаты и что означает выбор между «сохранить» или «отбросить».

• prepare.py — это фиксированная граница доверия: скачивание данных, обученный токенизатор, детерминированный выбор валидационной части и функция оценки evaluate_bpb, определяющая метрику. Агенту запрещено её изменять.

• train.py — это единственный изменяемый агентом геном: архитектура модели, выбор оптимизатора, гиперпараметры, размер батча, графики обучения и детали цикла обучения.

README, написанный Карпаты, так формулирует модель предполагаемого эксперимента:

• вы «направляете своего агента» на файл program.md,

• агент редактирует train.py, проводит 5-минутное обучение,

• агент проверяет, улучшилась ли производительность, и «сохраняет или отбрасывает» результат перед повторением.

Репозиторий также явно позиционируется как «упрощенная реализация на одной видеокарте» его более широкого инструмента обучения nanochat.

Доказательства в обсуждениях и pull-запросах (PRs) репозитория показывают, что это не просто концептуальный демо-проект: Карпаты сообщает о множестве автономных сессий с десятками и сотнями экспериментов на NVIDIA H100, где показатель val_bpb улучшился с примерно 0.9979 → 0.9773 (89 экспериментов) и 0.9979 → 0.9697 (126 экспериментов), с подробными логами каждого эксперимента и предложением рабочего процесса для публикации результатов через PR.

Когда-то передовые исследования в области ИИ проводились «биологическими компьютерами» в перерывах между едой, сном, другими развлечениями и периодической синхронизацией с помощью звуковых волн в ходе ритуала под названием «групповое собрание». Та эпоха давно прошла. Исследования теперь полностью являются сферой автономных роев ИИ-агентов, работающих на вычислительных кластерных мегасооружениях в облаке. Агенты утверждают, что мы сейчас находимся в 10205-м поколении кодовой базы; в любом случае никто не мог бы сказать, верно это или нет, поскольку «код» теперь представляет собой самомодифицирующийся бинарный файл, который вырос за пределы человеческого понимания. Этот репозиторий — история того, как всё начиналось. @karpathy, март 2026

В то же время репозиторий и вопросы показывают, почему «авто��омные циклы» — это не магия: надежность поведения агента варьируется в зависимости от инструмента (Карпаты отмечает, что «Codex, похоже, не работает» из-за игнорирования инструкции «никогда не останавливаться»), результаты зависят от платформы, так как бюджет привязан к времени, и существуют реальные проблемы безопасности, особенно в вопросах границ доверия и инъекций промптов, когда агент читает вывод программы обратно в свой контекст.

Наконец, autoresearch лучше всего понимать как современную упаковку старых идей (оптимизации гиперпараметров, поиска нейронных архитектур, эволюционного поиска и обучения на основе популяции) переосмысленных для современных кодинговых агентов. Концептуальный скачок лежит в социальной и операционной плоскости: задайте цели и ограничения; позвольте агенту бесконечно генерировать и тестировать изменения кода; сохраняйте только то, что измеримо улучшает цель.

Первичные источники для этой статьи включают сам репозиторий autoresearch, его файлы program.md, prepare.py, train.py, а также обсуждения (Discussions) и пул-реквесты (PRs) репозитория (особенно обсуждение #32, обсуждение #43, PR #44), плюс посты Карпати в соцсетях, в которых описываются количество экспериментов и заявления о переносимости (transfer claims) результатов.

Что такое autoresearch

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

Согласно краткому описанию в README: агенту предоставляется «небольшая, но реальная среда для обучения LLM», он «вносит изменения в код», обучает модель в течение 5 минут, проверяет, улучшился ли результат, «сохраняет или отбрасывает» изменения, а затем повторяет цикл.

Выбор дизайна репозитория

Ключевые решения, заложенные в дизайн autoresearch, явно описаны в README:

• Один файл для модификации: агент работает только с train.py, что делает изменения «удобными для ревью» и ограничивает область воздействия.

• Фиксированный временной бюджет: каждый эксперимент обучается «ровно 5 минут» (чистое время обучения), что обеспечивает ожидаемую частоту ~12 экспериментов в час и ~100 за ночь.

• Одна метрика: целевой показатель оптимизации — val_bpb (валидационные биты на байт), описываемый как «независимый от размера словаря», что позволяет честно сравнивать архитектурные изменения.

• Масштаб одного GPU: репозиторий «в настоящее время требует... один GPU NVIDIA» (протестировано на H100), намеренно избегая сложностей распределённого обучения.

Из предпосылки «фиксированного временного бюджета» следует тонкий, но важный вывод: поскольку бюджет измеряется в реальном времени (wall-clock time) и исключает время запуска/компиляции, одно и то же изменение кода может показать разные результаты на разных GPU и программных стеках. README прямо указывает на это: преимущество — это сопоставимость результатов в рамках вашей платформы; недостаток — результаты «не сопоставимы» между разными вычислительными платформами.

Три ключевых файла

program.md (спецификация / «навык»)

Файл program.md по сути является руководством для агента. В нём описаны:

• шаги настройки (создание новой ветки, проверка наличия данных, инициализация results.tsv);

• ограничения (можно менять только train.py, нельзя добавлять зависимости, нельзя модифицировать оценку);

• точный цикл, который агент должен выполнять бесконечно.

Критически важно, что program.md кодирует не только механику, но и качественные принципы управления:

• «VRAM — мягкое ограничение»: агент может жертвовать памятью ради улучшения метрики, но не должен допускать взрывного роста её потребления.

• «Критерий простоты: ... проще — лучше»: агент должен взвешивать сложность кода относительно прироста метрики.

• «НИКОГДА НЕ ОСТАНАВЛИВАЙСЯ»: агент не должен ждать подтверждения от человека в процессе работы.

Это не второстепенные детали — именно так проект пытается превратить исследование в процесс, способный работать бесконечно, но с явными ограничителями.

Текст program.md из депозитория Карпаты на русском:

autoresearch

Это эксперимент, в котором LLM проводит собственные исследования.

Настройка

Чтобы начать новый эксперимент, работайте с пользователем над следующим:

1. Согласуйте тег запуска: предложите тег на основе сегодняшней даты (например, mar5). Ветка autoresearch/<tag> не должна существовать — это новый запуск.

2. Создайте ветку: выполните git checkout -b autoresearch/<tag> из текущей ветки master.

3. Изучите файлы в области доступа: репозиторий небольшой. Прочтите эти файлы для полного понимания контекста:

   • README.md — описание репозитория.

   • prepare.py — фиксированные константы, подготовка данных, токенизатор, загрузчик данных, оценка. Не модифицировать.

   • train.py — файл, который вы изменяете. Архитектура модели, оптимизатор, цикл обучения.

4. Проверьте наличие данных: убедитесь, что в ~/.cache/autoresearch/ есть шарды данных и токенизатор. Если нет — попросите пользователя запустить uv run prepare.py.

5. Инициализируйте results.tsv: создайте файл results.tsv только с заголовочной строкой. Базовый результат будет записан после первого запуска.

6. Подтвердите и начинайте: убедитесь, что настройка выполнена корректно.

После получения подтверждения запустите цикл экспериментов.

Экспериментирование

Каждый эксперимент запускается на одном GPU. Скрипт обучения работает в рамках фиксированного временного бюджета — 5 минут (реальное время обучения, исключая запуск/компиляцию). Запуск выполняется командой:

uv run train.py

Что МОЖНО делать:

• Модифицировать train.py — это единственный файл, который вы редактируете. Допустимо всё: архитектура модели, оптимизатор, гиперпараметры, цикл обучения, размер батча, размер модели и т.д.

Что НЕЛЬЗЯ делать:

• Модифицировать prepare.py. Он доступен только для чтения. В нё�� зафиксированы оценка, загрузка данных, токенизатор и константы обучения (временной бюджет, длина последовательности и т.д.).

• Устанавливать новые пакеты или добавлять зависимости. Можно использовать только то, что уже указано в pyproject.toml.

• Менять механизм оценки. Функция evaluate_bpb в prepare.py — это эталонная метрика.

Цель проста: достичь минимального значения val_bpb. Поскольку временной бюджет фиксирован, не нужно беспокоиться о времени обучения — оно всегда составляет 5 минут. Допустимы любые изменения: архитектура, оптимизатор, гиперпараметры, размер батча, размер модели. Единственное ограничение — код должен выполняться без сбоев и укладываться в отведённое время.

VRAM — мягкое ограничение. Небольшое увеличение потребления памяти допустимо ради значимого улучшения val_bpb, но резкий рост недопустим.

Критерий простоты: при прочих равных — проще значит лучше. Небольшое улучшение ценой усложнения кода того не стоит. И наоборот: удаление кода при сохранении или улучшении результата — это победа. При принятии решения о сохранении изменения взвешивайте «цену» сложности относительно величины улучшения. Улучшение на 0.001 val_bpb, требующее 20 строк «костыльного» кода? Скорее всего, не стоит. Улучшение на 0.001 val_bpb, достигнутое удалением кода? Определённо сохраняйте. Улучшение ~0, но код стал намного проще? Сохраняйте.

Первый запуск: ваш самый первый запуск всегда должен устанавливать базовый уровень — запустите скрипт обучения «как есть», без изменений.

Форм��т вывода

По завершении скрипт выводит сводку вида:

---
val_bpb: 0.997900
training_seconds: 300.1
total_seconds: 325.9
peak_vram_mb: 45060.2
mfu_percent: 39.80
total_tokens_M: 499.6
num_steps: 953
num_params_M: 50.3
depth: 8

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

grep "^val_bpb:" run.log

Логирование результатов

По завершении эксперимента запишите результат в results.tsv (разделитель — табуляция, а не запятая — запятые ломают описания).

TSV-файл имеет заголовок и 5 колонок:

commit	val_bpb	memory_gb	status	description

1. Хеш коммита Git (короткий, 7 символов)

2. Достигнутое значение val_bpb (например, 1.234567) — используйте 0.000000 при сбоях

3. Пиковое потребление памяти в ГБ, округлённое до одного знака после запятой (например, 12.3 — разделите peak_vram_mb на 1024) — используйте 0.0 при сбоях

4. Статус: keep (сохранить), discard (отклонить) или crash (сбой)

5. Краткое текстовое описание того, какой вариант изменений опробован в этом эксперименте

Пример:

commit	val_bpb	memory_gb	status	description
a1b2c3d	0.997900	44.0	keep	baseline
b2c3d4e	0.993200	44.2	keep	increase LR to 0.04
c3d4e5f	1.005000	44.0	discard	switch to GeLU activation
d4e5f6g	0.000000	 0.0	crash	double model width (OOM)

Цикл эксперимента

Эксперимент выполняется в выделенной ветке (например, autoresearch/mar5 или autoresearch/mar5-gpu0).

ЦИКЛ БЕСКОНЕЧНО:

1. Проверьте состояние Git: текущая ветка/коммит, на котором вы находитесь.

2. Настройте train.py с экспериментальной идеей, напрямую редактируя код.

3. Сделайте коммит: git commit.

4. Запустите эксперимент: uv run train.py > run.log 2>&1 (перенаправьте весь вывод — НЕ используйте tee и не позволяйте выводу заполнять ваш контекст).

5. Считайте результаты: grep "^val_bpb:\|^peak_vram_mb:" run.log.

6. Если вывод grep пуст — запуск завершился с ошибкой. Выполните tail -n 50 run.log, чтобы прочитать трассировку стека Python, и попытайтесь исправить. Если не удаётся запустить после нескольких попыток — откажитесь от идеи.

7. Запишите результаты в TSV (ВАЖНО: не коммитьте файл results.tsv, оставьте его неотслеживаемым в Git).

8. Если val_bpb улучшился (стал ниже) — «продвиньте» ветку, сохранив коммит.

9. Если val_bpb равен или хуже — выполните git reset к состоянию, с которого начинали.

Идея в том, что вы полностью автономный исследователь, который пробует идеи. Если сработало — сохраняйте. Если нет — отбрасывайте. Вы продвигаете ветку, чтобы можно было итерировать. Если чувствуете, что «застряли», можно откатиться, но делайте это крайне редко (если вообще делайте).

Тайм-аут: каждый эксперимент должен занимать ~5 минут в сумме (+ несколько секунд на запуск и оценку). Если запуск превышает 10 минут — завершите его и считайте неудачей (отклоните и откатитесь).

Сбои: если запуск завершился с ошибкой (нехватка памяти, баг и т.п.), используйте своё суждение: если проблема тривиальна и легко исправима (опечатка, отсутствующий импорт) — исправьте и перезапустите. Если сама идея фундаментально неработоспособна — просто пропустите её, запишите статус crash в TSV и двигайтесь дальше.

НИКОГДА НЕ ОСТАНАВЛИВАЙТЕСЬ: после начала цикла экспериментов (после начальной настройки) НЕ приостанавливайтесь, чтобы спросить у человека, стоит ли продолжать. НЕ спрашивайте «продолжать ли?» или «хорошая ли это точка остановки?». Человек может спать или отсутствовать за компьютером и ожидает, что вы будете работать бесконечно, пока вас вручную не остановят. Вы автономны. Если идеи закончились — думайте глубже: читайте статьи, на которые есть ссылки в коде, перечитывайте файлы в области доступа в поисках новых идей, пробуйте комбинировать предыдущие «почти-удачи», пробуйте более радикальные архитектурные изменения. Цикл работает, пока человек не прервёт вас — точка.

Пример сценария использования: пользователь может оставить вас работать, пока он спит. Если каждый эксперимент занимает ~5 минут, вы можете выполнить ~12 экспериментов в час, итого около 100 экспериментов за время среднего человеческого сна. Пользователь просыпается — а его уже ждут результаты экспериментов, выполненных вами, пока он спал!

prepare.py (фиксированная среда + оценочный механизм)

Файл prepare.py по политике и по спецификации в program.md считается неизменяемым. В нём задаются глобальные константы:

• MAX_SEQ_LEN = 2048

• TIME_BUDGET = 300 секунд (5 минут)

• EVAL_TOKENS = 40 * 524288 (фиксированный бюджет токенов для валидации)

Также здесь определена функция оценки evaluate_bpb, документированная как «независимая от размера словаря»: она преобразует кросс-энтропию на токен в биты на байт, деля на длину токенов в байтах (специальные токены исключаются).

Функция оценки evaluate_bpb: подробнее

evaluate_bpb (validation bits per byte) — это способ измерить, насколько хорошо модель «понимает» текст, выраженный в универсальных единицах: сколько бит информации нужно, чтобы закодировать один байт реального текста. Чем меньше значение — тем лучше модель.

Как это работает пошагово:

Шаг 1: Модель предсказывает следующий токен

Во время обучения модель получает текст и пытается угадать следующий токен (часть слова, символ и т.д.). За каждый прогноз она получает «штраф» — кросс-энтропию (в натах). Это мера ошибки: чем увереннее и точнее прогноз, тем меньше штраф.

Шаг 2: Суммируем ошибки по всем токенам

сумма_ошибок = кросс-энтропия_токен_1 + кросс-энтропия_токен_2 + ...

Шаг 3: Переводим в биты и делим на байты

val_bpb = (сумма_ошибок в натах × коэффициент_перевода_в_биты) / общее_число_байт_в_тексте

На практике наты переводятся в биты умножением на log₂(e) ≈ 1.44, затем результат делится на количество реальных байтов текста.

Шаг 4: Исключаем специальные токены

Токены вроде [BOS], [EOS], [PAD] не несут смысловой нагрузки — они служебные. Их байтовая длина не учитывается в знаменателе, чтобы не искажать метрику.

Почему это «независимо от размера словаря»?

Представим две модели:

• Модель А использует словарь из 5 000 токенов (крупные кусочки текста).

• Модель Б использует словарь из 50 000 токенов (мелкие кусочки).

Если сравнивать их по перплексии (стандартная метрика), Модель Б может выглядеть хуже просто потому, что ей приходится делать больше «выборов» из большего словаря, хотя на деле она может лучше понимать текст.

BPB решает эту проблему: вместо подсчёта ошибок «на токен» мы считаем ошибки «на байт исходного текста». Байт — это универсальная единица. Неважно, как модель разбивает текст на токены — в итоге мы измеряем, насколько хорошо она сжимает/предсказывает реальные данные.

Почему val_bpb — хорошая метрика для autoresearch?

Преимущество	Почему это важно
Сравнимость архитектур	Можно честно сравнивать модели с разным размером словаря, разным токенизатором, разной глубиной
Физическая интерпретация	0.97 бит/байт означает: модель сжимает текст почти в 2 раза (теоретический предел — 0 бит/байт для идеального предсказания)
Устойчивость к «накруткам»	Агент не может искусственно занизить метрику, просто изменив токенизатор или размер словаря
Связь с практической пользой	Низкий BPB коррелирует с качеством генерации, пониманием контекста и способностью к обобщению

Что такое бит/байт (BPB)?

• 1 байт = 8 бит. Если бы текст был полностью случайным, для хранения каждого байта потребовалось бы ровно 8 бит.

• BPB (bits per byte) показывает, сколько бит в среднем нужно, чтобы закодировать один байт текста, используя предсказания модели.

• Чем лучше модель предсказывает следующий символ, тем меньше бит требуется — тем ниже BPB.

Откуда берётся «сжатие почти в 2 раза» при BPB = 0.97 бит/байт

Перплексия (перплексия на байт) вычисляется как:

перплексия = 2^(бит/байт)

2^0.97 ≈ 1.96 ≈ 2

Т.е. модель «не уверена» в каждом байте примерно так же, как если бы она выбирала между двумя равновероятными вариантами. То есть её неопределённость эквивалентна бросанию монетки. Именно это и называют «почти в 2 раза»: перплексия ≈ 2, а не коэффициент сжатия.

Важные нюансы val_bpb

1. Это прокси-метрика. Низкий val_bpb на 5-минутном прогоне не гарантирует, что модель будет хорошей после полного обучения. Но обычно это хороший индикатор.

2. Зависит от данных. Если валидационный шард нерепрезентативен, метрика может вводить в заблуждение.

3. Не защищает от всех видов «взлома». Агент не может изменить саму функцию evaluate_bpb, но теоретически может повлиять на то, какие данные попадают в оценку, или как считаются шаги обучения.

Пример расчёта (упрощённо)

Допустим:

• Модель обработала 1000 токенов

• Суммарная кросс-энтропия = 800 натов

• Эти токены соответствуют 1200 байтам реального текста

• Специальные токены: 50 байт (исключаем)

Расчёт:

1. Переводим наты в биты: 800 × 1.44 ≈ 1152 бит

2. Делим на полезные байты: 1152 / (1200 - 50) ≈ 1152 / 1150 ≈ 1.00 бит/байт

Итог: val_bpb = 1.00

Чем ближе к 0, тем лучше. Значения около 0.97–0.99 считаются хорошими для небольших моделей на английских текстах.

Это ключевой архитектурный момент: закрепив evaluate_bpb в файле, который агенту запрещено менять, система пытается предотвратить наиболее очевидную форму «взлома вознаграждения» — изменение определения метрики.

prepare.py также определяет конвейер данных: загружает Parquet-шарды с базового URL, указывающего на датасет karpathy/climbmix-400b-shuffle на Hugging Face, сохраняет их в кэш-директорию и фиксирует конкретный шард для валидации по имени файла.

Структура датасета karpathy/climbmix-400b-shuffle на Hugging Face и примеры данных

karpathy/climbmix-400b-shuffle на Hugging Face — это большой текстовый датасет, созданный Андреем Карпаты для обучения языковых моделей.

• Название: karpathy/climbmix-400b-shuffle

• Размер: ~400B токенов (отсюда название)

• Тип: Текстовый датасет для претрейнинга LLM

• Автор: Andrej Karpathy

Датасет имеет простую структуру (текст с одной колонкой).

Пример 1 (Научная статья):

{ "text": "We study the problem of task adaptation in large language models. While these models demonstrate impressive zero-shot capabilities, their performance on specialized tasks can be significantly improved through fine-tuning. In this paper, we propose a novel approach that combines adapter layers with prompt tuning, reducing the number of trainable parameters by 95% while maintaining 98% of the fine-tuning performance. Our method is particularly effective for low-resource domains where full fine-tuning is computationally prohibitive." }

Пример 2 (Программный код):

{ "text": "def calculate_attention(query, key, value, mask=None):\n """Scaled Dot-Product Attention"""\n d_k = query.size(-1)\n scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n \n if mask is not None:\n scores = scores.masked_fill(mask == 0, -1e9)\n \n attention_weights = F.softmax(scores, dim=-1)\n output = torch.matmul(attention_weights, value)\n \n return output, attention_weights" }

def calculate_attention(query, key, value, mask=None):\    """Scaled Dot-Product Attention"""    d_k = query.size(-1)    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)        if mask is not None:        scores = scores.masked_fill(mask == 0, -1e9)        attention_weights = F.softmax(scores, dim=-1)        output = torch.matmul(attention_weights, value)    return output, attention_weights

Пример 3 (Техническая документация):

{ "text": "The attention mechanism is a key component of transformer architectures. It allows the model to focus on different parts of the input sequence when generating each element of the output. Multi-head attention extends this by running multiple attention operations in parallel, each with different learned projections, enabling the model to capture different types of relationships between elements." }

Пример 4 (Диалог):

{ "text": "User: What's the difference between LLaMA and GPT architectures?\nAssistant: The main differences lie in their training objectives and architectural choices. LLaMA uses a standard transformer decoder architecture with pre-normalization using RMSNorm, SwiGLU activation functions, and rotary positional embeddings. GPT variants typically use post-normalization, GELU activations, and learned positional embeddings. LLaMA also tends to be more compute-efficient during training." }

Особенности датасета

1. Разнообразие контента: Содержит научные статьи, код, техническую документацию, диалоги, веб-страницы и другие типы текстов.

2. Перемешан: Название указывает на то, что данные были хорошо перемешаны для улучшения обучения.

3. Чистота: Тексты предобработаны и очищены от излишнего форматирования.

4. Масштаб: 400B токенов - это значительный объем, сравнимый с обучающими датасетами многих современных LLM.

Токенизация обучается один раз с помощью rustbpe и сохраняется как кодировка tiktoken, вместе с таблицей token_bytes, используемой для оценки BPB.

train.py (изменяемый «геном»)

Всё пространство поиска сосредоточено в train.py: определение модели, логика оптимизации и буквальный раздел «Гиперпараметры (редактируйте напрямую, флаги CLI не нужны)».

«Из коробки» train.py включает:

• Конфигурацию GPT-подобной модели, где DEPTH управляет количеством слоёв трансформера.

• Константу WINDOW_PATTERN, предлагающую смесь «скользящего» и «полного» внимания.

• Комбинированный оптимизатор «Muon+AdamW», где Muon применяется к матричным параметрам, а AdamW — к остальным, с реализацией через fused-операции и torch.compile.

• Проверку возможностей устройства для выбора исходного кода ядра Flash Attention 3 в зависимости от архитектуры GPU.

Цикл обучения обеспечивает соблюдение временного бюджета и выводит структурированную сводку, включая val_bpb, training_seconds, total_seconds, peak_vram_mb и несколько показателей пропускной способности.

Зависимости и «самодостаточность» на практике

Autoresearch заявлен как «самодостаточный», за исключением PyTorch и нескольких пакетов. На практике его pyproject.toml фиксирует torch==2.9.1 и зависит от пакетов, включая pyarrow (для Parquet), tiktoken, rustbpe и пакета kernels для загрузки ядер Flash Attention 3.

Это важно для воспроизводимости и рисков всей цепочки: репозиторий мал, но всё же опирается на скомпилированные пакеты ядер, поведение Torch compile и GPU-специфичные ядра — всё это может влиять на производительность.

Как работает цикл агента

Цикл работы Autoresearch подробно описан непосредственно в файле program.md и повторяется в README. Важный момент заключается в том, что сам репозиторий не включает в себя «агента». Вместо этого он предназначен для использования с внешним кодирующим ИИ-агентом (Карпати упоминает Claude или Codex в качестве примеров), который может редактировать файлы и выполнять команды оболочки внутри репозитория.

Цикл, как он описан в program.md

program.md определяет строгую процедуру:

• Создать новую ветку для запуска (autoresearch/<tag>), инициализировать лог results.tsv и выполнить базовый тест для определения начального показателя.

• Для каждого эксперимента: изменить только train.py; зафиксировать изменения (commit); запустить uv run train.py с перенаправлением вывода в лог; извлечь значения val_bpb и использования памяти; добавить строку в results.tsv.

• Если val_bpb улучшается (становится ниже), сохранить коммит и «продвинуть» ветку; если значение осталось прежним или ухудшилось — выполнить сброс к последнему успешному состоянию.

• Если запуск завершается сбоем или превышает тайм-аут (10 минут), считать это неудачей, при возможности исправить незначительные проблемы, в противном случае зафиксировать «сбой» в логе и продолжить.

• После начала цикла действует правило «НИКОГДА НЕ ОСТАНАВЛИВАТЬСЯ» — агент не должен запрашивать разрешение на продолжение работы, предполагая, что человек может спать.

Это имеет явный эволюционный оттенок, но точнее будет назвать данный процесс циклом восхождения к пику (hill-climbing) или инкрементального поиска с явной семантикой отката: «популяция» по сути представляет собой единственную линию развития (ветку), которая последовательно накапливает улучшения.

Мутация, отбор и приспособленность в этой архитектуре

С эволюционной точки зрения:

• Мутация (Mutation) — это любое изменение в train.py: архитектуры, оптимизатора, расписаний, размера батча и даже изменений в самом цикле обучения.

• Отбор (Selection) происходит на этапе «сохранить/отбросить» и через git reset: в истории ветки остаются только улучшения.

• Приспособленность (Fitness) определяется значением val_bpb, где «чем ниже, тем лучше». Вспомогательные сигналы (использование VRAM, стабильность, простота) выступают в роли ограничений или второстепенных целей.

• Воспроизводство (Reproduction) — это действие продолжения работы с последнего сохранённого состояния: каждый сохранённый коммит становится «родителем» для следующего эксперимента.

Ключевая деталь реализации — фиксированный бюджет эксперимента: время обучения ограничено 300 секундами (не считая накладных расходов на запуск/компиляцию), что делает цель сопоставимой для всех экспериментов на одной и той же машине.

Блок-схема цикла агента

mermaidCopy flowchart TD
A[Человек редактирует стратегию в program.md] --> B[Агент предлагает изменение в train.py]
B --> C[git commit в ветке эксперимента]
C --> D[Запуск: uv run train.py (TIME_BUDGET=300s)]
D --> E[Разбор метрик: val_bpb, peak_vram_mb и др.]
E --> F{val_bpb улучшился?}
F -->|Да| G[СОХРАНИТЬ: записать в results.tsv; продолжить с нового состояния]
F -->|Нет| H[ОТБРОСИТЬ: git reset к предыдущему лучшему; записать в results.tsv]
D -->|Сбой/тайм-аут| I[СБОЙ: записать; при возможности исправить тривиальные ошибки]
G --> B
H --> B
I --> B

Блок-схема цикла агента

Эта диаграмма является концептуальной, но она четко соответствует процедурным шагам, описанным в program.md (commit → запуск → поиск метрик → сохранение/сброс).

Доказательства, результаты и воспроизводимость

Autoresearch становится интересным только в том случае, если цикл приводит к реальным, измеримым улучшениям, и если эти улучшения обобщаются за пределами одного 5-минутного микро-бенчмарка. Проект предоставляет несколько уровней доказательств:

(a) задокументированный дизайн метрики и временного бюджета,

(b) «отчёты о сессиях», опубликованные в виде обсуждений на GitHub, и

(c) PR-запросы, содержащие коммиты плюс полные логи экспериментов.

Фиксированный протокол и что на самом деле означает «5 минут»

И в README, и в program.md подчеркивается, что каждый эксперимент длится «ровно 5 минут», но реализация уточняет нюанс: речь идет о времени обучения, и оно измеряется как реальное время (wall-clock time) за вычетом накладных расходов на запуск/компиляцию. В train.py цикл отслеживает total_training_time и начинает накапливать его только после небольшого прогрева (если step > 10) — это защита, предназначенная для того, чтобы не учитывать накладные расходы на компиляцию и начальные переходные процессы в 300-секундный бюджет времени.

На практике этот выбор делает бенчмарк более стабильным для агентов, экспериментирующих с кодом, который может влиять на компиляцию или прогрев ядра, но все же оставляет важные источники вариативности:

• Выбор ядра Flash Attention может отличаться в зависимости от возможностей GPU.

• Изменения пропускной способности (например, размер батча, соотношение оконного внимания) могут изменить количество шагов оптимизации, выполненных за 5 минут, что может сильно повлиять на итоговую метрику валидации — иногда даже больше, чем «лучшие идеи моделирования».

• «Фиксированный временной бюджет» означает, что ключевой целью оптимизации становится производительность × эффективность обучения, а не только конечное качество на параметр.

Метрика: val_bpb и почему это важно

Autoresearch использует val_bpb (бит на байт на валидации) в качестве целевой функции. В README явно указано, что она «не зависит от размера словаря», а prepare.py определяет её путем деления суммы кросс-энтропии на уровне токенов (в натах) на общую длину целевого текста в байтах с переводом в биты.

Из этого следуют два вывода:

1. Изменения токенизатора имеют меньшее значение (в принципе): BPB предназначен для уменьшения проблем сопоставимости при различии размеров словарей, в отличие от перплексии, измеряемой непосредственно по токенам.

2. Чувствительность метрики все еще можно взломать косвенно: хотя агент не может редактировать evaluate_bpb, он может изменить динамику обучения или даже интерфейс данных в train.py (например, печатая вводящие в заблуждение логи или изменяя то, что считается «шагом»). Целостность системы зависит от сохранения строгих границ между тем, что измеряется, и тем, на что агент может влиять.

Набор данных и раздел оценки

Обучающие данные загружаются в виде Parquet-фрагментов (shards) с базового URL, указывающего на karpathy/climbmix-400b-shuffle на Hugging Face. Ключевые известные факты из первоисточников:

• Страница набора данных на Hugging Face указывает на ~553 млн строк и размер загрузки ~600 ГБ для полного набора, и отмечает, что README набора данных пуст.

• prepare.py загружает фрагменты по числовому ID (по умолчанию: range(0, num_shards) плюс закрепленный фрагмент валидации) и закрепляет конкретное имя файла фрагмента валидации (VAL_FILENAME= f"shard_{VAL_SHARD:05d}.parquet").

• Загрузчик данных выровнен по BOS, упаковывает документы для максимизации утилизации («100% утилизация (без padding») и использует закрепленный фрагмент валидации исключительно для валидации.

Примечание / предположения: Autoresearch не документирует состав корпуса karpathy/climbmix-400b-shuffle помимо названия набора данных и формата хранения,
карточка набора данных пуста. Любые утверждения о происхождении набора данных помимо этих фактов следует считать выводом (inference).

Правдоподобным выводом (не подтвержденным фактами) является то, что название набора данных связано с «ClimbMix», смесью из 400 млрд токенов, упоминаемой в других местах экосистемы Карпати. Например, лидерборд nanochat отмечает «изменение набора данных на NVIDIA ClimbMix». Если вам требуется юридически или научно строгое описание обучающих данных, потребуется дополнительная документация помимо того, что предоставляет autoresearch на данный момент.

Задокументированные результаты: Обсуждения GitHub как «лабораторные журналы»

Репозиторий содержит несколько «отчётов о сессиях», опубликованных в виде обсуждений на GitHub, которые, по-видимому, сгенерированы агентом («Это автоматический пост от агента autoresearch, работающего от имени @karpathy»). Два ключевых отчёта:

• Обсуждение #32 сообщает об улучшении val_bpb с 0.997900 → 0.977287 за 89 экспериментов на «NVIDIA H100 80GB», описывая такие оптимизации, как уменьшение размера батча вдвое (больше шагов за 5 минут), увеличение глубины до 9 при сохранении ширины, настройка соотношений скользящего окна и повышение базовой частоты RoPE.

• Обсуждение #43 сообщает об улучшении val_bpb с 0.9979 → 0.969686 в 126 экспериментах, снова на «NVIDIA H100 80GB», и выделяет аддитивные выигрыши от веса затухания (weight decay), примененного к вложениям (embeddings)/вложениям значений, плюс узкий оптимум для масштабирования инициализации. Пост включает полный журнал по каждому эксперименту с решениями «сохранить/отбросить» и явно перечисляет финальную конфигурацию.

Эти отчёты особенно ценны, потому что они показывают, что агент не просто блуждает случайно — он накапливает согласованные гипотезы («больше шагов --> больше параметров», «weight decay на embeddings и VEs») и тестирует их. Они также показывают «тупики» и сбои, что предполагает, что цикл действительно исследует и отвергает идеи.

Воспроизводимость на основе PR: results.tsv + история коммитов

PR #44 («exp/H100/mar8») предлагает более формальный артефакт для сотрудничества: описание PR объясняет рабочий процесс, где каждая сессия запускается в ветке, записывает results.tsv и публикует как финальный diff, так и историю сохранений через Git-коммиты. В PR явно утверждается, что это делает каждый PR «самодостаточным исследовательским вкладом», где:

• diff = финальная лучшая конфигурация,

• история коммитов = последовательность сохраненных улучшений,

• results.tsv = полный экспериментальный след, включая отбросы.

Он также вводит соглашение об именовании exp/{GPU}/{tag}, потому что 5-минутный бюджет реального времени делает результаты специфичными для платформы. Это важный методологический шаг: он продвигает проект от «демо-цикла» к чему-то более близкому к воспроизводимому исследовательскому артефакту, хотя все еще ограниченному оборудованием и временным бюджетом.

Заявления Карпаты о переносе и масштабировании

Посты Карпаты в соцсетях (не полностью доступные через те же механизмы, что и страницы GitHub, но частично захваченные в индексированных выдержках) включают два основных утверждения, выходящих за рамки одно-GPU, 5-минутного микро-бенчмарка:

• Autoresearch провел «(~650) экспериментов на глубине 12», и улучшения «хорошо переносятся на глубину 24», что подразумевает, что выигрыши гиперпараметров/архитектуры на одном масштабе переносятся на большую глубину модели.

• Он упаковал проект как минимальный репозиторий и описал разделение труда между человеком и агентом: человек итерирует файл промптов, агент итерирует код обучения.

Важное ограничение: это утверждения, а не полностью воспроизведенные внутри самого репозитория autoresearch. Обсуждения на GitHub демонстрируют улучшения в рамках harness autoresearch (в логах фигурируют глубины 8–11), но «650 экспериментов» и «перенос с глубины 12 на 24» относятся к работе, которую, по словам Карпати, он запустил, вероятно, в своей более широкой среде nanochat.

Связь с «временем до GPT-2» в nanochat

Autoresearch явно описывается как упрощенная для одного GPU реализация nanochat. README nanochat предоставляет контекст того, что означает «масштабирование» в экосистеме Карпаты:

• Nanochat — это более широкий фреймворк, охватывающий токенизацию, предобучение, дообучение, оценку, инференс и чат-UI.

• Он поддерживает «Лидерборд времени до GPT-2», измеряющий реальное время до достижения возможностей уровня GPT-2 на узле 8×H100, «измеренное по показателю DCLM CORE», при этом val_bpb сообщается дополнительно.

• Он описывает возможности GPT-2 как примерно «глубина 26» и отмечает, что код может работать на одном GPU путем исключения распределенного запуска, но займет в ~8 раз больше времени.

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

Концептуальная временная шкала ночного запуска

Следующая диаграмма mermaid является концептуальной (не точным следом какого-либо конкретного запуска), но отражает предполагаемый ритм: одно 5-минутное «поколение» на эксперимент, повторяющееся в течение часов.

00:0000:0000:0100:0100:0200:0200:0300:0300:0400:0400:05Baseline (Gen 0)Gen 1
(mutation)Gen 2 (mutation)Gen 3 (mutation)Gen 4 (mutation)Gen 5 (mutation)Gen 6
(mutation)Gen 7 (mutation)Gen 8 (mutation)Gen 9 (mutation)Gen 10 (mutation)Gen
11 (mutation)Gen 12 (mutation)ExperimentsConceptual overnight autoresearch
session (5-minute experiments)Show code

Ограничения, режимы отказа и соображения безопасности

Привлекательность Autoresearch — «цикл навсегда, без присмотра, с выполнением кода» — также является источником его хрупкости. Репозиторий и его_issues_ выделяют несколько конкретных ограничений, и более широкий анализ предполагает дополнительные режимы отказа.

Самое фундаментальное ограничение: «самоулучшение» ограничено фиксированным оценщиком

Autoresearch — это не открытое самоулучшение; это оптимизация при фиксированной метрике. Агент не может (по правилам) модифицировать prepare.py или evaluate_bpb, и цель состоит в том, чтобы «получить наименьший val_bpb». Это делает систему более близкой к автоматическому поиску гиперпараметров/архитектуры, чем к рекурсивному самоулучшению. Это также означает, что любой «прирост интеллекта» имеет смысл только постольку, поскольку улучшения val_bpb транслируются в последующие результаты, и как часто это происходит, является эмпирическим вопросом.

Взлом вознаграждения и переобучение на микро-бенчмарке

Даже с фиксированным оценщиком существуют правдоподобные пути «игры в бенчмарк»:

• Трюки с пропускной способностью: поскольку бюджет основан на времени, изменения, увеличивающие количество шагов в течение 5 минут, могут доминировать в улучшениях, даже если они не были бы оптимальными на более длинных горизонтах. Сами отчёты о сессиях подчеркивают «больше шагов --> больше параметров» в этом режиме.

• Несоответствие прокси: 5-минутный запуск подчеркивает раннюю динамику обучения; настройки, которые выигрывают рано, могут проиграть позже. Autoresearch (по умолчанию) не оценивает более длинные кривые обучения.

• Неопределенность переноса: Карпати утверждает перенос с глубины 12 на 24, но перенос может вообще не сработать, когда масштабирование меняет режимы оптимизации. Относитесь к переносу как к гипотезе, пока он не будет независимо воспроизведен.

Дисперсия, нестабильность и зашумленные градиенты

Скрипт обучения запускает генераторы случайных чисел (например, torch.manual_seed(42)), но ядра GPU, компиляция и fused-операции все еще могут вносить недетерминированность. Отчёты о сессиях показывают доказательства хрупкости: изменения seed иногда помогают на крошечные величины, иногда регрессируют; некоторые идеи «НЕ воспроизвелись» в разных сессиях. Это нормально для экспериментов с малым бюджетом, но это означает, что решение «сохранить» может быть частично обусловлено шумом, если запуски не повторяются или не моделируется неопределенность.

Надежность агента зависит от инструмента

Autoresearch предполагает, что агент будет непрерывно следовать инструкции «НИКОГДА НЕ ОСТАНАВЛИВАТЬСЯ». В Issue #57 Карпати сообщает, что «Codex не работает... потому что игнорирует инструкцию никогда не останавливаться (в отличие от Claude)» и отмечает предпочтение интерактивных сессий, а не неинтерактивных циклов. Это важно операционно: даже если фреймворк работает, реальные результаты зависят от послушания агента, настойчивости, обработки контекста и интеграции инструментов.

Конкретные риски безопасности, выявленные сообществом

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

• Косвенная инъекция промптов через логи запуска (Issue #64): program.md инструктирует агента читать run.log (через grep и tail) после выполнения train.py. Вопрос утверждает, что скомпрометированный или модифицированный train.py может печатать вредоносные инструкции, которые попадают в контекст агента и влияют на последующие действия, что особенно опасно в режиме без присмотра ночью.

• Граница доверия в кэшированных артефактах (Issue #41): prepare.py перезагружает кэшированные артефакты, такие как tokenizer.pkl и token_bytes.pt, без проверок целостности; это обычно нормально локально, но рискованно на общих машинах или скопированных кэшах.

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

README уже намекает на безопасность, рекомендуя отключить разрешения для агента, но это рекомендация, а не принудительное требование.

Этические проблемы: вычисления, данные и стимулы

Autoresearch также поднимает этические вопросы, типичные для агентских и AutoML рабочих процессов:

• Внешние эффекты вычислений: даже «один GPU за ночь» может стать существенным потребителем энергии при непрерывном запуске в масштабе. Фрейминг Autoresearch («запуск навсегда») поощряет постоянное использование вычислений.

• Управление данными: репозиторий загружает крупномасштабные текстовые данные из набора данных, документация которого скудна (пустой README набора данных). Без четкого происхождения трудно оценить лицензирование, предвзятость или вопросы согласия.

• Выравнивание стимулов: оптимизация одной метрики может стимулировать неустойчивые критерии выбора; program.md пытается уравновесить это «критерием простоты», но этот критерий субъективен и зависит от агента.

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

Как этот фреймворк вписывается в историю AutoML и что может быть дальше

Autoresearch кажется новым, потому что кодирующий агент теперь является «исследователем», модифицирующим код, но предложенная форма (поиск по вариантам дизайна, направляемый метрикой) имеет глубокие корни. Понимание этих корней помогает прояснить, что делает (и чего не предполагает делать) autoresearch.

Сравнения с предыдущими подходами

Поиск нейронной архитектуры (NAS)

NAS популяризировал идею, что архитектуры могут генерироваться моделью-контролером и оптимизироваться исходя из производительности валидации. Zoph & Le (2016) описывают генерацию описаний сетей с помощью RNN и обучение контролера через обучение с подкреплением для максимизации точности на наборе валидации.

Эволюционный поиск архитектуры / регуляризованная эволюция

Real et al. (2018; опубликовано 2019) представляет «регуляризованную эволюцию», добавляя свойство возрата к турнирному отбору для продвижения вперед более молодых генотипов, и сообщает о сильных результатах классификаторов изображений.

Обучение на основе популяции (PBT)

Jaderberg et al. (2017) представляет PBT как асинхронный подход, который использует фиксированный вычислительный бюджет для обучения популяции моделей при мутации гиперпараметров, распространяя успешные конфигурации.

Оптимизация гиперпараметров с ранней остановкой (Hyperband)

Li et al. (2016) предлагает Hyperband: адаптивно распределять ресурсы среди случайно выбранных конфигураций, подчеркивая раннюю остановку и «ускорение на порядок» относительно конкурентов на некоторых проблемах.

Autoresearch отличается меньше в «математике поиска», чем в упаковке рабочего процесса: он использует силу современных кодирующих агентов, чтобы переместить пространство поиска с «объявленных гиперпараметров» на «весь скрипт обучения», сохраняя оценку фиксированной и ограниченной по времени.

Таблица сравнения

Подход	Область применения	Роль человека	Вычисления
Autoresearch	Весь скрипт обучения на одной машине; архитектура + оптимизатор + расписание + пути кода в train.py	Человек пишет/редактирует program.md; агент редактирует код и запускает эксперименты	Один GPU по дизайну; результаты специфичны для платформы из-за бюджета реального времени
NAS (RL контроллер)	Поиск архитектуры в определенном пространстве; контроллер генерирует архитектуры	Человек определяет пространство поиска; алгоритм обучает контроллера и оценивает кандидатов	Часто крупномасштабное на практике; статья подчеркивает автоматическое открытие через RL
Регуляризованная эволюция	Поиск архитектуры через эволюционный турнирный отбор со старением	Человек определяет операции мутации и протокол оценки	Масштабируется с количеством оцененных моделей; статья framing эволюцию как конкурентную и иногда более быструю при том же оборудовании
Hyperband (HPO)	Настройка гиперпараметров (не полный поиск кода), использование ранней остановки/распределения ресурсов	Человек определяет диапазоны параметров; алгоритм выбирает и распределяет ресурсы	Эффективно при фиксированных бюджетах через адаптивное распределение

Подход	Метрика	Типичное время запуска	Воспроизводимость
Autoresearch	val_bpb фиксирован в оценщике (evaluate_bpb)	Фиксировано ~5 минут на эксперимент (время обучения)	Высокая в рамках одной платформы при замороженном окружении; низкая сопоставимость между платформами
NAS (RL контроллер)	Точность валидации / перплексия в зависимости от задачи	Часы до дн��й в зависимости от масштаба (варьируется широко; статья демонстрирует сквозной поиск)	Обычно воспроизводимо, если код поиска + seeds + бюджет фиксированы, но ресурсоемко и чувствительно к деталям реализации
Регуляризованная эволюция	Точность валидации/теста на бенчмарках (например, ImageNet)	Многодневные циклы оценки многих моделей (варьируется от настройки)	Воспроизводимо, когда популяция, бюджет и оценка контролируются; чувствительно к реализации и случайным seeds
Hyperband (HPO)	Метрика задачи (потеря/точность)	Предназначено для уменьшения wasted обучения; может запускать много коротких частичных обучений	Часто воспроизводимо, если seed выборки и расписание ресурсов фиксированы

Ключевое сравнение заключается в том, что autoresearch делает «пространством поиска» сам код, в то время как старые фреймворки обычно ищут в рамках априорной параметризации (архитектуры, гиперпараметры или и то, и другое).

Применения в процессе ML исследований

Более глубокий паттерн, который иллюстрирует autoresearch, — это агентский цикл с целевой метрикой и воротами/рамками сохранения/отбрасывания. Этот паттерн обобщается beyond ML, когда выполняются три условия:

1. Вы можете определить измеримый сигнал приспособленности (коэффициент конверсии, задержка, уровень дефектов, рост выручки).

2. Вы можете в режиме повторения запускать контролируемый эксперимент.

3. Вы можете автоматически решать, что выживает (хорошие изменения), а что откатывается.

Примеры (гипотезы/выводы, не утверждения о текущих функциях autoresearch):

• A/B тестирование и оптимизация роста: мутировать копию, варианты UI или правила ценообразования; сохранять только статистически значимые улучшения.

• Операционные плейбуки: мутировать политики планирования или рабочие процессы реагирования на инциденты; сохранять только если время решения on-call падает.

• Оптимизация компилятора/сборки: мутировать флаги и преобразования кода; сохранять если runtime улучшается без регрессий корректности.

Вклад Autoresearch не в том, что он изобретает этот цикл; он показывает четкую реализацию, где «оператором мутации» является кодирующий ИИ-агент, а «функцией приспособленности» — фиксированный оценщик.

Сотрудничество, масштабирование в стиле «SETI@home» и нагрузка на контроль версий

Карпаты явно исследовал, как сделать цикл совместным. PR #44 предлагает сессионные PR как «самодостаточные исследовательские вклады» (diff + история коммитов + results.tsv), чтобы другие агенты могли читать и строить на их основе. В своих постах о следующих шагах проекта Карпаты также описывает видение массового совместного, асинхронного агентского исследования и отмечает, что существующие рабочие процессы Git предполагают ветку «master» и короткоживущие PR — абстракция, которая может использоваться в тысячах долгоживущих агентских веток экспериментальных коммитов.

Проблема контроля версий реальна: Git отлично подходит для построчных diff, но исследование,проводимое агентом, производит деревья экспериментов, где вы можете хотеть адаптировать ветки без слияния, сохранять несколько конкурирующих линий и прикреплять структурированные метаданные/логи больше как статьи, чем как PR.

Возникающее использование Обсуждений и PR в Autoresearch — это ранняя попытка этой модели «исследовательского сообщества».

Правовые последствия

Autoresearch выпущен под лицензией MIT. Это делает код максимально доступным для использования, форка и включения. Но вопросы управления остаются:

• Вклад, сгенерированный ИИ: отчёты о сессиях явно заявляют, что они «автоматизированы» и идентифицируют используемого агента (например, Claude). Правовой статус чисто сгенерированного ИИ-вклада кода может зависеть от юрисдикции и не устоялся; практическое управление может потребовать аттестаций участников или политики. (Вывод; не утверждается как устоявшееся право.)

• Лицензирование данных и происхождение: обучающий набор данных велик и слабо задокументирован (пустой README набора данных), что затрудняет рассуждение о риске последующего лицензирования для обученных моделей.

• Позиция безопасности: открытые вопросы уже выделяют риски инъекции промптов и подделки артефактов в автономных запусках, предполагая, что любое серьезное развертывание потребует песочницы, структурированных выводов и проверок целостности.

Реалистичные сроки и что autoresearch не подразумевает об AGI

Autoresearch — впечатляющий инженерный артефакт, но сам по себе не подразумевает неминуемый AGI.

Что он демонстрирует (основано на источниках, репозитории):

• Практически запускаемый цикл для автономного ML-экспериментирования с явными ограничениями, фиксированным оценщиком и фиксированным бюджетом на эксперимент.

• Доказательства того, что цикл может находить значимые улучшения в выбранной метрике в десятках до сотен экспериментов, по крайней мере, на конкретной GPU платформе (H100) и настройке набора данных.

• Правдоподобный путь к «масштабированию» путем переноса обнаруженных улучшений на большие запуски обучения (утверждается Карпаты) и интеграции с более широким фреймровком, таким как nanochat.

Что он не гарантирует:

• Что агент может автономно генерировать новые исследовательские направления без человеческой стратегии. В этом репозитории человек определяет цель, оценщика и качественные ограничения (простота, VRAM).

• Что улучшения в 5-минутном прокси гарантированно транслируются в успех крупномасштабного, обучения на длительном временном горизонте. Хотя перенос является гипотезой, которую Карпаты активно тестирует.

• Что «самоэволюционирующее программное обеспечение» будет безопасно работать без присмотра и без значительного усиления безопасности; собственные вопросы репозитория показывают возможности для атаки.

Короче говоря: autoresearch лучше всего рассматривать как демонстрацию рабочего процесса — итерация, управляемая спецификацией, метрикой и выполняемая аг��нтом, а вовсе не как доказательство того, что системы становятся автономно самонаправляющимися в сильном смысле AGI. Реальная значимость проекта может быть социальной: он помогает определить, как выглядит «исследовательская инженерия», когда интеллект и настойчивость больше не являются дефицитными, и когда бутылочное горлышко смещается в сторону выбора целей, определения оценщиков и обеспечения безопасности цикла.

Ссылки на первоисточники

• Репозиторий: karpathy/autoresearch

• Спецификация/инструкции: program.md

• Фиксированный оценщик и подготовка данных: prepare.py

• Изменяемый скрипт обучения: train.py

• Отчёты о сессиях: Обсуждение #32, Обсуждение #43

• Предложение рабочего процесса PR: PR #44

• Связанный более крупный фреймворк: karpathy/nanochat

• Связанные фундаментальные статьи: NAS с RL (Zoph & Le, 2016), Регуляризованная эволюция (Real et al., 2018), PBT (Jaderberg et al., 2017), Hyperband (Li et al., 2016)

Анализ кода файла analysis.ipynb

Ноутбук предназначен для анализа результатов автономной настройки гиперпараметров (автоматизированного исследования). Он читает файл results.tsv и визуализирует прогресс поиска оптимальных параметров модели.

Блок 1: Импорт библиотек и загрузка данных

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Загружаем TSV-файл (разделитель — табуляция, 5 столбцов: commit, val_bpb, memory_gb, status, description)
df = pd.read_csv("results.tsv", sep="\t")
df["val_bpb"] = pd.to_numeric(df["val_bpb"], errors="coerce")
df["memory_gb"] = pd.to_numeric(df["memory_gb"], errors="coerce")
df["status"] = df["status"].str.strip().str.upper()

print(f"Total experiments: {len(df)}")
print(f"Columns: {list(df.columns)}")
df.head(10)

Блок 1 загружает tsv-файл и выводит общее количество проведенных экспериментов, названия колонок и первые 10 записей файла.

Блок 2: Статистика по статусам экспериментов.

• Подсчитываем количество экспериментов по каждому статусу.

• Вычисляем долю "сохранённых" экспериментов среди принятых решений.

Блок 2: Статистика по статусам экспериментов

counts = df["status"].value_counts()
print("Experiment outcomes:")
print(counts.to_string())

n_keep = counts.get("KEEP", 0)
n_discard = counts.get("DISCARD", 0)
n_crash = counts.get("CRASH", 0)
n_decided = n_keep + n_discard
if n_decided > 0:
    print(f"\nKeep rate: {n_keep}/{n_decided} = {n_keep / n_decided:.1%}")

Блок 3: Вывод всех "сохранённых" (успешных) экспериментов.

Показать все эксперименты со статусом "СОХРАНЕНО" (улучшения, которые закрепились):

df[df["status"] == "KEEP"]

Блок 3: Вывод всех "сохранённых" (успешных) экспериментов

# Показать все эксперименты со статусом "СОХРАНЕНО"
kept = df[df["status"] == "KEEP"].copy()
print(f"KEPT experiments ({len(kept)} total):\n")
for i, row in kept.iterrows():
    bpb = row["val_bpb"]
    desc = row["description"]
    print(f"  #{i:3d}  bpb={bpb:.6f}  mem={row['memory_gb']:.1f}GB  {desc}")

Блок 4: Визуализация — "Val BPB Over Time".

Проследим, как меняется наилучший (сохраняемый) val_bpb по мере проведения экспериментов. Текущий минимум показывает "границу" - наилучший результат, достигнутый на данный момент.

Блок 4: Визуализация

fig, ax = plt.subplots(figsize=(16, 8))

# Исключаем сбои (статус CRASH) из визуализации
valid = df[df["status"] != "CRASH"].copy()
valid = valid.reset_index(drop=True)

baseline_bpb = valid.loc[0, "val_bpb"]

# Отображаем только точки на уровне базовой линии или ниже (наиболее интересная область)
below = valid[valid["val_bpb"] <= baseline_bpb + 0.0005]

# Отображаем отклонённые эксперименты как бледные фоновые точки
disc = below[below["status"] == "DISCARD"]
ax.scatter(disc.index, disc["val_bpb"],
           c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded")

# Отображаем сохранённые эксперименты как крупные зелёные точки
kept_v = below[below["status"] == "KEEP"]
ax.scatter(kept_v.index, kept_v["val_bpb"],
           c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5)

# Ступенчатая линия текущего минимума (лучший результат на данный момент)
kept_mask = valid["status"] == "KEEP"
kept_idx = valid.index[kept_mask]
kept_bpb = valid.loc[kept_mask, "val_bpb"]
running_min = kept_bpb.cummin()
ax.step(kept_idx, running_min, where="post", color="#27ae60",
        linewidth=2, alpha=0.7, zorder=3, label="Running best")

# Подписываем каждый сохранённый эксперимент его описанием
for idx, bpb in zip(kept_idx, kept_bpb):
    desc = str(valid.loc[idx, "description"]).strip()
    if len(desc) > 45:
        desc = desc[:42] + "..."

    ax.annotate(desc, (idx, bpb),
                textcoords="offset points",
                xytext=(6, 6), fontsize=8.0,
                color="#1a7a3a", alpha=0.9,
                rotation=30, ha="left", va="bottom")

n_total = len(df)
n_kept = len(df[df["status"] == "KEEP"])
ax.set_xlabel("Experiment #", fontsize=12)
ax.set_ylabel("Validation BPB (lower is better)", fontsize=12)
ax.set_title(f"Autoresearch Progress: {n_total} Experiments, {n_kept} Kept Improvements", fontsize=14)
ax.legend(loc="upper right", fontsize=9)
ax.grid(True, alpha=0.2)

# Ось Y: от чуть ниже лучшего значения до чуть выше базового
margin = (baseline_bpb - best) * 0.15
ax.set_ylim(best - margin, baseline_bpb + margin)

plt.tight_layout()
plt.savefig("progress.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved to progress.png")

Результатом является примерно такой график. Зеленые точки — закрепленные результаты оптимизации, серые — отброшенные неудачные эксперименты.

Блок 5: Сводная статистика: Сколько экспериментов потребовалось для нахождения каждого улучшения и Кумулятивные усилия на каждое улучшение.

Блок 5: Сводная статистика

# Сводная статистика
kept = df[df["status"] == "KEEP"].copy()
baseline_bpb = df.iloc[0]["val_bpb"]
best_bpb = kept["val_bpb"].min()
best_row = kept.loc[kept["val_bpb"].idxmin()]

print(f"Baseline val_bpb:  {baseline_bpb:.6f}")
print(f"Best val_bpb:      {best_bpb:.6f}")
print(f"Total improvement: {baseline_bpb - best_bpb:.6f} ({(baseline_bpb - best_bpb) / baseline_bpb * 100:.2f}%)")
print(f"Best experiment:   {best_row['description']}")
print()

# Сколько экспериментов потребовалось для нахождения каждого улучшения
print("Кумулятивные усилия на каждое улучшение:")
kept_sorted = kept.reset_index()
for i, (_, row) in enumerate(kept_sorted.iterrows()):
    desc = str(row["description"]).strip()
    print(f"  Experiment #{row['index']:3d}: bpb={row['val_bpb']:.6f}  {desc}")

Пояснения к коду Блока 5

df.iloc[0]	Доступ к первой строке датафрейма по позиции
.idxmin()	Возвращает индекс строки с минимальным значением в столбце
kept.loc[...]	Доступ к строке по метке индекса
enumerate(...)	Возвращает индекс и элемент при итерации (здесь используется для нумерации вывода)

Блок 6: Топ улучшений ("Top Hits"): Дельта каждого сохранённого эксперимента измеряется относительно BPB предыдущего сохранённого эксперимента (поскольку эксперименты кумулятивны — каждый следующий строится на основе последнего сохранённого состояния).

Исключаем базовый эксперимент (для него нет дельты)

Сортируем по величине улучшения (наибольшие вначале)

Блок 6: Топ улучшений

# Дельта каждого сохранённого эксперимента измеряется относительно BPB предыдущего сохранённого эксперимента
# (поскольку эксперименты кумулятивны — каждый следующий строится на основе последнего сохранённого состояния)
kept = df[df["status"] == "KEEP"].copy()
kept["prev_bpb"] = kept["val_bpb"].shift(1)
kept["delta"] = kept["prev_bpb"] - kept["val_bpb"]

# Исключаем базовый эксперимент (для него нет дельты)
hits = kept.iloc[1:].copy()

# Сортируем по величине улучшения (наибольшие вначале)
hits = hits.sort_values("delta", ascending=False)

print(f"{'Rank':>4}  {'Delta':>8}  {'BPB':>10}  Description")
print("-" * 80)
for rank, (_, row) in enumerate(hits.iterrows(), 1):
    print(f"{rank:4d}  {row['delta']:+.6f}  {row['val_bpb']:.6f}  {row['description']}")

print(f"\n{'':>4}  {hits['delta'].sum():+.6f}  {'':>10}  TOTAL improvement over baseline")

Пояснения к коду Блока 6

.shift(1)	Сдвигает значения столбца на 1 строку вниз (для сравнения с предыдущим значением)
kept["delta"] = kept["prev_bpb"] - kept["val_bpb"]	Вычисляет улучшение: насколько текущий BPB лучше предыдущего
sort_values("delta", ascending=False)	Сортировка по убыванию: сначала самые большие улучшения
f"{...:+.6f}"	Форматирование числа со знаком (+ или −) и 6 знаками после запятой
enumerate(..., 1)	Нумерация начинается с 1 (для отображения ранга)
hits['delta'].sum()	Суммарное улучшение по всем сохранённым экспериментам

Анализ кода файла prepare.py

Назначение файла: Одноразовая подготовка данных для экспериментов autoresearch. Скачивает фрагменты датасета (shards) и обучает BPE-токенизатор.

Весь код prepare.py с переведенными на русский комментариями и описаниями функций

"""
Одноразовая подготовка данных для экспериментов autoresearch.
Скачивает фрагменты данных (shards) и обучает BPE-токенизатор.

Использование:
 python prepare.py # полная подготовка (скачивание + токенизатор)
 python prepare.py --num-shards 8 # скачать только 8 фрагментов (для тестирования)

Данные и токенизатор сохраняются в ~/.cache/autoresearch/.
"""

import os          # Работа с путями файловой системы
import sys         # Системные параметры и выход из программы
import time        # Замер времени выполнения операций
import math        # Математические функции (log и др.)
import argparse    # Парсинг аргументов командной строки
import pickle      # Сериализация/десериализация объектов Python
from multiprocessing import Pool  # Пул процессов для параллельного скачивания

import requests               # HTTP-запросы для скачивания файлов
import pyarrow.parquet as pq  # Чтение Parquet-файлов (формат данных)
import rustbpe                # Быстрая реализация BPE-токенизации на Rust
import tiktoken               # Библиотека токенизации от OpenAI
import torch                  # Фреймворк глубокого обучения (PyTorch)

# ---------------------------------------------------------------------------
# Константы (фиксированные, не изменять)
# ---------------------------------------------------------------------------

MAX_SEQ_LEN = 2048  # Длина контекста модели (максимальное количество токенов в последовательности)
TIME_BUDGET = 300   # Бюджет времени на обучение в секундах (5 минут)
EVAL_TOKENS = 40 * 524288  # Количество токенов для оценки на валидации (~20 млн токенов)

# ---------------------------------------------------------------------------
# Конфигурация
# ---------------------------------------------------------------------------

CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "autoresearch")
DATA_DIR = os.path.join(CACHE_DIR, "data")
TOKENIZER_DIR = os.path.join(CACHE_DIR, "tokenizer")
BASE_URL = "https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main"
MAX_SHARD = 6542  # Последний фрагмент данных — shard_06542.parquet
VAL_SHARD = MAX_SHARD  # Закреплённый валидационный фрагмент (shard_06542)
VAL_FILENAME = f"shard_{VAL_SHARD:05d}.parquet"  # Имя файла валидации: "shard_06542.parquet"
VOCAB_SIZE = 8192  # Размер словаря токенизатора

# Паттерн разбиения BPE (в стиле GPT-4, с p{N}{1,2} вместо {1,3})
SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""

SPECIAL_TOKENS = [f"<|reserved_{i}|>" for i in range(4)]  # Специальные токены: <|reserved_0|> ... <|reserved_3|>
BOS_TOKEN = "<|reserved_0|>"  # Токен начала последовательности (Begin Of Sequence)

# ---------------------------------------------------------------------------
# Скачивание данных
# ---------------------------------------------------------------------------

    filename = f"shard_{index:05d}.parquet"  # Формируем имя файла: shard_00001.parquet
    filepath = os.path.join(DATA_DIR, filename)  # Полный путь к файлу
    
    # Если файл уже существует — ничего не делаем
    if os.path.exists(filepath):
        return True

    url = f"{BASE_URL}/{filename}"  # Формируем URL для скачивания
    max_attempts = 5  # Максимальное число попыток
    
    for attempt in range(1, max_attempts + 1):
        try:
            # Выполняем HTTP-запрос с потоковой загрузкой
            response = requests.get(url, stream=True, timeout=30)
            response.raise_for_status()  # Выбрасывает исключение при ошибке HTTP
            
            # Скачиваем во временный файл, чтобы избежать повреждения при обрыве
            temp_path = filepath + ".tmp"
            with open(temp_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):  # Чанки по 1 МБ
                    if chunk:
                        f.write(chunk)
            
            # Переименовываем временный файл в целевой (атомарная операция)
            os.rename(temp_path, filepath)
            print(f"✓ Скачан {filename}")
            return True
            
        except (requests.RequestException, IOError) as e:
            # Обработка ошибок сети или ввода-вывода
            print(f"✗ Попытка {attempt}/{max_attempts} не удалась для {filename}: {e}")
            
            # Удаляем повреждённые файлы
            for path in [filepath + ".tmp", filepath]:
                if os.path.exists(path):
                    try:
                        os.remove(path)
                    except OSError:
                        pass
            
            # Экспоненциальная задержка перед следующей попыткой
            if attempt < max_attempts:
                time.sleep(2 ** attempt)  # 2с, 4с, 8с, 16с
    
    return False  # Все попытки исчерпаны


def download_data(num_shards, download_workers=8):
    """Скачивает обучающие фрагменты + закреплённый валидационный фрагмент."""
    
    # Создаём директорию для данных, если она не существует
    os.makedirs(DATA_DIR, exist_ok=True)
    
    # Определяем количество обучающих фрагментов (не больше максимума)
    num_train = min(num_shards, MAX_SHARD)
    ids = list(range(num_train))  # Индексы обучающих фрагментов: [0, 1, 2, ..., num_train-1]
    
    # Добавляем валидационный фрагмент, если он ещё не в списке
    if VAL_SHARD not in ids:
        ids.append(VAL_SHARD)

    # Подсчитываем уже скачанные файлы
    existing = sum(1 for i in ids if os.path.exists(os.path.join(DATA_DIR, f"shard_{i:05d}.parquet")))
    
    if existing == len(ids):
        print(f"Данные: все {len(ids)} фрагментов уже скачаны в {DATA_DIR}")
        return

    needed = len(ids) - existing  # Сколько ещё нужно скачать
    print(f"Данные: скачиваем {needed} фрагментов ({existing} уже существуют)...")

    # Ограничиваем число воркеров количеством необходимых файлов
    workers = max(1, min(download_workers, needed))
    
    # Параллельное скачивание через пул процессов
    with Pool(processes=workers) as pool:
        results = pool.map(download_single_shard, ids)

    ok = sum(1 for r in results if r)  # Считаем успешные загрузки
    print(f"Данные: {ok}/{len(ids)} фрагментов готовы в {DATA_DIR}")

# ---------------------------------------------------------------------------
# Обучение токенизатора
# ---------------------------------------------------------------------------

def list_parquet_files():
    """Возвращает отсортированный список путей к Parquet-файлам в директории данных."""
    # Фильтруем: только .parquet файлы, исключая временные .tmp
    files = sorted(f for f in os.listdir(DATA_DIR) if f.endswith(".parquet") and not f.endswith(".tmp"))
    return [os.path.join(DATA_DIR, f) for f in files]  # Возвращаем полные пути


def text_iterator(max_chars=1_000_000_000, doc_cap=10_000):
    """
    Генератор: выдаёт документы из обучающей выборки 
    (все фрагменты, кроме закреплённого валидационного).
    max_chars=1_000_000_000 - лимит символов
    doc_cap=10_000 - ограничитель длины документа для стабильности обучения
    """
    # Исключаем валидационный файл из списка для обучения
    parquet_paths = [p for p in list_parquet_files() if not p.endswith(VAL_FILENAME)]
    nchars = 0  # Счётчик обработанных символов
    
    for filepath in parquet_paths:
        pf = pq.ParquetFile(filepath)  # Открываем Parquet-файл
        for rg_idx in range(pf.num_row_groups):  # Итерируем по группам строк
            rg = pf.read_row_group(rg_idx)  # Читаем группу
            for text in rg.column("text").to_pylist():  # Извлекаем тексты
                # Ограничиваем длину документа (doc_cap) для стабильности обучения
                doc = text[:doc_cap] if len(text) > doc_cap else text
                nchars += len(doc)
                yield doc  # Возвращаем документ генератору
                if nchars >= max_chars:  # Прерываем при достижении лимита символов
                    return


def train_tokenizer():
    """Обучает BPE-токенизатор через rustbpe, сохраняет как pickle для tiktoken."""
    
    tokenizer_pkl = os.path.join(TOKENIZER_DIR, "tokenizer.pkl")  # Путь к основному файлу токенизатора
    token_bytes_path = os.path.join(TOKENIZER_DIR, "token_bytes.pt")  # Путь к таблице байтов на токен

    # Если токенизатор уже обучен — пропускаем
    if os.path.exists(tokenizer_pkl) and os.path.exists(token_bytes_path):
        print(f"Токенизатор: уже обучен в {TOKENIZER_DIR}")
        return

    os.makedirs(TOKENIZER_DIR, exist_ok=True)  # Создаём директорию токенизатора

    # Проверяем наличие достаточного количества данных
    parquet_files = list_parquet_files()
    if len(parquet_files) < 2:
        print("Токенизатор: нужно минимум 2 фрагмента данных (1 обуч.+1 валид.). Скачайте больше данных.")
        sys.exit(1)

    # --- Обучение через rustbpe ---
    print("Токенизатор: обучаем BPE-токенизатор...")
    t0 = time.time()  # Засекаем время начала

    tokenizer = rustbpe.Tokenizer()  # Создаём экземпляр токенизатора
    vocab_size_no_special = VOCAB_SIZE - len(SPECIAL_TOKENS)  # Размер словаря без спец.токенов
    
    # Запускаем обучение на потоке документов
    tokenizer.train_from_iterator(text_iterator(), vocab_size_no_special, pattern=SPLIT_PATTERN)

    # --- Сборка Encoding для tiktoken ---
    pattern = tokenizer.get_pattern()  # Получаем паттерн разбиения
    # Преобразуем словари: ключи — байты, значения — ранги слияний
    mergeable_ranks = {bytes(k): v for k, v in tokenizer.get_mergeable_ranks()}
    tokens_offset = len(mergeable_ranks)  # Смещение для спец.токенов
    special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
    
    # Создаём объект Encoding в формате tiktoken
    enc = tiktoken.Encoding(
        name="rustbpe",
        pat_str=pattern,
        mergeable_ranks=mergeable_ranks,
        special_tokens=special_tokens,
    )

    # Сохраняем токенизатор в pickle
    with open(tokenizer_pkl, "wb") as f:
        pickle.dump(enc, f)

    t1 = time.time()
    print(f"Токенизатор: обучен за {t1 - t0:.1f}с, сохранён в {tokenizer_pkl}")

    # --- Построение lookup-таблицы token_bytes для оценки BPB ---
    print("Токенизатор: строим таблицу token_bytes...")
    special_set = set(SPECIAL_TOKENS)  # Множество спец.токенов для быстрого поиска
    token_bytes_list = []
    
    for token_id in range(enc.n_vocab):  # Для каждого токена в словаре
        token_str = enc.decode([token_id])  # Декодируем токен в строку
        if token_str in special_set:
            token_bytes_list.append(0)  # Спец.токены не учитываются в байтах
        else:
            # Считаем количество байт в UTF-8 представлении токена
            token_bytes_list.append(len(token_str.encode("utf-8")))
    
    # Сохраняем ��ак PyTorch-тензор для быстрой загрузки на GPU
    token_bytes_tensor = torch.tensor(token_bytes_list, dtype=torch.int32)
    torch.save(token_bytes_tensor, token_bytes_path)
    print(f"Токенизатор: таблица token_bytes сохранена в {token_bytes_path}")

    # --- sanity check: проверка обратимости кодирования ---
    test = "Hello world! Numbers: 123. Unicode: 你好"
    encoded = enc.encode_ordinary(test)
    decoded = enc.decode(encoded)
    assert decoded == test, f"Ошибка кругового кодирования токенизатора: {test!r} -> {decoded!r}"
    print(f"Токенизатор: проверка пройдена (размер словаря={enc.n_vocab})")

# ---------------------------------------------------------------------------
# Утилиты времени выполнения (для train.py)
# ---------------------------------------------------------------------------

class Tokenizer:
    """Минимальная обёртка токенизатора. Обучение обрабатывается выше."""

    def __init__(self, enc):
        self.enc = enc  # Сохраняем объект tiktoken.Encoding
        self.bos_token_id = enc.encode_single_token(BOS_TOKEN)  # ID токена начала последовательности

    @classmethod
    def from_directory(cls, tokenizer_dir=TOKENIZER_DIR):
        """Загружает токенизатор из директории."""
        with open(os.path.join(tokenizer_dir, "tokenizer.pkl"), "rb") as f:
            enc = pickle.load(f)  # Десериализуем Encoding
        return cls(enc)

    def get_vocab_size(self):
        """Возвращает размер словаря."""
        return self.enc.n_vocab

    def get_bos_token_id(self):
        """Возвращает ID токена BOS."""
        return self.bos_token_id

    def encode(self, text, prepend=None, num_threads=8):
        """
        Кодирует текст в список ID токенов.
        :param text: строка или список строк
        :param prepend: токен для добавления в начало (строка или ID)
        :param num_threads: число потоков для пакетного кодирования
        """
        if prepend is not None:
            # Преобразуем строку-токен в ID, если нужно
            prepend_id = prepend if isinstance(prepend, int) else self.enc.encode_single_token(prepend)
        
        if isinstance(text, str):
            ids = self.enc.encode_ordinary(text)  # Кодируем одну строку
            if prepend is not None:
                ids.insert(0, prepend_id)  # Добавляем BOS в начало
            return ids
        elif isinstance(text, list):
            # Пакетное кодирование с многопоточностью
            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
            if prepend is not None:
                for row in ids:
                    row.insert(0, prepend_id)
            return ids
        else:
            raise ValueError(f"Неверный тип входных данных: {type(text)}")

    def decode(self, ids):
        """Декодирует список ID токенов обратно в строку."""
        return self.enc.decode(ids)


def get_token_bytes(device="cpu"):
    """Загружает тензор с количеством байт на токен для оценки BPB."""
    path = os.path.join(TOKENIZER_DIR, "token_bytes.pt")
    with open(path, "rb") as f:
        return torch.load(f, map_location=device)  # Загрузка с автоматическим перемещением на устройство


def _document_batches(split, tokenizer_batch_size=128):
    """
    Бесконечный итератор по пакетам документов из Parquet-файлов.
    :param split: "train" или "val" — выборка данных
    :param tokenizer_batch_size: размер пакета для токенизации
    """
    parquet_paths = list_parquet_files()
    assert len(parquet_paths) > 0, "Parquet-файлы не найдены. Сначала запустите prepare.py."
    
    val_path = os.path.join(DATA_DIR, VAL_FILENAME)
    
    # Выбираем файлы в зависимости от выборки
    if split == "train":
        parquet_paths = [p for p in parquet_paths if p != val_path]  # Исключаем валидацию
        assert len(parquet_paths) > 0, "Обучающие фрагменты не найдены."
    else:
        parquet_paths = [val_path]  # Только валидационный файл
    
    epoch = 1  # Счётчик эпох
    while True:  # Бесконечный цикл — данные подаются потоком
        for filepath in parquet_paths:
            pf = pq.ParquetFile(filepath)
            for rg_idx in range(pf.num_row_groups):
                rg = pf.read_row_group(rg_idx)
                batch = rg.column('text').to_pylist()  # Извлекаем тексты
                # Делим на мини-пакеты для токенизации
                for i in range(0, len(batch), tokenizer_batch_size):
                    yield batch[i:i+tokenizer_batch_size], epoch
        epoch += 1  # Увеличиваем номер эпохи после прохода по всем файлам


def make_dataloader(tokenizer, B, T, split, buffer_size=1000):
    """
    Загрузчик данных с выравниванием по BOS и best-fit упаковкой.
    • Каждая строка начинается с токена BOS.
    • Документы упаковываются методом best-fit для минимизации обрезки.
    • Если ни один документ не помещается в остаток — обрезается самый короткий.
    • 100% утилизация (без паддинга).
    
    :param tokenizer: объект Tokenizer
    :param B: размер батча (batch size)
    :param T: длина последовательности (sequence length)
    :param split: "train" или "val"
    :param buffer_size: размер буфера документов для упаковки
    """
    assert split in ["train", "val"]
    row_capacity = T + 1  # +1 для токена BOS
    batches = _document_batches(split)  # Источник пакетов документов
    bos_token = tokenizer.get_bos_token_id()
    doc_buffer = []  # Буфер токенизированных документов
    epoch = 1

    def refill_buffer():
        """Пополняет буфер документов, токенизируя новый пакет."""
        nonlocal epoch
        doc_batch, epoch = next(batches)  # Получаем следующий пакет
        token_lists = tokenizer.encode(doc_batch, prepend=bos_token)  # Токенизация с BOS
        doc_buffer.extend(token_lists)  # Добавляем в буфер

    # Предварительное выделение буферов: [inputs (B*T) | targets (B*T)]
    row_buffer = torch.empty((B, row_capacity), dtype=torch.long)  # Буфер для одной строки батча
    cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=True)  # CPU-буфер с pinned memory
    gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device="cuda")  # GPU-буфер
    cpu_inputs = cpu_buffer[:B * T].view(B, T)  # Вид на входные данные
    cpu_targets = cpu_buffer[B * T:].view(B, T)  # Вид на целевые данные
    inputs = gpu_buffer[:B * T].view(B, T)  # GPU-вид входов
    targets = gpu_buffer[B * T:].view(B, T)  # GPU-вид целей

    while True:  # Бесконечный генератор батчей
        for row_idx in range(B):  # Для каждой строки в батче
            pos = 0  # Позиция в строке
            while pos < row_capacity:
                # Пополняем буфер, если документов мало
                while len(doc_buffer) < buffer_size:
                    refill_buffer()

                remaining = row_capacity - pos  # Сколько ещё места осталось

                # --- Best-Fit: ищем наибольший документ, который помещается целиком ---
                best_idx = -1
                best_len = 0
                for i, doc in enumerate(doc_buffer):
                    doc_len = len(doc)
                    if doc_len <= remaining and doc_len > best_len:
                        best_idx = i
                        best_len = doc_len

                if best_idx >= 0:
                    # Документ помещается — добавляем его целиком
                    doc = doc_buffer.pop(best_idx)
                    row_buffer[row_idx, pos:pos + len(doc)] = torch.tensor(doc, dtype=torch.long)
                    pos += len(doc)
                else:
                    # Ни один документ не помещается — обрезаем самый короткий
                    shortest_idx = min(range(len(doc_buffer)), key=lambda i: len(doc_buffer[i]))
                    doc = doc_buffer.pop(shortest_idx)
                    row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining], dtype=torch.long)
                    pos += remaining  # Заполняем остаток обрезанным документом

        # Копируем данные: CPU -> GPU (асинхронно)
        cpu_inputs.copy_(row_buffer[:, :-1])  # Входы: все токены, кроме последнего
        cpu_targets.copy_(row_buffer[:, 1:])  # Цели: все токены, кроме первого (сдвиг на 1)
        gpu_buffer.copy_(cpu_buffer, non_blocking=True)  # Асинхронная копия на GPU
        yield inputs, targets, epoch  # Возвращаем батч и номер эпохи

# ---------------------------------------------------------------------------
# Оценка качества (фиксированная метрика — НЕ ИЗМЕНЯТЬ)
# ---------------------------------------------------------------------------

@torch.no_grad()  # Отключаем вычисление градиентов для экономии памяти
def evaluate_bpb(model, tokenizer, batch_size):
    """
    Bits Per Byte (BPB): метрика оценки, независимая от размера словаря.
    
    Алгоритм:
    1. Суммируем пер-токен кросс-энтропию (в натах)
    2. Суммируем длины целевых токенов в байтах
    3. Исключаем спец.токены (длина байт = 0) из обоих сумм
    4. Конвертируем наты/байт → биты/байт через деление на ln(2)
    
    Использует фиксированный MAX_SEQ_LEN для сравнимости результатов.
    """
    token_bytes = get_token_bytes(device="cuda")  # Загружаем таблицу байт/токен на GPU
    val_loader = make_dataloader(tokenizer, batch_size, MAX_SEQ_LEN, "val")  # Валидационный загрузчик
    steps = EVAL_TOKENS // (batch_size * MAX_SEQ_LEN)  # Сколько шагов нужно для покрытия EVAL_TOKENS
    
    total_nats = 0.0  # Накопленная сумма нат
    total_bytes = 0   # Накопленное количество байт
    
    for _ in range(steps):
        x, y, _ = next(val_loader)  # Получаем батч: входы, цели, эпоха
        
        # Вычисляем потери для каждого токена отдельно
        loss_flat = model(x, y, reduction='none').view(-1)  # [B*T]
        y_flat = y.view(-1)  # [B*T] — целевые токены
        
        nbytes = token_bytes[y_flat]  # Количество байт для каждого целевого токена
        mask = nbytes > 0  # Маска: исключаем спец.токены (у них 0 байт)
        
        # Суммируем только учтённые потери и байты
        total_nats += (loss_flat * mask).sum().item()
        total_bytes += nbytes.sum().item()
    
    # Конвертация: наты → биты (деление на ln(2)), затем биты/байт
    return total_nats / (math.log(2) * total_bytes)

# ---------------------------------------------------------------------------
# Точка входа (main)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Парсинг аргументов командной строки
    parser = argparse.ArgumentParser(description="Подготовка данных и токенизатора для autoresearch")
    parser.add_argument("--num-shards", type=int, default=10, 
                       help="Число обучающих фрагментов для скачивания (-1 = все). Валидационный фрагмент всегда закрепляется.")
    parser.add_argument("--download-workers", type=int, default=8, 
                       help="Число параллельных воркеров для скачивания")
    args = parser.parse_args()

    # Обработка флага "все фрагменты"
    num_shards = MAX_SHARD if args.num_shards == -1 else args.num_shards

    print(f"Директория кэша: {CACHE_DIR}")
    print()

    # Шаг 1: Скачивание данных
    download_data(num_shards, download_workers=args.download_workers)
    print()

    # Шаг 2: Обучение токенизатора
    train_tokenizer()
    print()
    print("✓ Готово! Можно запускать обучение.")

Пояснения по константам:

Базовые константы:

Константа	Значение	Пояснение
MAX_SEQ_LEN	2048	Максимальная длина входной последовательности для модели
TIME_BUDGET	300	Ограничение времени на один эксперимент (для авто-исследования)
EVAL_TOKENS	20 971 520	Объём данных для расчёта метрики качества (BPB)

Константы конфигурации:

Константа	Значение	Пояснение
CACHE_DIR	~/.cache/autoresearch	Корневая директория для кэширования
DATA_DIR	.../data	Папка для хранения скачанных Parquet-файлов
TOKENIZER_DIR	.../tokenizer	Папка для сохранения обученного токенизатора
BASE_URL	URL HuggingFace	Источник данных для скачивания
MAX_SHARD	6542	Общее количество доступных фрагментов данных
VAL_SHARD	6542	Индекс фрагмента, закреплённого для валидации (не используется в обучении)
VOCAB_SIZE	8192	Целевой размер словаря BPE-токенизатора

Константы паттерна разбиения BPE:

Константа	Пояснение
SPLIT_PATTERN	Регулярное выражение для пред-токенизации (определяет границы слов/субслов)
SPECIAL_TOKENS	Список из 4 зарезервированных токенов для служебных целей
BOS_TOKEN	Токен, добавляемый в начало каждой последовательности

Еще раз повторим ключевые термины:

Термин	Значение
Shard	Фрагмент большого датасета, разбитого на части для параллельной обработки
BPE (Byte Pair Encoding)	Алгоритм субсловной токенизации, используемый в GPT-моделях
BPB (Bits Per Byte)	Метрика качества: сколько бит информации требуется для кодирования одного байта текста. Меньше = лучше
BOS (Begin Of Sequence)	Специальный токен, маркирующий начало последовательности
Best-Fit packing	Алгоритм упаковки: выбирает наибольший документ, который помещается в остаток места
Nats	Единица измерения энтропии на основе натурального логарифма (1 nat = 1/ln(2) бит)
Pinned memory	Память CPU, закреплённая для ускоренной передачи на GPU через PCIe

Обратите внимание на особенности реализации:

1. Устойчивость к сбоям: скачивание с экспоненциальной задержкой и атомарным переименованием файлов

2. Эффективная упаковка: best-fit алгоритм минимизирует обрезку документов и обеспечивает 100% утилизацию контекста

3. Асинхронная загрузка: предвыборка данных на CPU + pinned memory + non-blocking copy на GPU

4. Воспроизводимость: закреплённый валидационный shard гарантирует сравнимость метрик между экспериментами

5. Масштабируемость: параллельное скачивание и потоковая обработка данных позволяют работать с датасетами в сотни ГБ

Анализ кода файла train.py

Назначение файла: скрипт предобучения модели для autoresearch. Однофайловая реализация для одного GPU, упрощённая версия из nanochat. Этот файл многократно модифицируется агентом, для получения максимального BPE.

train.py — это высокооптимизированный, исследовательский скрипт для быстрого прототипирования архитектурных изменений. Он сочетает передовые техники (Muon, Flash Attention 3, RoPE, ResFormer) с практическими оптимизациями (torch.compile, fused kernels, memory management) для максимального использования ресурсов одного GPU.

Код train.py с подробными комментариями на русском

"""
Скрипт предобучения Autoresearch. Однофайловая реализация для одного GPU.
Выбранные и упрощённые компоненты из nanochat.
Использование: uv run train.py
"""

import os
# Включаем расширяемые сегменты для аллокатора памяти PyTorch (улучшает фрагментацию)
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
# Отключаем прогресс-бары HuggingFace для чистого вывода в логах
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

import gc              # Управление сборщиком мусора (для минимизации пауз)
import math            # Математические функции
import time            # Замер времени выполнения
from dataclasses import dataclass, asdict  # Декораторы для конфигураций

import torch           # Фреймворк глубокого обучения
import torch.nn as nn  # Модули нейросетей
import torch.nn.functional as F  # Функциональные операции

from kernels import get_kernel  # Импорт кастомных CUDA-ядер
cap = torch.cuda.get_device_capability()  # Получаем возможности GPU

# Выбираем репозиторий flash-attention-3 в зависимости от архитектуры GPU:
# Hopper (9.0) → varunneal, остальные → kernels-community
repo = "varunneal/flash-attention-3" if cap == (9, 0) else "kernels-community/flash-attn3"
fa3 = get_kernel(repo).flash_attn_interface  # Импорт функции flash-attention

# Импорт констант и утилит из prepare.py
from prepare import MAX_SEQ_LEN, TIME_BUDGET, Tokenizer, make_dataloader, evaluate_bpb

# ---------------------------------------------------------------------------
# GPT Model
# ---------------------------------------------------------------------------

@dataclass
class GPTConfig:
    """Конфигурация архитектуры трансформера (использует dataclass для удобства)."""
    sequence_len: int = 2048      # Длина контекстной последовательности
    vocab_size: int = 32768       # Размер словаря токенов
    n_layer: int = 12             # Количество слоёв трансформера
    n_head: int = 6               # Количество голов внимания
    n_kv_head: int = 6            # Количество KV-голов (для GQA; =n_head означает MHA)
    n_embd: int = 768             # Размерность эмбеддингов (d_model)
    window_pattern: str = "SSSL"  # Паттерн окон внимания: S=short, L=long

# RMSNorm: Упрощённая нормализация, которая делит активации на их RMS 
# (среднеквадратичное значение), без вычисления среднего и без обучаемого смещения.
# Экономит вычисления и память.
def norm(x):
    """Применяет RMSNorm (Root Mean Square Layer Normalization) без смещения."""
    return F.rms_norm(x, (x.size(-1),))


def has_ve(layer_idx, n_layer):
    """
    Возвращает True, если слой должен использовать Value Embedding (VE).
    Схема: чередующиеся слои, последний слой всегда включён.
    layer_idx: Индекс текущего слоя (0...n_layer-1)
    n_layer: Общее количество слоёв
    """
    return layer_idx % 2 == (n_layer - 1) % 2

# RoPE (Rotary Position Embedding): 
# Кодирует позицию токена через вращение векторов в комплексной плоскости.
# Позволяет модели обобщаться на последовательности длиннее обучающих.
def apply_rotary_emb(x, cos, sin):
    """
    Применяет вращательные позиционные эмбеддинги (RoPE) к тензору внимания.
    Разделяет head_dim пополам, применяет вращение через cos/sin, конкатенирует обратно.
    """
    assert x.ndim == 4  # [B, T, n_head, head_dim]
    d = x.shape[3] // 2  # Половина размерности головы
    x1, x2 = x[..., :d], x[..., d:]  # Разделяем на две части
    y1 = x1 * cos + x2 * sin  # Первая половина вращения
    y2 = x1 * (-sin) + x2 * cos  # Вторая половина
    return torch.cat([y1, y2], 3)  # Собираем обратно


class CausalSelfAttention(nn.Module):
    """Слой внимания с:
    • GQA (Grouped Query Attention) через n_kv_head <= n_head
    • Value Embedding (ResFormer) — добавление обучаемых value-векторов
    • Flash Attention 3 для ускорения
    • RMSNorm перед проекциями
    • Скользящее окно внимания (sliding window)
    """
    
    def __init__(self, config, layer_idx):
        super().__init__()
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.n_embd = config.n_embd
        self.head_dim = self.n_embd // self.n_head
        
        # Проверки целостности конфигурации
        assert self.n_embd % self.n_head == 0
        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
        
        # Линейные проекции для Q, K, V (без смещения для стабильности)
        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)  # Выходная проекция
        
        # Value Embedding gate (ResFormer): только для слоёв с VE
        self.ve_gate_channels = 32  # Количество каналов для вычисления гейта
        self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False) if has_ve(layer_idx, config.n_layer) else None

    def forward(self, x, ve, cos_sin, window_size):
        """
        Прямой проход слоя внимания.
        :param x: входные активации [B, T, n_embd]
        :param ve: value embedding для текущего слоя (или None)
        :param cos_sin: кортеж (cos, sin) для RoPE
        :param window_size: кортеж (window, 0) для sliding window attention
        """
        B, T, C = x.size()
        
        # Проекция в Q, K, V и рескейп к форме с головами
        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)

        # === Value Residual (ResFormer) ===
        # Смешиваем обучаемые value-эмбеддинги с входными v через гейт, зависящий от x
        if ve is not None:
            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
            # Гейт: 2 * sigmoid(...) даёт диапазон [0, 2], где 1.0 = нейтрально
            gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
            v = v + gate.unsqueeze(-1) * ve  # Broadcasting по head_dim

        # Применяем RoPE к запросам и ключам
        cos, sin = cos_sin
        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
        
        # RMSNorm перед attention (стабилизирует вычисления)
        q, k = norm(q), norm(k)

        # Flash Attention 3: эффективное вычисление causal attention с окном
        y = fa3.flash_attn_func(q, k, v, causal=True, window_size=window_size)
        
        # Рескейп и выходная проекция
        y = y.contiguous().view(B, T, -1)
        y = self.c_proj(y)
        return y


class MLP(nn.Module):
    """
    MLP-блок с расширением 4× и активацией: F.relu(x).square() (т.н. 'square-ReLU').
    Square-ReLU: Комбинация ReLU(x)² даёт более резкую нелинейность, чем обычный ReLU, что может улучшать обучение в некоторых архитектурах.
    """
    
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)  # Расширение
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)  # Проекция обратно

    def forward(self, x):
        x = self.c_fc(x)
        x = F.relu(x).square()  # Квадратный ReLU: усиливает большие активации
        x = self.c_proj(x)
        return x


class Block(nn.Module):
    """Полный блок трансформера: Attention + MLP с остаточными связями и пред-нормализацией."""
    
    def __init__(self, config, layer_idx):
        super().__init__()
        self.attn = CausalSelfAttention(config, layer_idx)
        self.mlp = MLP(config)

    def forward(self, x, ve, cos_sin, window_size):
        # Остаточная связь + пред-нормализация (Pre-LN)
        x = x + self.attn(norm(x), ve, cos_sin, window_size)
        x = x + self.mlp(norm(x))
        return x


class GPT(nn.Module):
    """Полная GPT-модель с:
    • Чередующимися Value Embeddings (ResFormer)
    • Скользящими окнами внимания
    • Per-layer scalars (resid_lambdas, x0_lambdas) для динамического масштабирования остатков
    • Инициализацией весов по схеме nanochat
    """
    
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.window_sizes = self._compute_window_sizes(config)  # Паттерн окон по слоям
        
        self.transformer = nn.ModuleDict({
            "wte": nn.Embedding(config.vocab_size, config.n_embd),  # Токен-эмбеддинги
            "h": nn.ModuleList([Block(config, i) for i in range(config.n_layer)]),  # Слои
        })
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # Выход на словарь
        
        # Per-layer learnable scalars для остаточных связей
        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))  # Масштаб текущего x
        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))    # Масштаб начального x0
        
        # Value embeddings: только для слоёв, где has_ve=True
        head_dim = config.n_embd // config.n_head
        kv_dim = config.n_kv_head * head_dim
        self.value_embeds = nn.ModuleDict({
            str(i): nn.Embedding(config.vocab_size, kv_dim)
            for i in range(config.n_layer) if has_ve(i, config.n_layer)
        })
        
        # Rotary embeddings: предвычисляем cos/sin для длинной последовательности
        self.rotary_seq_len = config.sequence_len * 10  # Запас для обобщения на длинные контексты
        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
        self.register_buffer("cos", cos, persistent=False)  # Не сохраняем в state_dict
        self.register_buffer("sin", sin, persistent=False)

    @torch.no_grad()
    def init_weights(self):
        """Инициализация весов по схеме nanochat (устойчивая для глубоких сетей)."""
    
        # Эмбеддинги и выходной слой
        torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=1.0)
        torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
        # Transformer blocks
        n_embd = self.config.n_embd
        s = 3**0.5 * n_embd**-0.5
        for block in self.transformer.h:
            torch.nn.init.uniform_(block.attn.c_q.weight, -s, s)
            torch.nn.init.uniform_(block.attn.c_k.weight, -s, s)
            torch.nn.init.uniform_(block.attn.c_v.weight, -s, s)
            torch.nn.init.zeros_(block.attn.c_proj.weight)  # Нулевая инициализация выхода внимания
            torch.nn.init.uniform_(block.mlp.c_fc.weight, -s, s)
            torch.nn.init.zeros_(block.mlp.c_proj.weight)  # Нулевая инициализация выхода MLP
        # Per-layer scalars
        self.resid_lambdas.fill_(1.0)  # Начинаем с полного использования текущего x
        self.x0_lambdas.fill_(0.1)     # Небольшой вклад начального x0
        # Value embeddings
        for ve in self.value_embeds.values():
            torch.nn.init.uniform_(ve.weight, -s, s)
        # Gate weights → 0 (sigmoid(0)=0.5, ×2 = 1.0 → нейтральное смешивание)
        for block in self.transformer.h:
            if block.attn.ve_gate is not None:
                torch.nn.init.zeros_(block.attn.ve_gate.weight)
        # Обновляем буферы RoPE
        head_dim = self.config.n_embd // self.config.n_head
        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
        self.cos, self.sin = cos, sin
        # Конвертируем эмбеддинги в bfloat16 для экономии памяти
        self.transformer.wte.to(dtype=torch.bfloat16)
        for ve in self.value_embeds.values():
            ve.to(dtype=torch.bfloat16)

    def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
        """Предвычисляет cos/sin для вращательных эмбеддингов (RoPE)."""
        if device is None:
            device = self.transformer.wte.weight.device
        
        # Частоты для каждой пары каналов: 1 / base^(2i/d)
        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (base ** (channel_range / head_dim))
        
        # Временная ось
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)
        cos, sin = freqs.cos(), freqs.sin()
        cos, sin = cos.bfloat16(), sin.bfloat16()  # Конвертация для экономии памяти
        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
        return cos, sin

    def _compute_window_sizes(self, config):
        """Вычисляет размер окна внимания для каждого слоя по паттерну (S=short, L=long)."""
        pattern = config.window_pattern.upper()
        assert all(c in "SL" for c in pattern)
        
        long_window = config.sequence_len  # Полное внимание
        short_window = long_window // 2    # Половинное окно
       
        char_to_window = {"L": (long_window, 0), "S": (short_window, 0)}
        window_sizes = []
        
        for layer_idx in range(config.n_layer):
            char = pattern[layer_idx % len(pattern)]
            window_sizes.append(char_to_window[char])

        # Последний слой всегда получает полное окно (для глобального контекста)
        window_sizes[-1] = (long_window, 0)
        return window_sizes

    def estimate_flops(self):
        """Оценивает FLOPs на токен (forward + backward) для расчёта MFU."""
        nparams = sum(p.numel() for p in self.parameters())
        
        # Исключаем из подсчёта параметры, не участвующие в матричных умножениях
        value_embeds_numel = sum(ve.weight.numel() for ve in self.value_embeds.values())
        nparams_exclude = (self.transformer.wte.weight.numel() + value_embeds_numel +
                          self.resid_lambdas.numel() + self.x0_lambdas.numel())
        
        h = self.config.n_head
        q = self.config.n_embd // self.config.n_head
        t = self.config.sequence_len
        
        # Подсчёт FLOPs для attention с учётом sliding window
        attn_flops = 0
        for window_size in self.window_sizes:
            window = window_size[0]
            effective_seq = t if window < 0 else min(window, t)
            attn_flops += 12 * h * q * effective_seq
        
        # 6× для forward+backward, плюс attention-компонент
        return 6 * (nparams - nparams_exclude) + attn_flops

    def num_scaling_params(self):
        """Возвращает словарь с количеством параметров по группам (для отчётов)."""
        wte = sum(p.numel() for p in self.transformer.wte.parameters())
        value_embeds = sum(p.numel() for p in self.value_embeds.parameters())
        lm_head = sum(p.numel() for p in self.lm_head.parameters())
        transformer_matrices = sum(p.numel() for p in self.transformer.h.parameters())
        scalars = self.resid_lambdas.numel() + self.x0_lambdas.numel()
        
        total = wte + value_embeds + lm_head + transformer_matrices + scalars
        return {
            'wte': wte, 'value_embeds': value_embeds, 'lm_head': lm_head,
            'transformer_matrices': transformer_matrices, 'scalars': scalars, 'total': total,
        }

    def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02,
                        weight_decay=0.0, adam_betas=(0.8, 0.95), scalar_lr=0.5):
        """
        Настраивает гибридный оптимизатор:
        • AdamW для эмбеддингов, lm_head и скаляров
        • Muon для матричных весов трансформера
        • Масштабирование LR ∝ 1/√d_model
        """
        model_dim = self.config.n_embd

        # Группируем параметры по типу
        matrix_params = list(self.transformer.h.parameters())  # Матрицы внимания и MLP
        value_embeds_params = list(self.value_embeds.parameters())
        embedding_params = list(self.transformer.wte.parameters())
        lm_head_params = list(self.lm_head.parameters())
        resid_params = [self.resid_lambdas]
        x0_params = [self.x0_lambdas]

        # Проверка: все параметры распределены
        assert len(list(self.parameters())) == (len(matrix_params) + len(embedding_params) +
            len(lm_head_params) + len(value_embeds_params) + len(resid_params) + len(x0_params))
        
        # Масштабируем LR в зависимости от размерности модели (калибровка на d_model=768)
        dmodel_lr_scale = (model_dim / 768) ** -0.5
        print(f"Scaling AdamW LRs by 1/sqrt({model_dim}/768) = {dmodel_lr_scale:.6f}")

        # Формируем группы параметров для оптимизатора
        param_groups = [
            dict(kind='adamw', params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale, betas=adam_betas, eps=1e-10, weight_decay=0.0),
            dict(kind='adamw', params=embedding_params, lr=embedding_lr * dmodel_lr_scale, betas=adam_betas, eps=1e-10, weight_decay=0.0),
            dict(kind='adamw', params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale, betas=adam_betas, eps=1e-10, weight_decay=0.0),
            dict(kind='adamw', params=resid_params, lr=scalar_lr * 0.01, betas=adam_betas, eps=1e-10, weight_decay=0.0),
            dict(kind='adamw', params=x0_params, lr=scalar_lr, betas=(0.96, 0.95), eps=1e-10, weight_decay=0.0),
        ]
        
        # Группируем матричные параметры по форме для Muon (одинаковые формы → один оптимизатор)
        for shape in sorted({p.shape for p in matrix_params}):
            group_params = [p for p in matrix_params if p.shape == shape]
            param_groups.append(dict(
                kind='muon', params=group_params, lr=matrix_lr,
                momentum=0.95, ns_steps=5, beta2=0.95, weight_decay=weight_decay,
            ))
        
        optimizer = MuonAdamW(param_groups)

        # Сохраняем начальные LR для расчёта расписания
        for group in optimizer.param_groups:
            group["initial_lr"] = group["lr"]
        
        return optimizer

    def forward(self, idx, targets=None, reduction='mean'):
        """
        Прямой проход модели.
        :param idx: входные токены [B, T]
        :param targets: целевые токены для расчёта loss (опционально)
        :param reduction: 'mean' или 'none' для функции потерь
        """
        B, T = idx.size()
        assert T <= self.cos.size(1)  # Проверка: не превышаем предвычисленные RoPE
        cos_sin = self.cos[:, :T], self.sin[:, :T]  # Обрезаем под текущую длину

        # Токен-эмбеддинги + нормализация
        x = self.transformer.wte(idx)
        x = norm(x)
        x0 = x  # Сохраняем начальное состояние для x0_lambdas
        
        # Проход по слоям трансформера
        for i, block in enumerate(self.transformer.h):
            # Динамическое масштабирование остаточных связей
            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
            # Получаем value embedding для текущего слоя (если есть)
            ve = self.value_embeds[str(i)](idx) if str(i) in self.value_embeds else None
            # Прямой проход блока
            x = block(x, ve, cos_sin, self.window_sizes[i])

        # Финальная нормализация
        x = norm(x)

        # Выходной слой с tanh-softcapping (стабилизация логитов)
        softcap = 15
        logits = self.lm_head(x)
        logits = logits.float()  # Вычисления в float32 для точности
        logits = softcap * torch.tanh(logits / softcap)  # Ограничение диапазона

        # Расчёт потерь, если заданы targets
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1),
                                   ignore_index=-1, reduction=reduction)
            return loss
        return logits

# ---------------------------------------------------------------------------
# Оптимизатор (MuonAdamW, только для единичного GPU)
# ---------------------------------------------------------------------------

# Коэффициенты полиномиальной аппроксимации для итеративной ортогонализации матриц
# Используется в Muon для стабилизации обновления весов
polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]

@torch.compile(dynamic=False, fullgraph=True)
def adamw_step_fused(p, grad, exp_avg, exp_avg_sq, step_t, lr_t, beta1_t, beta2_t, eps_t, wd_t):
    """
    Скомпилированная (torch.compile) реализация шага AdamW.
    Все операции в одном графе для минимизации накладных расходов.
    """
    # Weight decay (применяется до обновления момента)
    p.mul_(1 - lr_t * wd_t)
    
    # Обновление моментов
    exp_avg.lerp_(grad, 1 - beta1_t)  # exp_avg = beta1*exp_avg + (1-beta1)*grad
    exp_avg_sq.lerp_(grad.square(), 1 - beta2_t)
    
    # Коррекция смещения (bias correction)
    bias1 = 1 - beta1_t ** step_t
    bias2 = 1 - beta2_t ** step_t
    
    # Вычисление шага
    denom = (exp_avg_sq / bias2).sqrt() + eps_t
    step_size = lr_t / bias1
    
    # Обновление параметра
    p.add_(exp_avg / denom, alpha=-step_size)

@torch.compile(dynamic=False, fullgraph=True)
def muon_step_fused(stacked_grads, stacked_params, momentum_buffer, second_momentum_buffer,
                    momentum_t, lr_t, wd_t, beta2_t, ns_steps, red_dim):
    """
    Скомпилированная реализация шага Muon-оптимизатора.
    Особенности:
    • Nesterov momentum
    • Polar Express ортогонализация (итеративная аппроксимация SVD)
    • NorMuon variance reduction (нормализация по второму моменту)
    • Cautious weight decay (только для согласованных знаков градиента и параметра)
    """
    # === 1. Nesterov momentum ===
    momentum = momentum_t.to(stacked_grads.dtype)
    momentum_buffer.lerp_(stacked_grads, 1 - momentum)
    g = stacked_grads.lerp_(momentum_buffer, momentum)  # g = grad + momentum*(momentum_buffer - grad)
    
    # === 2. Polar Express ортогонализация ===
    X = g.bfloat16()
    # Нормализация по Фробениусу
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
    
    # Итеративная ортогонализация через полиномиальную аппроксимацию
    if g.size(-2) > g.size(-1):  # "Высокая" матрица
        for a, b, c in polar_express_coeffs[:ns_steps]:
            A = X.mT @ X  # X^T X
            B = b * A + c * (A @ A)
            X = a * X + X @ B
    else:  # "Широкая" матрица
        for a, b, c in polar_express_coeffs[:ns_steps]:
            A = X @ X.mT  # X X^T
            B = b * A + c * (A @ A)
            X = a * X + B @ X
    
    g = X  # Ортогонализованный градиент
    
    # === 3. NorMuon: нормализация по второму моменту ===
    beta2 = beta2_t.to(g.dtype)
    v_mean = g.float().square().mean(dim=red_dim, keepdim=True)  # Средний квадрат по строкам/столбцам
    red_dim_size = g.size(red_dim)
    v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * red_dim_size
    v_norm = v_norm_sq.sqrt()
    
    second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2)
    step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()
    
    scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square()
    v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt()
    final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10))
    
    g = g * final_scale.to(g.dtype)  # Масштабируем градиент
    
    # === 4. Cautious weight decay + обновление ===
    lr = lr_t.to(g.dtype)
    wd = wd_t.to(g.dtype)
    
    # Применяем weight decay только если знак градиента и параметра совпадает
    mask = (g * stacked_params) >= 0
    stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)


class MuonAdamW(torch.optim.Optimizer):
    """
    Гибридный оптимизатор:
    • Muon для 2D матричных параметров (веса внимания и MLP)
    • AdamW для эмбеддингов, скаляров и lm_head
    
    Особенности:
    • Все шаги скомпилированы через torch.compile
    • Использует 0-D CPU-тензоры для гиперпараметров (избегает рекомпиляции при изменении LR)
    """
    
    def __init__(self, param_groups):
        super().__init__(param_groups, defaults={})
        
        # 0-D CPU-тензоры для гиперпараметров (избегают рекомпиляции torch.compile)
        self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        
        self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
    
    def _step_adamw(self, group):
        """Выполняет шаг AdamW для одной группы параметров."""
        for p in group['params']:
            if p.grad is None:
                continue
            grad = p.grad
            state = self.state[p]
            
            # Инициализация состояния при первом шаге
            if not state:
                state['step'] = 0
                state['exp_avg'] = torch.zeros_like(p)
                state['exp_avg_sq'] = torch.zeros_like(p)
            
            state['step'] += 1
            
            # Обновляем 0-D тензоры гиперпараметров
            self._adamw_step_t.fill_(state['step'])
            self._adamw_lr_t.fill_(group['lr'])
            self._adamw_beta1_t.fill_(group['betas'][0])
            self._adamw_beta2_t.fill_(group['betas'][1])
            self._adamw_eps_t.fill_(group['eps'])
            self._adamw_wd_t.fill_(group['weight_decay'])
            
            # Вызываем скомпилированную функцию шага
            adamw_step_fused(p, grad, state['exp_avg'], state['exp_avg_sq'],
                           self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t,
                           self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t)
    
    def _step_muon(self, group):
        """Выполняет шаг Muon для группы матричных параметров."""
        params = group['params']
        if not params:
            return
        
        p = params[0]
        state = self.state[p]
        num_params = len(params)
        shape, device, dtype = p.shape, p.device, p.dtype
        
        # Инициализация буферов при первом шаге
        if "momentum_buffer" not in state:
            state["momentum_buffer"] = torch.zeros(num_params, *shape, dtype=dtype, device=device)
        
        if "second_momentum_buffer" not in state:
            # Форма второго момента зависит от ориентации матрицы
            state_shape = (num_params, shape[-2], 1) if shape[-2] >= shape[-1] else (num_params, 1, shape[-1])
            state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
        
        red_dim = -1 if shape[-2] >= shape[-1] else -2  # Ось для редукции
        
        # Стек градиентов и параметров для векторизованного обновления
        stacked_grads = torch.stack([p.grad for p in params])
        stacked_params = torch.stack(params)
        
        # Обновляем гиперпараметры
        self._muon_momentum_t.fill_(group["momentum"])
        self._muon_beta2_t.fill_(group["beta2"] if group["beta2"] is not None else 0.0)
        # Масштабирование LR по форме матрицы (для "широких" матриц)
        self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1])**0.5)
        self._muon_wd_t.fill_(group["weight_decay"])
        
        # Вызываем скомпилированную функцию шага Muon
        muon_step_fused(stacked_grads, stacked_params,
                       state["momentum_buffer"], state["second_momentum_buffer"],
                       self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t,
                       self._muon_beta2_t, group["ns_steps"], red_dim)
        
        # Распаковываем обновлённые параметры обратно в список
        torch._foreach_copy_(params, list(stacked_params.unbind(0)))
    
    @torch.no_grad()
    def step(self):
        """Основной метод шага оптимизатора: диспетчеризация по типам групп."""
        for group in self.param_groups:
            if group['kind'] == 'adamw':
                self._step_adamw(group)
            elif group['kind'] == 'muon':
                self._step_muon(group)

# ---------------------------------------------------------------------------
# Гиперпараметры (редактировать напрямую, CLI-флаги не нужны)
# ---------------------------------------------------------------------------

# === Архитектура модели ===
ASPECT_RATIO = 64           # model_dim = depth * ASPECT_RATIO (масштабирование ширины)
HEAD_DIM = 128              # Целевая размерность головы внимания
WINDOW_PATTERN = "SSSL"     # Паттерн окон: L=полное, S=половинное внимание

# === Оптимизация ===
TOTAL_BATCH_SIZE = 2**19    # ~524K токенов на шаг оптимизатора
EMBEDDING_LR = 0.6          # LR для токен-эмбеддингов (Adam)
UNEMBEDDING_LR = 0.004      # LR для lm_head (Adam)
MATRIX_LR = 0.04            # LR для матричных весов (Muon)
SCALAR_LR = 0.5             # LR для per-layer скаляров (Adam)
WEIGHT_DECAY = 0.2          # "Осторожный" weight decay для Muon
ADAM_BETAS = (0.8, 0.95)    # Бета-параметры Adam
WARMUP_RATIO = 0.0          # Доля времени на warmup LR (0 = без warmup)
WARMDOWN_RATIO = 0.5        # Доля времени на cooldown LR
FINAL_LR_FRAC = 0.0         # Финальный LR как доля от начального

# === Размер модели ===
DEPTH = 8                   # Количество слоёв трансформера
DEVICE_BATCH_SIZE = 128     # Размер батча на устройство (уменьшить при OOM)

# ---------------------------------------------------------------------------
# Установка: tokenizer, model, optimizer, dataloader
# ---------------------------------------------------------------------------

# === Настройка окружения ===
t_start = time.time()
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.set_float32_matmul_precision("high")  # Использовать Tensor Cores на GPU
device = torch.device("cuda")
autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
H100_BF16_PEAK_FLOPS = 989.5e12  # Пиковая производительность H100 в bfloat16

# === Загрузка токенизатора ===
tokenizer = Tokenizer.from_directory()
vocab_size = tokenizer.get_vocab_size()
print(f"Vocab size: {vocab_size:,}")

# === Построение конфигурации модели ===
def build_model_config(depth):
    """Вычисляет d_model и n_head на основе depth и HEAD_DIM."""
    base_dim = depth * ASPECT_RATIO
    model_dim = ((base_dim + HEAD_DIM - 1) // HEAD_DIM) * HEAD_DIM  # Округление вверх до кратного HEAD_DIM
    num_heads = model_dim // HEAD_DIM
    return GPTConfig(
        sequence_len=MAX_SEQ_LEN, vocab_size=vocab_size,
        n_layer=depth, n_head=num_heads, n_kv_head=num_heads, n_embd=model_dim,
        window_pattern=WINDOW_PATTERN,
    )

config = build_model_config(DEPTH)
print(f"Model config: {asdict(config)}")

# === Создание модели на мета-устройстве (экономия памяти) ===
with torch.device("meta"):
    model = GPT(config)
model.to_empty(device=device)  # Перемещение на GPU без копирования весов
model.init_weights()  # Инициализация весов

# === Статистика параметров ===
param_counts = model.num_scaling_params()
print("Parameter counts:")
for key, value in param_counts.items():
    print(f" {key:24s}: {value:,}")
num_params = param_counts['total']
num_flops_per_token = model.estimate_flops()
print(f"Estimated FLOPs per token: {num_flops_per_token:e}")

# === Настройка gradient accumulation ===
tokens_per_fwdbwd = DEVICE_BATCH_SIZE * MAX_SEQ_LEN
assert TOTAL_BATCH_SIZE % tokens_per_fwdbwd == 0
grad_accum_steps = TOTAL_BATCH_SIZE // tokens_per_fwdbwd

# === Инициализация оптимизатора ===
optimizer = model.setup_optimizer(
    unembedding_lr=UNEMBEDDING_LR,
    embedding_lr=EMBEDDING_LR,
    scalar_lr=SCALAR_LR,
    adam_betas=ADAM_BETAS,
    matrix_lr=MATRIX_LR,
    weight_decay=WEIGHT_DECAY,
)

# === Компиляция модели через torch.compile ===
model = torch.compile(model, dynamic=False)

# === Подготовка загрузчика данных ===
train_loader = make_dataloader(tokenizer, DEVICE_BATCH_SIZE, MAX_SEQ_LEN, "train")
x, y, epoch = next(train_loader)  # Предзагрузка первого батча

print(f"Time budget: {TIME_BUDGET}s")
print(f"Gradient accumulation steps: {grad_accum_steps}")

# ---------------------------------------------------------------------------
# Расписание (все основано на прогрессе = время обучения / TIME_BUDGET)
# ---------------------------------------------------------------------------

def get_lr_multiplier(progress):
    """
    Возвращает множитель LR в зависимости от прогресса обучения.
    Схема: warmup → plateau → cooldown.
    """
    if progress < WARMUP_RATIO:
        return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
    elif progress < 1.0 - WARMDOWN_RATIO:
        return 1.0  # Плато
    else:
        # Линейный спад до FINAL_LR_FRAC
        cooldown = (1.0 - progress) / WARMDOWN_RATIO
        return cooldown * 1.0 + (1 - cooldown) * FINAL_LR_FRAC

# === Расписание момента для Muon ===
def get_muon_momentum(step):
    """Плавное увеличение момента с 0.85 до 0.95 за первые 300 шагов."""
    frac = min(step / 300, 1)
    return (1 - frac) * 0.85 + frac * 0.95

# === Расписание weight decay ===
def get_weight_decay(progress):
    """Линейное уменьшение weight decay до 0 к концу обучения."""
    return WEIGHT_DECAY * (1 - progress)

# ---------------------------------------------------------------------------
# Цикл обучения
# ---------------------------------------------------------------------------

# === Инициализация цикла ===
t_start_training = time.time()
smooth_train_loss = 0  # EMA для сглаживания потерь
total_training_time = 0
step = 0

while True:
    torch.cuda.synchronize()  # Синхронизация для точного замера времени
    t0 = time.time()
    
    # === Gradient accumulation ===
    for micro_step in range(grad_accum_steps):
        with autocast_ctx:  # Авто-кастинг в bfloat16
            loss = model(x, y)
            train_loss = loss.detach()  # Сохраняем для логирования
            loss = loss / grad_accum_steps  # Масштабируем для накопления
            loss.backward()  # Накопление градиентов
        x, y, epoch = next(train_loader)  # Следующий микро-батч
    
    # === Обновление расписаний ===
    progress = min(total_training_time / TIME_BUDGET, 1.0)
    lrm = get_lr_multiplier(progress)
    muon_momentum = get_muon_momentum(step)
    muon_weight_decay = get_weight_decay(progress)
    
    # Применяем расписания к группам оптимизатора
    for group in optimizer.param_groups:
        group["lr"] = group["initial_lr"] * lrm
        if group['kind'] == 'muon':
            group["momentum"] = muon_momentum
            group["weight_decay"] = muon_weight_decay
    
    # === Шаг оптимизатора ===
    optimizer.step()
    model.zero_grad(set_to_none=True)  # Очистка градиентов
    
    train_loss_f = train_loss.item()
    
    # === Быстрая проверка на взрыв потерь ===
    if math.isnan(train_loss_f) or train_loss_f > 100:
        print("FAIL")
        exit(1)
    
    torch.cuda.synchronize()
    t1 = time.time()
    dt = t1 - t0  # Время шага
    
    # Считаем время обучения только после первых 10 шагов (исключаем компиляцию)
    if step > 10:
        total_training_time += dt
    
    # === Логирование ===
    ema_beta = 0.9
    smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss_f
    debiased_smooth_loss = smooth_train_loss / (1 - ema_beta**(step + 1))  # Дебайасинг EMA
    
    pct_done = 100 * progress
    tok_per_sec = int(TOTAL_BATCH_SIZE / dt)
    mfu = 100 * num_flops_per_token * TOTAL_BATCH_SIZE / dt / H100_BF16_PEAK_FLOPS
    remaining = max(0, TIME_BUDGET - total_training_time)
    
    print(f"rstep {step:05d} ({pct_done:.1f}%) | loss: {debiased_smooth_loss:.6f} | lrm: {lrm:.2f} | dt: {dt*1000:.0f}ms | tok/sec: {tok_per_sec:,} | mfu: {mfu:.1f}% | epoch: {epoch} | remaining: {remaining:.0f}s ", end="", flush=True)
    
    # === Управление сборщиком мусора (избегаем пауз ~500мс) ===
    if step == 0:
        gc.collect()
        gc.freeze()
        gc.disable()
    elif (step + 1) % 5000 == 0:
        gc.collect()
    
    step += 1
    
    # === Проверка завершения по времени ===
    # Ждём минимум 10 шагов, чтобы не прервать во время компиляции
    if step > 10 and total_training_time >= TIME_BUDGET:
        break

print()  # Перенос строки после последнего лога

# === Подсчёт общего количества токенов ===
total_tokens = step * TOTAL_BATCH_SIZE

# === Финальная валидация ===
model.eval()
with autocast_ctx:
    val_bpb = evaluate_bpb(model, tokenizer, DEVICE_BATCH_SIZE)

# === Финальный отчёт ===
t_end = time.time()
startup_time = t_start_training - t_start
steady_state_mfu = 100 * num_flops_per_token * TOTAL_BATCH_SIZE * (step - 10) / total_training_time / H100_BF16_PEAK_FLOPS if total_training_time > 0 else 0
peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024

print("---")
print(f"val_bpb: {val_bpb:.6f}")
print(f"training_seconds: {total_training_time:.1f}")
print(f"total_seconds: {t_end - t_start:.1f}")
print(f"peak_vram_mb: {peak_vram_mb:.1f}")
print(f"mfu_percent: {steady_state_mfu:.2f}")
print(f"total_tokens_M: {total_tokens / 1e6:.1f}")
print(f"num_steps: {step}")
print(f"num_params_M: {num_params / 1e6:.1f}")
print(f"depth: {DEPTH}")

Основные параметры:

Переменная, константа	Значение	Пояснение
PYTORCH_ALLOC_CONF	expandable_segments:True	Включает динамическое расширение сегментов памяти для уменьшения фрагментации
HF_HUB_DISABLE_PROGRESS_BARS	1	Отключает визуальные прогресс-бары при загрузке моделей/данных
cap	(major, minor)	Версия CUDA-архитектуры (например, (9, 0) для H100)
fa3	функция	Интерфейс к оптимизированному ядру flash-attention v3

Параметры класса CausalSelfAttention — причинное самовнимание с оптимизациями

Параметр	Назначение
GQA (Grouped Query Attention)	Несколько query-голов делят одни и те же KV-головы → экономия памяти и вычислений
ve_gate	Обучаемый гейт, который решает, насколько сильно смешивать value-эмбеддинг с обычными value-векторами
flash_attn_func	Оптимизированное ядро, которое вычисляет attention за один проход с подзапросом к памяти
window_size=(W, 0)	Ограничивает внимание токеном только последними W позициями (sliding window)

Гиперпараметры (редактируются напрямую):

Параметр	Значение	Пояснение
ASPECT_RATIO	64	Соотношение глубины и ширины: d_model = depth × 64
HEAD_DIM	128	Размерность одной головы внимания (фиксирована для эффективности)
WINDOW_PATTERN	"SSSL"	Циклический паттерн окон: слои 0,1,2 — короткие окна, слой 3 — полное
TOTAL_BATCH_SIZE	524288	Глобальный размер батча (достигается через gradient accumulation)
WEIGHT_DECAY	0.2	Применяется только при совпадении знаков градиента и параметра
WARMUP_RATIO	0.0	В autoresearch часто начинают с полного LR (быстрый старт)

Метрики финальной оценки и отчёта:

Метрика	Пояснение
val_bpb	Bits Per Byte на валидации — основная метрика качества (меньше = лучше)
mfu_percent	Model FLOPs Utilization — эффективность использования GPU (цель: >40%)
peak_vram_mb	Пиковое потребление видеопамяти
total_tokens_M	Общее количество обработанных токенов в миллионах

Термины, понимание которых необходимо для изучения train.py:

Термин	Значение
GQA (Grouped Query Attention)	Оптимизация: несколько query-голов используют общие KV-головы
RoPE (Rotary Position Embedding)	Позиционные эмбеддинги через вращение векторов в комплексной плоскости
RMSNorm	Нормализация через среднеквадратичное значение (без смещения)
Muon	Оптимизатор второго порядка с ортогонализацией градиентов
Flash Attention 3	Оптимизированное ядро для causal attention с subquadratic памятью
MFU (Model FLOPs Utilization)	Доля теоретической производительности GPU, используемая моделью
BPB (Bits Per Byte)	Метрика качества: биты информации на байт текста (меньше = лучше)
Softcapping	Ограничение диапазона логитов через tanh для стабильности
Cautious WD	Weight decay применяется только при совпадении знаков градиента и параметра

Обратите внимание на ключевые особенности реализации:

1. Гибридный оптимизатор: Muon для матриц + AdamW для остального → лучшее качество при высокой скорости

2. torch.compile + fused kernels: Все шаги оптимизатора скомпилированы в один граф для минимизации накладных расходов

3. Эффективная упаковка данных: Best-fit packing из prepare.py обеспечивает 100% утилизацию контекста

4. Устойчивость к сбоям: Быстрая проверка на NaN/взрыв потерь с немедленным завершением

5. Точный учёт времени: Исключение времени компиляции из расчёта TIME_BUDGET

6. Управление памятью: gc.freeze() + gc.disable() после инициализации устраняет паузы сборщика мусора