Wonjae (Dan) Kim김원재

Delete your account?계정을 지울까요?

This signs you out and permanently removes your account, sign-in sessions and coffee chat history.로그아웃되고, 계정과 로그인 기록, 커피챗 기록이 영구히 지워집니다.

Your comments are deleted too. A comment that already has replies is emptied and kept as a placeholder so the replies still make sense. Backups may keep a copy for up to 14 days.쓰신 댓글도 지워집니다. 이미 답글이 달린 댓글은 답글이 어색해지지 않도록 내용만 비우고 자리를 남깁니다. 백업에는 최대 14일 동안 사본이 남을 수 있습니다.

Christian Mills: 컨퍼런스 발표 12: PyTorch FSDP와 torchao로 메모리 부족(OOM) 잡기Christian Mills: Conference Talk 12: Slaying OOMs with PyTorch FSDP and torchao

모델 메모리가 파라미터, 그래디언트, 옵티마이저 상태, 활성값으로 어떻게 나뉘는지부터 4비트 양자화, QLoRA, 텐서 서브클래스, FSDP1과 FSDP2의 차이, 그리고 FSDP2 벤치마크에서 찾은 성능 격차와 해법까지 Mark Saroufim과 Jane Xu의 발표를 정리한 Christian Mills의 노트In this talk, Mark Saroufim and Jane Xu, discuss techniques and tools for mitigating Out of Memory (OOM) errors in PyTorch, specifically when working with LLMs.

· 10 min read10분 분량KO · translation한국어 · 번역

Contents목차

번역문 · 원저자: Christian Mills (발표: Mark Saroufim, Jane Xu) · 원문: Conference Talk 12: Slaying OOMs with PyTorch FSDP and torchao · 원문 라이선스: CC BY-NC-SA 4.0. 이 번역문도 원문과 같은 CC BY-NC-SA 4.0을 따릅니다. 원문의 그림 5개는 모두 발표자(Mark Saroufim, Jane Xu)가 만든 발표 슬라이드 화면입니다. 이 글의 라이선스가 슬라이드까지 덮는지 확인하지 못해 싣지 않았고, 그림 자리에는 설명만 남겼습니다. 그림은 원문에서 볼 수 있습니다.

원문: 「Conference Talk 12: Slaying OOMs with PyTorch FSDP and torchao」, Christian Mills, https://christianjmills.com/posts/mastering-llms-course-notes/conference-talk-012/, 2024년 7월 24일 게시. Mastering LLMs 강좌에서 Mark Saroufim과 Jane Xu가 한 발표를 정리한 노트입니다. 이 글은 원문을 한국어로 옮긴 번역문입니다.

라이선스: 원문 사이트 바닥글에 “콘텐츠는 CC BY-NC-SA 4.0, 코드 예시는 MIT 라이선스를 따른다”고 적혀 있습니다. 이 번역도 같은 CC BY-NC-SA 4.0 조건(저작자 표시, 비영리, 동일 조건 변경 허락)을 따르며, 코드는 원문 그대로 둡니다. 글 끝에 붙은 사이트 공통 자기소개와 상담 안내 문단은 본문이 아니어서 옮기지 않았습니다.

이 발표에서 Mark Saroufim과 Jane Xu는 PyTorch에서, 특히 LLM을 다룰 때 메모리 부족OOM 오류를 줄이는 기법과 도구를 이야기합니다.

발표 자료: Slaying OOMs

들어가며

  • 큰 PyTorch 모델을 다룰 때 OOM(메모리 부족) 오류는 흔히 부딪히는 문제입니다.
  • 배치 크기나 모델 크기를 줄이는 전통적인 해법은 효과에 한계가 있습니다.

모델 메모리 이해하기

  • VRAM 제약: 요즘 GPU의 VRAM은 한정되어 있어서(예: 3090과 4090은 24GB), 큰 모델을 학습할 때 어려움이 생깁니다.
  • 모델 메모리의 구성 요소:
    • 파라미터: 모델 가중치로, 보통 FP16(파라미터당 2바이트)으로 저장합니다. 파라미터 70억 개짜리 LLaMa 모델이라면 14GB입니다.
    • 그래디언트: 역전파 중에 계산되며, 파라미터와 같은 크기의 저장 공간이 필요합니다(LLaMa 7B라면 14GB가 더 듭니다).
    • 옵티마이저 상태: Adam 같은 옵티마이저는 정보를 따로 더 저장하는데, 흔히 파라미터의 두 배 크기입니다(LLaMa 7B라면 28GB).
    • 활성값: 모델 층들의 중간 출력입니다.
      • 크기를 어림하기가 더 어렵고, 배치 크기나 컨텍스트 길이 같은 요인에 따라 달라집니다.
      • 배치 크기와 컨텍스트 길이가 커지면 메모리 사용량의 대부분을 활성값이 차지하는 경향이 있습니다.
  • 예: 파라미터 70억 개짜리 LLaMa 모델을 전체 미세 조정하면 56GB(파라미터 14GB + 그래디언트 14GB + 옵티마이저 상태 28GB)를 쉽게 넘을 수 있어, 대부분의 소비자용 GPU의 VRAM 용량을 넘깁니다.
  • 포럼 글: How to measure memory usage from your model without running it? (모델을 돌려 보지 않고 메모리 사용량을 재는 방법)

메모리 사용량 줄이기

옵티마이저 메모리

  • 논문: Adam: A Method for Stochastic Optimization
  • SGD 같은 다른 옵티마이저는 메모리 부담이 적어 끌릴 수 있지만, 실제로는 효과가 좋아서 여전히 Adam이 압도적으로 많이 쓰입니다.
  • Adam을 바꾸기는 어렵고, 바꾼다고 해서 메모리를 크게 아끼지 못할 수도 있습니다.

파라미터 양자화

  • 4비트 양자화: 모델 파라미터의 정밀도를 FP16에서 INT4(파라미터당 반 바이트)로 낮춥니다.
  • 이 기법으로 메모리 사용량을 크게 줄일 수 있습니다. 예를 들어 파라미터 70억 개짜리 LLaMa 모델의 크기가 14GB에서 3.5GB로 줄어듭니다.
  • Torch Compile: 이 도구를 쓰면 파이썬 코드에서 바로 효율적인 양자화 커널을 만들 수 있어서, CUDA 커널을 직접 짤 필요가 없어집니다.

    • quantize_tensor 함수와 dequantize_tensor 함수에 @torch.compile() 데코레이터를 붙입니다.

    • 예:

      import os
os.environ["TORCH_LOGS"] = "output_code"
import torch

@torch.compile()
def quantize_tensor(x_fp32):
    absmax = torch.max(torch.abs(x_fp32))
    c = 127.0 / absmax
    x_int8 = torch.round(c * x_fp32).to(torch.int8)
    return x_int8, c

@torch.compile()
def dequantize_tensor(x_int8, c):
    x_fp32 = x_int8.to(torch.float32) / c
    return x_fp32

x_int8, c = quantize_tensor(torch.randn(10, device="cuda"))
x_fp32 = dequantize_tensor(x_int8, c)

    • 문서: https://pytorch.org/docs/stable/generated/torch.compile.html

    • 튜토리얼: Introduction to torch.compile

LoRA로 그래디언트 줄이기

  • 그래디언트를 곧바로 4비트로 양자화하면 수렴이 나빠집니다.
  • LoRA(Low-Rank Adaptation, 저랭크 적응): 대부분의 파라미터는 얼려 두고 작은 파라미터 묶음(어댑터)만 학습해서 이 문제를 비켜 갑니다.
  • QLoRA: LoRA와 파라미터 양자화를 함께 써서, 정확도를 해치지 않고 메모리를 크게 아끼는 기법입니다.

QLoRA의 어려움과 해법

  • QLoRA는 구현이 복잡해서 흔히 CUDA 커널을 직접 짜야 합니다(예: Tim Dettmers의 원래 구현은 CUDA 코드 4,000줄입니다).
    • 가중치가 int4가 아니라, 정규 분포에 더 가까운 NF4입니다.
    • NF4 텐서끼리는 행렬 곱을 할 수 없습니다.
      • 역양자화한 다음 행렬 곱을 해야 합니다.
    • 모든 값에 같은 최댓값을 쓸 수는 없습니다. 그러면 이상치에 너무 민감해집니다.
    • 양자화는 보통 블록 단위로 하고, 블록마다 따로 스케일을 둡니다.
    • QLoRA는 그 스케일까지 양자화합니다(이중 양자화).
  • Torch Compile로 단순하게 구현하기: PyTorch 개발자 Driss Guessous는 Torch Compile을 써서 파이썬 코드 약 900줄로 QLoRA를 구현했습니다.
  • 텐서 서브클래스: PyTorch의 텐서 서브클래스 기능을 쓰면 NF4(QLoRA에서 쓰는 형식) 같은 사용자 정의 데이터 타입을 만들 수 있어서, 양자화된 텐서를 더 효율적으로 표현하고 다룰 수 있습니다.

    • 예:

      return F.linear(input, weight.to(input.dtype))

    • 문서: torch.Tensor

    • subclass_zoo: PyTorch 텐서 서브클래스 예제를 여럿 모아 둔 저장소입니다.

FSDP로 모델 병렬화하기

데이터 병렬화

  • 데이터 병렬화: 배치를 여러 장치에 나눠 주고, 각 장치에 그래디언트, 모델 파라미터, 옵티마이저 상태의 사본을 하나씩 둡니다.
  • 데이터 병렬화만으로도 도움은 되지만, 아주 큰 모델에는 모자랄 수 있습니다.

[그림: 데이터 병렬화에서 장치마다 모델 파라미터, 그래디언트, 옵티마이저 상태 사본을 모두 들고 있어 메모리가 넘치는 모습을 그린 슬라이드]

데이터 병렬화에서의 OOM, 「Slaying OOMs」 슬라이드 19

FSDP

  • FSDP(Fully Sharded Data Parallel, 완전 샤딩 데이터 병렬): 모델 파라미터를(그리고 그에 따라 그래디언트와 옵티마이저 상태도) 여러 GPU에 나눠 두어서, 장치 하나의 메모리 용량을 넘는 모델도 학습할 수 있게 합니다.
    • 지금 처리하는 층에 해당하는 메모리는 그 층의 처리가 끝나면 해제됩니다.

[그림: FSDP에서 파라미터와 AdamW 상태를 GPU 0과 GPU 1에 나눠 두어(PARAMS 0, PARAMS 1) GPU마다 올라가는 몫이 데이터 병렬화 때보다 줄어든 모습을 그린 손그림 슬라이드]

FSDP 메모리, 「Slaying OOMs」 슬라이드 21

  • 샤딩: 모델 구성 요소를 작은 조각(샤드)으로 나눠 서로 다른 GPU에 둘 수 있게 하는 것입니다.
  • 올게더All Gather 연산: 모델 학습 중에 여러 GPU에서 필요한 데이터를 모아 올 때 씁니다.
  • 층:
    • PyTorch의 nn 모듈은 모두 더 작은 nn 모듈들로 이루어진 트리입니다.
    • 무엇을 하나의 “층”으로 다룰지는 사용자의 래핑 정책이 정합니다.
    • 무엇을 래핑하기로 했는지가 메모리 사용량에 영향을 줍니다.
      • 덩어리가 작을수록 한 번에 올게더해야 하는 메모리가 줄어듭니다.
  • CPU 오프로딩:

    • 파라미터를 CPU에 두었다가 순전파와 역전파를 계산할 때 GPU로 옮겨 올 수 있습니다.

    • 옵티마이저 업데이트는 CPU에서 하므로, 옵티마이저 상태도 CPU에 둡니다.

      [그림: FSDP에서 CPU 오프로딩을 켜면 파라미터와 옵티마이저 상태를 CPU 메모리에 두고 계산할 때만 GPU로 옮겨 GPU 메모리 사용량이 줄어드는 모습을 그린 슬라이드]

      FSDP CPU 오프로딩, 「Slaying OOMs」 슬라이드 25

FSDP1과 FSDP2

  • 목표: 올게더를 효율적으로 만드는 것
  • 제약: NCCL(NVIDIA Collective Communications Library)에서는 GPU마다 같은 크기의 텐서를 내놓아야 합니다.
  • FSDP1:

    • 샤딩하기 전에 모든 텐서를 평탄화해 이어 붙이기 때문에, 데이터 타입과 메타데이터가 서로 부딪칠 수 있습니다.

      [그림: FSDP1의 FlatParameter 방식에서 텐서 t1, t2, t3를 하나로 평탄화해 이어 붙인 뒤 GPU 수만큼 같은 크기로 잘라, t2가 두 GPU에 걸쳐 나뉘는 모습을 그린 슬라이드]

      FlatParam FSDP, 「Slaying OOMs」 슬라이드 28

      • t1과 t2의 앞부분이 텐서 하나로 합쳐집니다.
      • t2의 뒷부분은 t3와 함께 텐서 하나로 합쳐집니다.
      • t2가 두 GPU에 나뉩니다.
      • 그래서 t1, t2, t3가 모두 같은 dtype과 같은 메타데이터를 갖도록 강제됩니다.

    • 메모리가 비결정적으로 튈 수 있어서 메모리 관리가 어려워집니다.

  • FSDP2(파라미터별 FSDP):

    • 분산 텐서D-tensor를 도입해 텐서를 하나하나 따로 샤딩할 수 있게 하므로, 텐서마다 원래 데이터 타입과 메타데이터가 그대로 유지됩니다.

      [그림: FSDP2의 파라미터별 방식에서 텐서 t1, t2, t3를 각각 따로 GPU 수만큼 나눠 각 GPU가 모든 텐서의 한 조각씩을 들고 있는 모습을 그린 슬라이드]

      파라미터별 FSDP, 「Slaying OOMs」 슬라이드 30

    • 메모리 사용이 더 결정적이어서, 메모리 사용량을 예측하고 관리하기가 더 쉽습니다.

    • 특히 QLoRA 같은 기법과 함께 쓸 때, 더 유연하고 효율적인 학습 방식이 가능해집니다.

    • FSDP1에 비해 올게더할 때 복사가 더 필요합니다.

    • 포럼 글: FSDP & CUDACachingAllocator: an outsider newb perspective (FSDP와 CUDA 캐싱 할당자를 바깥 초보자의 눈으로 본 글)

CPU 오프로딩

  • CPU 메모리에 파라미터와 옵티마이저 상태를 두어 GPU의 메모리 부담을 더 줄입니다.
  • 옵티마이저 단계를 CPU로 넘겨서, GPU는 계산량이 많은 순전파와 역전파에 집중할 수 있게 합니다.

FSDP2 벤치마크

벤치마크 계획

  • 목표: 벤치마크를 돌려 FSDP2와 기준선(FSDP1과 bnb) 사이의 성능 격차를 찾고, FSDP2를 더 빠르게 할 방법을 살펴봅니다.
  • 하드웨어: Vast AI에서 빌린 NVIDIA 3090 GPU 두 장(소비자용, 각각 VRAM 24GB)
  • 기준선: Answer.ai의 train.py에 FSDP1과 bnb를 쓰고, 배치 크기는 8입니다.

첫 벤치마크와 어긋난 점

  • 벤치마크 환경:
    • Answer.ai의 train.py(명령줄로 설정) 
      • python train.py --model_name meta-llama/Llama-2-7b-hf --batch_size 8 --context_length 2048 --train_type
qLoRA --use_gradient_checkpointing True --reentrant_checkpointing True --dataset dummy --dataset_samples
48
    • TorchTune 레시피(YAML로 설정)
  • 첫 결과:
    • 배치 크기최대 메모리한 단계 실행 시간
      train.py812.88 GiB14.0s
      torchtune812.60 GiB16.5s
    • FSDP2가 최대 메모리 사용량은 더 나았습니다.
    • FSDP2가 실행 시간은 더 느렸습니다.
  • 트레이스 분석(Perfetto):

근본 원인 찾기

  • 근본 원인 1: 설정 차이
    • TorchTune 레시피는 출력 프로젝션에도 LoRA를 붙여서(층마다 저랭크 어댑터 두 개를 더함), train.py보다 파라미터가 64개 더 많았습니다.
  • 근본 원인 2: 텐서 서브클래스 디스패치 오버헤드
    • 분산 텐서(D-tensor)는 커널을 호출할 때마다 메타데이터를 벗겨 내야 해서 오버헤드가 생깁니다.
    • 이 오버헤드가 늘어난 파라미터 수 때문에 더 커져서, FSDP2에서 옵티마이저 단계가 눈에 띄게 느려졌습니다.

어긋난 점을 바로잡고 다시 벤치마크하기

  • 같은 조건으로 비교하기:
    • 벤치마크마다 데이터셋, 파라미터 수, 래핑 정책을 똑같이 맞췄습니다.
  • 조건을 맞춘 뒤의 결과:
    • FSDP2의 실행 시간은 여전히 더 느렸습니다.
    • 최대 메모리 사용량은 크게 나아졌습니다.
  • 추가 분석:
    • 남은 성능 격차의 주범으로 순전파, 역전파, 옵티마이저에 초점을 맞췄습니다.

성능 격차 깊이 들여다보기

  • 격차 1: 옵티마이저 단계가 느려짐
    • 문제: D-tensor 오버헤드 때문에 옵티마이저 단계가 3배 느려졌습니다.
    • 해법: Intel과 함께 디스패치 한 번에 벡터화까지 하는 퓨즈드 Adam 옵티마이저 커널을 만들어 8배 빨라졌습니다.
  • 격차 2: 올게더 연산이 더 커짐
    • 문제 1: FSDP2는 기준선과 달리 올게더 연산에 데이터(스칼라와 양자화 계수)를 더 많이 담았습니다.
    • 문제 2: LoRA를 쓰지 않기로 하면 FSDP2에서는 출력 프로젝션이 양자화되지 않아, 기준선보다 크기가 4배 컸습니다.
    • 해법:
      • 문제 1은 예상된 동작입니다.
      • 문제 2는 TorchTune이 다음 릴리스에서 고칠 예정입니다.
  • 격차 3: 역양자화 오버헤드
    • 문제: CUDA 커널에 넘기려고 NF4를 BF16으로 역양자화하는 FSDP2의 과정이, 기준선이 쓰는 전용 구현보다 덜 효율적이었습니다.
    • 해법: 더 빠른 역양자화를 위해 Torch Compile, 사용자 정의 Torch 커널, Triton 커널을 써 보는 방안을 살펴봅니다.
  • 격차 4: 서로 다른 RoPE 알고리즘
    • 문제: TorchTune과 기준선이 서로 다른 RoPE 알고리즘을 써서, SDPA 앞에서 하는 연산이 달랐습니다.
    • 해법: TorchTune이 더 다양한 RoPE 알고리즘을 고를 수 있게 할 예정입니다.
  • 격차 5: 통신과 계산 겹치기
    • 문제: FSDP2는 메모리 관리가 더 엄격하다 보니, CPU 오프로딩과 계산을 겹치는 부분의 비효율이 드러났습니다. 계산 작업이 통신 작업보다 훨씬 작아서 노는 시간이 생겼습니다.
    • 해법: 래핑 정책을 바꿔 층을 다르게 묶어서, 겹치기를 개선하고 노는 시간을 줄였습니다. 이 해법은 한 층 안에서 서로 다른 정밀도를 섞어 다룰 수 있는 FSDP2에서만 쓸 수 있습니다.