서진호

4 분 소요

모델 훈련 속도와 메모리 사용 효율성을 향상시키기 위한 성능 최적화 기법을 이해하려면, 훈련 중 GPU가 어떻게 활용되는지와 연산 종류에 따라 계산 집약도가 어떻게 달라지는지를 이해하는 것이 도움이 된다. GPU 활용 예시와 모델 훈련 과정을 통해 동기를 부여하는 사례를 살펴보겠다.

1. 모델 학습 구조의 시작

  • 메모리 및 계산 효율성을 높이기 위해 큰 행렬이나 데이터 블록을 작은 조각(타일, tile)으로 나누어 처리하는 최적화 기법
pip install transformers datasets accelerate nvidia-ml-py3

  • nvidia-ml-py3 라이브러리는 Python 내부에서 모델의 메모리 사용량을 모니터링할 수 있게 해준다.

  • 터미널에서 사용하는 nvidia-smi 명령어에 익숙할 수도 있는데, 이 라이브러리를 사용하면 동일한 정보를 Python 코드에서 직접 확인할 수 있다.

import numpy as np
from datasets import Dataset


seq_len, dataset_size = 512, 512
dummy_data = {
    "input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),
    "labels": np.random.randint(0, 2, (dataset_size)),
}
ds = Dataset.from_dict(dummy_data)
ds.set_format("pt")
  • GPU 활용도 및 Trainer를 사용한 훈련 과정에 대한 요약 통계를 출력하기 위해, 두 개의 ㅗHelper 함수를 정의해야 함
from pynvml import *


def print_gpu_utilization():
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU memory occupied: {info.used//1024**2} MB.")


def print_summary(result):
    print(f"Time: {result.metrics['train_runtime']:.2f}")
    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
    print_gpu_utilization()
  • 먼저 GPU 메모리가 비어 있는 상태인지 확인해보자!
print_gpu_utilization()
  • 아직 어떤 모델도 로드하지 않았기 때문에 GPU 메모리가 예상대로 점유되지 않은 상태. 만약 사용 중인 시스템에서 그렇지 않다면, GPU 메모리를 사용하는 모든 프로세스를 종료했는지 확인하라.
  • GPU 메모리가 모두 사용자에게 할당 가능한 것은 아님. 모델이 GPU에 로드될 때, 커널도 함께 로드되며 이 과정에서 보통 1~2GB 정도의 메모리를 차지할 수 있음. 이 메모리 사용량을 확인하기 위해, 아주 작은 텐서를 GPU에 로드하여 커널 로딩을 유도함.
import torch


torch.ones((1, 1)).to("cuda")
print_gpu_utilization()

2. 모델 불러오기

  • 먼저 google-bert/bert-large-uncased 모델을 불러오기함. 모델 가중치를 GPU에 직접 불러와서, 가중치만이 차지하는 메모리 용량이 얼마나 되는지 확인함.
from transformers import AutoModelForSequenceClassification


model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-large-uncased").to("cuda")
print_gpu_utilization()
  • 모델의 가중치만으로도 약 1.3GB의 GPU 메모리를 차지하는 것을 확인함 (이 수치는 사용 중인 GPU의 종류에 따라 달라짐)
  • 최신 GPU에서는 모델이 더 빠르게 동작하도록 최적화된 방식으로 가중치를 로드하기 때문에, 더 많은 메모리를 사용할 수도 있다.
  • nvidia-smi 명령어를 사용할 때와 동일한 결과가 나오는지도 빠르게 확인해 볼 수 있음
nvidia-smi

Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P0    39W / 300W |   2631MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      3721      C   ...nvs/codeparrot/bin/python     2629MiB |
+-----------------------------------------------------------------------------+
  • 앞서 확인한 것과 동일한 수치를 얻었고, 현재 사용 중인 GPU가 메모리 16GB를 가진 V100이라는 것도 확인할 수 있음.
  • 이제 모델 훈련을 시작하고 GPU 메모리 사용량이 어떻게 변하는지 살펴보겠다.
  • 먼저, 몇 가지 표준 훈련 파라미터를 설정함:
default_args = {
    "output_dir": "tmp",
    "eval_strategy": "steps",
    "num_train_epochs": 1,
    "log_level": "error",
    "report_to": "none",
}
  • 여러 실험을 실행할 계획이라면, 실험 간에 메모리를 제대로 정리하기 위해 각 실험 사이마다 Python 커널을 재시작함

3. 기본 훈련 시 메모리 사용량

  • GPU 성능 최적화 기법을 전혀 사용하지 않고, 배치 크기(batch size)를 4로 설정한 상태에서 Trainer를 사용해 모델을 훈련시켜 보겠다.
from transformers import TrainingArguments, Trainer, logging

logging.set_verbosity_error()


training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)

Time: 57.82
Samples/second: 8.86
GPU memory occupied: 14949 MB.
  • 비교적 작은 배치 크기임에도 불구하고 GPU의 메모리가 거의 꽉 차는 것을 확인함
  • 일반적으로 더 큰 배치 크기는 모델의 수렴 속도를 높이거나 더 나은 최종 성능을 가져오는 경우가 많음
  • GPU의 제약이 아닌, 모델의 요구에 맞춰 배치 크기를 조정하는 것이 바람직함
  • 모델 자체의 크기보다 훨씬 더 많은 메모리를 사용하고 있다는 것
  • 모델의 연산 방식과 메모리 요구 사항

4. 모델 연산의 구조

  • 텐서 수축 (Tensor Contractions): 선형 계층(Linear layers)과 다중 헤드 어텐션(Multi-Head Attention)의 구성 요소들은 모두 배치된 행렬-행렬 곱셈(Batched Matrix-Matrix Multiplication)을 수행함. 이 연산들은 Transformer 훈련에서 가장 연산 집약적인(compute-intensive) 부분.
  • 통계적 정규화 (Statistical Normalizations): Softmax와 Layer Normalization은 텐서 수축보다는 덜 연산 집약적이며, 하나 이상의 축에 대해 값을 줄이는 reduction 연산을 포함하며, 이 결과는 이후 map 연산을 통해 적용됨.
  • 원소별 연산 (Element-wise Operators): Bias 추가, Dropout, 활성화 함수(activation), 잔차 연결(residual connections) 등이 여기에 해당함. 이 연산들은 연산량이 가장 적으며 compute cost가 낮다.
  • 모델의 성능 병목 현상(performance bottlenecks)을 분석할 때 매우 유용: Data Movement Is All You Need: A Case Study on Optimizing Transformers

5. 모델 메모리의 구조

  • GPU 메모리 사용 구성 요소

    • 모델 가중치 (Model Weights)
    • 옵티마이저 상태 (Optimizer States)
    • 그래디언트 (Gradients)
    • 그래디언트 계산을 위한 Forward Activation 값
    • 임시 버퍼 (Temporary Buffers)
    • 기능 특화 메모리 (Functionality-specific Memory)

  • 메모리 요구량 (혼합 정밀도 기준 - float16과 float32)

    • 혼합 정밀도 + AdamW 옵티마이저로 훈련 시: → 모델 파라미터당 18 바이트 + 활성화 메모리
    • 혼합 정밀도 추론 시:→ 모델 파라미터당 6바이트 + 활성화 메모리 (옵티마이저 상태 및 그래디언트 없음)

  • 세부 메모리 요구량

    • 모델 가중치 (Model Weights)
      • fp32 훈련: 4 bytes * 파라미터 수
      • 혼합 정밀도 훈련: 6 bytes * 파라미터 수 (fp32와 fp16 버전을 동시에 유지)
    • 옵티마이저 상태 (Optimizer States)
      • AdamW (일반): 8 bytes * 파라미터 수 (2개의 상태 유지)
      • 8-bit 옵티마이저 (bitsandbytes 등): 2 bytes * 파라미터 수
      • SGD + Momentum: 4 bytes * 파라미터 수 (1개의 상태만 유지)
    • 그래디언트 (Gradients)
      • fp32 또는 혼합 정밀도 공통: 4 bytes * 파라미터 수 (그래디언트는 항상 fp32)
    • Forward 활성화 값 (Activations)
      • 크기는 다음 요인에 따라 결정됨: → 시퀀스 길이, hidden size, 배치 크기
      • forward/backward 함수에서 입출력 및 그래디언트 계산을 위해 저장됨.

  • 임시 메모리 (Temporary Memory)

    • 계산 도중에만 존재하는 다양한 임시 변수들이 존재함
    • 계산이 끝나면 해제되지만, 그 순간에는 OOM(메모리 초과)를 일으킬 수 있음
    • 코딩 시 불필요한 변수는 가능한 한 빨리 del하거나 torch.cuda.empty_cache()를 사용하는 등 전략적 메모리 관리가 중요함

  • 기능 특화 메모리 (Functionality-specific Memory)

    • 소프트웨어 기능에 따라 추가적인 메모리가 필요 ->
    • 예: 빔 서치(Beam Search)를 활용한 텍스트 생성 시, 여러 복사본의 입력 및 출력 상태를 동시에 유지해야 함.

  • Forward vs Backward 속도 차이

    • 선형/합성곱 계층: → Backward는 Forward보다 약 2배 FLOPs (연산량), 속도도 일반적으로 2배 느림
    • Activation 연산: → 보통 메모리 대역폭에 의존. → Backward에서는 더 많은 데이터를 읽음 예:
      • Forward: 1회 읽기, 1회 쓰기
      • Backward: 2회 읽기 (gradOutput, forward output), 1회 쓰기 (gradInput)

댓글남기기