FSDP2

Fully Sharded Data Parallel (FSDP2)는 모델, 그래디언트, 옵티마이저 상태를 여러 GPU에 걸쳐 샤딩합니다. 계산 전에 각 GPU는 모든 샤드에서 완전한 파라미터 세트를 수집한 뒤, 계산이 끝나면 이를 해제합니다. 샤딩은 단일 GPU의 메모리보다 큰 모델을 학습할 수 있게 해주지만, DDP보다 통신량이 더 많이 발생합니다. 모델이나 옵티마이저 상태가 단일 GPU에 맞지 않을 때 FSDP를 사용하세요.

                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step

Sharding strategies

FSDP2는 ~TrainingArguments.fsdp_config로 샤딩을 제어합니다. FSDP를 활성화하려면 fsdp=True로 설정하고, FSDP 설정에서 reshard_after_forward를 설정하여 메모리와 처리량 간의 트레이드오프를 선택합니다.

auto_wrap_policy는 모듈이 FSDP 단위로 래핑되는 방식을 제어합니다. 기본값은 "TRANSFORMER_BASED_WRAP"이며, 이는 모델의 transformer 레이어를 래핑합니다. 래핑 없이("NO_WRAP") 사용할 경우, 전체 모델이 하나의 FSDP 단위가 되어 샤딩의 메모리 이점을 잃게 됩니다.

FSDP 구성

이 필드들은 FSDP2가 모델을 래핑, 샤딩 및 로드하는 방식을 제어합니다. reshard_after_forward와 auto_wrap_policy는 샤딩 전략에서 다룹니다.

• cpu_offload는 사용되지 않는 파라미터와 그래디언트를 CPU로 오프로드하여 GPU 메모리를 절약합니다.

• transformer_layer_cls_to_wrap은 auto_wrap_policy가 "TRANSFORMER_BASED_WRAP"인 경우 FSDP unit으로 래핑할 transformer layer를 정의합니다. 각 unit은 자신의 gather 및 scatter 연산을 관리합니다. forward pass 동안에는 현재 unit의 파라미터만 수집됩니다. 이전 unit의 파라미터는 메모리를 절약하기 위해 해제됩니다.

  최상위 모델만 래핑하면 GPU 메모리 절감 효과가 없습니다. 개별 Linear 레이어를 모두 래핑하면 단위 간 통신 비용이 매우 높아집니다. 이 필드를 비워 두면 FSDP가 모델 정의에서 값을 읽습니다.

• min_num_params는 크기 기반 래핑을 위한 모듈별 최소 매개변수 수를 설정합니다. auto_wrap_policy가 "SIZE_BASED_WRAP"인 경우에만 사용됩니다.

• state_dict_type는 체크포인트 형식을 제어합니다. 기본값은 "FULL_STATE_DICT"이며, 이는 단일 Transformers 호환 체크포인트를 생성합니다. 대형 모델의 경우 더 빠른 "SHARDED_STATE_DICT"를 사용하면 랭크별로 하나의 체크포인트 파일이 생성됩니다. 샤드된 체크포인트는 FSDP로만 다시 로드할 수 있으므로, FSDP 외부에서 공유하거나 로드하려는 최종 체크포인트는 "FULL_STATE_DICT"로 저장하세요.

• cpu_ram_efficient_loading은 rank 0에서만 디스크에서 체크포인트를 로드합니다. 다른 GPU는 빈 모델을 초기화한 후 브로드캐스트를 통해 가중치를 받아, 여러 프로세스가 큰 모델을 CPU RAM에 로드하는 것을 방지합니다.

• activation_checkpointing는 저장 대신 역전파 과정에서 activations를 다시 계산합니다. TrainingArguments에서 gradient checkpointing 대신 이 옵션을 사용하세요. 두 옵션을 모두 설정하면 오류가 발생합니다.

FSDP 학습은 Accelerate 설정 파일 또는 fsdp_config에 전달되는 FSDP 설정 파일 중 하나로 구성할 수 있습니다.

accelerate launch train.py

{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

mixture-of-experts 모델의 경우, fsdp_size를 tp_size 및 ep_size와 결합하여 전문가(experts)를 하나의 메시(mesh) 차원에, 나머지는 다른 차원에 샤딩(shard)할 수 있습니다. expert parallelism을 참조하세요.

컨텍스트 병렬성을 사용하여 평가

FSDP2를 사용한 컨텍스트 병렬성(CP)은 Accelerate의 ParallelismConfig cp_size를 통해 설정됩니다. 실행 시 --use_parallelism_config와 --parallelism_config_cp_size를 모두 전달해야 합니다. Accelerate YAML에 비어 있지 않은 parallelism_config가 있으면 use_parallelism_config도 자동으로 켜집니다. --parallelism_config_cp_size만 단독으로 전달하면 무시되고 cp_size=1이 유지됩니다.

FSDP2에서 CP는 시퀀스를 랭크별로 샤딩하여 각 GPU가 컨텍스트의 슬라이스를 보유하도록 합니다(allgather 또는 alltoall). 이는 DeepSpeed ALST를 통해 실행되며 동일한 설정에서 sp_size를 사용하는 Ulysses 시퀀스 병렬성과 다릅니다. CP 또는 Ulysses SP 중 하나를 선택하십시오. cp_size와 sp_size를 모두 1보다 크게 설정하지 마십시오.

evaluate()는 prediction_step()를 통해 train()과 동일한 context-parallel 컨텍스트를 사용합니다. 해당 컨텍스트가 활성화된 동안 각 랭크는 자신의 시퀀스 샤드를 실행합니다. eval_loss는 올바른 손실 값입니다(cp_size로 곱해지지 않음). 조기 종료와 load_best_model_at_end는 이 메트릭을 사용합니다.

Logits는 시퀀스 방향으로 샤딩된 상태로 유지되며, 이를 전체 길이의 예측으로 재조립하는 과정은 없습니다. CP 환경에서는 Logits에 대한 예측 및 compute_metrics가 지원되지 않습니다.

eval_loss가 필요할 때 기존 FSDP2 + cp_size 실행을 재사용하세요.

prediction_step()를 서브클래스로 구현하고 토큰을 직접 카운팅하는 경우, 컨텍스트-파럴렐 컨텍스트에 진입하기 전에 num_items_in_batch를 카운팅해야 합니다. 해당 컨텍스트가 활성화된 동안 시퀀스 버퍼는 랭크 간에 인 플레이스로 샤딩됩니다. 이후에 카운팅하면 배치 수가 과소 집계됩니다. 기본 Trainer는 이미 먼저 카운팅하므로, prediction_step를 오버라이드하는 경우에만 이 작업이 필요합니다.

다음 단계

• 모델이 GPU 하나에 들어가는 경우의 데이터 병렬 학습에 대해서는 DDP를 참조하세요.
• ZeRO 최적화 및 NVMe 오프로딩에 대해 DeepSpeed를 참조하세요.
• PyTorch/XLA를 사용하는 TPU에서 FSDP를 사용하려면 ~TrainingArguments.fsdp_config에 xla, xla_fsdp_settings, xla_fsdp_grad_ckpt를 설정하세요.
• Read the FSDP chapter from The Ultra-Scale Playbook for more information about how FSDP works.