[CS336] Lecture2. PyTorch, Resource Accounting

딥러닝 모델을 학습할 때 보통 모델 구조나 성능에 집중하지만, 실제로는 메모리와 연산 자원(resource)이 훨씬 더 중요한 제약 조건이 됩니다. 특히 대규모 모델을 다루게 되면 이 모델이 주어진 자원으로 학습 가능한지 자체가 가장 먼저 해결해야 할 이슈입니다. 이번 글에서는 Stanford CS336 강의를 기반으로, PyTorch 관점에서 딥러닝 학습 시 어떤 요소들이 메모리와 연산량을 결정하는지를 정리해보았습니다. 단순한 개념 설명이 아니라, 실제로 모델을 학습할 때 반드시 알아야 하는 resource accounting 관점에 초점을 맞추려고 합니다.

 

- 딥러닝 학습 = 텐서가 차지하는 메모리 + 그 텐서들에 대해 수행하는 FLOPs 계산
- 학습 전에 메모리·연산량·시간을 계산하는 것 중요

 

- Tensor
    - Memory = 개수 × dtype 크기
    - tensor는 딥러닝의 기본 단위이며 모든 것은 tensor로 표현됨
    - 메모리는 값의 개수와 dtype에 의해 결정됨
    - dtype(float32, bf16 등)는 메모리와 연산 성능에 직접 영향
    - view는 메모리를 공유하고, clone/contiguous는 새로운 메모리를 생성

- dtype: float32, float16, bfloat16(Brain Floating Point), fp8
    - dtype은 각 값이 몇 byte로 저장되는지를 의미
    - float32는 안정적이지만 메모리 사용량이 큼
    - float16, bf16, fp8은 메모리를 줄이지만 정밀도와 안정성 trade-off 존재
    - precision을 낮출수록 메모리와 속도는 좋아지지만 학습 안정성은 떨어질 수 있음
    - 해결 방법은 mixed precision을 사용해 안정성과 효율을 동시에 확보하는 것

- Model
    - 학습 메모리는 파라미터만 있는 것이 아님
    - 학습 시에는 모델 파라미터 외에도 여러 텐서가 더 붙는다는 것이 중요
    - 구성요소: 파라미터, activation, gradient, optimizer state
    - Parameters: 모델 가중치
    - Activations: forward 중간 결과, backward 때 필요해서 저장
    - Gradients: 각 파라미터에 대한 미분값
    - Optimizer states: Adam/AdaGrad 같은 옵티마이저가 들고 있는 추가 통계량

- Optimizer
    - 옵티마이저는 파라미터만 업데이트하는 것이 아님
    - 각 파라미터마다 추가 state(m, v 등)를 저장
    - 따라서 모델이 커질수록 optimizer 메모리도 같이 증가
    - Adam 계열은 특히 메모리를 많이 사용하는 대표적인 경우

- Compute (FLOPs)
    - 딥러닝 연산량은 FLOPs로 측정됨
    - forward보다 backward가 더 많은 연산을 수행
    - 전체 학습 연산량은 대략 forward의 약 3배 수준
    - 학습 시간은 FLOPs와 GPU 성능, MFU에 의해 결정됨

- MFU (Model FLOPs Utilization)
    - GPU의 최대 연산 성능 대비 실제 사용 비율
    - MFU가 낮으면 GPU를 제대로 활용하지 못하는 상태
    - 데이터 이동, 작은 batch, 비효율적인 연산 등이 원인

- Memory Bottleneck
    - 실제 메모리 병목은 파라미터보다 activation에서 발생하는 경우가 많음
    - activation은 layer마다 저장되기 때문에 빠르게 증가
    - batch size, sequence length가 커질수록 메모리 사용량 급증

- Training Loop
    - 흐름: forward → loss → backward → optimizer step → zero_grad
    - forward에서 activation 저장
    - backward에서 gradient 계산
    - optimizer step에서 state 활용
    - 각 단계마다 메모리 사용 패턴이 다름

- Mixed Precision
    - 하나의 dtype만 사용하는 것이 아니라 여러 dtype을 혼합하여 사용
    - forward/activation은 저정밀(fp16/bf16), 파라미터/optimizer는 고정밀(fp32) 유지
    - 메모리 절약과 학습 안정성을 동시에 확보하는 방법