[LLM] 모델 학습 with HuggingFace (Accelerate) (6)

1. 개념

• 멀티 GPU/TPU 환경 등에서 파이토치(PyTorch) 기반 훈련 스크립트를 간편하게 확장할 수 있도록 도와주는 경량 라이브러리
• 전통적으로 PyTorch 분산 학습(DDP)나 TPU/XLA를 직접 설정하려면, torch.distributed API나 TPU용 코드를 별도로 작성
• 사용자 정의 훈련 루프를 최대한 그대로 두면서 분산 학습 설정과 관련된 복잡도를 줄여주는 것이 핵심
• 다양한 하드웨어 가속 환경을 지원

 

출처: Accelerate 링크

Accelerate

 

설치 방법

pip install accelerate

 

---

2. 실행 방법

2.1. accelerate config (설정 세팅)

• 사용 목적: 사용자의 환경(로컬 머신/멀티 GPU/TPU/멀티 노드 등)에 맞추어, 분산 학습 설정을 빠르게 상호작용형(Interactive) 모드로 구성
• 해당 명령어를 입력하면, 터미널에서 여러 질문이 차례로 나옴.
• GPU 장치 개수, 혼합 정밀도(FP16/BF16), DeepSpeed 사용 여부, gradient_accumulation_steps 등의 설정을 yaml 파일로 저장

$ accelerate config

 

설정 테스트

accelerate test

 

---

2.2. accelerate launch (설정 실행)

• 설정된 config 정보를 기반으로, 분산 학습 프로세스(DDP, DeepSpeed, TPU 등)를 자동으로 띄워 주는 명령어
• 일반적인 Python 스크립트(train.py)를 분산 실행 형태로 간단히 확장

$ accelerate launch python train.py --config accelerate_config.yaml

 

파라미터 예시

• --multi_gpu : 로컬 머신에 여러 GPU가 있는 경우

• --num_processes 4 : 프로세스를 4개 띄우고 싶은 경우 (직접 지정)

• --main_process_port 29500 : 분산 프로세스들 간 통신 포트

• --mixed_precision fp16 : 혼합 정밀도 적용(Direct CLI옵션; config.yaml 대신 CLI로 즉시 지정할 수도 있음)

 

---

3.  Pytorch 코드 반영

pytorch 반영 방법은 아래와 같다.

 

3.1. 학습

기본 방법

• 먼저 Accelerator() 객체를 생성한다.
• device를 위 객체를 이용해 만들어 준다.
• 모델, 옵티마이저, 데이터 등을 accelerator에 등록해 준다.
• 위에서 생성한 Accelerator에 등록한 모델을 학습한다.
• 미분은 accelerator.backward(loss) 을 이용해준다.

 

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ device = accelerator.device
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
-     # inputs = inputs.to(device)
-     # targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()

 

Gradient Accumulation (그레디언트 모으기)

• with accelerator.accumulate(model): 블럭을 만들어 준다. 이용함
• 해당 블록 내에서 그래디언트를 몇 스텝 누적하고 언제 step을 해야 하는지를 자동으로 관리
• optimizer.step(), ptimizer.zero_grad() 호출을 지연시키는 등의 동작을 함
• 즉, 코드상으로는 매 스텝마다 step/zero_grad가 보이지만, 실질적으로는 설정된 accumulation_steps만큼 그래디언트를 누적한 뒤에 한 번씩 동작

accelerator = Accelerator(gradient_accumulation_steps=2)
model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

for input, label in training_dataloader:
    # ↓ 이 컨텍스트 안에서 그래디언트 누적 로직이 자동 제어됨
    with accelerator.accumulate(model):
        predictions = model(input)
        loss = loss_function(predictions, label)
        accelerator.backward(loss)
        
        # 실제로는 특정 스텝(accumulation_steps를 만족할 때)에만 optimizer.step()을 호출
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

---

3.2. 분산 평가

•  gather_for_metrics((predictions, targets)) : 각 장치의 예측한 결과값을 모아줌
• metric.add_batch(predictions, references) : 한 번의 미니배치(혹은 여러 스텝)에 대한 예측값(predictions)과 실제 정답값(references)을 한 번에 모아서 metric 객체에 저장

from datasets import load_metric
metric = load_metric("accuracy")

validation_dataloader = accelerator.prepare(validation_dataloader)

for inputs, targets in validation_dataloader:
    predictions = model(inputs)
    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    # Example of use with a *Datasets.Metric*
    metric.add_batch(all_predictions, all_targets)

 

---

3.3. 모델 불러오기 (빈 가중치 할당)

• init_empty_weights() : 내부에서는 파이토치의 meta 장치를 활용하여 “빈 텐서”로 모델 파라미터를 생성합니다. 이렇게 하면 실제 데이터(가중치)가 할당되지 않으므로 메모리를 거의 사용하지 않음

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

 

---

3.4. 모델 불러오기 (가중치 불러오기)

• Accelerate.load_checkpoint_and_dispatch() : 로컬 체크포인트에서 모델 가중치를 불러온 뒤, 가용 디바이스로 자동 분산
• model_checkpoint 경로에 저장된 체크포인트(가중치)를 불러온다.
• device_map="auto" 옵션을 통해, 가능한 하드웨어 자원(GPU/CPU)을 자동 탐색해 모델의 레이어(파라미터)를 분산 배치
• no_split_module_classes=['Block']는 특정 모듈(Block) 단위로는 쪼개지 않고 통째로 할당

from accelerate import load_checkpoint_and_dispatch

model_checkpoint = "your-local-model-folder"
model = load_checkpoint_and_dispatch(
    model, 
    checkpoint=model_checkpoint, 
    device_map="auto",
    no_split_module_classes=['Block']
)

 

---

4. Transformers.Trainer 와 함께 사용

Transformers의 Trainer 내부에서 이미 Accelerate가 기본적으로 적용되어 있어서,

사용자가 별도로 코드를 많이 수정할 필요가 없다.