Paralelismo de Dados Totalmente Fragmentado

Fully sharded data parallel (FSDP) foi desenvolvido para o treinamento distribuído de grandes modelos pré-treinados com até 1T de parâmetros. O FSDP atinge isso ao dividir (sharding) os parâmetros do modelo, os gradientes e os estados do otimizador entre processos de paralelismo de dados, e também pode descarregar (offload) os parâmetros do modelo divididos para a CPU. A eficiência de memória proporcionada pelo FSDP permite escalar o treinamento para tamanhos maiores de lote ou de modelo.

Ambas essas funcionalidades são suportadas no 🤗 Accelerate e você pode usá-las com o 🤗 PEFT.

Use PEFT e FSDP

Esta seção do guia ajudará você a aprender como usar nosso script de treinamento do DeepSpeed para realizar SFT. Você configurará o script para executar SFT (ajuste fino supervisionado) do modelo Llama-70B com LoRA e FSDP em 8 GPUs H100 de 80GB em uma única máquina. Você pode configurá-lo para escalar para múltiplas máquinas alterando a configuração do accelerate.

Configuração

Comece executando o seguinte comando para criar um arquivo de configuração do FSDP com o 🤗 Accelerate. A flag --config_file permite salvar o arquivo de configuração em um local específico; caso contrário, ele é salvo como um arquivo default_config.yaml no cache do 🤗 Accelerate.

O arquivo de configuração é usado para definir as opções padrão ao iniciar o script de treinamento.

accelerate config --config_file fsdp_config.yaml

Você será solicitado a responder algumas perguntas sobre sua configuração e configurar os seguintes argumentos. Neste exemplo, você responderá ao questionário conforme mostrado na imagem abaixo.

Creating Accelerate's config to use FSDP

Uma vez concluído, a configuração correspondente deve ter o formato abaixo e pode ser encontrada na pasta de configuração em 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

Comando de inicialização

O comando de inicialização está disponível em run_peft_fsdp.sh e também é mostrado abaixo:

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

Observe que estamos usando LoRA com rank=8, alpha=16 e direcionando para todas as camadas lineares. Estamos passando o arquivo de configuração do FSDP e fazendo o ajuste fino do modelo Llama de 70B em um subconjunto do conjunto de dados ultrachat.

As partes importantes

Vamos mergulhar um pouco mais no script para que você possa ver o que está acontecendo e entender como ele funciona.

A primeira coisa a saber é que o script usa FSDP para treinamento distribuído, pois a configuração do FSDP foi passada. A classe SFTTrainer cuida de todo o trabalho pesado de criar o modelo PEFT usando a configuração peft que é passada. Depois disso, quando você chama trainer.train(), o Trainer usa internamente o 🤗 Accelerate para preparar o modelo, o otimizador e o trainer usando a configuração do FSDP para criar um modelo envolto em FSDP, que então é treinado. O trecho principal de código está abaixo:

# 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()

Aqui, um ponto principal a observar atualmente ao usar FSDP com PEFT é que use_orig_params precisa ser False para realizar economias de memória de GPU. Devido a use_orig_params=False, a política de auto wrap para FSDP precisa mudar para que os parâmetros treináveis e não treináveis sejam embrulhados separadamente. Isso é feito pelo trecho de código abaixo, que usa a função utilitária fsdp_auto_wrap_policy do 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)

Uso de memória

No exemplo acima, a memória consumida por GPU é de 72-80 GB (90-98%), conforme visto no screenshot abaixo. O leve aumento na memória da GPU no final ocorre ao salvar o modelo usando o tipo de state dict FULL_STATE_DICT em vez do SHARDED_STATE_DICT, para que o modelo tenha pesos de adaptador que podem ser carregados normalmente com o método from_pretrained durante a inferência:

GPU memory usage for the training run

Use PEFT QLoRA e FSDP para ajuste fino de modelos grandes em várias GPUs

Nesta seção, veremos como usar QLoRA e FSDP para fazer o ajuste fino do modelo llama de 70B em GPUs 2X24GB. O Answer.AI, em colaboração com bitsandbytes e Hugging Face 🤗, disponibilizou em código aberto o código que permite o uso de FSDP+QLoRA e explicou todo o processo em seu blogpost informativo You can now train a 70b language model at home. Isso agora está integrado no ecossistema do Hugging Face.

Para isso, primeiro precisamos de bitsandbytes>=0.43.3, accelerate>=1.0.1, transformers>4.44.2, trl>0.11.4 e peft>0.13.0. Precisamos definir fsdp_cpu_ram_efficient_loading=true, fsdp_use_orig_params=false e fsdp_offload_params=true (offloading para CPU) ao usar a configuração do Accelerate. Quando não estiver usando o launcher do accelerate, você pode alternativamente definir a variável de ambiente export FSDP_CPU_RAM_EFFICIENT_LOADING=true. Aqui, estaremos usando a configuração do accelerate e a configuração abaixo pode ser encontrada em 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

O comando de lançamento é apresentado abaixo e está disponível em 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"

Observe o novo argumento sendo passado, bnb_4bit_quant_storage_dtype, que denota o tipo de dados para empacotar os parâmetros de 4 bits. Por exemplo, quando definido como bfloat16, 16/4 = 4 parâmetros de 4 bits são empacotados juntos após a quantização. Ao usar treinamento com precisão mista com bfloat16, bnb_4bit_quant_storage_dtype pode ser bfloat16 para finetuning puro em bfloat16, ou float32 para precisão mista automática (isso consome mais memória de GPU). Ao usar treinamento com precisão mista com float16, bnb_4bit_quant_storage_dtype deve ser definido como float32 para um treinamento estável com precisão mista automática.

Em termos de código de treinamento, as alterações importantes no código são:

...

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,
)

Observe que o dtype de AutoModelForCausalLM é o mesmo que o tipo de dados bnb_4bit_quant_storage. É isso. Tudo o mais é tratado pelo Trainer e pelo TRL.

Uso de memória

No exemplo acima, a memória consumida por GPU é de 19,6 GB, enquanto o uso de RAM da CPU fica em torno de 107 GB. Ao desativar o offloading para a CPU, o uso de memória de GPU é de 35,6 GB/ GPU. Portanto, o que exigia 16 GPUs de 80GB para finetuning completo, 8 GPUs de 80GB com FSDP+LoRA e algumas GPUs de 80GB com DDP+QLoRA, agora requer apenas 2 GPUs de 24GB. Isso torna o finetuning de modelos grandes mais acessível.

Mais recursos

Você também pode consultar o repositório llama-recipes e o guia Getting started with Llama sobre como fazer fine-tuning usando FSDP e PEFT.

Observações

A fusão ao usar PEFT e FSDP não é suportada no momento e resultará em erro.
Passar o parâmetro de configuração modules_to_save não é testado no momento.
A economia de memória de GPU ao usar CPU Offloading ainda não foi testada no momento.
Ao usar FSDP+QLoRA, paged_adamw_8bit atualmente resulta em um erro ao salvar um checkpoint.
O treinamento com DoRA usando FSDP deve funcionar (embora com velocidade menor do que a do LoRA). Se combinado com bitsandbytes (QDoRA), a quantização de 4 bits também deve funcionar, mas a quantização de 8 bits tem problemas conhecidos e não é recomendada.
Com o FSDP2, set_adapter() e set_requires_grad() re-shardam os módulos FSDP2 dentro do PeftModel, incluindo o próprio PeftModel, para que as alterações em requires_grad alcancem os parâmetros particionados. A saída de disable_adapter() só faz isso se o FSDP2 substituiu os objetos de parâmetro registrados dentro dele, como um forward faz com parâmetros particionados, ou se a reativação altera requires_grad em relação ao valor que ele tinha na entrada. Módulos FSDP2 que particionam parâmetros de uma camada de adaptador mesclado não são re-shardados, portanto, desmescle antes de chamar esses métodos. Com reshard_after_forward=False, ou para o módulo raiz com o valor padrão None, o próximo forward faz o all-gather novamente. Módulos FSDP2 que envolvem o PeftModel não são re-shardados; chame reshard() neles antes desses métodos. Como essas chamadas podem deixar os pesos base particionados, mesclar imediatamente depois delas não funciona corretamente (veja 1). Os métodos de aprendizado de prompt e adaption prompt ainda não são cobertos.

Update on GitHub