PyTorch 훈련 중 GPU 메모리 부족 문제 해결을 위한 빌드즈-GPU 활용법

소개

buildz-gpu는 Python 기반의 GPU 메모리 최적화 도구로, 특히 PyTorch 모델 훈련 시 발생하는 GPU 캐시 메모리 과다 사용 문제를 완화하기 위해 설계되었습니다. 이 도구는 일부 계산 결과를 저장하지 않고 필요 시 재계산함으로써 메모리 사용량을 줄이는 방식인 recomputation (재계산) 전략을 자동화하여, 제한된 VRAM 환경에서도 큰 모델을 실행할 수 있게 해줍니다.

설치 방법

pip install buildz-gpu

핵심 기능: recals_with_rngs

주요 인터페이스는 buildz.gpu.torch.recals_with_rngs 함수입니다. 이 함수는 지정된 메모리 한도 내에서 모델의 활성화(activation) 캐싱을 관리합니다.

파라미터 설명

  • cache_limit_size: 허용할 최대 캐시 크기(byte 단위). 예: 4 * (1024**3)은 4GB에 해당.
  • gpu_devices: 난수 상태(rnd state)를 보존할 GPU 장치 지정. 다음 값이 가능:
    • "all": 기본값. 모든 GPU 장치의 난수 상태를 추적.
    • "none": 난수 상태를 저장하지 않음. Dropout 등 확률적 연산이 없을 경우 사용 가능.
    • int: 특정 GPU 인덱스 (예: 0).
    • List[int]: 여러 GPU 인덱스 목록 (예: [0, 1]).
    • torch.Tensor: 해당 텐서가 위치한 장치를 자동 감지.

이 함수는 Recals 객체를 반환하며, 이를 통해 모델 호출 시 메모리 관리를 수행합니다.

작동 원리

  1. 정방향 전파(forward pass) 중, 캐시 용량이 설정된 한도를 초과하면 가장 오래된 캐시 항목이 삭제됩니다.
  2. 역방향 전파(backward pass) 시, 필요한 캐시가 존재하지 않으면 해당 시점까지 정방향 계산을 다시 수행(recompute)하여 누락된 값을 복원합니다.
  3. 이때 PyTorch의 난수 상태(torch.cuda.get_rng_state())를 사전에 저장하고 복원함으로써, dropout 등 무작위성이 포함된 연산의 일관성을 유지합니다.

※ 참고: PyTorch 자체에도 torch.utils.checkpoint 기능이 유사한 목적을 제공하지만, 수동으로 체크포인트 레이어를 지정해야 하며 사용 난이도가 높습니다. 반면 buildz-gpu는 자동화 및 용량 기반 제어가 가능해 접근성이 더 높습니다.

사용 예시

1. 메모리 요구량 분석 (훈련 전)

실제 훈련 없이 모델의 캐시 사용 패턴을 분석할 수 있습니다.

from buildz.gpu.torch import recals_with_rngs

# 캐시 제한 없이 분석 모드 실행 (-1)
recals = recals_with_rngs(-1)

model.train()
total_size, max_unit_size = recals.analyze(model, inputs)

print(f"총 캐시 크기: {total_size / (1024**3):.2f} GB")
print(f"단일 레이어 최대 크기: {max_unit_size / (1024**3):.2f} GB")

이 정보를 바탕으로 적절한 cache_limit_size 값을 설정할 수 있습니다.

2. 기본 훈련 코드 적용

기존의 model(input) 호출을 recals(model, input)으로 대체합니다.

from buildz.gpu.torch import recals_with_rngs

# 4GB 메모리 제한 설정
recals = recals_with_rngs(4 * (1024**3))
model.train()

for data, label in dataloader:
    optimizer.zero_grad()
    
    # 기존: output = model(data)
    output = recals(model, data)
    
    loss = criterion(output, label)
    loss.backward()
    optimizer.step()

3. 멀티 디바이스 환경 적용

여러 GPU 또는 CPU와 혼합된 파이프라인에서 각 GPU별로 독립적인 캐시 제한을 설정할 수 있습니다.

# 각 디바이스별 Recals 인스턴스 생성
recals_a = recals_with_rngs(6 * (1024**3), device_a)  # GPU A: 6GB 제한
recals_b = recals_with_rngs(4 * (1024**3), device_b)  # GPU B: 4GB 제한

for data, target in dataloader:
    data = data.to(device_a)
    target = target.to(device_c)

    # GPU A에서 두 번 호출 → 동일한 recals_a 사용
    mid_a0 = recals_a(model_cuda_a0, data)
    mid_cpu = model_cpu(mid_a0.cpu())
    mid_a1 = recals_a(model_cuda_a1, mid_cpu.to(device_a))

    # GPU B에서 호출
    mid_b = recals_b(model_cuda_b, mid_a1.to(device_b))

    # 최종 출력은 일반 호출 (또는 다른 제어 필요 시 추가 설정)
    output = model_cuda_c(mid_b.to(device_c))

    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

각 Recals 객체는 자신이 호출한 연산들에 대해서만 캐시를 집계하고 관리하므로, 디바이스 간 분리된 제어가 가능합니다.

성능 고려사항

  • 캐시 한도를 너무 작게 설정하면 재계산 빈도가 증가하여 훈련 시간이 크게 늘어납니다.
  • 이론적으로, 전체 캐시의 50%만 허용하면 약 17% 계산 증가, 25% 허용 시 약 50% 계산 증가가 발생합니다.
  • 실제 네트워크는 각 레이어의 계산 비용이 다르므로, 최적의 균형점은 실험을 통해 결정해야 합니다.

태그: PyTorch gpu-memory-optimization recomputation checkpointing deep-learning-training

7월 19일 07:36에 게시됨

인기 태그

©2026 괴물 클럽