Настройте Falcon 7B и другие LLM на Amazon SageMaker с помощью декоратора @remote

Настройте Falcon 7B и другие LLM с помощью @remote на Amazon SageMaker.

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

Теперь с помощью Amazon SageMaker вы можете запускать задачу обучения SageMaker, просто аннотируя свой код Python с помощью декоратора @remote. С помощью SageMaker Python SDK ваша существующая рабочая среда, а также любой связанный код обработки данных и наборы данных автоматически преобразуются в задачу обучения SageMaker, которая выполняется на платформе обучения. Это имеет преимущество написания кода более естественным, объектно-ориентированным способом и все же использования возможностей SageMaker для запуска задач обучения на удаленном кластере с минимальными изменениями.

В этом посте мы покажем, как настроить модели Foundation Models (FM) Falcon-7B с помощью декоратора @remote из SageMaker Python SDK. Он также использует библиотеку параметрической эффективной настройки Hugging Face (PEFT) и техники квантования через bitsandbytes для поддержки настройки. Представленный в этом блоге код также можно использовать для настройки других моделей FM, таких как Llama-2 13b.

Полные представления этой модели с полной точностью могут иметь проблемы с помещением в память одного или даже нескольких графических процессоров (GPUs) — или даже может потребоваться более крупный экземпляр. Поэтому для настройки этой модели без увеличения затрат мы используем известную технику, называемую квантованными LLM с низкоранговыми адаптерами (QLoRA). QLoRA — это эффективный подход к настройке, который снижает использование памяти LLM, сохраняя при этом очень хорошую производительность.

Преимущества использования декоратора @remote

Прежде чем двигаться дальше, давайте разберемся, как декоратор @remote повышает производительность разработчика при работе с SageMaker:

  • Декоратор @remote запускает задачу обучения непосредственно с использованием нативного кода Python, без явного вызова оценщиков SageMaker и каналов ввода SageMaker
  • Низкий порог вхождения для разработчиков, обучающих модели на SageMaker.
  • Не нужно переключаться между интегрированными средами разработки (IDE). Продолжайте писать код в выбранной вами IDE и вызывайте задачи обучения SageMaker.
  • Не нужно изучать контейнеры. Продолжайте предоставлять зависимости в файле requirements.txt и передавайте его в декоратор @remote.

Требования

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

В этом посте мы используем Amazon SageMaker Studio с образом Data Science 3.0 и быстрым запуском экземпляра ml.t3.medium. Однако вы можете использовать любую интегрированную среду разработки (IDE) по вашему выбору. Просто убедитесь, что у вас правильно настроены учетные данные AWS Command Line Interface (AWS CLI). Дополнительные сведения см. в разделе Настройка AWS CLI.

Для настройки модели Falcon-7B в этом посте используется экземпляр ml.g5.12xlarge. Пожалуйста, убедитесь в наличии достаточной емкости для этого экземпляра в учетной записи AWS.

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

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

  1. Установка предварительных требований для настройки модели Falcon-7B
  2. Настройка конфигурации декоратора @remote
  3. Предварительная обработка набора данных, содержащего вопросы и ответы о службах AWS
  4. Настройка Falcon-7B на вопросы и ответы о службах AWS
  5. Тестирование моделей с настройкой на примерах вопросов, связанных со службами AWS

1. Установка предварительных требований для настройки модели Falcon-7B

Запустите блокнот falcon-7b-qlora-remote-decorator_qa.ipynb в SageMaker Studio, выбрав образ Data Science и ядро Python 3. Установите все необходимые библиотеки, указанные в файле requirements.txt. Некоторые из библиотек должны быть установлены непосредственно на экземпляр блокнота. Выполните другие операции, необходимые для обработки набора данных и запуска задачи обучения SageMaker.

%pip install -r requirements.txt

%pip install -q -U transformers==4.31.0
%pip install -q -U datasets==2.13.1
%pip install -q -U peft==0.4.0
%pip install -q -U accelerate==0.21.0
%pip install -q -U bitsandbytes==0.40.2
%pip install -q -U boto3
%pip install -q -U sagemaker==2.154.0
%pip install -q -U scikit-learn

2. Настройка конфигураций удаленного декоратора

Создайте файл конфигурации, в котором указаны все конфигурации, связанные с заданием обучения Amazon SageMaker. Этот файл считывается декоратором @remote при выполнении задания обучения. В файле содержатся настройки, такие как зависимости, образ обучения, экземпляр и роль выполнения, используемые для задания обучения. Для подробной информации о всех поддерживаемых настройках файла конфигурации ознакомьтесь с разделом “Настройка и использование значений по умолчанию с помощью SageMaker Python SDK”.

SchemaVersion: '1.0'
SageMaker:
  PythonSDK:
    Modules:
      RemoteFunction:
        Dependencies: ./requirements.txt
        ImageUri: '{aws_account_id}.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04'
        InstanceType: ml.g5.12xlarge
        RoleArn: arn:aws:iam::111122223333:role/ExampleSageMakerRole

Использование файла config.yaml необязательно для работы с декоратором @remote. Это просто более удобный способ предоставления всех конфигураций декоратору @remote. Это позволяет хранить параметры SageMaker и AWS вне кода с минимальными усилиями по настройке конфигурационного файла, который используется всеми участниками команды. Все конфигурации также могут быть указаны непосредственно в аргументах декоратора, но это ухудшает читаемость и поддерживаемость изменений в долгосрочной перспективе. Кроме того, файл конфигурации может быть создан администратором и распространен среди всех пользователей в рабочей среде.

Предварительная обработка набора данных, содержащего вопросы и ответы по AWS-сервисам

Следующий шаг – загрузить и предварительно обработать набор данных, чтобы он был готов для задания обучения. Сначала давайте посмотрим на набор данных:

Показаны часто задаваемые вопросы по одной из услуг AWS. Помимо QLoRA, используется bitsanbytes для преобразования в 4-битную точность для квантования замороженного LLM до 4 бит и прикрепления адаптеров LoRA к нему.

Создайте шаблон запроса для преобразования каждого образца вопроса и ответа в формат запроса:

from random import randint

# начало пользовательской инструкции запроса
prompt_template = f"{{question}}\n---\nОтвет:\n{{answer}}{{eos_token}}"

# шаблон набора данных для добавления запроса к каждому образцу
def template_dataset(sample):
    sample["text"] = prompt_template.format(question=sample["question"],
                                            answer=sample["answers"],
                                            eos_token=tokenizer.eos_token)
    return sample

Следующий шаг – преобразовать входы (текст) в идентификаторы токенов. Это делается с помощью токенизатора Hugging Face Transformers.

from transformers import AutoTokenizer

model_id = "tiiuae/falcon-7b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
# Установка токенизатора Falcon
tokenizer.pad_token = tokenizer.eos_token

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

4. Точная настройка Falcon-7B по вопросам и ответам по AWS-сервисам

Теперь вы можете подготовить скрипт обучения, определить функцию обучения train_fn и добавить декоратор @remote к функции.

Функция обучения выполняет следующие действия:

  • токенизирует и разбивает набор данных на части
  • настраивает BitsAndBytesConfig, который указывает, что модель должна быть загружена в 4-битном формате, но вычисления должны быть преобразованы в bfloat16.
  • Загружает модель
  • Находит целевые модули и обновляет необходимые матрицы с помощью вспомогательного метода find_all_linear_names
  • Создает конфигурации LoRA, которые определяют ранжирование матриц обновления (s), коэффициент масштабирования (lora_alpha), модули, к которым применяются матрицы обновления LoRA (target_modules), вероятность отсева для слоев Lora (lora_dropout), task_type и т. д.
  • Запускает обучение и оценку
import bitsandbytes as bnb

def find_all_linear_names(hf_model):
    lora_module_names = set()
    for name, module in hf_model.named_modules():
        if isinstance(module, bnb.nn.Linear4bit):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if "lm_head" in lora_module_names:
        lora_module_names.remove("lm_head")
    return list(lora_module_names)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from sagemaker.remote_function import remote
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import transformers

# Начать обучение
@remote(volume_size=50)
def train_fn(
        model_name,
        train_ds,
        test_ds,
        lora_r=8,
        lora_alpha=32,
        lora_dropout=0.05,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        learning_rate=2e-4,
        num_train_epochs=1
):
    # токенизировать и разделить набор данных
    lm_train_dataset = train_ds.map(
        lambda sample: tokenizer(sample["text"]), batched=True, batch_size=24, remove_columns=list(train_dataset.features)
    )


    lm_test_dataset = test_ds.map(
        lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(test_dataset.features)
    )

    # Вывести общее количество обучающих образцов
    print(f"Общее количество обучающих образцов: {len(lm_train_dataset)}")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    # Falcon требует разрешить удаленное выполнение кода. Это связано с тем, что модель использует новую архитектуру, которая еще не является частью библиотеки transformers.
    # Код предоставлен авторами модели в репозитории.
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        quantization_config=bnb_config,
        device_map="auto")

    model.gradient_checkpointing_enable()
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

    # получить целевые модули lora
    modules = find_all_linear_names(model)
    print(f"Найдено {len(modules)} модулей для квантования: {modules}")

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM"
    )

    model = get_peft_model(model, config)
    print_trainable_parameters(model)

    trainer = transformers.Trainer(
        model=model,
        train_dataset=lm_train_dataset,
        eval_dataset=lm_test_dataset,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=per_device_train_batch_size,
            per_device_eval_batch_size=per_device_eval_batch_size,
            logging_steps=2,
            num_train_epochs=num_train_epochs,
            learning_rate=learning_rate,
            bf16=True,
            save_strategy="no",
            output_dir="outputs"
        ),
        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )
    model.config.use_cache = False

    trainer.train()
    trainer.evaluate()

    model.save_pretrained("/opt/ml/model")

И вызовите train_fn()

train_fn(model_id, train_dataset, test_dataset)

Настройка будет выполняться на кластере обучения Amazon SageMaker. Дождитесь завершения настройки.

Теперь пришло время запустить несколько тестов на модели. Сначала давайте загрузим модель:

from peft import PeftModel, PeftConfig
import torch
from transformers import AutoModelForCausalLM

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'

config = PeftConfig.from_pretrained("./model")
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
model = PeftModel.from_pretrained(model, "./model")
model.to(device)

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

Вот пример вопроса из обучающего набора и оригинальный ответ:

Теперь этот же вопрос задается модели Falcon-7B, которая прошла настройку:

Это завершает реализацию мелкой настройки модели Falcon-7B на наборе данных вопросов и ответов по службам AWS FAQ с использованием декоратора @remote из Amazon SageMaker Python SDK.

Очистка

Выполните следующие шаги для очистки ресурсов:

  • Выключите экземпляры Amazon SageMaker Studio, чтобы избежать дополнительных затрат.

  • Очистите каталог Amazon Elastic File System (Amazon EFS), удалив каталог кэша Hugging Face:

    rm -R ~/.cache/huggingface/hub

Вывод

В этом посте мы показали, как эффективно использовать возможности декоратора @remote для мелкой настройки модели Falcon-7B с использованием QLoRA, Hugging Face PEFT с bitsandbtyes без значительных изменений в блокноте для обучения и использовать возможности Amazon SageMaker для запуска заданий по обучению на удаленном кластере.

Весь показанный в этом посте код для мелкой настройки модели Falcon-7B доступен в репозитории GitHub. Репозиторий также содержит блокнот, демонстрирующий, как провести мелкую настройку модели Llama-13B.

Как следующий шаг, мы призываем вас ознакомиться с функциональностью декоратора @remote и API Python SDK и использовать их в выбранной вами среде и IDE. Дополнительные примеры доступны в репозитории amazon-sagemaker-examples, чтобы вы могли быстро начать. Вы также можете ознакомиться с следующими постами:

  • Запуск вашего локального кода машинного обучения в качестве заданий обучения Amazon SageMaker с минимальными изменениями кода
  • Доступ к закрытым репозиториям с использованием декоратора @remote для рабочих нагрузок обучения Amazon SageMaker
  • Интерактивная мелкая настройка модели Falcon-40B и других LLM на блокнотах Amazon SageMaker Studio с использованием QLoRA