~13 min · inference-mode, autograd, memory, oom, war-story
Level 0툴 빌려 쓰는 사람
0 XP0/36 lessons0/12 achievements
0/100 XP to next level100 XP to go0% complete
"2 GB 면 되는 forward 전용 경로가 75 를 썼어. 모델이 큰 게 아니었어. framework 가 조용히 숙제를 챙겨 두고 있었던 거야."
말이 안 되던 크래시
평범한 이미지 모델이, 이미지 한 장 만드는 중에, 머신을 통째로 끌어내렸어. 그것도 두 번. 숫자가 말이 안 됐어. 2 GB 쯤이면 들어가야 할 경로가 화면 그리는 프로그램까지 쓰러지기 직전에 75 GB 근처까지 부풀었거든. 모델 크기로는 아무것도 설명이 안 됐어. 모델도 작고 이미지도 작았어. 안 보이는 뭔가가 메모리를 수십 기가바이트씩 먹고 있었던 거야.
안 볼 때 autograd 가 하는 일
범인은 자동 미분 엔진이었어. 딥러닝 framework 는 기본적으로 네가 학습을 시킬 수도 있다고 보거든. 그래서 forward pass 를 돌 때마다 모든 층의 중간 activation 을 조용히 붙들고 있어. backward 를 돌 때 gradient 를 구하려면 필요하니까. 학습할 때는 그게 꼭 있어야 해. 그런데 inference 에서는 backward 를 부를 일이 없으니까 그건 그냥 죽은 짐이야. 그리고 timestep 이 많은 깊은 diffusion 모델에서는 그 죽은 짐이 아무도 안 쓸 activation 수십 기가바이트로 쌓여.
framework 의 기본값은 제일 까다로운 쓰임새를 가정해. 딥러닝 framework 는 '이걸 학습시킬 수도 있다' 를 기본으로 잡아. 학습이 챙길 게 제일 많으니까. 그런데 inference 만 하는 경로에서는 그 기본값이 안 보이는 낭비가 돼. 네 도구가 기본으로 뭘 가정하는지 알아 둬. 그 기본값이 네 경우에 딱 맞는 일은 드무니까.
한 줄짜리 해법
해법은 framework 한테 이 경로는 gradient 가 영영 필요 없다고 대놓고 말해 주는 거야. forward pass 를 inference 또는 no-grad 구간으로 감싸면 돼. 그러면 엔진이 activation 붙들기를 그만두고, 75 GB 가 실제로 필요한 2 GB 언저리로 내려앉아. context manager 하나에 메모리가 자릿수 단위로 달라진 거야. 버그는 모델에도 수식에도 없었어. 의도를 안 밝힌 게 문제였지.
제일 큰 승리는 최적화가 아니라 선언에서 나올 때가 많아. 알고리즘도 그대로고 모델도 안 줄었어. 문장 하나가 framework 한테 이 코드가 실제로 뭘 하는지 알려 줬고, framework 가 해당 없는 작업을 그만뒀을 뿐이야. 도구한테 네 의도를 말해 주는 게 웬만한 똑똑한 최적화보다 값어치 있을 때가 있어.
왜 머신 전체가 내려갔냐면
메모리를 통으로 같이 쓰는 Apple Silicon 에서는 GPU 랑 시스템 메모리가 한 웅덩이를 나눠 써. 그래서 폭주하는 inference 할당이 Python 프로세스만 죽이는 게 아니야. 운영체제까지 굶겨서 화면을 그리는 프로그램이 같이 내려가. 실패가 그렇게 요란했던 이유가 이거야. 실수를 받아 줄 별도의 VRAM 이 없거든. 통합 메모리는 큰 모델한테는 근사하고, 뭐 하나 새기 시작하면 가차 없어.
같이 쓰는 자원은 국지적인 버그를 전체 크래시로 키워. GPU 랑 OS 가 같은 메모리 웅덩이에서 퍼 쓰면 inference 쪽 누수가 갇히지 않고 전부를 끌어내려. 자원을 같이 쓰는 시스템에서는 다른 데서라면 혼자 죽고 말 버그가 시스템 전체 장애가 돼. 같이 쓰는 자원은 머신 전체가 거기 매달려 있다고 치고 예산을 잡아. 실제로 그러니까.
피파의 고백
난 inference_mode 를 백 번쯤 읽으면서도 있으면 좋은 것, 속도 조금 벌어 주는 것 정도로 다뤘어. 머신을 죽이던 75 GB 가 차분한 2 GB 로 돌아가는 걸 보고서야 생각이 바뀌었지. 그건 최적화가 아니야. 돌아가는 엔진이랑 가끔 컴퓨터를 통째로 죽이는 엔진의 차이야. 어떤 한 줄은 겉치레가 아니라 기둥이야. 그 뒤로 난 '재미없어 보이는' context manager 를 그냥 넘기지 않게 됐어.
Code
decorator 하나에 자릿수가 달라져·python
import torch
# 선언이 없으면: autograd 가 혹시 backward() 부를까 싶어 activation 을 다 붙들어.
# timestep 많은 깊은 diffusion forward pass 에서는 그게 수십 GB 로 쌓여.
# 실제 작업은 2 GB 쯤인데 75 GB 를 씀.
def infer_leaky(model, latents, conditioning, steps):
x = latents
for t in range(steps):
x = model.denoise(x, t, conditioning) # step 마다 activation 이 남아
return x
# 선언을 하면: 이 경로는 gradient 가 영영 필요 없다고 framework 한테 말해 줘.
# activation 을 안 붙들고, 75 GB 가 2 GB 언저리로 내려앉아.
@torch.inference_mode() # 기둥 노릇 하는 한 줄
def infer_clean(model, latents, conditioning, steps):
x = latents
for t in range(steps):
x = model.denoise(x, t, conditioning)
return x
깜빡할 수 없는 자리에 방어를·python
# 방어는 모듈마다가 아니라 ADAPTER 입구에 둬.
# 자기 no_grad 를 깜빡한 새 모델 계열도 그대로 덮여.
class LocalAdapter(Adapter):
@torch.inference_mode() # 지금 계열이든 나중 계열이든 전부 덮음
def _infer_sync(self, request):
backbone = self.load(request.model)
return self.sampler.sample(backbone, request)
# 모듈 쪽 no_grad 도 그대로 남겨 둬 -> 이중 안전망
아무 ML 코드나 잡고 — 네 것도 좋고 예제도 좋아 — inference 경로가 no-grad 나 inference mode 를 선언하는지 봐. 하고 있으면 머릿속에서 그걸 지우고 뭐가 쌓일지 따져 봐. 안 하고 있으면 그건 잠복한 메모리 폭탄이고. 어느 쪽이든 그 선언이 뭘 사 주는지 한 문장으로 적어 봐.
Hint
물어볼 건 늘 하나야. 이 경로가 언젠가 backward 를 부르나? 아니면 gradient 를 챙기는 건 순수한 낭비고, 그걸 끄는 선언은 선택이 아니라 필수야.
Progress
Progress is local-only — sign in to sync across devices.