Настройка Llama 2 70B с использованием PyTorch FSDP

Настройка Llama 2 70B с PyTorch FSDP

Введение

В этой статье блога мы рассмотрим, как настроить Llama 2 70B с использованием PyTorch FSDP и связанных лучших практик. Мы будем использовать библиотеки Hugging Face Transformers, Accelerate и TRL. Мы также узнаем, как использовать Accelerate с SLURM.

Полностью общая параллельность данных (FSDP) – это парадигма, в которой состояния оптимизатора, градиенты и параметры разделены между устройствами. Во время прямого прохода каждая единица FSDP выполняет операцию all-gather, чтобы получить полные веса, выполняется вычисление, а затем отбрасываются фрагменты с других устройств. После прямого прохода вычисляется потеря, а затем выполняется обратный проход. В обратном проходе каждая единица FSDP выполняет операцию all-gather, чтобы получить полные веса, выполняется вычисление для получения локальных градиентов. Эти локальные градиенты усредняются и разделены между устройствами с помощью операции reduce-scatter, чтобы каждое устройство могло обновить параметры своего фрагмента. Для получения дополнительной информации о том, что такое PyTorch FSDP, обратитесь к этой статье блога: Ускорение обучения больших моделей с помощью полностью общей параллельности данных в PyTorch.

(Источник: ссылка)

Используемое оборудование

Количество узлов: 2. Минимально необходимое количество – 1. Количество GPU на узел: 8. Тип GPU: A100. Память GPU: 80 ГБ. Внутриузловое соединение: NVLink. Оперативная память на узел: 1 ТБ. Количество ядер ЦП на узел: 96. Межузловое соединение: Elastic Fabric Adapter.

Проблемы с настройкой LLaMa 70B

Мы столкнулись с тремя основными проблемами при попытке настроить LLaMa 70B с использованием FSDP:

  1. FSDP оборачивает модель после загрузки предобученной модели. Если каждый процесс/ранг внутри узла загружает модель Llama-70B, то это потребует 70*4*8 ГБ ~ 2 ТБ оперативной памяти ЦП, где 4 – количество байтов на параметр, а 8 – количество GPU на каждом узле. Это приведет к нехватке оперативной памяти ЦП и прекращению работы процессов.

  2. Сохранение полных промежуточных контрольных точек с использованием FULL_STATE_DICT с выгрузкой на ЦП на ранге 0 занимает много времени и часто приводит к ошибкам времени ожидания NCCL из-за неопределенного зависания во время трансляции. Однако в конце обучения мы хотим получить полный словарь состояния модели, а не словарь состояния, разделенный на фрагменты, который совместим только с FSDP.

  3. Нам необходимо улучшить скорость и снизить использование VRAM для более быстрого обучения и снижения затрат на вычисления.

Давайте рассмотрим, как решить вышеуказанные проблемы и настроить модель 70B!

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

  1. Кодовая база: https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/training с патчем мартышки flash-attn V2

  2. Конфигурация FSDP: https://github.com/pacman100/DHS-LLM-Workshop/blob/main/chat_assistant/training/configs/fsdp_config.yaml

  3. Скрипт SLURM launch.slurm: https://gist.github.com/pacman100/1cb1f17b2f1b3139a63b764263e70b25

  4. Модель: meta-llama/Llama-2-70b-chat-hf

  5. Набор данных: smangrul/code-chat-assistant-v1 (смесь LIMA+GUANACO с правильным форматированием в готовом к обучению формате)

Предварительные требования

Сначала выполните следующие шаги для установки Flash Attention V2: Dao-AILab/flash-attention: Fast and memory-efficient exact attention (github.com). Установите последние версии PyTorch с CUDA ≥11.8. Установите остальные требования в соответствии с файлом requirements.txt в DHS-LLM-Workshop/code_assistant/training. Здесь мы устанавливаем 🤗 Accelerate и 🤗 Transformers из основной ветки.

Настройка

Решение вызова 1

PRs huggingface/transformers#25107 и huggingface/accelerate#1777 решают первую проблему и не требуют изменений кода со стороны пользователя. Они выполняют следующие действия:

  1. Создание модели без весов на всех рангах (с использованием устройства meta).
  2. Загрузка состояния словаря только на ранге==0 и установка весов модели с использованием этого состояния словаря на ранге 0.
  3. Для всех остальных рангов выполните torch.empty(*param.size(), dtype=dtype) для каждого параметра на устройстве meta.
  4. Таким образом, ранг==0 загрузит модель с правильным состоянием словаря, в то время как все остальные ранги будут иметь случайные веса.
  5. Установите sync_module_states=True, чтобы объект FSDP заботился о трансляции их на все ранги перед началом обучения.

Ниже приведен фрагмент вывода для модели 7B на 2-х GPU, измеряющий объем памяти и параметры модели на различных этапах. Мы можем заметить, что при загрузке предобученной модели ранг 0 и ранг 1 имеют пиковое объем памяти CPU 32744 MB и 1506 MB соответственно. Следовательно, только ранг 0 загружает предобученную модель, что приводит к эффективному использованию оперативной памяти CPU. Полные журналы можно найти здесь

accelerator.process_index=0 Оперативная память GPU до загрузки: 0
accelerator.process_index=0 Оперативная память GPU, потребляемая к концу загрузки (конец-начало): 0
accelerator.process_index=0 Максимальное потребление оперативной памяти GPU во время загрузки (максимум-начало): 0
accelerator.process_index=0 Общее максимальное потребление оперативной памяти GPU во время загрузки (максимум): 0
accelerator.process_index=0 Оперативная память CPU до загрузки: 926
accelerator.process_index=0 Оперативная память CPU, потребляемая к концу загрузки (конец-начало): 26415
accelerator.process_index=0 Максимальное потребление оперативной памяти CPU во время загрузки (максимум-начало): 31818
accelerator.process_index=0 Общее максимальное потребление оперативной памяти CPU во время загрузки (максимум): 32744

accelerator.process_index=1 Оперативная память GPU до загрузки: 0
accelerator.process_index=1 Оперативная память GPU, потребляемая к концу загрузки (конец-начало): 0
accelerator.process_index=1 Максимальное потребление оперативной памяти GPU во время загрузки (максимум-начало): 0
accelerator.process_index=1 Общее максимальное потребление оперативной памяти GPU во время загрузки (максимум): 0
accelerator.process_index=1 Оперативная память CPU до загрузки: 933
accelerator.process_index=1 Оперативная память CPU, потребляемая к концу загрузки (конец-начало): 10
accelerator.process_index=1 Максимальное потребление оперативной памяти CPU во время загрузки (максимум-начало): 573
accelerator.process_index=1 Общее максимальное потребление оперативной памяти CPU во время загрузки (максимум): 1506

Решение вызова 2

Он решается выбором типа состояния словаря SHARDED_STATE_DICT при создании конфигурации FSDP. SHARDED_STATE_DICT сохраняет отдельные куски для каждого GPU, что делает быстрым сохранение или возобновление обучения из промежуточной контрольной точки. Когда используется FULL_STATE_DICT, первый процесс (ранг 0) собирает всю модель на CPU, а затем сохраняет ее в стандартном формате.

Создадим конфигурацию ускорения с помощью следующей команды:

accelerate config --config_file "fsdp_config.yaml"

Результирующая конфигурация доступна здесь: fsdp_config.yaml. Здесь стратегия шардинга – FULL_SHARD. Мы используем TRANSFORMER_BASED_WRAP для автоматической обертки политики и он использует _no_split_module для поиска имени блока Transformer для вложенной автоматической обертки FSDP. Мы используем SHARDED_STATE_DICT для сохранения промежуточных контрольных точек и состояний оптимизатора в этом формате, рекомендованном командой PyTorch. Обязательно включите передачу параметров модуля с ранга 0 в начале, как указано в предыдущем абзаце, решающем вызов 1. Мы включаем тренировку смешанной точности bf16.

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

если trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(script_args.output_dir) # альтернативно, trainer.push_to_hub(), если весь ckpt ниже 50 ГБ, поскольку лимит LFS на файл составляет 50 ГБ 

Решение проблемы 3

Для более быстрой тренировки и сокращения использования VRAM требуется включение Flash Attention и градиентного чекпоинтинга, чтобы обеспечить возможность проведения доводки и экономить вычислительные ресурсы. В данный момент кодовая база использует манки-патчинг, и реализация находится в chat_assistant/training/llama_flash_attn_monkey_patch.py.

FlashAttention: Быстрое и энергосберегающее точное внимание с учетом ввода-вывода предлагает способ вычисления точного внимания, который одновременно быстр и энергосберегающ, используя знания о памятиерархии аппаратного обеспечения/графических процессоров – Чем выше пропускная способность/скорость памяти, тем меньше ее емкость, так как она становится более дорогой.

Если мы следуем посту блога “Сделаем Deep Learning быстрее на основе первых принципов”, мы можем понять, что модуль Attention на текущем оборудовании является ограниченным памятью/пропускной способностью. Причина в том, что внимание в основном состоит из поэлементных операций, как показано на рисунке слева. Мы можем заметить, что операции маскировки, softmax и отсева занимают большую часть времени, вместо матричных умножений, которые составляют основную часть FLOPs.

(Источник: ссылка)

Именно такую проблему решает Flash Attention. Идея заключается в удалении избыточных операций чтения/записи HBM. Это достигается путем хранения всего в SRAM, выполнения всех промежуточных шагов и только после этого записи окончательного результата обратно в HBM, также известного как слияние ядра. Ниже показано, как это преодолевает ограничение памяти.

(Источник: ссылка)

Во время прямого и обратного проходов используется тайлинг, чтобы разбить вычисление NxN softmax/оценок на блоки и преодолеть ограничение размера памяти SRAM. Для активации тайлинга используется алгоритм онлайн-softmax. Во время обратного прохода используется повторное вычисление, чтобы избежать сохранения всей матрицы NxN softmax/оценок во время прямого прохода. Это существенно снижает потребление памяти.

Для упрощенного и глубокого понимания Flash Attention, пожалуйста, ознакомьтесь с блог-постами ELI5: FlashAttention и Сделаем Deep Learning быстрее на основе первых принципов, а также с оригинальной статьей FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.

Сводим все воедино

Чтобы запустить тренировку с использованием запуска Accelerate с помощью SLURM, см. этот файл gist launch.slurm. Ниже приведена эквивалентная команда, демонстрирующая использование запуска Accelerate для запуска тренировки. Обратите внимание, что мы переопределяем значения main_process_ip, main_process_port, machine_rank, num_processes и num_machines в файле fsdp_config.yaml. Здесь также важно отметить, что хранилище хранится между всеми узлами.

accelerate launch \
    --config_file configs/fsdp_config.yaml \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --machine_rank \$MACHINE_RANK \
    --num_processes 16 \
    --num_machines 2 \
    train.py \
    --model_name "meta-llama/Llama-2-70b-chat-hf" \
    --dataset_name "smangrul/code-chat-assistant-v1" \
    --max_seq_len 2048 \
    --max_steps 500 \
    --logging_steps 25 \
    --eval_steps 100 \
    --save_steps 250 \
    --bf16 True \
    --packing True \
    --output_dir "/shared_storage/sourab/experiments/full-finetune-llama-chat-asst" \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --dataset_text_field "content" \
    --use_gradient_checkpointing True \
    --learning_rate 5e-5  \
    --lr_scheduler_type "cosine" \
    --weight_decay 0.01 \
    --warmup_ratio 0.03 \
    --use_flash_attn True

Настройка завершена примерно за 13,5 часа, и ниже представлена диаграмма потерь обучения. Давайте рассчитаем использование модельных операций FLOPS (Model Flops Utilization, MFU) для обучения.

  1. Графические процессоры A100 выполняют примерно 3,12 * 10^14 операций FLOPS в секунду (в формате float32 или bfloat16)
  2. Количество токенов, на которых было проведено обучение в вышеуказанных экспериментах, равно произведению длины последовательности, размера пакета и числа шагов обучения: (2048 * 16 *
    1. = 16,384,000 = 1,64 * 10^7
  3. Примерное количество вычислений для вышеуказанных экспериментов равно 6 * P (число параметров) * D (число токенов) = 6 * 7 * 10^10 * 1,64 * 10^7 = 6,89 * 10^18 операций FLOPS
  4. Скорость обучения в операциях FLOPS в секунду = Примерное количество вычислений / время обучения = 6,89 * 10^17 / (13,5 * 3600) = 6,89 * 10^18 / 4,86 * 10^4 = 1,42 * 10^14 операций FLOPS в секунду
  5. MFU = Скорость обучения в операциях FLOPS в секунду / Пиковая производительность A100 = 1,42 * 10^14 / 3,12 * 10^14 = 0,4551 = 45,51% от пика производительности

Ниже приведен пример разговора, используя указанную модель:

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

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

- Человек: Что такое глубокое обучение? Объясни, будто я пират.

+ Помощник: Аррр, грумет! Тебе нужно объяснение о глубоком обучении? 
+ Хорошо, слушай внимательно, мой верный! Глубокое обучение - это тип машинного обучения, который использует искусственные нейронные сети для анализа данных. 
+ Ай, ты слышал правильно, матросы! Искусственные нейронные сети - это модели человеческого мозга, но вместо использования нейронов они используют взаимосвязанные узлы или "искусственные нейроны" для обработки информации. 
+ Эти узлы работают вместе, чтобы формировать слои обработки, которые могут учиться на больших объемах данных. 
+ Чем больше данных получают эти системы, тем лучше они становятся в распознавании образов и принятии решений самостоятельно. 
+ Так что поднимай паруса и ставь курс на приключения с глубоким обучением, мои друзья! Здесь тебя ждет кладезь знаний, ожидающий своего открытия. 
+ Счастливого плавания!

- Человек: Теперь объясни это, будто я шеф-повар.

+ Помощник: Конечно! Вот объяснение глубокого обучения с точки зрения кулинарии:

+ Представь рецепт как последовательность шагов, используемых для превращения сырых ингредиентов в вкусное блюдо. 
+ Аналогично в глубоком обучении есть несколько слоев "ингредиентов" (или функций), которые объединяются и преобразуются с помощью различных операций для получения конечного результата или прогноза. 
+ Каждый слой основан на предыдущем, улучшая и углубляя представление входных данных до достижения желаемого результата. 
+ Точно так же, как разные специи усиливают вкус пищи, каждый слой добавляет глубину и сложность понимания модели входных данных, позволяя ей делать более точные прогнозы. 
+ И, подобно тому, как шеф-повары должны тщательно сбалансировать количество соли, сахара и других приправ в рецепте, процесс обучения в глубоком обучении включает настройку весов и смещений сети для достижения идеального сочетания точности и эффективности. 
+ С достаточной практикой и экспериментами каждый может стать мастером кулинарии... Имею в виду, мастером глубокого обучения! Приятного аппетита!

Вся беседа форматируется следующим образом:

<|system|> Системное сообщение <|endoftext|> <|prompter|> Вопрос 1 <|endoftext|> <|assistant|> Ответ 1 <|endoftext|> ...

Вывод

Мы успешно провели настройку модели Llama 70B с использованием PyTorch FSDP в многонодовой многографической среде, учитывая различные проблемы. Мы увидели, как 🤗 Transformers и 🤗 Accelerates теперь поддерживают эффективный способ инициализации больших моделей при использовании FSDP для преодоления нехватки оперативной памяти процессора. Затем мы рассмотрели рекомендуемые практики сохранения/загрузки промежуточных контрольных точек и способы сохранения конечной модели для ее последующего использования. Чтобы обеспечить более быстрое обучение и снижение использования памяти GPU, мы описали важность Flash Attention и Gradient Checkpointing. В целом, мы видим, как простая конфигурация с использованием 🤗 Accelerate позволяет настраивать такие большие модели в многонодовых многографических средах.