Hugging Face Trainer를 활용한 LLM 학습 시 평가 단계에서 GPU 메모리 부족 해결 방법

Hugging Face의 Trainer 클래스로 LLM(Large Language Model)을 학습할 때 자주 발생하는 문제 중 하나는 GPU 메모리 부족 현상입니다. 특히 학습 단계보다 평가(Evaluation) 단계에서 갑자기 GPU 메모리 사용량이 급격히 증가하며, 평가가 끝난 이후에도 메모리가 줄어들지 않고 계속 유지되어 결국 메모리 부족(OOM, Out of Memory)이 발생할 수 있습니다.

이러한 문제를 해결하기 위한 두 가지 주요 방법을 소개하겠습니다.

1. preprocess_logits_for_metrics 활용

preprocess_logits_for_metrics는 평가 단계에서 모델의 예측 결과인 logits를 사전 처리할 수 있도록 제공되는 파라미터입니다. 이를 통해 GPU 메모리를 효과적으로 관리할 수 있습니다.

아래와 같이 logits를 CPU로 이동시키도록 설정하면 GPU 메모리에 불필요한 텐서가 남지 않게 됩니다.

def preprocess_logits_for_metrics(logits, labels):
    return logits.detach().cpu()

이렇게 설정하면 평가 과정에서 logits가 GPU에서 즉시 CPU로 이동되어 GPU 메모리 사용을 줄일 수 있으며, 특히 compute_metrics 함수에서 logits를 사용하는 경우 GPU가 아닌 CPU에서 연산을 수행하게 되어 GPU 메모리 부담을 크게 줄일 수 있습니다.

2. eval_accumulation_steps 활용

평가 시 모든 예측 결과를 GPU에 쌓아두지 않고 일정 간격으로 CPU로 이동시키는 방법입니다. 이 방법은 특히 평가 데이터셋이 크거나, 모델이 크고 GPU 메모리가 제한적일 때 유용합니다.

예를 들어, 배치를 4로 설정하면 GPU는 평가 데이터를 4개의 미니배치 단위로 처리하게 됩니다. 각 미니배치에서 생성된 logits는 GPU 메모리에 임시로 저장되며, 4개의 미니배치 처리가 끝나면 logits는 GPU에서 CPU로 이동됩니다. 이렇게 하면 GPU 메모리에 과도한 데이터가 계속 쌓이는 것을 방지하고, 메모리 사용량을 효율적으로 관리할 수 있습니다.

 

 

 

구체적으로는 다음과 같이 설정할 수 있습니다:

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir='./results',
    eval_accumulation_steps=4,  # GPU에서 4개의 미니배치마다 CPU로 이동
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)

compute_metrics 함수는 평가 결과를 계산할 때 logits와 labels를 사용합니다. 이때 logits가 CPU로 옮겨진 상태에서 계산을 수행하면 GPU 메모리를 절약할 수 있습니다. 또한, eval_accumulation_steps를 통해 메모리 사용을 효과적으로 관리하면 GPU 메모리가 부족해 평가가 중단되는 현상을 예방할 수 있습니다.

결론

이 두 가지 방법을 조합하여 사용하면 Hugging Face의 Trainer로 LLM을 학습하거나 평가할 때 GPU 메모리를 효율적으로 관리할 수 있습니다. 특히 GPU 메모리 제한이 있는 환경에서는 이러한 설정이 필수적일 수 있습니다.