GigaChat Pro. Технические детали, сдача ЕГЭ и сравнение с конкурентами

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

За последнее время нами были выпущены новые модели и добавлены новые функциональные возможности в сервис. А теперь мы представляем новую нейросетевую модель, у которой более 29 миллиардов параметров. Она успешно прошла тест на ЕГЭ и показала отличные результаты в сравнении с другими системами. Если вы только начинаете интересоваться новыми технологиями в области искусственного интеллекта, то имейте в виду, именно они в ближайшие годы будут формировать наше с вами будущее и давать преимущество тем, кто ими владеет. Давайте познакомимся с ними поближе.

Не ещё один чат

GigaChat и подобные системы, такие, как ChatGPT, являются универсальным инструментом, который позволяет решать творческие, повседневные и аналитические задачи.

Эти системы обычно состоят из нескольких специализированных нейросетевых моделей. Основная модель в этом ансамбле играет роль дирижера, она генерирует креативные ответы на запросы пользователей, взаимодействует с другими моделями и требует особого внимания при обучении и настройке. Помимо основной модели есть и другие, которые расширяют функциональность системы. Например, есть модель Kandinsky, которая может создавать изображения по вашему запросу. Есть модели SaluteSpeech, способные распознавать и синтезировать речь, а также другие, предназначенные для поиска информации во внешних источниках или анализа загружаемых документов.

Приведем пару примеров работы с GigaChat. Запросы могут быть самые разные:

Технические

Даны списки:
a = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]; b = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13].

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

Ответ:

Для решения этой задачи можно использовать встроенную функцию Python set(), которая возвращает уникальные элементы списка. После этого можно преобразовать полученное множество обратно в список с помощью функции list(). В итоге получим список, состоящий из элементов, которые присутствуют и в первом, и во втором списке.

Вот как это можно сделать:
a = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89];
b = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
c = list(set(a) & set(b)).

После выполнения этого кода переменная c будет содержать список [1, 2, 3, 5, 8, 13].

Креативные

Давно хочу продать свой гараж в Химках. Напиши продающий текст, чтобы каждый захотел купить его. Не жалей эпитетов.

Ответ:

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

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

Могу сказать наверняка: ваш автомобиль заслуживает лучшего, и этот гараж станет для него настоящим домом. Не упустите свой шанс обзавестись недвижимостью в экологически чистом районе Химок. Жду ваших обращений!

Развлекательные

Сочини считалку про программиста

Ответ:

В офисе скрестили шпаги
Наши программисты:
Разбирали они баги,
Фича-рецидивисты.

Что под капотом

Процесс создания GigaChat, как и всех других AI-сервисов на основе LLM от ChatGPT до Claude, состоит из двух этапов:

• Unsupervised Pretraining — учим нейронную сеть на большом количестве легкодоступных данных. GigaChat учится предсказывать следующее слово в текстах из Интернета, книгах и т. д.

• Alignment — учим нейронную сеть с предыдущего шага следовать инструкциям пользователя, при этом не нарушая никакие законы и не будучи предвзятым. Эта стадия называется alignment (выравнивание), потому что происходит выравнивание желаний и ценностей людей с тем, что генерирует нейронная сеть.

В этой статье мы расскажем про нашу новую pretrain модель, которая легла в основу GigaChat Pro версии.

Возможности новой модели

Самое большое событие релиза — это достижение паритета с ChatGPT-3.5-turbo на метрике SBS (Side-By-Side), где мы просим людей сравнить качество ответов двух моделей.

У SBS есть минусы и поэтому мы дополняем его с помощью автоматических метрик — задач, ответы на которые можно проверить без участия людей. Такие метрики позволяют получать оценку модели очень быстро, а также оценивать Pretrain'ы, которые еще не научились следовать инструкциям.

Для оценки моделей после alignment в автоматических метриках мы используем продуктовые задачи из банка и от внешних клиентов, которые передали нам корзины запросов для оценки, а для оценки pretrain моделей замеряем их на открытых бенчмарках.

SBS (Side-By-Side)

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

• Лучше ответила 1-ая модель (всего таких ответов good_a);

• Лучше ответила 2-ая модель (всего таких ответов good_b);

• Обе модели ответили хорошо (всего таких ответов both);

• Обе модели ответили плохо (всего таких ответов none).

Результат SBS считается по формуле:

Результаты SBS (50%) GigaChatPro vs (50%) ChatGPT

Результаты SBS (62%) GigaChatPro vs (38%) YandexGPT

Автоматические метрики pretrain модели

Метрики для таблицы были замерены нами, чтобы сравнивать все модели в одинаковых условиях.

Основной акцент GigaChat’а — это русский язык, но наш pretrain показывает результаты, сравнимые с лучшими Open Source моделями и на английском языке.
На английском MMLU GigaChat Pro до alignment’а дает более высокие результаты, чем LLama-2-34B (64,7% против 62,6%). К сожалению, LLama-2-34B не была выпущена в Open Source, поэтому мы не можем замерить другие метрики вроде MMLU-RU.

На MMLU-RU GigaChat Pro побеждает одну из самых сильных Open Source моделей на данный момент Mistral-7B (57,8% против 50,5%).

Автоматические метрики после alignment

Для получения метрики в «продуктовых задачах» мы усреднили результаты замеров на разных датасетах.

В «Ответы по документам» необходимо указать процент аванса и число дней отсрочки, если они указаны в документе.

Сдача ЕГЭ по обществознанию

GigaChat смог выполнить все задания единого государственного экзамена по обществознанию и набрал 67 баллов, что превышает минимальный балл для подачи в вуз (45 баллов) и средний балл по предмету в 2023 году (56,4 балла). Ответы модели были проверены независимыми экспертами НИУ ВШЭ.

Чтобы получить такой высокий результат, были использованы различные подходы. На итоговое качество влияли инструкции, параметры генерации, размер модели и пайплайн формирования конечного ответа. Так, модели 29b достаточно было просто задать вопрос из ЕГЭ, в то время как 7b конфигурация демонстрировала схожее качество работы уже с помощью цепочки рассуждений, включающей генерацию размышления модели о теме вопроса, критику ответа, исправление ошибок и приведение к требуемому формату.

Результаты экзамена говорят о том, что GigaChat хорошо эрудирован в области социально-юридических наук и способен корректно решить задачи разного типа.

Выводы

GigaChat Pro достиг паритета по способностям с ChatGPT, а на продуктовых задачах в банке обходит его. Сама pretrain модель находится на уровне сильнейших Open Source на английском языке сравнимого размера, а на русском является лидером.

API

Чтобы попросить GigaChat перевести текст, придумать анекдот или сгенерировать картинку с котиком, отлично подходит веб интерфейс, но для более серьезных применений нужен удобный программный интерфейс — API. ChatGPT от OpenAI на днях запретил обращение к своей системе с российских IP-адресов и GigaChat может стать хорошей альтернативой, достигая такого же или лучшего качества на многих задачах, особенно на русском языке.

Мы предоставляем индивидуальным разработчикам специальный freemium тариф на 1 миллион бесплатных токенов. Если их не хватает, можно докупить нужный пакет токенов. Как по мне, так это отличная возможность интегрировать нейросетевые технологии в свой проект или облегчить свои рутинные задачи на работе.

Сценарии работы с API подробно описаны в нашей документации, но в целом алгоритм следующий:

1. Получить в личном кабинете авторизационные данные, которые будут использоваться для получения временного токена доступа.

curl --location --request POST 'https://ngw.devices.sberbank.ru:9443/api/v2/oauth' \
     --header 'Authorization: Bearer <auth_token>' \
     --header 'RqUID:f2911be0-8c91-4576-bbd4-5238d834c17a' \
     --header 'Content-Type: application/x-www-form-urlencoded' \
     --data-urlencode 'scope=GIGACHAT_API_PERS'
Ответ:
{
  "access_token":"eyJjdHkiOiJqd3Q...iLCJlbmMiOiJBMjU2Q0JDLU",
  "expires_at":1700396614784
}

2. Использовать токен при обращении к необходимой модели.

   Полный контекст диалога можно передавать в виде массива объектов с ролями user и assistant.

curl -k https://gigachat.devices.sberbank.ru/api/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -H 'Authorization: Bearer eyJjdHkiOiJqd3Q...iLCJlbmMiOiJBMjU2Q0JDLU' \
  -d '{
  "model": "GigaChat-Pro",
  "messages": [
    {
      "role": "user",
      "content": "Придумай альтернативное название для шоколадки Алёнка, 
            оно должно быть необычным и футуристичным"
    }
  ],
  "temperature": 0.65
}'
Ответ:
{
  "choices": [
    {
      "finish_reason": "stop",
      "index": 0,
      "message": {
        "content": "Как насчёт \"Квантовой сладости\"?
                    Это название звучит необычно и подразумевает использование
                    передовых технологий в производстве шоколада.",
        "role": "assistant"
      }
    }
  ],
  "created": 1700395635,
  "model": "GigaChat:latest",
  "object": "chat.completion",
  "usage": {
    "completion_tokens": 30,
    "prompt_tokens": 191,
    "total_tokens": 221
  }
}

Отмечу также и еще одно преимущество работы с API. Вы можете самостоятельно управлять семплингом и регулировать параметры top_p, top_k, temperature и repetition_penalty. Это дает вам контроль над тем, будет ли модель генерировать самый вероятный ответ или будет варьировать распределение токенов, чтобы каждый раз выдавать что-то новое. Первый вариант больше подходит для решения задач на классификацию текста или ответа на экзаменационный вопрос, а последний — для генерации креативных материалов, типа маркетинговых текстов.

Техническая часть

Чтобы внести свой вклад в развитие LLM, мы хотим поделиться подробностями нашего пайплайна для pretraining’а. Мы надеемся, что эта информация поможет сообществу развиваться и поддерживать конкурентную среду.

Начнем с общей технической информации о наших моделях. Как и предыдущая модель, новый GigaChat Pro — это декодер трансформер 29 миллиардов параметров с улучшениями для ускорения сходимости и обучения. Для кодирования позиционной информации используются Rotary эмбеддинги. В качестве оптимизатора используется Adam с .

Основная часть pretraining’а происходит с контекстом в 4 тысячи токенов. Ближе к концу обучения мы расширяем контекст моделей, используя свойства Rotary эмбеддингов, и заканчиваем обучение с большим контекстом (24-32 тысячи токенов).

Для обучения моделей мы используем фреймворк FSDP (Fully Sharded Data Parallel). По функциональности он очень похож на DeepSpeed и позволяет шардировать модели при обучении, т.е. хранить на каждой GPU только часть весов и состояний оптимизатора.

Большинство Open Source LLM сейчас используют по сути одинаковую архитектуру и процесс обучения нейронной сети. Направления, в которых идет основная конкуренция — это качество pretrain данных, эффективность обучения и небольшие детали в процедуре обучений. Именно про продвижение в этих направлениях мы расскажем далее.

Данные

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

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

Во-вторых, самые качественные Open Source датасеты вроде FRW (Falcon Redefined Web) отфильтровывали русские данные, а мультиязычные датасеты не делали акцент на русском языке и в целом слабее FRW.

Для сбора качественного датасета мы решили самостоятельно очистить Common Crawl, сделав акцент на двух языках — русском и английском. Извлечение текста из HTML страниц проводилось с помощью библиотеки trafilatura, а для дальнейшей очистки мы дополнили классические фильтрации на основе эвристик с помощью двухстадийного нейронного классификатора. После очисток применяется дедупликация с помощью SimHash.

Нейронные классификаторы

Основной тип ошибок, которые мы ищем — некорректное извлечение текста из HTML trafilatur’ой. Первая стадия очистки разделяет документы на три цвета — зеленый, желтый и красный.

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

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

Отладка процесса обучения

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

GigaChat Pro — трансформер, содержащий 29 миллиардов параметров, и его обучение потребовало от нас решения многих инфраструктурных и архитектурных проблем.

Инициализация модели

Текущие модели настолько большие, что для инференса порой не хватает 80 GB GPU памяти даже на Nvidia-A100. Обучение модели требует еще больше ресурсов, поэтому используется техника шардирования (разбиения) модели между вычислительными узлами, как, например, показано на картинке.

Шардировать модель можно тремя способами (рассмотрим случай, когда инициализация слоев происходит независимо друг от друга):

1. Инициализировать ее полностью на GPU и разослать соответствующие куски модели на другие GPU. Этот подход не сработает, если модель достаточно большая, чтобы не поместиться на одну GPU.

2. Инициализировать ее на CPU, а потом передавать куски модели на соответствующие GPU. Это намного дольше, чем пункт 1, но зато позволяет использовать оперативную память вычислительного узла, которой зачастую в десятки раз больше, чем на видеокарте.

3. Если известен алгоритм шардирования модели, то можно сразу произвести инициализацию шардов на нужных GPU без полной инициализации модели в одном месте. Это самый быстрый и самый эффективный по памяти метод (meta init).

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

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

Токенизатор

Токен — некоторая смысловая единица текста (например, подстрока или последовательность байт). Задача токенизатора состоит в том, чтобы эффективно разбивать поступающий в модель текст на подстроки или последовательности байт и сопоставлять им соответствующие токены.

Изначально мы использовали токенизатор со словарем в 132 тысяч токенов и byte-level токенизаций. Это приводило к ряду проблем:

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

   

2. Неверная токенизация математических формул. Исходный токенизатор сопоставлял отдельный токен не каждой цифре, а отдельным числам или последовательностям цифр.

3. Развал модели при более агрессивных параметрах оптимизации (уменьшении Adam epsilon).

Мы попробовали решить их комплексно и обучили новый мультиязычный токенизатор на 42 тысячи токенов. У нас получилась достаточно высокая степень сжатия: в среднем на один токен приходится 3.9 символа текста. Также в новом токенизаторе каждая цифра — отдельный токен, что позволило GigaChat’у лучше понимать математические задачи.

Оптимизатор

В своих обучениях мы используем Decoupled AdamW оптимизатор. Проведя ряд экспериментов с другими оптимизаторами, мы выбрали Adam. Он дает самые низкие значения loss функции, а также достаточно прост во внедрении.

Рассмотрим гиперпараметр Epsilon: самая популярная Open Source модель LLaMa-2 обучалась с . Наши эксперименты показывают, что:

1. При большом размере токенизатора уменьшение epsilon может вызывать развал обучения.

2. Более низкий epsilon с первых шагов и до конца обучения демонстрирует наименьшие значения loss функции, то есть модель обучается лучше.

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

Рестарты обучения

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

1. Каждые N часов производится сохранение чекпоинта обучения, и, если падение случилось ровно перед сохранением модели, мы теряем N часов работы кластера.

2. Каждый рестарт также занимает некоторое время.

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

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

NaN в обучении

Во время pretrain’а GigaChat’а мы столкнулись с тем, что неожиданно число батчей, на которых возникает NaN, начало резко расти. А обновлять NaN’ом веса нельзя, так как это ведет к развалу модели.

поиске проблемы нам помогла система мониторинга обучения модели. Она собирает информацию о:

• нормах, минимумах и максимумах весов, их градиентов и соответствующих моментов оптимизатора;

• нормах, минимумах и максимумах активаций;

• номерах слоев на forward и backward, где возник NaN;

• скорости обучения, использования памяти, метриках качества и прочих деталях.

По этим показателям мы заметили, что NaN возникал на backward сразу перед последним слоем, и придумали решение. Обновление весов модели при шардированном обучении обычно реализовано следующим образом:

1. на каждой GPU считается loss по своему батчу;

2. с каждой карты аккумулируются и усредняются градиенты;

3. происходит backward и обновление параметров.

Был модифицирован пункт 2, и те градиенты, где встречаются NaN’ы, мы обнулили. Плюсы и минусы этого подхода:

Плюсы:

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

Минусы:

• Тратим лишние вычислительные ресурсы на проверку на NaN’ы и зануление градиентов. Снижается скорость обучения.

• Снижаем норму общего градиента с учетом зануления.

Этот трюк помог «вылечить» нашу модель от вычислительной нестабильности. Включив его, число NaN’ов начало падать и в итоге сократилось до нуля.

Позже мы отключили его, и модель продолжила стабильно обучаться.

BF16 и FP32

Активации наших моделей для экономии памяти и ускорения подсчетов на обучении хранятся в типе данных bf16 (bfloat16). Этот тип данных имеет тот же range, что и fp32, но меньшую точность. Для аккумуляции градиентов и на этапе вычисления backward мы используем тип fp32, чтобы не допустить потерь в точности и ошибок вычислений.

Ускорение обучения трансформеров

А теперь давайте рассмотрим методы, которые позволили нам ускорить обучение в несколько раз.

Flash Attention 2

FlashAttention реализует механизм внимания (attention is all you need), который оптимизирован для работы с устройствами, имеющими высокую пропускную способность памяти (HBM) и быстродействующую оперативную память (SRAM).

Чем отличается HBM память от SRAM?

SRAM (Static Random Access Memory) и HBM (High Bandwidth Memory) — это два разных типа памяти, которые используются в компьютерных системах. Они отличаются по многим параметрам, включая скорость работы, емкость, стоимость и конструктивные особенности.

SRAM — это тип статической оперативной памяти. Она быстрее и надежнее, чем обычная оперативная память (RAM), но более дорогая и имеет меньшую емкость. SRAM используется в качестве кэш-памяти процессора, так как кэш требует высокой скорости работы.

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

Картинка из официального репозитория.

Ускорение достигается путем разбиения входного потока данных на более мелкие блоки и обработки этих блоков параллельно. Это позволяет уменьшить количество обращений к «медленной» памяти, что ускоряет выполнение операций.

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

Небольшой технический совет

В реализации из flash-attention обратите внимание на переменную causal.

Во многих реализациях блока внимания на основе flash-attention, подача непустой attention_mask (в том числе состоящей из единиц) сопровождается использованием функции unpad_input на каждом слое. Эта функция вызывает множество последовательных синхронизаций CPU и GPU. Чтобы найти подобные узкие места, можно воспользоваться PyTorch Profiler.

Чем плохи синхронизации между CPU и GPU?

Синхронизация между CPU и GPU может стать узким местом в процессе обработки данных. Это связано с тем, что передача данных между CPU и GPU требует времени и ресурсов, что может замедлить выполнение операций. Кроме того, синхронизация может приводить к простою одного из устройств во время работы другого, что также снижает общую производительность системы.

Пример синхронизации CPU & GPU от использования unpad_input. На больших обучениях такие простои могут стоить много миллионов рублей.

Другие fused слои

Что такое fused слои?

Fused слой — это слой, в котором операции выполняются с использованием «слияния» или совместной обработки. В контексте глубокого обучения, «слияние» обычно относится к оптимизации, при которой несколько операций объединяются в одну, чтобы улучшить эффективность вычислений. Например, вместо того, чтобы выполнять операцию свертки, а затем операцию активации, можно объединить их в одну операцию, которая выполняет обе функции одновременно. Это может ускорить выполнение модели и снизить использование памяти.

К примеру, FlashAttention тоже относится к fused слоям. Но он особенный, поэтому мы написали о нем отдельно :).

Мы используем множество фьюзов, вот некоторые из них: Fused Cross-Entropy, Fused LayerNorm, Fused Rotary Embeddings. Каждый дает большой прирост к скорости, а некоторые — огромный выигрыш в видеопамяти.

Подробнее можно почитать в уже классическом блог-посте Making Deep Learning go Brrrr From First Principles

Selective Activation Checkpointing

Что такое Activation Checkpointing?

Activation checkpointing — это техника оптимизации, используемая в глубоком обучении для уменьшения количества памяти, необходимой для хранения активаций между операциями. Вместо того, чтобы сохранять активации для каждого отдельного примера во время обучения, activation checkpointing сохраняет в памяти только часть активаций (например, для каждого второго слоя), а остальные пересчитывает при вычислении backward. Это может существенно снизить использование памяти и дает возможность обучать модели большего размера.

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

При достаточно большом контексте (или модели), мы сталкиваемся с тем, что активации всей модели уже не помещаются в видеопамять GPU. Тем не менее, нам нужны активации для подсчета backward pass. Решение — можем хранить только активации для N слоев, а остальные пересчитывать по необходимости вычисления градиентов.

Но вместо того, чтобы хранить все активации для N слоев, мы могли бы сохранять только самые вычислительно-затратные активации на единицу памяти. Это позволит нам сохранять «дорогие» активации для дополнительных M слоев.

Реализации подобных слоев с гибкой настройкой под модель позволили хранить в 5 раз больше вычислительно-затратных на единицу памяти активаций и тем самым увеличить скорость обучения на 10%!

В дополнение, если речь идет о меньших моделях, то снижение потребляемой активациями памяти позволяет отказаться от полного разделения параметров модели по всем GPU (ZeRO 3), оставив только шардирование состояний оптимизатора и градиентов (ZeRO 2), что дает около 20% в производительности.

Синхронный GarbageCollector

Как PyTorch управляет памятью?

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

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

При работе на границе видеопамяти GPU кэширующий аллокатор в PyTorch работает на пределе. Это происходит из-за недетерминированного перекрытия computation и communication, что приводит к увеличению доли памяти, недоступной из-за фрагментации, и в конечном итоге может привести к ошибке OOM (out of memory).

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

Подробнее про работу аллокатора и FSDP можно почитать тут.

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

Чтобы минимизировать простой и суммарное время простоя, можно синхронизировать чистку мусора. Таким образом, мы смогли эффективно использовать ресурсы GPU и избежать проблем с памятью.

Улучшаем общение между GPU

Когда у Вас есть кластер с тысячей видеокарт, выстроить эффективное общение между ними — это сложная задача. GPU от Nvidia для этого используют библиотеку NCCL (Никель).

Простыми словами про NCCL.

Nvidia NCCL— это библиотека для организации параллельных вычислений на графических процессорах (GPU) от компании Nvidia.

NCCL позволяет организовать взаимодействие между GPU, которые подключены к одному компьютеру или находятся в одной сети. Библиотека обеспечивает быстрый обмен данными между GPU, что в разы увеличивает скорость работы алгоритмов, которые взаимодействуют с большими объёмами данных.

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

Так мы смогли сократить передаваемый объем данных по сети, улучшить топологию общения и увеличить скорость обучения на 5%!

Итоговые ускорения

Тесты проводились на модели 29B, 128GPU A100. *Эффект от настройки NCCL приведен для 1024 GPU, так как на маленьком числе узлов эффект не так заметен.

Благодарности

Задача по осуществлению такого проекта требует больших усилий, внимательности и вовлеченности. Участвуя в этом процессе, я не перестаю удивляться энтузиазму коллег, генерирующих идеи и готовых воплощать их в жизнь в любое время. Девушки и парни из SberDevices, Sber AI и Научно-исследовательского института искусственного интеллекта AIRI, вы лучшие.

Ждём вашу обратную связь

Также приглашаю вас в наш Telegram-канал Salute AI, в котором мы с коллегами начали делиться наработками в области машинного обучения и другими рабочими моментами. А в соответствующем чатике Salute AI Community можно напрямую поспрашивать про всё это и просто пообщаться.

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

Авторы: Сергей Аверкиев (@averkij), Валентин Мамедов (@vltnmmdv), Евгений Косарев (@evgenijkkk), Григорий Лелейтнер (@runfme), Эльдар Дамиров (@edmr_x5720), Дарья Латорцева (@latortseva)