PEFT documentation
Fully Sharded Data Parallel
Fully Sharded Data Parallel
Fully sharded data parallel(FSDP)은 최대 1T 파라미터 규모의 대규모 사전 학습 모델의 분산 학습을 위해 개발되었습니다. FSDP는 모델 파라미터, 그래디언트, 옵티마이저 상태를 데이터 병렬 프로세스에 걸쳐 샤딩(sharding)하여 이를 달성하며, 샤딩된 모델 파라미터를 CPU로 오프로드(offload)할 수도 있습니다. FSDP가 제공하는 메모리 효율성 덕분에 더 큰 배치 또는 모델 크기로 학습을 확장할 수 있습니다.
이 두 가지 기능은 모두 🤗 Accelerate에서 지원되며, 🤗 PEFT와 함께 사용할 수 있습니다.
PEFT와 FSDP 사용
이 가이드 섹션에서는 DeepSpeed 학습 스크립트를 사용하여 SFT를 수행하는 방법을 배울 수 있습니다. 단일 머신에서 8xH100 80GB GPU를 활용하여 LoRA와 FSDP를 적용해 Llama-70B 모델의 SFT(지도 기반 미세 조정)를 수행하도록 스크립트를 구성합니다. accelerate config를 변경하여 여러 머신으로 확장하도록 구성할 수 있습니다.
구성
🤗 Accelerate를 사용하여 FSDP 구성 파일을 생성하는 다음 명령을 실행하여 시작합니다. --config_file 플래그를 사용하면 구성 파일을 특정 위치에 저장할 수 있으며, 그렇지 않으면 🤗 Accelerate 캐시에 default_config.yaml 파일로 저장됩니다.
설정 파일은 학습 스크립트를 실행할 때 기본 옵션을 설정하는 데 사용됩니다.
accelerate config --config_file fsdp_config.yaml
설정에 대한 몇 가지 질문이 표시됩니다. 아래 이미지에 표시된 대로 설문지를 응답하여 다음 인수를 구성합니다.
이 작업이 완료되면 해당 config는 다음과 같이 보이며 fsdp_config.yaml 파일은 config 폴더에서 찾을 수 있습니다:
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 Falserank=8, alpha=16을 사용하고 모든 linear layer를 대상으로 하는 LoRA를 적용하고 있음을 주목하세요. FSDP 설정 파일을 전달하고 ultrachat dataset의 하위 집합으로 70B Llama 모델을 fine-tuning하고 있습니다.
중요한 부분
스크립트를 조금 더 자세히 살펴보겠습니다. 어떻게 동작하는지 이해할 수 있도록 내부 구조를 확인해 보겠습니다.
가장 먼저 알아둘 점은, 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와 함께 사용할 때 주목해야 할 주요 사항은 GPU 메모리 절약을 실현하려면 use_orig_params이 False여야 한다는 것입니다. use_orig_params=False로 인해 FSDP의 자동 래핑 정책은 학습 가능 파라미터와 학습 불가능 파라미터를 각각 별도로 래핑하도록 변경되어야 합니다. 이는 아래 코드 스니펫에서 PEFT의 유틸 함수 fsdp_auto_wrap_policy를 사용하여 수행됩니다:
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 GB(90-98%)입니다. 끝부분에서 GPU 메모리가 약간 증가하는 이유는 SHARDED_STATE_DICT 대신 FULL_STATE_DICT 상태 딕트 유형으로 모델을 저장하기 때문이며, 이를 통해 추론 시 from_pretrained 메서드로 정상적으로 로드할 수 있는 어댑터 가중치를 모델이 갖게 됩니다:
PEFT QLoRA와 FSDP를 사용하여 여러 GPU에서 대규모 모델을 파인튜닝합니다
이 섹션에서는 2X24GB GPU에서 70B llama 모델을 파인튜닝하기 위해 QLoRA와 FSDP를 사용하는 방법을 살펴봅니다. 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 offloading)을 설정해야 합니다. accelerate launcher를 사용하지 않는 경우, 대신 환경 변수 export FSDP_CPU_RAM_EFFICIENT_LOADING=true를 설정할 수 있습니다. 여기서는 accelerate config를 사용하며, 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"4-bit 파라미터를 패킹하는 데 사용되는 데이터 타입을 나타내는 새로운 인자 bnb_4bit_quant_storage_dtype가 전달되는 것을 확인하세요. 예를 들어, bfloat16으로 설정하면 양자화 후 16/4 = 4개의 4-bit 파라미터가 함께 패킹됩니다. 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,
)AutoModelForCausalLM의 dtype이 bnb_4bit_quant_storage 데이터 타입과 동일하다는 점에 유의하세요. 이 정도면 충분합니다. 나머지는 Trainer와 TRL이 처리합니다.
메모리 사용량
위 예시에서 GPU당 소비되는 메모리는 19.6 GB이고, CPU RAM 사용량은 약 107 GB입니다. CPU 오프로딩을 비활성화하면 GPU 메모리 사용량은 35.6 GB/ GPU가 됩니다. 따라서 전체 파인튜닝에는 16X80GB GPU, FSDP+LoRA에는 8X80GB GPU, DDP+QLoRA에는 80GB GPU 몇 개가 필요했던 작업이 이제는 2X24GB GPU로 가능합니다. 이를 통해 대규모 모델의 파인튜닝이 더 쉽게 접근할 수 있게 되었습니다.
추가 리소스
FSDP와 PEFT를 사용하여 파인튜닝하는 방법에 대해서는 llama-recipes 저장소와 Getting started with Llama 가이드를 참고할 수 있습니다.
주의 사항
- PEFT와 FSDP를 함께 사용할 때의 병합은 현재 지원되지 않으며 오류가 발생합니다.
modules_to_save구성 매개변수를 전달하는 기능은 현재 테스트되지 않았습니다.- CPU Offloading을 사용할 때의 GPU Memory 절감 효과는 현재까지 테스트되지 않았습니다.
- FSDP+QLoRA를 사용할 때
paged_adamw_8bit는 현재 체크포인트를 저장할 때 오류가 발생합니다. - FSDP를 사용한 DoRA 학습은 작동해야 합니다(LoRA보다 속도는 느릴 수 있음). bitsandbytes(QDoRA)와 결합할 경우 4-bit 양자화도 작동해야 하지만, 8-bit 양자화는 알려진 문제가 있어 권장되지 않습니다.
- FSDP2에서
set_adapter()와set_requires_grad()는 PeftModel 내부의 FSDP2 모듈,PeftModel자체를 포함하여,requires_grad변경 사항이 샤드된 매개변수에 도달하도록 리샤드(reshard)합니다.disable_adapter()를 종료할 때는, FSDP2가 내부에 등록된 매개변수 객체를 샤드된 매개변수에 대한 포워드와 같이 교체했거나, 재활성화 시requires_grad가 진입 시점의 값과 달라지는 경우에만 리샤드합니다. 병합된 어댑터 레이어의 매개변수를 샤드하는 FSDP2 모듈은 리샤드되지 않으므로, 이 메서드를 호출하기 전에 언머지(unmerge)를 수행해야 합니다.reshard_after_forward=False인 경우 또는 루트 모듈의 기본값인None인 경우, 다음 포워드에서 다시 올게더(all-gather)가 수행됩니다.PeftModel를 래핑하는 FSDP2 모듈은 리샤드되지 않으므로, 이 메서드들을 호출하기 전에 해당 모듈에reshard()를 호출해야 합니다. 이러한 호출로 인해 기본 가중치가 샤드된 상태로 남을 수 있으므로, 바로 이후에 병합을 수행하면 올바르게 작동하지 않습니다(1 참조). 프롬프트 학습 방법과 어댑션 프롬프트는 아직 지원되지 않습니다.