# Fully Sharded Data Parallel

[Fully sharded data parallel](https://pytorch.org/docs/stable/fsdp.html)(FSDP)은 최대 1T 파라미터 규모의 대규모 사전 학습 모델의 분산 학습을 위해 개발되었습니다. FSDP는 모델 파라미터, 그래디언트, 옵티마이저 상태를 데이터 병렬 프로세스에 걸쳐 샤딩(sharding)하여 이를 달성하며, 샤딩된 모델 파라미터를 CPU로 오프로드(offload)할 수도 있습니다. FSDP가 제공하는 메모리 효율성 덕분에 더 큰 배치 또는 모델 크기로 학습을 확장할 수 있습니다.

이 두 가지 기능은 모두 🤗 Accelerate에서 지원되며, 🤗 PEFT와 함께 사용할 수 있습니다.

# PEFT와 FSDP 사용
이 가이드 섹션에서는 DeepSpeed [학습 스크립트](https://github.com/huggingface/peft/blob/main/examples/sft/train.py)를 사용하여 SFT를 수행하는 방법을 배울 수 있습니다. 단일 머신에서 8xH100 80GB GPU를 활용하여 LoRA와 FSDP를 적용해 Llama-70B 모델의 SFT(지도 기반 미세 조정)를 수행하도록 스크립트를 구성합니다. accelerate config를 변경하여 여러 머신으로 확장하도록 구성할 수 있습니다.

## 구성

🤗 Accelerate를 사용하여 [FSDP 구성 파일을 생성](https://huggingface.co/docs/accelerate/quicktour#launching-your-distributed-script)하는 다음 명령을 실행하여 시작합니다. `--config_file` 플래그를 사용하면 구성 파일을 특정 위치에 저장할 수 있으며, 그렇지 않으면 🤗 Accelerate 캐시에 `default_config.yaml` 파일로 저장됩니다.

설정 파일은 학습 스크립트를 실행할 때 기본 옵션을 설정하는 데 사용됩니다.

```bash
accelerate config --config_file fsdp_config.yaml
```

설정에 대한 몇 가지 질문이 표시됩니다. 아래 이미지에 표시된 대로 설문지를 응답하여 다음 인수를 구성합니다.

    

FSDP를 사용하도록 Accelerate의 config를 생성

이 작업이 완료되면 해당 config는 다음과 같이 보이며 [fsdp_config.yaml](https://github.com/huggingface/peft/blob/main/examples/sft/configs/fsdp_config.yaml) 파일은 config 폴더에서 찾을 수 있습니다:

```yml
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](https://github.com/huggingface/peft/blob/main/examples/sft/run_peft_fsdp.sh)에서 사용할 수 있으며, 아래에도 표시됩니다:
```bash
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
```

rank=8, alpha=16을 사용하고 모든 linear layer를 대상으로 하는 LoRA를 적용하고 있음을 주목하세요. FSDP 설정 파일을 전달하고 [ultrachat dataset](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k)의 하위 집합으로 70B Llama 모델을 fine-tuning하고 있습니다.

## 중요한 부분

스크립트를 조금 더 자세히 살펴보고, 어떤 일이 일어나고 있는지, 그리고 어떻게 작동하는지 이해해 보겠습니다.

가장 먼저 알아둘 점은, FSDP 설정이 전달되었으므로 스크립트가 분산 학습에 FSDP를 사용한다는 것입니다. [SFTTrainer](https://huggingface.co/docs/trl/main/en/sft_trainer#trl.SFTTrainer) 클래스는 전달된 peft 설정을 사용하여 PEFT 모델을 생성하는 모든 핵심 작업을 처리합니다. 이후 `trainer.train()`을 호출하면, Trainer는 내부적으로 🤗 Accelerate를 사용하여 FSDP 설정에 따라 FSDP로 래핑된 모델을 생성하고, 이 모델을 학습하기 위해 모델, 옵티마이저, 트레이너를 준비합니다. 주요 코드 스니펫은 다음과 같습니다:

```python
# 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` 메서드로 정상적으로 로드할 수 있는 어댑터 가중치를 모델이 갖게 됩니다:

    

학습 실행의 GPU 메모리 사용량

# PEFT QLoRA와 FSDP를 사용하여 여러 GPU에서 대규모 모델을 파인튜닝합니다

이 섹션에서는 2X24GB GPU에서 70B llama 모델을 파인튜닝하기 위해 QLoRA와 FSDP를 사용하는 방법을 살펴봅니다. [Answer.AI](https://www.answer.ai/)는 bitsandbytes 및 Hugging Face 🤗와 협력하여 FSDP+QLoRA 사용을 가능하게 하는 코드를 오픈소스로 공개했으며, 통찰력 있는 블로그 포스트 [You can now train a 70b language model at home](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html)에서 전체 과정을 설명했습니다. 이 기능은 이제 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](https://github.com/huggingface/peft/blob/main/examples/sft/configs/fsdp_config_qlora.yaml)에서 찾을 수 있는 설정은 다음과 같습니다:

```yml
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](https://github.com/huggingface/peft/blob/main/examples/sft/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`로 설정해야 합니다.

학습 코드 측면에서 중요한 코드 변경 사항은 다음과 같습니다:

```diff
...

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](https://github.com/facebookresearch/llama-recipes/?tab=readme-ov-file#fine-tuning) 저장소와 [Getting started with Llama](https://llama.meta.com/get-started/#fine-tuning) 가이드를 참고할 수 있습니다.

## 주의 사항
1. PEFT와 FSDP를 함께 사용할 때의 병합은 현재 지원되지 않으며 오류가 발생합니다.
2. `modules_to_save` 구성 매개변수를 전달하는 기능은 현재 테스트되지 않았습니다.
3. CPU Offloading을 사용할 때의 GPU Memory 절감 효과는 현재까지 테스트되지 않았습니다.
4. FSDP+QLoRA를 사용할 때 `paged_adamw_8bit`는 현재 체크포인트를 저장할 때 오류가 발생합니다.
5. FSDP를 사용한 DoRA 학습은 작동해야 합니다(LoRA보다 속도는 느릴 수 있음). bitsandbytes(QDoRA)와 결합할 경우 4-bit 양자화도 작동해야 하지만, 8-bit 양자화는 알려진 문제가 있어 권장되지 않습니다.
6. FSDP2에서 `set_adapter()`와 `set_requires_grad()`는 [PeftModel](/docs/peft/main/ko/package_reference/peft_model#peft.PeftModel) 내부의 FSDP2 모듈, `PeftModel` 자체를 포함하여, `requires_grad` 변경 사항이 샤드된 매개변수에 도달하도록 리샤드(reshard)합니다. `disable_adapter()`를 종료할 때는, FSDP2가 내부에 등록된 매개변수 객체를 샤드된 매개변수에 대한 포워드처럼 교체했거나, 재활성화 시 `requires_grad`가 진입 시점의 값과 달라지는 경우에만 리샤드합니다. 병합된 어댑터 레이어의 매개변수를 샤드하는 FSDP2 모듈은 리샤드되지 않으므로, 이 메서드를 호출하기 전에 언머지(unmerge)를 수행해야 합니다. `reshard_after_forward=False`인 경우 또는 루트 모듈의 기본값인 `None`인 경우, 다음 포워드에서 다시 올게더(all-gather)가 수행됩니다. `PeftModel`를 래핑하는 FSDP2 모듈은 리샤드되지 않으므로, 이 메서드들을 호출하기 전에 해당 모듈에 `reshard()`를 호출해야 합니다. 이러한 호출로 인해 기본 가중치가 샤드된 상태로 남을 수 있으므로, 바로 이후에 병합을 수행하면 올바르게 작동하지 않습니다(1 참조). 프롬프트 학습 방법과 어댑션 프롬프트는 아직 지원되지 않습니다.

