DeepSpeed Gradient Accumulation 메모리 최적화 심층 분석: 초거대 모델 학습을 위한 실전 전략

초거대 모델 학습 시 메모리 부족 문제를 겪고 계신가요? DeepSpeed의 Gradient Accumulation을 활용하면 배치 사이즈를 늘리는 효과를 내면서도 메모리 사용량을 최적화할 수 있습니다. 본 포스트에서는 Gradient Accumulation의 작동 원리를 파악하고, 실제 설정 방법과 메모리 최적화 전략을 자세히 설명하여 초거대 모델 학습 효율을 극대화하는 방법을 제시합니다.

1. The Challenge / Context

초거대 언어 모델 (LLM)과 같은 대규모 모델을 학습하는 것은 상당한 컴퓨팅 자원을 필요로 합니다. 가장 큰 제약 중 하나는 GPU 메모리 용량입니다. 모델 크기가 커질수록, 단일 GPU에 적합한 배치 크기는 작아지고, 이는 훈련 속도를 저하시키고 GPU 활용률을 떨어뜨립니다. 데이터 병렬 처리 (Data Parallelism)는 여러 GPU에 데이터를 분산하여 문제를 해결하지만, 통신 오버헤드가 발생하고, 여전히 개별 GPU의 메모리 제약은 존재합니다. Gradient Accumulation은 배치 크기를 늘리지 않고도 더 큰 배치 크기의 효과를 얻을 수 있는 효과적인 기술이지만, 단순히 적용하는 것만으로는 메모리 효율성을 극대화할 수 없습니다. 특히, 초거대 모델의 경우 Gradient Accumulation 적용 시 발생할 수 있는 메모리 폭증 문제를 해결하고 최적의 성능을 확보하는 것이 중요합니다.

2. Deep Dive: DeepSpeed Gradient Accumulation

DeepSpeed의 Gradient Accumulation은 여러 미니 배치의 기울기를 누적하여 단일 업데이트를 수행하는 기술입니다. 기본적으로 배치 크기를 효과적으로 늘리는 것과 같습니다. 예를 들어, accumulation steps가 4라면, 4개의 미니 배치의 기울기를 누적한 후에 모델의 가중치를 업데이트합니다. 이렇게 하면 GPU 메모리에 한 번에 로드해야 하는 데이터 양이 줄어들기 때문에 더 큰 모델을 학습할 수 있습니다. DeepSpeed는 이 과정을 효율적으로 관리하여 메모리 사용량을 최적화합니다.

DeepSpeed의 Gradient Accumulation은 PyTorch의 autograd 엔진과 긴밀하게 통합되어 작동합니다. 각 미니 배치의 forward pass 이후, backward pass를 통해 계산된 기울기는 즉시 가중치에 적용되는 대신, accumulate됩니다. 지정된 accumulation steps 수만큼 기울기가 누적되면, Optimizer가 누적된 기울기를 사용하여 모델 가중치를 업데이트합니다. 이 과정에서 DeepSpeed는 다양한 메모리 최적화 기술을 활용하여 메모리 footprint를 줄입니다.

3. Step-by-Step Guide / Implementation

다음은 DeepSpeed를 사용하여 Gradient Accumulation을 설정하고 메모리를 최적화하는 단계별 가이드입니다.

Step 1: DeepSpeed 설치 및 환경 설정

먼저 DeepSpeed를 설치합니다. pip를 사용하여 간단하게 설치할 수 있습니다.

pip install deepspeed

DeepSpeed를 사용하려면 `ds_config.json` 파일을 통해 설정을 지정해야 합니다. 이 파일은 데이터 병렬 처리, optimizer 구성, gradient accumulation steps 등 다양한 설정을 포함합니다.

Step 2: DeepSpeed 설정 파일 (ds_config.json) 작성

다음은 Gradient Accumulation을 위한 `ds_config.json` 파일의 예시입니다. 다른 최적화 옵션(예: ZeRO)도 포함되어 있습니다.

{
  "train_batch_size": 32,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 0.0001,
      "weight_decay": 0.01
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "reduce_scatter": true,
    "contiguous_gradients": true
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 32,
    "hysteresis": 2,
    "min_loss_scale": 1
  }
}

주요 파라미터는 다음과 같습니다.

  • `train_batch_size`: 전체 배치 크기.
  • `train_micro_batch_size_per_gpu`: 각 GPU 당 미니 배치 크기.
  • `gradient_accumulation_steps`: 기울기 누적 단계 수. `train_batch_size`는 `train_micro_batch_size_per_gpu * gradient_accumulation_steps`와 같아야 합니다.
  • `zero_optimization`: ZeRO 최적화 레벨을 설정합니다. stage 2 또는 stage 3를 사용하면 메모리 사용량을 크게 줄일 수 있습니다.
  • `offload_optimizer`, `offload_param`: optimizer 상태 및 모델 파라미터를 CPU로 오프로드하여 GPU 메모리를 확보합니다.
  • `fp16`: FP16 정밀도를 사용하여 메모리 사용량을 줄입니다.

Step 3: DeepSpeed 엔진 초기화 및 모델 훈련

PyTorch 훈련 스크립트에서 DeepSpeed 엔진을 초기화하고 모델을 훈련합니다. 아래는 기본적인 코드 스니펫입니다.

import deepspeed
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

# 간단한 모델 정의
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x)

# 더미 데이터셋 정의
class DummyDataset(Dataset):
    def __init__(self, size):
        self.size = size
        self.data = torch.randn(size, 10)
        self.labels = torch.randn(size, 1)

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]


# DeepSpeed 설정 파일 경로
config_path = 'ds_config.json'

# 모델, 옵티마이저, 데이터 로더 초기화
model = SimpleModel()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) # 실제로는 DeepSpeed에 의해 덮어씌워짐
dataset = DummyDataset(size=1000)
dataloader = DataLoader(dataset, batch_size=4)  # micro_batch_size_per_gpu와 일치해야 함

# DeepSpeed 엔진 초기화
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config=config_path
)


# 훈련 루프
for epoch in range(2):
    for step, (data, labels) in enumerate(dataloader):
        data = data.to(model_engine.device) # 데이터를 DeepSpeed가 관리하는 device로 옮김
        labels = labels.to(model_engine.device)

        outputs = model_engine(data)
        loss = torch.nn.functional.mse_loss(outputs, labels)

        # 역전파 및 가중치 업데이트 (gradient_accumulation_steps에 따라 자동으로 누적 후 업데이트됨)
        model_engine.backward(loss)
        model_engine.step()

        if step % 10 == 0:
            print(f"Epoch: {epoch}, Step: {step}, Loss: {loss.item()}")

print("훈련 완료!")

이 코드는 DeepSpeed 엔진을 초기화하고, 각 미니 배치에 대해 forward pass, backward pass, 가중치 업데이트를 수행합니다. `gradient_accumulation_steps`에 지정된 횟수만큼 기울기가 누적된 후 가중치가 업데이트됩니다.

Step 4: 메모리 프로파일링 및 최적화

훈련 중 메모리 사용량을 모니터링하여 추가적인 최적화가 필요한지 확인합니다. `torch.cuda.memory_allocated()` 및 `torch.cuda.memory_cached()`와 같은 PyTorch 유틸리티를 사용하여 메모리 사용량을 추적할 수 있습니다. DeepSpeed Profiler를 사용하면 더욱 상세한 메모리 프로파일링이 가능합니다. 설정 파일(`ds_config.json`)에 프로파일링 관련 설정을 추가하고 DeepSpeed 엔진 초기화 시 활성화할 수 있습니다.

{
  ...,
  "profiling": {
    "enabled": true,
    "profile_steps": [1, 10],  // 프로파일링할 step 범위
    "output_folder": "profiling_output"
  }
}

프로파일링 결과를 분석하여 메모리 병목 현상을 식별하고, ZeRO stage 조정, offload 옵션 조정, activation checkpointing 등의 기법을 통해 메모리 사용량을 추가적으로 최적화할 수 있습니다.

4. Real-world Use Case / Example

실제 LLM 학습 시나리오에서, 초기에는 배치 사이즈를 16으로 설정했을 때 OOM (Out of Memory) 오류가 발생했습니다. Gradient Accumulation을 4로 설정하여 효과적으로 배치 사이즈를 64로 늘린 결과, OOM 오류를 해결하고 훈련 속도를 15% 향상시킬 수 있었습니다. ZeRO stage 2를 활성화하고 optimizer를 CPU로 오프로드하여 GPU 메모리 사용량을 추가적으로 줄일 수 있었습니다. Activation checkpointing 기법을 통해 메모리 footprint를 더욱 줄여 더 큰 모델을 학습할 수 있게 되었습니다.

5. Pros & Cons / Critical Analysis

  • Pros:
    • GPU 메모리 제한 극복: 작은 메모리 용량으로도 큰 배치 사이즈의 효과를 얻을 수 있습니다.
    • 훈련 속도 향상: 작은 배치 사이즈로 인한 GPU 활용률 저하를 방지하고 훈련 속도를 향상시킵니다.
    • 코드 변경 최소화: 기존 코드를 크게 수정하지 않고 DeepSpeed를 통합할 수 있습니다.
  • Cons:
    • Hyperparameter 튜닝 필요: `gradient_accumulation_steps` 값을 적절하게 설정해야 최적의 성능을 얻을 수 있습니다.
    • 훈련 시간 증가 가능성: 과도한 Gradient Accumulation은 업데이트 빈도를 줄여 수렴 속도를 늦출 수 있습니다.
    • 메모리 최적화 추가 작업 필요: 단순히 Gradient Accumulation만 사용하는 것만으로는 충분하지 않을 수 있으며, ZeRO, Offload 등의 추가적인 메모리 최적화 기술을 함께 적용해야 합니다.

6. FAQ

  • Q: Gradient Accumulation은 항상 훈련 속도를 향상시키나요?
    A: 항상 그런 것은 아닙니다. `gradient_accumulation_steps`가 너무 크면 업데이트 빈도가 줄어 수렴 속도가 늦어질 수 있습니다. 적절한 값을 찾아야 합니다.
  • Q: ZeRO stage는 어떤 것을 선택해야 하나요?
    A: stage 1은 optimizer 상태를 분산하고, stage 2는 optimizer 상태와 기울기를 분산하며, stage 3는 optimizer 상태, 기울기, 모델 파라미터를 모두 분산합니다. 메모리 용량이 매우 제한적인 경우 stage 3를 사용하는 것이 좋지만, 통신 오버헤드가 증가할 수 있습니다. 일반적으로 stage 2가 좋은 절충안입니다.
  • Q: DeepSpeed Profiler는 어떻게 사용하나요?
    A: `ds_config.json` 파일에 프로파일링 설정을 추가하고, DeepSpeed 엔진 초기화 시 `profiling=True` 옵션을 설정합니다. 훈련이 완료되면 지정된 `output_folder`에 프로파일링 결과가 저장됩니다. TensorBoard와 같은 도구를 사용하여 결과를 분석할 수 있습니다.

7. Conclusion

DeepSpeed의 Gradient Accumulation은 초거대 모델 학습 시 메모리 제약을 극복하고 훈련 효율성을 높이는 강력한 기술입니다. 하지만 단순히 적용하는 것만으로는 충분하지 않으며, ZeRO, Offload, Activation Checkpointing 등 다양한 메모리 최적화 기법과 함께 사용해야 합니다. 제시된 단계별 가이드와 팁을 활용하여 여러분의 모델 학습 파이프라인을 최적화하고, 더 크고 강력한 모델을 학습해 보세요. 지금 바로 DeepSpeed를 설치하고, `ds_config.json` 파일을 조정하여 여러분의 환경에 맞는 최적의 설정을 찾아보십시오!