Как Azercell обучила азербайджанскую LLM на Amazon SageMaker AI

Решение построено на open source-инструментах, включая PyTorch, Hugging Face Transformers и Liger Kernels. Авторы также благодарят Aiham Taleb, Arefeh Ghahvechi, Manav Choudhary, Rohit Thekkanal, Daz Akbarov, Jamila Jamilova, Ross Povelikin, Almas Moldakanov, Christelle Xu и Ivan Khvostishkov за вклад в реализацию проекта.

Azercell Telecom LLC, ведущий телеком-провайдер Азербайджана, хотел создать азербайджанскую большую языковую модель (LLM) на Amazon SageMaker AI для телеком-сценариев и клиентского чат-бота. Задача заключалась в адаптации foundation models (FM) к морфологически богатому языку с ограниченным объемом обучающих данных и без готового шаблона эффективного обучения LLM на азербайджанском. За шесть недель Azercell вместе с AWS Generative AI Innovation Center выстроила production-ready фреймворк на Amazon SageMaker AI, который обеспечил на 23% более высокую пропускную способность обучения и на 58% более низкое пиковое потребление GPU-памяти благодаря оптимизациям на уровне kernel на экземпляре ml.p5.48xlarge. Также фреймворк добился 2× улучшения по показателю tokens per word с помощью собственного tokenizer, фактически удвоив объем азербайджанского текста, помещающийся в контекстное окно модели. Если вы работаете с языками с низкой ресурсной обеспеченностью или сложной морфологией, этот материал поможет оценить аналогичные подходы.

Обзор решения

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

• Этап 1: разработка tokenizer создает эффективный tokenizer для азербайджанского языка. Мы оценили три подхода: базовые tokenizer, оптимизированные под английский, расширение словаря и собственные одноязычные tokenizer, измеряя эффективность кодирования по стандартным метрикам. Лучший результат показал собственный одноязычный tokenizer, сократив tokens per word вдвое по сравнению с базовой версией.
• Этап 2: Continued pre-training (CPT) адаптирует FM (Llama 3.2 1B) к пониманию азербайджанского языка с помощью distributed training и оптимизаций Liger Kernel в training jobs Amazon SageMaker AI. Это позволяет использовать большие batch size и добиться более высокой пропускной способности на том же оборудовании. Хотя distributed training не был необходим для proof-of-concept уровня 1B, он станет критически важным по мере перехода Azercell к более крупным моделям.
• Этап 3: Supervised fine-tuning с Low-Rank Adaptation (LoRA) превращает дообученную модель в conversational assistant. После CPT модель умеет предсказывать азербайджанские токены, но не умеет вести диалог. На третьем этапе применяется LoRA — метод parameter-efficient fine-tuning, который значительно сокращает число обучаемых параметров.

Этапы обучения (CPT и LoRA fine-tuning) выполнялись как training jobs Amazon SageMaker AI, запущенные из Amazon SageMaker Unified Studio, и каждый из них ссылался на собственный training script. Каждый job поднимает свежие экземпляры Amazon Elastic Compute Cloud (Amazon EC2) и завершает работу после окончания, поэтому вы платите только за фактическое время вычислений без затрат на простой кластера.

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

Рисунок 1. Архитектура pipeline обучения. Операторы запускают training jobs из Amazon SageMaker AI Notebook Instances. Обучающие данные и артефакты модели хранятся в Amazon Simple Storage Service (Amazon S3). Метрики обучения отслеживаются в TensorBoard в Amazon SageMaker AI, а системные метрики собираются через Amazon CloudWatch.

Разработка азербайджанского tokenizer

Такие языки, как азербайджанский, имеют богатую морфологию: одно слово может нести грамматическое значение с помощью суффиксов, которое в английском языке выразилось бы несколькими словами. Однако стандартные tokenizer, оптимизированные под английский, дробят такие сложные словоформы. Например, слово «kitablardan» (в значении «из книг») может разбиваться на несколько subword tokens, как показано на рисунке 2, из-за чего в фиксированное context window помещается меньше фактического содержания.

Рисунок 2. Сравнение базовой и пользовательской tokenization для азербайджанского текста, показывающее меньшее дробление токенов.

Чтобы решить эту проблему, мы обучили собственный tokenizer на азербайджанском тексте с использованием алгоритма Byte-Level Byte-Pair Encoding (BBPE), который итеративно объединяет наиболее частые пары байтов в элементы словаря. Использование необработанных байтов вместо заранее заданных наборов символов обеспечивает полное покрытие специфических для азербайджанского языка символов без ручного определения алфавита. Мы экспериментировали с размерами словаря от 50k до 100k tokens, чтобы найти баланс: слишком маленький словарь приводит к избыточному дроблению слов, слишком большой — к нехватке обучающего сигнала для редких токенов.

Мы обучали пользовательские tokenizer с помощью библиотеки Hugging Face tokenizers с той же конфигурацией, что и у нативного tokenizer Llama 3.2, меняя только размер словаря. После обучения и оценки нескольких tokenizer с разными размерами словаря мы выбрали финальный словарь на 100k tokens. Чтобы убедиться, что пользовательский tokenizer не ухудшил качество моделирования, мы сравнили модели после continued pre-training по метрике Bits-Per-Byte (BPB), а не perplexity, потому что BPB нормализует различия в словаре и измеряет качество предсказания на уровне байтов. Модель с пользовательским tokenizer показала BPB 0.5795 на валидационном наборе против 0.6830 у базовой версии, что подтвердило отсутствие компромисса в качестве при лучшей эффективности кодирования.

Помимо сохранения качества моделирования, пользовательский tokenizer дает заметный практический выигрыш в эффективности. Эффективность кодирования можно измерять через fertility score — среднее число токенов на слово, где меньшие значения означают более эффективное кодирование. Базовый tokenizer Llama 3.2 в среднем использовал 3.22 токена на азербайджанское слово, тогда как собственный одноязычный tokenizer — 1.59, то есть в 2 раза эффективнее. В сочетании с context window Llama 3.2 на 128k tokens это дает ощутимую разницу в емкости: примерно 40k слов для базового tokenizer против 80k для оптимизированного, фактически удваивая объем текста, который модель учитывает одновременно.

Continued pre-training

Continued pre-training адаптирует FM (Llama 3.2 1B) к пониманию азербайджанского языка. Основное узкое место на этом этапе — GPU memory: оптимизация использования памяти напрямую определяет, какая часть инвестиций в оборудование превращается в training throughput. Мы провели бенчмарки на экземплярах ml.p4d.24xlarge (8× NVIDIA A100 GPUs) и ml.p5.48xlarge (8× NVIDIA H100 GPUs). В следующих разделах описаны две оптимизации, которые мы сравнивали: distributed training с Fully Sharded Data Parallel (FSDP) в PyTorch и интеграция Liger Kernel.

Distributed training с Fully Sharded Data Parallel (FSDP)

Памятный след модели включает не только веса, но и gradients, optimizer states и activations. Для более крупных моделей, таких как Llama 3.1 8B, эти компоненты в смешанной точности могут превышать 100 GB. Мы разработали и проверили схему distributed training на модели 1B, чтобы при переходе к более крупным архитектурам требовалось только изменение конфигурации, а не переработка pipeline. Стандартный Distributed Data Parallel (DDP) реплицирует всю модель на каждом GPU, что ограничивает достижимые batch size и масштаб модели. FSDP делит параметры, gradients и optimizer states между GPU и динамически собирает только то, что нужно на каждом шаге вычислений. Это снизило память состояния модели на одном GPU с 9.23 GB до 1.17 GB на ml.p4d.24xlarge, освободив запас для больших batch size.

Интеграция Liger Kernel

Liger Kernels — это экономящие память реализации типовых операций LLM на основе Triton, которые объединяют несколько операций в один запуск GPU kernel, сокращая промежуточные выделения памяти при численно эквивалентных результатах. Они поддерживают несколько популярных архитектур, включая Llama. Мы рекомендуем проверить совместимость с вашей архитектурой перед внедрением.

Интеграция требует минимальных изменений кода: один вызов функции патчит модель оптимизированными kernel перед инициализацией, а Liger Kernels работают с PyTorch FSDP без изменений в distributed training setup. Корректность выполнения мы проверили с помощью PyTorch Profiler, подтвердив наличие fused operations в трассировке. Следующая таблица суммирует совокупный эффект каждой оптимизации на обоих типах экземпляров. Обратите внимание, что память и throughput DDP на p5 не измерялись, поскольку целевой конфигурацией был FSDP.

На ml.p4d.24xlarge полный стек оптимизаций увеличил максимальный batch size в 7 раз по сравнению с DDP. На ml.p5.48xlarge пиковое потребление GPU memory снизилось на 58%, а throughput на GPU вырос на 23% после добавления Liger Kernels к FSDP.

Настройка pre-training

Каждая конфигурация tokenizer из этапа 1 была перенесена через CPT end-to-end, чтобы сравнить поведение сходимости и качество на downstream-задачах. При использовании собственного азербайджанского tokenizer (словарь 100k) объем обучающего корпуса составляет примерно 2.5B tokens.

Пользовательский training script поддерживает настраиваемое context window, смешанную точность BFloat16, cosine learning rate scheduling с AdamW и автоматическое checkpointing в Amazon S3 для отказоустойчивости. Мы установили context window в 2,048 tokens, потому что после tokenization более 90% обучающих примеров были короче этого значения, хотя конфигурация поддерживает до нативного лимита модели в 128k tokens.

При добавлении новых токенов в словарь CPT выполняется в два этапа. На первом этапе backbone модели замораживается, и обучается только embedding layer. Это адаптирует представления новых токенов к уже существующему внутреннему пространству модели, не разрушая предварительно обученные знания. На втором этапе параметры размораживаются для полного обучения, что позволяет модели глубже усвоить шаблоны азербайджанского языка. Следующая таблица показывает конфигурацию обучения с пользовательским азербайджанским tokenizer (100k vocabulary). Обучение выполнялось на двух экземплярах ml.p4d.24xlarge (всего 16 NVIDIA A100 GPUs) с оптимизациями FSDP и Liger Kernel.

Более низкий learning rate на этапе полного обучения сохраняет знания, полученные во время адаптации embedding. При эффективном batch size 224 (14 на GPU × 16 GPU) и context window 2,048 tokens каждый шаг обучения обрабатывает примерно 450k tokens, что дает расчетное время на эпоху около 4.3 часа в этой конфигурации. На ml.p5.48xlarge более высокая пропускная способность на GPU и больший batch size дополнительно сократили бы время на эпоху.

Supervised fine-tuning с LoRA

После CPT модель умеет бегло предсказывать следующий азербайджанский token, но не понимает структуру диалога. На вопрос она выдает правдоподобное продолжение, а не полезный ответ. LoRA эффективно закрывает этот разрыв: pre-trained weights замораживаются, а в attention- и feed-forward layers модели обучаются небольшие матрицы низкого ранга. Вместо обновления полной weight matrix LoRA обучает две меньшие матрицы, произведение которых аппроксимирует полный update, сокращая число обучаемых параметров до небольшой доли от общего объема. Следующая таблица суммирует конфигурацию fine-tuning с LoRA.

Компактный размер решения позволил выполнить fine-tuning на одном экземпляре ml.g5.8xlarge (1× NVIDIA A10G GPU) и завершить его за минуты. Для fine-tuning использовалось примерно 2,000 одноходовых пар вопросов и ответов на азербайджанском языке с Hugging Face SFTTrainer и learning rate 1e-4 — выше, чем в CPT, потому что адаптеры LoRA инициализируются случайно и выигрывают от более сильных градиентных обновлений.

Обучение использовало Llama-style chat template с assistant-only loss masking: модель штрафуется только за предсказание токенов ответа assistant и токена окончания реплики (<|eot_id|>), тогда как пользовательские запросы и разделители шаблона исключаются из loss. В результате модель сосредотачивает вычислительную емкость на генерации уместных ответов, а не на запоминании паттернов пользовательского ввода.

Результаты и проверка

Continued pre-training использовал примерно 2.5B tokens с пользовательским азербайджанским tokenizer, а fine-tuning — 2,000 пар вопросов и ответов. Фреймворк дал измеримые улучшения по четырем направлениям:

• 2× повышение эффективности кодирования благодаря custom tokenization. Собственный одноязычный tokenizer снизил fertility score вдвое — с 3.22 до 1.59 токена на слово, фактически удвоив объем азербайджанского текста, помещающегося в 128k-token context window модели. Показатель BPB 0.5795 против 0.6830 у базовой версии подтвердил, что это улучшение не ухудшило качество моделирования.
• Существенная оптимизация памяти и пропускной способности с Fully Sharded Data Parallel (FSDP). Шардирование FSDP и интеграция Liger Kernel позволили увеличить batch size на том же оборудовании: до 7× на ml.p4d.24xlarge и 4.5× на ml.p5.48xlarge относительно соответствующих базовых конфигураций DDP, при этом пиковое потребление GPU memory снизилось на 58%, а throughput на GPU вырос на 23%.
• Production-ready и масштабируемая инфраструктура. Проверенные конфигурации для экземпляров ml.p4d.24xlarge и ml.p5.48xlarge дают Azercell гибкую основу, которая масштабируется по мере роста требований к обучению, а distributed setup спроектирован так, что переход к более крупным моделям требует только изменения конфигурации.
• Согласованная генерация текста на азербайджанском языке. Готовая модель Llama 3.2 1B без дообучения при запросах на азербайджанском языке уходит в повторяющиеся и бессмысленные ответы — типичная проблема для недопредставленных языков. Дообученная модель выдает короткие, семантически связные ответы, демонстрируя уверенное понимание азербайджанского языка. Это показано в таблице ниже на примере запроса, предоставленного Azercell.

Вывод

В этом материале мы показали, как Azercell и AWS Generative AI Innovation Center построили фреймворк для обучения азербайджанских language models на Amazon SageMaker AI. Трехэтапный pipeline — custom tokenization, continued pre-training с оптимизациями FSDP и Liger Kernel, а также fine-tuning с LoRA — превращает универсальную foundation model в азербайджанского conversational assistant и одновременно максимально эффективно использует GPU. Теперь Azercell эксплуатирует этот фреймворк самостоятельно, а методология масштабируется на более крупные корпусы, архитектуры и новые сценарии использования. Чтобы узнать больше, изучите следующие ресурсы:

• Amazon SageMaker AI — managed machine learning service для обучения и развертывания
• AWS Generative AI Innovation Center — программы совместной разработки AI-решений
• Distributed training в Amazon SageMaker AI — принципы distributed training
• PyTorch FSDP tutorial — руководство по реализации
• Liger Kernels — memory-efficient реализации kernel
• Hugging Face Tokenizers — библиотека для обучения tokenizer
• Llama 3.2 1B — базовая модель

Если вы хотите реализовать похожее решение, обратитесь к своей AWS account team или посетите AWS Generative AI Innovation Center. Если вы обучаете LLM для языков с низкой ресурсной обеспеченностью или оптимизируете использование GPU на SageMaker AI, нам будет интересно узнать о вашем опыте. Делитесь мыслями и вопросами в комментариях.

Материал — перевод статьи с английского.

Оригинал: Training Azerbaijani language models on Amazon SageMaker AI