Полностью шардированная параллельная обработка данных

Fully sharded data parallel (FSDP) разработан для распределённого обучения крупных предобученных моделей объёмом до 1T параметров. FSDP достигает этого, разбивая параметры модели, градиенты и состояния оптимизатора между процессами, работающими в режиме параллелизма данных, а также позволяет выгружать разбитые параметры модели в оперативную память CPU. Эффективность использования памяти, обеспечиваемая FSDP, позволяет масштабировать обучение до более крупных размеров батчей или моделей.

Обе эти функции поддерживаются в 🤗 Accelerate, и вы можете использовать их вместе с 🤗 PEFT.

Использование PEFT и FSDP

Этот раздел руководства поможет вам освоить использование нашего скрипта обучения DeepSpeed для выполнения SFT. Вы настроите скрипт для SFT (супервизированной дообучения) модели Llama-70B с использованием LoRA и FSDP на 8 GPU H100 80GB на одном компьютере. Вы можете настроить масштабирование на несколько машин, изменив конфигурацию accelerate.

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

Начните с выполнения следующей команды, чтобы создать файл конфигурации FSDP с помощью 🤗 Accelerate. Флаг --config_file позволяет сохранить файл конфигурации в указанное место, в противном случае он сохраняется как файл default_config.yaml в кэше 🤗 Accelerate.

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

accelerate config --config_file fsdp_config.yaml

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

После выполнения этих шагов соответствующий конфигурационный файл должен выглядеть следующим образом. Его можно найти в папке config по пути fsdp_config.yaml:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Команда запуска

Команда запуска доступна в run_peft_fsdp.sh и также приведена ниже:

accelerate launch --config_file "configs/fsdp_config.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-lora-fsdp" \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 8 \
--gradient_accumulation_steps 4 \
--gradient_checkpointing True \
--use_reentrant False \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization False

Обратите внимание, что мы используем LoRA с rank=8, alpha=16 и нацеливаем его на все линейные слои. Мы передаём конфигурационный файл FSDP и дообучаем модель Llama 70B на подмножестве набора данных ultrachat.

Ключевые части

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

Первое, что нужно знать: скрипт использует FSDP для распределённого обучения, так как была передана конфигурация FSDP. Класс SFTTrainer выполняет всю основную работу по созданию PEFT-модели с использованием переданной конфигурации peft. После этого, когда вы вызываете trainer.train(), Trainer внутренне использует 🤗 Accelerate, чтобы подготовить модель, оптимизатор и тренаер с помощью конфигурации FSDP, создавая модель, обёрнутую в FSDP, которая затем обучается. Основной фрагмент кода приведён ниже:

# trainer
trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
)
trainer.accelerator.print(f"{trainer.model}")
if model_args.use_peft_lora:
    # handle PEFT+FSDP case
    trainer.model.print_healthcheck()
    if getattr(trainer.accelerator.state, "fsdp_plugin", None):
        from peft.utils.other import fsdp_auto_wrap_policy

        fsdp_plugin = trainer.accelerator.state.fsdp_plugin
        fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(trainer.model)

# train
checkpoint = None
if training_args.resume_from_checkpoint is not None:
    checkpoint = training_args.resume_from_checkpoint
trainer.train(resume_from_checkpoint=checkpoint)

# saving final model
if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
trainer.save_model()

Здесь важно отметить, что при использовании FSDP с PEFT параметр use_orig_params должен быть установлен в False, чтобы обеспечить экономию памяти GPU. Из-за use_orig_params=False политику автоматического обёртывания для FSDP необходимо изменить так, чтобы обучаемые и необучаемые параметры обёртывались отдельно. Это реализовано в приведённом ниже фрагменте кода, который использует утилиту fsdp_auto_wrap_policy из PEFT:

if getattr(trainer.accelerator.state, "fsdp_plugin", None):
    from peft.utils.other import fsdp_auto_wrap_policy

    fsdp_plugin = trainer.accelerator.state.fsdp_plugin
    fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(trainer.model)

Использование памяти

В приведённом выше примере объём памяти, потребляемой на каждый GPU, составляет 72–80 ГБ (90–98%), как показано на скриншоте ниже. Незначительное увеличение объёма памяти GPU в конце происходит при сохранении модели с использованием типа state dict FULL_STATE_DICT вместо SHARDED_STATE_DICT, чтобы у модели были веса адаптеров, которые можно было бы загрузить обычным способом с помощью метода from_pretrained во время инференса:

Используйте PEFT QLoRA и FSDP для дообучения больших моделей на нескольких GPU

В этом разделе мы рассмотрим, как использовать QLoRA и FSDP для дообучения модели llama 70B на 2X24GB GPU. Answer.AI совместно с bitsandbytes и Hugging Face 🤗 открыли исходный код, позволяющий использовать FSDP+QLoRA, и подробно описали весь процесс в своей информативной статье You can now train a 70b language model at home. Теперь это интегрировано в экосистему Hugging Face.

Для этого нам в первую очередь нужны bitsandbytes>=0.43.3, accelerate>=1.0.1, transformers>4.44.2, trl>0.11.4 и peft>0.13.0. При использовании конфигурации Accelerate необходимо установить fsdp_cpu_ram_efficient_loading=true, fsdp_use_orig_params=false и fsdp_offload_params=true (оффлоад на CPU). Если вы не используете запускатель accelerate, можно вместо этого задать переменную окружения export FSDP_CPU_RAM_EFFICIENT_LOADING=true. Здесь мы будем использовать конфигурацию accelerate, и ниже приведена конфигурация, которую можно найти в fsdp_config_qlora.yaml:

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false                                                                                                                                                                 
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Команда запуска приведена ниже и доступна в файле run_peft_qlora_fsdp.sh:

accelerate launch --config_file "configs/fsdp_config_qlora.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-qlora-fsdp" \
--per_device_train_batch_size 2 \
--per_device_eval_batch_size 2 \
--gradient_accumulation_steps 2 \
--gradient_checkpointing True \
--use_reentrant True \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization True \
--use_nested_quant True \
--bnb_4bit_compute_dtype "bfloat16" \
--bnb_4bit_quant_storage_dtype "bfloat16"

Обратите внимание на новый передаваемый аргумент bnb_4bit_quant_storage_dtype, который определяет тип данных для упаковки 4-битных параметров. Например, при установке значения bfloat16 после квантования вместе упаковываются 16/4 = 4 4-битных параметра. При использовании обучения с смешанной точностью на bfloat16 значение bnb_4bit_quant_storage_dtype может быть установлено в bfloat16 для чистого дообучения в bfloat16 или в float32 для автоматической смешанной точности (это потребляет больше памяти GPU). При использовании обучения с смешанной точностью на float16 значение bnb_4bit_quant_storage_dtype следует установить в float32 для стабильного обучения с автоматической смешанной точностью.

Что касается кода обучения, важные изменения в коде:

...

bnb_config = BitsAndBytesConfig(
    load_in_4bit=args.use_4bit_quantization,
    bnb_4bit_quant_type=args.bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=args.use_nested_quant,
+   bnb_4bit_quant_storage=quant_storage_dtype,
)

...

model = AutoModelForCausalLM.from_pretrained(
    args.model_name_or_path,
    quantization_config=bnb_config,
    trust_remote_code=True,
    attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
+   dtype=quant_storage_dtype or torch.float32,
)

Обратите внимание, что dtype для AutoModelForCausalLM совпадает с типом данных bnb_4bit_quant_storage. Всё. Остальное обрабатывается Trainer и TRL.

Использование памяти

В приведённом выше примере объём памяти, потребляемой на каждый GPU, составляет 19,6 ГБ, а использование оперативной памяти CPU составляет около 107 ГБ. При отключении выгрузки на CPU использование памяти GPU составляет 35,6 ГБ/ GPU. Таким образом, то, что для полной дообучения требовало 16 GPU по 80 ГБ, 8 GPU по 80 ГБ с FSDP+LoRA и пару GPU по 80 ГБ с DDP+QLoRA, теперь требует всего 2 GPU по 24 ГБ. Это делает дообучение больших моделей более доступным.

Дополнительные ресурсы

Вы также можете обратиться к репозиторию llama-recipes и руководству Getting started with Llama, чтобы узнать, как выполнять дообучение с использованием FSDP и PEFT.

Ограничения

1. Слияние при использовании PEFT и FSDP в настоящее время не поддерживается и приводит к ошибке.
2. Передача параметра конфигурации modules_to_save в настоящее время не протестирована.
3. Экономия памяти GPU при использовании CPU Offloading на данный момент не протестирована.
4. При использовании FSDP+QLoRA paged_adamw_8bit в настоящее время приводит к ошибке при сохранении контрольной точки.
5. Обучение с DoRA в сочетании с FSDP должно работать (хотя и с меньшей скоростью, чем LoRA). При использовании вместе с bitsandbytes (QDoRA) 4-битное квантование также должно работать, однако 8-битное квантование имеет известные проблемы и не рекомендуется.
6. В FSDP2 методы set_adapter() и set_requires_grad() выполняют ресхардирование модулей FSDP2 внутри PeftModel, включая сам PeftModel, чтобы изменения requires_grad распространялись на шардированные параметры. Выход из disable_adapter() выполняет ресхардирование только в том случае, если FSDP2 заменил зарегистрированные объекты параметров внутри него, как это происходит с шардированными параметрами при forward, или если повторное включение изменяет requires_grad по сравнению со значением на входе. Модули FSDP2, шардирующие параметры слоя объединённого адаптера, не ресхардируются, поэтому перед вызовом этих методов необходимо выполнить unmerge. При reshard_after_forward=False или для корневого модуля со значением по умолчанию None следующий forward снова выполняет all-gather. Модули FSDP2, оборачивающие PeftModel, не ресхардируются; вызовите reshard() для них перед этими методами. Поскольку эти вызовы могут оставить базовые веса в шардированном состоянии, слияние сразу после них работает некорректно (см. 1). Методы prompt learning и adaption prompt пока не поддерживаются.