Hugging Face Trainer와 함께하는 DeepSpeed ZeRO 가이드
GPU를 최대한 활용하기 위한 가이드! 이 글은 AI로 번역되었습니다. 오역이 있을 수 있으니 댓글로 알려주세요.
Created on September 12|Last edited on September 12
Comment
소개
대규모 모델을 효율적이고 효과적으로 학습시키는 일은 연구자와 개발자 모두에게 핵심 과제가 되었습니다. 모델의 복잡도와 규모가 커질수록, 더 정교한 최적화 기법에 대한 수요가 혁신적인 해결책의 개발로 이어지고 있습니다.
그런 혁신적인 해결책 중 하나는 DeepSpeed ZeRODeepSpeed 라이브러리의 구성 요소로서, 사용 가능한 하드웨어에서 속도를 희생하거나 막대한 리소스를 요구하지 않고도 대규모 모델을 학습할 수 있게 하여 AI의 대중화를 목표로 합니다. 이 가이드는 DeepSpeed ZeRO가 무엇인지, 어떤 이점이 있는지, 그리고 어떻게 동작하는지 자세히 다룹니다.
DeepSpeed로 만들 수 있는 구성은 매우 다양하며, 이를 모두 다루려면 여러 튜토리얼이 필요할 것입니다. 따라서 이 가이드에서는 GPU 메모리 사용을 가장 크게 줄여 줄 핵심 구성 요소와 설정 위주로 살펴보겠습니다.
목차
소개 목차DeepSpeed ZeRO 이해하기단일 GPU 전략 CPU 오프로딩그래디언트 체크포인팅 저정밀도 데이터 타입 메모리 효율적 옵티마이저 멀티 GPU 전술 스테이지 1: 옵티마이저 상태 샤딩스테이지 2: 그래디언트 분할스테이지 3: 모델 파라미터 분할Hugging Face와 DeepSpeed로 학습하기 전속력으로 앞으로 관련 문서
DeepSpeed ZeRO 이해하기
DeepSpeed ZeRO는 오픈 소스 딥러닝 최적화 라이브러리인 DeepSpeed의 구성 요소로, 계산 부담을 줄이고 학습 효율을 높이도록 설계되었습니다. Microsoft가 개발한 DeepSpeed는 대규모 모델 학습에서 발생하는 다양한 과제를 해결하는 여러 기능을 제공합니다. 그중에서도 Zero Redundancy Optimizer(ZeRO)는 메모리 사용을 최적화하고 학습을 전례 없는 규모로 확장하는 독창적인 접근 방식으로 두드러집니다.
모델 규모가 커질수록 파라미터, 그래디언트, 옵티마이저 상태에 필요한 메모리가 사용 가능한 GPU의 용량을 넘어설 수 있어, 대규모 모델을 효율적으로 학습하기 어려워집니다. ZeRO는 분산 학습에서 전통적인 데이터 병렬성과 모델 병렬성 기법의 한계를 해결하도록 설계되었습니다. ZeRO는 사용 가능한 디바이스 전반에 모델 상태를 분할해 각 디바이스당 메모리 사용량을 크게 줄이고, 이전에는 불가능했던 규모의 모델을 학습할 수 있도록 합니다.
전통적인 데이터 병렬성(DP)은 연산과 통신 효율은 높지만, 모든 프로세스에 전체 모델을 복제하기 때문에 메모리 효율이 낮아 중복된 메모리 사용이 발생합니다. 반면 모델 병렬성(MP)은 모델을 분할해 메모리 효율을 높이지만, 연산 및 통신 효율이 떨어져 미세한 단위의 연산과 비용이 큰 통신이 늘어나 확장성이 저해되는 문제가 있습니다.
두 접근 방식 모두 학습 내내 모든 모델 상태를 정적으로 유지하는데, 모든 상태가 항상 필요한 것은 아니므로 비효율적입니다. 이러한 문제를 해결하기 위해 Microsoft는 ZeRO-DP(Zero Redundancy Optimizer for Data Parallelism)를 개발했습니다. ZeRO-DP는 메모리 효율을 위해 MP처럼 모델 상태를 프로세스 간에 분할하는 동시에, DP의 연산 세분성과 통신 효율을 그대로 유지함으로써 두 방식의 장점을 결합합니다. 이는 동적인 통신 스케줄을 통해 이루어지며, DP의 이점을 희생하지 않고 중복된 메모리 사용을 효과적으로 제거합니다.
단일 GPU 전략
DeepSpeed의 많은 기능은 멀티 GPU 환경을 염두에 두고 설계되었지만, 단일 GPU만으로 라이브러리를 활용하는 방법도 있습니다. 아래에서는 단일 GPU 환경에서 적용할 수 있는 몇 가지 전략을 먼저 살펴보겠습니다.
CPU 오프로딩
단일 GPU 환경에서 주목할 만한 한 가지 전략은 CPU 오프로딩이 방법은 네트워크 파라미터와 옵티마이저 상태를 CPU RAM으로 오프로딩해 GPU의 VRAM 요구량을 줄여 줍니다. GPU 메모리가 제한된 시스템에서 대규모 모델을 다루기에 특히 유용하며, 하드웨어 업그레이드 없이도 더 복잡한 연산을 가능하게 합니다. 다만 CPU RAM을 충분히 확보해야 한다는 점을 유의하세요. CPU 오프로딩을 사용할 때 이 점을 반드시 고려하시기 바랍니다.
그래디언트 체크포인팅
별도로, 단일 GPU 환경에서 효과적인 또 다른 전략은 활성화 체크포인팅이라고도 불리는 그래디언트 체크포인팅입니다. 그래디언트 체크포인팅이 기법은 순전파 단계에서 중간 활성화의 일부만 저장하고 역전파 단계에서 나��지를 재계산함으로써 VRAM 요구량을 줄여 줍니다.
이 방법은 계산 오버헤드를 늘릴 수 있지만 메모리 사용량을 크게 줄여 단일 GPU에서도 더 큰 모델을 학습하거나 더 큰 배치 크기를 사용할 수 있게 합니다. 보통 저는 DeepSpeed 설정 파일 바깥에서 이를 활성화하고 ‘ 설정을 합니다.gradient_checkpointing’로 true 에서 TrainingArguments Hugging Face용 클래스
저정밀도 데이터 타입
DeepSpeed에서는 FP16과 BF16이 학습 속도와 효율을 높이기 위한 정밀도 형식으로 사용됩니다. FP16은 16비트 부동소수점을 사용해 메모리 사용량을 줄이고 연산을 가속하여, 지원되는 하드웨어에서 처리 속도를 사실상 두 배로 높여 줍니다. BF16은 FP32에 가까운 정확도를 유지하면서 더 넓은 표현 범위를 제공하여, 최신 CPU와 GPU에서의 학습을 최적화합니다.
두 형식 모두 최대 효율 향상을 위해 하드웨어 지원이 필요하며, DeepSpeed는 성능과 정밀도의 균형을 맞추도록 구현을 자동화해 더 빠르고 대규모의 모델 학습을 가능하게 합니다.
메모리 효율적 옵티마이저
DeepSpeed는 하드웨어와 모델 요구 사항에 맞춘 다양한 고성능 옵티마이저를 제공합니다:
Adam (CPU)매개변수별로 학습률을 적응적으로 조정하는 범용 옵티마이저로, CPU 환경에서 다양한 딥러닝 작업에 적합합니다.
AdamW (CPU): 가중치 감쇠 정규화를 분리한 Adam의 변형으로, 과적합에 취약한 모델에서 특히 학습 안정성과 성능을 개선하는 경향이 있습니다.
FusedAdam (GPU): GPU에서 성능을 높이기 위해 연산을 융합해 Adam을 개선한 옵티마이저로, 최적 활용을 위해 Apex 설치가 필요합니다.
FusedLamb (GPU): 대규모 배치 학습을 위한 LAMB 알고리즘을 적용하여, 적응형 학습률을 지원하면서 GPU에서의 학습 효율을 최적화합니다.
OneBitAdam (GPU): 분산 학습의 통신 오버헤드를 줄이기 위해 1비트 압축을 도입하여, 다중 GPU로 확장하는 데 최적입니다.
ZeroOneAdam (GPU): 제로 및 1비트 양자화의 장점을 결합하여 분산 학습 효율과 통신을 한층 더 최적화합니다.
원비트LAMB (GPU): 1비트 압축을 적용한 LAMB 알고리즘을 구현하여, 모델 학습의 확장성을 높이고 GPU 메모리 사용량을 줄입니다.
멀티 GPU 전술
여러 GPU를 사용할 때는 DeepSpeed가 VRAM 요구량을 줄이기 위한 더 많은 도구를 제공합니다. 일반적으로 사용되는 3가지 주요 “스테이지”가 있으며, 아래에서 각 스테이��를 설명하겠습니다. 참고로, CPU 오프로딩을 함께 사용하는 경우 단일 GPU에서도 이 스테이지들을 적용할 수 있습니다.
스테이지 1: 옵티마이저 상태 샤딩
DeepSpeed ZeRO의 스테이지 1은 옵티마이저 상태의 메모리 사용량을 해결합니다. 특히 Adam이나 RMSprop과 같은 옵티마이저의 상태는 모델 파라미터만큼 많은 메모리를 차지할 수 있습니다. 전통적인 설정에서는 각 GPU가 옵티마이저 상태의 사본을 개별적으로 유지하기 때문에 상당한 메모리 오버헤드가 발생합니다.
스테이지 1은 이전 단계에서 파라미터와 그래디언트를 분할했던 방식과 유사하게, 옵티마이저 상태를 GPUs 전체에 샤딩하여 이 문제를 해결합니다. 각 GPU는 옵티마이저 상태의 일부만 보유하며, 업데이트 단계에서는 분산된 파라미터와 그래디언트를 기준으로 해당되는 옵티마이저 상태만 업데이트됩니다. 아래는 그래디언트를 유지하면서 보여주는 예시입니다.
{"zero_optimization": {"stage": 1},"gradient_accumulation_steps": 1,"train_micro_batch_size_per_gpu": 1,"gradient_clipping": 1.0,"fp16": {"enabled": true}}
앞에서 언급했듯이 옵티마이저 상태를 CPU로 오프로딩하는 것도 가능하며, 다음 설정에 간단한 항목을 추가하면 쉽게 구현할 수 있습니다. offload_optimizer 플래그입니다. 또한 옵티마이저를 CPU로 오프로딩하는 것은 ZeRO 스테이지 1, 2, 3 모두에서 유효하다는 점도 중요합니다!
{"zero_optimization": {"stage": 1},"offload_optimizer": {"device": "cpu"}"gradient_accumulation_steps": 1,"train_micro_batch_size_per_gpu": 1,"gradient_clipping": 1.0,"fp16": {"enabled": true}}
스테이지 2: 그래디언트 분할
DeepSpeed ZeRO의 두 번째 스테이지는 학습에 참여하는 모든 GPU에 그래디언트를 분할해 저장함으로써 메모리 사용을 최적화하는 데 초점을 맞춥니다. 전통적인 분산 학습에서는 각 GPU가 전체 그래디언트 집합을 보관하기 때문에, 모델 크기가 커질수록 메모리 병목이 빠르게 발생할 수 있습니다.
스테이지 2는 각 GPU가 그래디언트의 일부만 보유하도록 함으로써 이 문제를 해결합니다. 즉, 역전파 동안 그래디언트가 계산된 뒤 GPU들 사이에 분산되어 어느 한 GPU도 전체 그래디언트를 모두 보유하지 않게 됩니다. 이 방식은 학습에 필요한 메모리를 크게 줄여, 더 큰 배치 크기나 모델 크기를 사용할 수 있게 합니다.
스테이지 1 구성 JSON 파일 예시는 다음과 같습니다:
{"zero_optimization": {"stage": 2},"gradient_accumulation_steps": 1,"train_micro_batch_size_per_gpu": 1,"gradient_clipping": 1.0,"fp16": {"enabled": true}}
스테이지 3: 모델 파라미터 분할
첫 두 스테이지에서 마련된 기반 위에, 스테이지 3는 분할 전략을 모델 파라미터로까지 확장합니다. 이 단계에서는 모델의 파라미터를 여러 GPU에 나누어 저장하며, 각 GPU는 모델 파라미터의 일부만 저장하고 업데이트하는 책임을 집니다.
이러한 분할 덕분에 각 GPU가 모델을 저장하는 데 필요한 메모리가 크게 줄어들어, 기존 하드웨어로는 불가능했던 훨씬 더 큰 모델도 학습할 수 있게 됩니다.
학습 중에는 특정 GPU에서 연산에 필요한 파라미터가 GPU들 사이에서 동적으로 통신되어, 각 장치가 자신에게 할당된 연산에 필요한 파라미터에 접근할 수 있도록 합니다. 또한 모델의 파라미터를 CPU로 오프로딩하여 GPU VRAM 요구량을 줄일 수도 있습니다. 다만 파라미터를 저장하려면 충분한 CPU RAM이 필요하다는 점을 유의해야 하므로, 사용 중인 머신이 이러한 요구 사항을 충족하는지 반드시 확인하세요.
{"zero_optimization": {"stage": 3,"offload_param": {"device": "cpu","pin_memory": true},"offload_optimizer": {"device": "cpu","pin_memory": true}},"gradient_accumulation_steps": 1,"train_micro_batch_size_per_gpu": 1,"gradient_clipping": 1.0,"fp16": {"enabled": true}}
스테이지 3에서는 VRAM 사용량을 줄이기 위해 모델의 파라미터를 CPU로 오프로딩할 수도 있습니다. 아래는 예시입니다:
{"zero_optimization": {"stage": 3,"offload_param": {"device": "cpu"},"offload_optimizer": {"device": "cpu"}},"gradient_accumulation_steps": 1,"train_micro_batch_size_per_gpu": 1,"gradient_clipping": 1.0,"fp16": {"enabled": true}}
Hugging Face와 DeepSpeed로 학습하기
이들 다양한 설정을 시험해 보고 싶다면, 아래에 다양한 DeepSpeed 설정을 테스트할 수 있는 학습 스크립트를 추가하겠습니다. 이 스크립트를 실행하려면 일반적으로 사용하는 ‘python’ 명령을 ‘ 로 바꿔야 한다는 점을 유의하세요.deepspeed그래서 스크립트 이름이 train.py, 명령을 실행하면 됩니다 deepspeed train.py.
이 스크립트의 핵심은 DeepSpeed 설정 파일을 TrainingArguments에 전달하여 Hugging Face 트레이너가 해당 설정을 활용하도록 하는 것입니다.
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArgumentsfrom datasets import load_datasetfrom transformers import DataCollatorForLanguageModelingfrom transformers import TrainerCallbackimport torchimport transformersdef main():# Load the tokenizertokenizer = AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-step-50K-105b')tokenizer.add_special_tokens({'pad_token': '[PAD]'})# Load dataset from the Hugging Face datasets librarydataset = load_dataset("wikitext", "wikitext-2-raw-v1")# Tokenize the textsdef tokenize_function(examples):return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)tokenized_datasets = dataset.map(tokenize_function, batched=True)# Load the data collatordata_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,mlm=False, # TinyLlama uses a causal (not masked) language model, similar to GPT-2)# Load the modelmodel = AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-step-50K-105b')# Define the training argumentstraining_args = TrainingArguments(output_dir='./results',overwrite_output_dir=True,num_train_epochs=3,per_device_train_batch_size=1,save_steps=10_000,save_total_limit=2,fp16=True,deepspeed="path_to_deepspeed_config.json", # Path to DeepSpeed config filegradient_checkpointing=True,report_to='wandb')# Initialize the Trainertrainer = Trainer(model=model,args=training_args,data_collator=data_collator,train_dataset=tokenized_datasets["train"],eval_dataset=tokenized_datasets["validation"])# Start the trainingtrainer.train()# Save the final model and tokenizermodel.save_pretrained('./final_model')tokenizer.save_pretrained('./final_model')if __name__ == "__main__":main()
참고로 저는 … 설정했습니다 gradient_checkpointing 플래그를 …하도록 true 학습 인자에서 설정하여 앞서 언급한 gradient checkpointing을 활성화합니다. 추가로, 저는 report_to 인자에서 …로 'wandb' 결과를 Weights & Biases에 기록하기 위해서입니다.
전속력으로 앞으로
마무리하자면, DeepSpeed ZeRO는 대규모 AI 모델 학습의 판도를 크게 바꿉니다. 메모리를 효율적으로 관리해, 일반적인 하드웨어 구성에서도 모델 크기의 한계를 넓힐 수 있게 해줍니다. DeepSpeed ZeRO를 사용하면 하드웨어 제약에서 벗어나 모델로 무엇을 해낼 수 있는지에 집중할 수 있어, 복잡한 AI 문제를 더 단순한 과정으로 풀어갈 수 있습니다. 이 혁신은 모델을 효율적으로 확장하려는 모든 실무자에게 매우 중요한 의미가 있습니다. 질문이나 의견이 있으시면 아래에 남겨 주세요. 튜토리얼이 유익하셨길 바랍니다!
관련 문서
Building a RAG-Based Digital Restaurant Menu with LlamaIndex and W&B Weave
Powered by RAG, we will transform the traditional restaurant PDF menu into an AI powered interactive menu!
Fine-Tuning Mistral7B on Python Code With A Single GPU!
A tutorial for fine-tuning Mistral7B on Python Code using a single GPU!
How to Fine-Tune LLaVA on a Custom Dataset
A tutorial for fine-tuning LLaVA on your own data!
Skin Lesion Classification on HAM10000 with HuggingFace using PyTorch and W&B
Explore the use of HuggingFace, PyTorch, and W&B for classifying skin lesions with the HAM10000 dataset. We will build, train, and evaluate models for medical diagnostics!
출처:
Add a comment