FSDP과 DeepSpeed

Accelerate는 Pytorch FSDP와 Microsoft DeepSpeed라는 두 가지 강력한 분산 학습 도구를 통합하여 학습 프레임워크의 유연성을 제공합니다. 이 튜토리얼의 목표는 두 프레임워크 간의 유사점과 잠재적 차이를 비교하여 사용자가 이 두 프레임워크를 원활하게 전환할 수 있도록 하는 것입니다.

프레임워크를 전환하려면 accelerate launch 코드를 실행할 때 --config_file로 올바른 설정 파일을 전달하거나, FSDP 및 DeepSpeed의 경우 각각의 인수를 직접 전달하는 것을 권장합니다.

DeepSpeed와 FSDP의 Accelerate 설정 예제는 “Launch Configurations” 섹션의 example zoo에서 확인할 수 있습니다.

이 튜토리얼은 단일 노드, 다중 GPU 시나리오에만 해당합니다.

기능 구성

모델 텐서는 모델 크기를 확장하기 위해 서로 다른 GPU로 분할됩니다. 이는 FSDP에서는 sharding, DeepSpeed에서는 partitioning이라고 불립니다. FSDP sharding과 DeepSpeed ZeRO(partitioning) 단계는 각각 --fsdp_sharding_strategy와 --zero_stage로 구성됩니다. 특히, FSDP FULL_SHARD는 DeepSpeed ZeRO 단계 3에 매핑됩니다. FSDP sharding과 DeepSpeed ZeRO 설정 간의 포괄적인 매핑을 참조하세요. 아래 표는 유사한 설정을 요약하고 그룹화합니다:

위 항목에 대한 자세한 설명은 Accelerate 런칭 문서를 참고하세요.

혼합 정밀도 설정과 같은 다른 DeepSpeed 구성에 액세스하려면 --deepspeed_config_file를 전달해야 하며, 문서를 참조하세요.

DeepSpeed는 DeepSpeedPlugin을 통해서도 구성할 수 있습니다. 예를 들어, DeepSpeedPlugin.zero_stage은 --zero_stage와 동등하며, DeepSpeedPlugin.hf_ds_config는 --deepeed_config_file.을 전달하는 데 사용할 수 있습니다.

FSDP는 FullyShardedDataParallelPlugin를 통해서도 구성할 수 있습니다. 예를 들어, FullyShardedDataParallelPlugin.sharding_strategy은 --fsdp_sharding_strategy와 동등합니다.

Checkpointing

FSDP는 --fsdp_state_dict_type를 통해 전체 / 샤드된 체크포인트를 저장하도록 구성할 수 있다는 점에 유의하세요.

DeepSpeed Zero3의 경우 --zero3_save_16bit_model true를 전달할 수 있으며, 이는 모델을 단일 rank로 통합하여 저장하는 편리한 방법입니다. 이는 fsdp_state_dict_type: FULL_STATE_DICT의 FSDP에 해당하는 기능입니다.

대형 모델의 경우 모델을 단일 rank로 통합하는 과정이 매우 느릴 수 있습니다.

더 빠른 체크포인트 저장을 위해, FSDP의 경우 fsdp_state_dict_type: SHARDED_STATE_DICT를 사용하고, DeepSpeed Zero3의 경우 zero_to_fp32.py 스크립트를 사용하여 샤드된 체크포인트를 사후 변환하세요.

오프로딩

FSDP는 all-or-nothing 오프로드만 허용합니다(즉, 매개변수, 그래디언트, 옵티마이저를 모두 오프로드하거나 모두 GPU에 유지해야 합니다). 반면 DeepSpeed는 매개변수와 옵티마이저를 다르게 오프로드할 수 있습니다. 또한 DeepSpeed는 NVME로의 오프로드도 지원합니다.

Prefetching

FSDP는 통신/계산의 오버랩을 개선하기 위해 추가 메모리 비용으로 --fsdp_forward_prefetch와 --fsdp_backward_prefetch라는 두 가지 프리페칭 구성을 허용합니다. 자세한 내용은 FSDP 문서를 참조하세요. DeepSpeed의 경우, 필요할 때 프리페칭이 활성화되며, stage3_param_persistence_threshold, stage3_max_reuse_distance와 같은 특정 하이퍼파라미터에 따라 활성화됩니다. 이러한 하이퍼파라미터는 Zero3에 대해 구성할 수 있습니다. deepspeed 구성 파일에서 이러한 값을 명시적으로 설정하지 않으면 accelerate가 자동으로 설정할 수 있습니다.

메모리가 허용하는 경우 FSDP의 경우 fsdp_backward_prefetch: BACKWARD_PRE를 설정하여 처리량을 개선할 수 있습니다.

모델 로딩

FSDP의 경우 효율적인 모델 로딩을 활성화하려면 --fsdp_cpu_ram_efficient_loading true를 명시적으로 지정해야 하지만, transformers는 DeepSpeed Zero3가 사용될 때 유사한 기능을 자동으로 활성화합니다.

FSDP의 경우 --fsdp_cpu_ram_efficient_loading true를 설정할 때마다 accelerate는 sync_module_states를 자동으로 true로 설정합니다. RAM 효율적 로딩을 위해 가중치는 단일 rank에서만 로드되며, 따라서 다른 rank로 가중치를 브로드캐스트하려면 sync_module_states가 필요합니다.

모델

FSDP는 all-gather 및 reduce-scatter 연산의 스케줄링 방식을 알고리즘이 결정하도록 --fsdp_auto_wrap_policy를 명시적으로 요구합니다. 하지만 DeepSpeed의 경우 이는 사용자에게 투명하게 처리됩니다.

FSDP의 경우, fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP만 설정하면 됩니다. 최신 transformers 버전에서는 HF transformers 모델에 적합한 fsdp_transformer_layer_cls_to_wrap를 자동으로 파악하도록 최선을 다하고 있습니다. 하지만 해당 항목과 관련된 오류가 발생하면, 해당 값을 명시적으로 지정해 주세요.

Parameters Summoning

torch.compile를 사용하는 경우 FSDP에는 --fsdp_use_orig_params 플래그가 명시적으로 필요합니다. pytorch 문서를 참조하세요. DeepSpeed의 경우 사용자에게 투명하게 처리됩니다.

FSDP를 사용할 때 torch.compile을 적용하려면 fsdp_use_orig_params: True를 설정하세요.

학습

DeepSpeed는 --gradient_accumulation_steps 및 --gradient_clipping 플래그를 명시적으로 요구합니다. FSDP의 경우 이는 사용자에게 투명하게 처리됩니다.

DeepSpeed를 사용할 때 gradient_accumulation_steps: "auto"와 gradient_clipping: "auto"를 설정하면 Accelerator 또는 TrainingArguments(transformers를 사용하는 경우)에 설정된 값을 자동으로 가져옵니다.

데이터 정밀도 처리의 차이

FSDP와 DeepSpeed에서 데이터 정밀도가 어떻게 처리되는지 논의하려면, 먼저 이러한 프레임워크에서 모델 파라미터가 어떻게 처리되는지 개요를 살펴보는 것이 도움이 됩니다. 모델/옵티마이저 파라미터가 GPU에 분배되기 전에, 파라미터 준비 과정이 개입되어 먼저 1차원 torch.Tensor로 “flatten”됩니다. FSDP/DeepSpeed의 구현은 이러한 “flatten”된 파라미터가 저장되는 dtype 측면에서 차이가 있으며, 이는 torch.Optimizer가 dtype을 할당하는 방식에 영향을 미칩니다. 아래 표는 두 프레임워크의 프로세스를 요약하며, “Local” 열은 GPU별 수준에서 발생하는 프로세스를 나타내므로, 업캐스팅으로 인한 메모리 오버헤드는 사용된 GPU 수로 상쇄된다고 이해해야 합니다.

일반적으로 자동 혼합 정밀도로 안정적인 학습을 하려면 모든 학습 가능한 파라미터가 torch.float32 형식이어야 합니다.

따라서 DeepSpeed를 사용할 때 GPU 수가 적은 경우, 준비 단계에서 발생하는 업캐스팅으로 인해 상당한 메모리 오버헤드가 발생할 수 있음을 유의하세요.

FSDP를 사용할 때 혼합 정밀도가 없는 경우, GPU 수가 적은 환경에서 유용할 수 있는 저정밀도 torch_dtype로 torch.Optimizer를 동작시킬 수 있습니다.

혼합 정밀도를 사용할 경우, FSDP와 DeepSpeed는 모델 준비 단계에서 업캐스트를 수행합니다(위 표 참조). 단, FSDP는 업캐스트된 정밀도로 체크포인트를 저장합니다. DeepSpeed는 --zero3_save_16bit_model이 지정된 경우에도 여전히 낮은 정밀도의 체크포인트를 저장할 수 있습니다.

위 표를 명확히 하기 위해 아래 구체적인 예시를 살펴보세요. 간결성을 위해 옵티마이저의 사전 단계와 실제 단계를 함께 표시했습니다. FSDP를 사용하면 아래에 표시된 두 가지 모드로 작동할 수 있지만, DeepSpeed는 하나의 모드로만 작동할 수 있습니다.